diff options
| -rw-r--r-- | users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs | 33 |
1 files changed, 29 insertions, 4 deletions
diff --git a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs index 57daeb3aceef..a542f8c7b899 100644 --- a/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs +++ b/users/Profpatsch/my-prelude/src/Postgres/MonadPostgres.hs @@ -5,6 +5,7 @@ module Postgres.MonadPostgres where +import Arg import AtLeast (AtLeast) import Control.Exception ( Exception (displayException), @@ -26,7 +27,6 @@ import Data.List qualified as List import Data.Pool (Pool) import Data.Pool qualified as Pool import Data.Text qualified as Text -import Data.Time (getCurrentTime) import Data.Typeable (Typeable) import Database.PostgreSQL.Simple (Connection, FormatError, FromRow, Query, QueryError, ResultError, SqlError, ToRow) import Database.PostgreSQL.Simple qualified as PG @@ -48,7 +48,7 @@ import Pretty (showPretty) import Seconds import System.Exit (ExitCode (..)) import Tool -import UnliftIO (MonadUnliftIO (withRunInIO), bracket, hClose) +import UnliftIO (MonadUnliftIO (withRunInIO), bracket, hClose, mask_) import UnliftIO.Concurrent (forkIO) import UnliftIO.Process (ProcessHandle) import UnliftIO.Process qualified as Process @@ -429,7 +429,8 @@ data PgFormatProcess = PgFormatProcess { stdinHdl :: Handle, stdoutHdl :: Handle, stderrHdl :: Handle, - procHdl :: ProcessHandle + procHdl :: ProcessHandle, + startedAt :: Otel.Timestamp } initPgFormatPool :: (HasField "pgFormat" tools Tool) => tools -> IO PgFormatPool @@ -463,11 +464,34 @@ initPgFormatPool tools = do destroyPgFormatPool :: PgFormatPool -> IO () destroyPgFormatPool pool = Pool.destroyAllResources pool.pool +-- | Get the oldest resource from the pool, or stop if you find a resource that’s older than `cutoffPointMs`. +takeOldestResource :: PgFormatPool -> Arg "cutoffPointMs" Integer -> IO (PgFormatProcess, Pool.LocalPool PgFormatProcess) +takeOldestResource pool cutoffPointMs = do + now <- Otel.getTimestamp + mask_ $ do + a <- Pool.takeResource pool.pool + (putBack, res) <- go now [] a + -- make sure we don’t leak any resources we didn’t use in the end + for_ putBack $ \(x, xLocal) -> Pool.putResource xLocal x + pure res + where + mkMs ts = (ts & Otel.timestampNanoseconds & toInteger) `div` 1000_000 + go now putBack a@(a', _) = + if abs (mkMs now - mkMs a'.startedAt) > cutoffPointMs.unArg + then pure (putBack, a) + else + Pool.tryTakeResource pool.pool >>= \case + Nothing -> pure (putBack, a) + Just b@(b', _) -> do + if a'.startedAt < b'.startedAt + then go now (b : putBack) a + else go now (a : putBack) b + -- | Format the given SQL with pg_formatter. Will use the pool of already running formatters to speed up execution. runPgFormat :: PgFormatPool -> ByteString -> IO (T3 "exitCode" ExitCode "formatted" ByteString "stderr" ByteString) runPgFormat pool sqlStatement = do bracket - (Pool.takeResource pool.pool) + (takeOldestResource pool 200) ( \(a, localPool) -> do -- we always destroy the resource, because the process exited Pool.destroyResource pool.pool localPool a @@ -772,6 +796,7 @@ pgFormatStartCommandWaitForInput :: m PgFormatProcess pgFormatStartCommandWaitForInput tools = do do + startedAt <- Otel.getTimestamp (Just stdinHdl, Just stdoutHdl, Just stderrHdl, procHdl) <- Process.createProcess ( ( Process.proc |