From f721f42bb9f5edf01db4953d15caf27f2f27274a Mon Sep 17 00:00:00 2001
From: Rory Tyler Hayford <52039264+ngua@users.noreply.github.com>
Date: Fri, 29 Nov 2024 14:57:51 +0700
Subject: [PATCH] [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                | 27 ++++++++++---------
 inferno-ml-server/inferno-ml-server.cabal     |  3 ++-
 inferno-ml-server/src/Inferno/ML/Server.hs    | 23 ++++++++--------
 .../src/Inferno/ML/Server/Inference/Model.hs  |  4 +--
 .../src/Inferno/ML/Server/Types.hs            | 26 ++++++++++++++++--
 .../src/Inferno/ML/Server/Utils.hs            | 13 ++++++---
 6 files changed, 65 insertions(+), 31 deletions(-)

diff --git a/inferno-ml-server/CHANGELOG.md b/inferno-ml-server/CHANGELOG.md
index 267d3e0..b1a7d93 100644
--- a/inferno-ml-server/CHANGELOG.md
+++ b/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
diff --git a/inferno-ml-server/inferno-ml-server.cabal b/inferno-ml-server/inferno-ml-server.cabal
index 4f8fc7a..11de856 100644
--- a/inferno-ml-server/inferno-ml-server.cabal
+++ b/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
diff --git a/inferno-ml-server/src/Inferno/ML/Server.hs b/inferno-ml-server/src/Inferno/ML/Server.hs
index b93629d..0b6126a 100644
--- a/inferno-ml-server/src/Inferno/ML/Server.hs
+++ b/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
diff --git a/inferno-ml-server/src/Inferno/ML/Server/Inference/Model.hs b/inferno-ml-server/src/Inferno/ML/Server/Inference/Model.hs
index 0256836..f7e306f 100644
--- a/inferno-ml-server/src/Inferno/ML/Server/Inference/Model.hs
+++ b/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 <-
diff --git a/inferno-ml-server/src/Inferno/ML/Server/Types.hs b/inferno-ml-server/src/Inferno/ML/Server/Types.hs
index ff7eda1..f0797a8 100644
--- a/inferno-ml-server/src/Inferno/ML/Server/Types.hs
+++ b/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
diff --git a/inferno-ml-server/src/Inferno/ML/Server/Utils.hs b/inferno-ml-server/src/Inferno/ML/Server/Utils.hs
index 7d0c9a1..0f6acf5 100644
--- a/inferno-ml-server/src/Inferno/ML/Server/Utils.hs
+++ b/inferno-ml-server/src/Inferno/ML/Server/Utils.hs
@@ -5,6 +5,7 @@ module Inferno.ML.Server.Utils
     firstOrThrow,
     queryStore,
     executeStore,
+    withConns,
   )
 where
 
@@ -12,9 +13,11 @@ 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