[inferno-ml] Use connection pool (#149)

Currently, the entire server runs within `withConnect` and holds a
single `Connection` within its `Env`. If Postgres is unreachable for
some reason when the server starts, the process will hang until it's
killed by an exception (which takes quite some time). If the connection
fails while the server is running, the entire process will be brought
down. This switches the server to use a `Pool` from `resource-pool`
instead to create/hold connections on demand

inferno-ml-server/CHANGELOG.md

@@ -1,39 +1,42 @@
 # Revision History for `inferno-ml-server`
 
-## 2023.11.26
+## 2024.11.29
+* Use `Pool` to hold Postgres connections
+
+## 2024.11.26
 * Add creation date to models and versions
 
-## 2023.10.18
+## 2024.10.18
 * Add new testing route
 * Some improvements to model caching
 * Make `/status` not awful and confusing
 
-## 2023.9.27
+## 2024.9.27
 * Change entity DB representation to `numeric`
 
-## 2023.7.2
+## 2024.7.2
 * Use new `loadModel` primitive and pass model names to script evaluator
 
-## 2023.6.19
+## 2024.6.19
 * Save `BridgeInfo` to DB
 
-## 2023.6.5
+## 2024.6.5
 * Support linking multiple models to inference parameters
 
-## 2023.6.1
+## 2024.6.1
 * Add `resolution` to `InferenceParam`
 
-## 2023.5.29
+## 2024.5.29
 * Change representation of script inputs/outputs
 
-## 2023.5.22
+## 2024.5.22
 * Add support for tracking evaluation info
 
-## 2023.4.3
+## 2024.4.3
 * Add `terminated` column to DB types
 
-## 2023.3.26
+## 2024.3.26
 * Move to `inferno` repo
 
-## 2023.3.6
+## 2024.3.6
 * Initial release

inferno-ml-server/inferno-ml-server.cabal

@@ -1,6 +1,6 @@
 cabal-version:      2.4
 name:               inferno-ml-server
-version:            2023.11.26
+version:            2024.11.29
 synopsis:           Server for Inferno ML
 description:        Server for Inferno ML
 homepage:           https://github.com/plow-technologies/inferno.git#readme
@@ -78,6 +78,7 @@ library
     , plow-log-async
     , postgresql-simple
     , prettyprinter
+    , resource-pool
     , scientific
     , servant-client
     , servant-server

inferno-ml-server/src/Inferno/ML/Server.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE OverloadedRecordDot #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
 module Inferno.ML.Server
@@ -14,7 +15,6 @@ import Control.Monad.Except (ExceptT (ExceptT))
 import Control.Monad.Reader (ReaderT (runReaderT))
 import qualified Data.ByteString.Lazy.Char8 as ByteString.Lazy.Char8
 import Data.Proxy (Proxy (Proxy))
-import Database.PostgreSQL.Simple (withConnect)
 import Inferno.ML.Server.Inference
 import Inferno.ML.Server.Log
 import Inferno.ML.Server.Types
@@ -72,19 +72,20 @@ main = runServer =<< mkOptions
         mkSettings :: (Request -> Status -> Maybe Integer -> IO ()) -> Settings
         mkSettings logger =
           defaultSettings
-            & setPort (view #port cfg & fromIntegral)
+            & setPort (fromIntegral cfg.port)
             & setLogger logger
 
 runInEnv :: Config -> (Env -> IO ()) -> IO ()
-runInEnv cfg f = withRemoteTracer $ \tracer -> do
-  traceWith tracer $ InfoTrace StartingServer
-  withConnect (view #store cfg) $ \conn ->
-    f
-      =<< Env cfg conn tracer
-        <$> newMVar ()
-        <*> newEmptyMVar
-        <*> newManager defaultManagerSettings
-        <*> newIORef Nothing
+runInEnv cfg f =
+  withConnectionPool cfg.store $ \pool ->
+    withRemoteTracer $ \tracer -> do
+      traceWith tracer $ InfoTrace StartingServer
+      f
+        =<< Env cfg tracer pool
+          <$> newMVar ()
+          <*> newEmptyMVar
+          <*> newManager defaultManagerSettings
+          <*> newIORef Nothing
 
 infernoMlRemote :: Env -> Application
 infernoMlRemote env = serve api $ hoistServer api (`toHandler` env) server

inferno-ml-server/src/Inferno/ML/Server/Inference/Model.hs

@@ -29,8 +29,8 @@ import Inferno.ML.Server.Types
 import Inferno.ML.Server.Utils
   ( firstOrThrow,
     queryStore,
+    withConns,
   )
-import Lens.Micro.Platform
 import UnliftIO (MonadUnliftIO (withRunInIO))
 import UnliftIO.Exception (bracket)
 
@@ -62,7 +62,7 @@ getModelsAndVersions =
 -- with the number of bytes
 getModelVersionSizeAndContents :: Oid -> RemoteM (Integer, ByteString)
 getModelVersionSizeAndContents m =
-  view #store >>= \conn -> withRunInIO $ \r ->
+  withConns $ \conn -> withRunInIO $ \r ->
     withTransaction conn . r $ do
       size <- getModelVersionSize m
       bs <-

inferno-ml-server/src/Inferno/ML/Server/Types.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DeriveGeneric #-}
@@ -23,6 +24,7 @@ where
 import Control.Applicative (Alternative ((<|>)), asum, (<**>))
 import Control.Exception (Exception (displayException))
 import Control.Monad.Extra (whenM)
+import Control.Monad.IO.Class (liftIO)
 import Control.Monad.Reader (ReaderT)
 import Data.Aeson
   ( FromJSON (parseJSON),
@@ -47,6 +49,8 @@ import Data.Data (Typeable)
 import Data.Generics.Labels ()
 import Data.Generics.Wrapped (wrappedTo)
 import Data.Map.Strict (Map)
+import Data.Pool (Pool)
+import qualified Data.Pool as Pool
 import Data.Scientific (Scientific)
 import Data.Text (Text)
 import qualified Data.Text as Text
@@ -59,6 +63,8 @@ import Database.PostgreSQL.Simple
   ( ConnectInfo (ConnectInfo),
     Connection,
     ResultError (ConversionFailed, UnexpectedNull),
+    close,
+    connect,
     (:.) ((:.)),
   )
 import Database.PostgreSQL.Simple.FromField
@@ -104,7 +110,8 @@ import Plow.Logging.Message
 import Servant.Client.Streaming (ClientError)
 import System.Posix.Types (EpochTime)
 import Text.Read (readMaybe)
-import UnliftIO (Async)
+import UnliftIO (Async, MonadUnliftIO)
+import UnliftIO.Exception (bracket)
 import UnliftIO.IORef (IORef)
 import UnliftIO.MVar (MVar)
 import Web.HttpApiData (FromHttpApiData, ToHttpApiData)
@@ -113,8 +120,8 @@ type RemoteM = ReaderT Env IO
 
 data Env = Env
   { config :: Config,
-    store :: Connection,
     tracer :: IOTracer RemoteTrace,
+    store :: Pool Connection,
     -- Lock for starting inference evaluation
     lock :: MVar (),
     -- The current inference evaluation job, if any
@@ -500,3 +507,18 @@ traceLevel = \case
   InfoTrace {} -> LevelInfo
   WarnTrace {} -> LevelWarn
   ErrorTrace {} -> LevelError
+
+-- | Create the connection pool for the DB
+newConnectionPool :: ConnectInfo -> IO (Pool Connection)
+#if MIN_VERSION_resource_pool(0,4,0)
+newConnectionPool ci = Pool.newPool $ Pool.defaultPoolConfig (connect ci) close 60 10
+#else
+newConnectionPool ci = Pool.newPool $ Pool.PoolConfig (connect ci) close 60 10
+#endif
+
+withConnectionPool ::
+  forall m a. MonadUnliftIO m => ConnectInfo -> (Pool Connection -> m a) -> m a
+withConnectionPool = flip bracket destroyPool . liftIO . newConnectionPool
+  where
+    destroyPool :: Pool Connection -> m ()
+    destroyPool = liftIO . Pool.destroyAllResources

inferno-ml-server/src/Inferno/ML/Server/Utils.hs

@@ -5,16 +5,19 @@ module Inferno.ML.Server.Utils
     firstOrThrow,
     queryStore,
     executeStore,
+    withConns,
   )
 where
 
 import Control.Monad (void)
 import Control.Monad.Catch (Exception, MonadThrow (throwM))
 import Control.Monad.IO.Class (liftIO)
 import Data.Generics.Labels ()
+import Data.Pool (withResource)
 import Data.Vector (Vector, (!?))
 import Database.PostgreSQL.Simple
-  ( FromRow,
+  ( Connection,
+    FromRow,
     Query,
     ToRow,
     execute,
@@ -23,18 +26,22 @@ import Database.PostgreSQL.Simple
 import Database.PostgreSQL.Simple.Vector (query)
 import Inferno.ML.Server.Types
 import Lens.Micro.Platform (view)
+import UnliftIO (MonadUnliftIO (withRunInIO))
 
 throwInfernoError :: forall e a. Exception e => Either e a -> RemoteM a
 throwInfernoError = either (throwM . InfernoError . SomeInfernoError) pure
 
 queryStore :: (ToRow b, FromRow a) => Query -> b -> RemoteM (Vector a)
-queryStore q x = view #store >>= \conn -> liftIO $ query conn q x
+queryStore q x = withConns $ \conn -> liftIO $ query conn q x
 
 executeStore :: ToRow a => Query -> a -> RemoteM ()
 executeStore q x =
-  view #store >>= \conn ->
+  withConns $ \conn ->
     liftIO . withTransaction conn . void $
       execute conn q x
 
 firstOrThrow :: (MonadThrow m, Exception e) => e -> Vector a -> m a
 firstOrThrow e = maybe (throwM e) pure . (!? 0)
+
+withConns :: (Connection -> RemoteM b) -> RemoteM b
+withConns f = view #store >>= \cs -> withRunInIO $ \r -> withResource cs $ r . f