Add pool stats to /metrics admin endpoint

nix/overlays/haskell-packages.nix

@@ -50,6 +50,15 @@ let
           }
           { };
 
+      resource-pool-fork-avanov =
+        prev.callHackageDirect
+          {
+            pkg = "resource-pool-fork-avanov";
+            ver = "0.2.4.0";
+            sha256 = "0y5hk4wi2n5xzdb11jvb9f8mh3lmycjfyxii81kl6s412ir5gpm5";
+          }
+          { };
+
       hasql-dynamic-statements =
         lib.dontCheck (lib.unmarkBroken prev.hasql-dynamic-statements);
 

postgrest.cabal

@@ -56,6 +56,7 @@ library
                       PostgREST.Logger
                       PostgREST.Middleware
                       PostgREST.OpenAPI
+                      PostgREST.Pool
                       PostgREST.Query.QueryBuilder
                       PostgREST.Query.SqlFragment
                       PostgREST.Query.Statements
@@ -86,7 +87,6 @@ library
                     , hasql                     >= 1.4 && < 1.5
                     , hasql-dynamic-statements  == 0.3.1
                     , hasql-notifications       >= 0.1 && < 0.3
-                    , hasql-pool                >= 0.5 && < 0.6
                     , hasql-transaction         >= 1.0.1 && < 1.1
                     , heredoc                   >= 0.2 && < 0.3
                     , http-types                >= 0.12.2 && < 0.13
@@ -103,6 +103,7 @@ library
                     , protolude                 >= 0.3 && < 0.4
                     , regex-tdfa                >= 1.2.2 && < 1.4
                     , retry                     >= 0.7.4 && < 0.10
+                    , resource-pool-fork-avanov >= 0.2.4 && < 0.2.5
                     , scientific                >= 0.3.4 && < 0.4
                     , swagger2                  >= 2.4 && < 2.7
                     , text                      >= 1.2.2 && < 1.3

src/PostgREST/Admin.hs

@@ -2,16 +2,18 @@ module PostgREST.Admin
   ( postgrestAdmin
   ) where
 
-import qualified Data.Text as T
+import           Data.Aeson as JSON
+import qualified Data.Text  as T
 
 import Network.Socket
 import Network.Socket.ByteString
 
 import qualified Network.HTTP.Types.Status as HTTP
 import qualified Network.Wai               as Wai
 
-import qualified Hasql.Pool    as SQL
-import qualified Hasql.Session as SQL
+import qualified Data.Pool      as ResourcePool
+import qualified Hasql.Session  as SQL
+import qualified PostgREST.Pool as SQL
 
 import qualified PostgREST.AppState as AppState
 import           PostgREST.Config   (AppConfig (..))
@@ -33,6 +35,9 @@ postgrestAdmin appState appConfig req respond  = do
       respond $ Wai.responseLBS (if isMainAppReachable && isConnectionUp && isSchemaCacheLoaded then HTTP.status200 else HTTP.status503) [] mempty
     ["live"] ->
       respond $ Wai.responseLBS (if isMainAppReachable then HTTP.status200 else HTTP.status503) [] mempty
+    ["metrics"] -> do
+      pStats <- SQL.stats (AppState.getPool appState)
+      respond $ Wai.responseLBS HTTP.status200 [] $ JSON.encode $ Metrics pStats
     _ ->
       respond $ Wai.responseLBS HTTP.status404 [] mempty
 
@@ -54,3 +59,16 @@ reachMainApp appConfig =
         connect sock $ addrAddress addr
         return sock
     sendEmpty sock = void $ send sock mempty
+
+newtype Metrics = Metrics ResourcePool.PoolStats
+
+instance JSON.ToJSON Metrics where
+    toJSON (Metrics (ResourcePool.PoolStats highwaterUsage currentUsage takes creates createFailures)) =
+      JSON.object [
+        "dbPoolStats" .= JSON.object
+          ["highwaterUsage"  .= highwaterUsage
+          ,"currentUsage"   .= currentUsage
+          ,"takes"          .= takes
+          ,"creates"        .= creates
+          ,"createFailures" .= createFailures]
+      ]

src/PostgREST/App.hs

@@ -31,14 +31,14 @@ import qualified Data.ByteString.Lazy            as LBS
 import qualified Data.HashMap.Strict             as M
 import qualified Data.Set                        as S
 import qualified Hasql.DynamicStatements.Snippet as SQL (Snippet)
-import qualified Hasql.Pool                      as SQL
 import qualified Hasql.Transaction               as SQL
 import qualified Hasql.Transaction.Sessions      as SQL
 import qualified Network.HTTP.Types.Header       as HTTP
 import qualified Network.HTTP.Types.Status       as HTTP
 import qualified Network.HTTP.Types.URI          as HTTP
 import qualified Network.Wai                     as Wai
 import qualified Network.Wai.Handler.Warp        as Warp
+import qualified PostgREST.Pool                  as SQL
 
 import qualified PostgREST.Admin                    as Admin
 import qualified PostgREST.AppState                 as AppState

src/PostgREST/AppState.hs

@@ -27,7 +27,7 @@ module PostgREST.AppState
   , waitListener
   ) where
 
-import qualified Hasql.Pool as SQL
+import qualified PostgREST.Pool as SQL
 
 import Control.AutoUpdate (defaultUpdateSettings, mkAutoUpdate,
                            updateAction)

src/PostgREST/CLI.hs

@@ -11,9 +11,9 @@ module PostgREST.CLI
 import qualified Data.Aeson                 as JSON
 import qualified Data.ByteString.Char8      as BS
 import qualified Data.ByteString.Lazy       as LBS
-import qualified Hasql.Pool                 as SQL
 import qualified Hasql.Transaction.Sessions as SQL
 import qualified Options.Applicative        as O
+import qualified PostgREST.Pool             as SQL
 
 import Data.Text.IO (hPutStrLn)
 import Text.Heredoc (str)

src/PostgREST/Config/Database.hs

@@ -9,11 +9,11 @@ import PostgREST.Config.PgVersion (PgVersion (..))
 
 import qualified Hasql.Decoders             as HD
 import qualified Hasql.Encoders             as HE
-import qualified Hasql.Pool                 as SQL
 import           Hasql.Session              (Session, statement)
 import qualified Hasql.Statement            as SQL
 import qualified Hasql.Transaction          as SQL
 import qualified Hasql.Transaction.Sessions as SQL
+import qualified PostgREST.Pool             as SQL
 
 import Text.InterpolatedString.Perl6 (q)
 

src/PostgREST/Error.hs

@@ -20,9 +20,9 @@ import qualified Data.ByteString.Char8     as BS
 import qualified Data.Text                 as T
 import qualified Data.Text.Encoding        as T
 import qualified Data.Text.Encoding.Error  as T
-import qualified Hasql.Pool                as SQL
 import qualified Hasql.Session             as SQL
 import qualified Network.HTTP.Types.Status as HTTP
+import qualified PostgREST.Pool            as SQL
 
 import Data.Aeson  ((.=))
 import Network.Wai (Response, responseLBS)

src/PostgREST/Pool.hs

@@ -0,0 +1,91 @@
+module PostgREST.Pool
+(
+  Pool,
+  Settings,
+  acquire,
+  release,
+  UsageError(..),
+  use,
+  stats,
+)
+where
+
+import qualified Data.Pool as ResourcePool
+import           Data.Time (NominalDiffTime)
+
+import qualified Hasql.Connection
+import qualified Hasql.Session
+
+import Protolude
+
+-- |
+-- A pool of connections to DB.
+newtype Pool =
+  Pool (ResourcePool.Pool (Either Hasql.Connection.ConnectionError Hasql.Connection.Connection))
+  deriving (Show)
+
+stats :: Pool -> IO ResourcePool.PoolStats
+stats (Pool p) = ResourcePool.poolStats <$> ResourcePool.stats p False
+
+-- |
+-- Settings of the connection pool. Consist of:
+--
+-- * Pool-size.
+--
+-- * Timeout.
+-- An amount of time for which an unused resource is kept open.
+-- The smallest acceptable value is 0.5 seconds.
+--
+-- * Connection settings.
+--
+type Settings =
+  (Int, NominalDiffTime, Hasql.Connection.Settings)
+
+-- |
+-- Given the pool-size, timeout and connection settings
+-- create a connection-pool.
+acquire :: Settings -> IO Pool
+acquire (size, timeout, connectionSettings) =
+  Pool <$> ResourcePool.createPool acq rel stripes timeout size
+  where
+    acq =
+      Hasql.Connection.acquire connectionSettings
+    rel =
+      either (const (pure ())) Hasql.Connection.release
+    stripes =
+      1
+
+-- |
+-- Release the connection-pool.
+release :: Pool -> IO ()
+release (Pool pool) =
+  ResourcePool.destroyAllResources pool
+
+-- |
+-- A union over the connection establishment error and the session error.
+data UsageError =
+  ConnectionError Hasql.Connection.ConnectionError |
+  SessionError Hasql.Session.QueryError
+  deriving (Show, Eq)
+
+-- |
+-- Use a connection from the pool to run a session and
+-- return the connection to the pool, when finished.
+use :: Pool -> Hasql.Session.Session a -> IO (Either UsageError a)
+use (Pool pool) session =
+  fmap (either (Left . ConnectionError) (either (Left . SessionError) Right)) $
+  withResourceOnEither pool $
+  traverse $
+  Hasql.Session.run session
+
+withResourceOnEither :: ResourcePool.Pool resource -> (resource -> IO (Either failure success)) -> IO (Either failure success)
+withResourceOnEither pool act = mask_ $ do
+  (resource, localPool) <- ResourcePool.takeResource pool
+  failureOrSuccess <- act resource `onException` ResourcePool.destroyResource pool localPool resource
+  case failureOrSuccess of
+    Right success -> do
+      ResourcePool.putResource localPool resource
+      return (Right success)
+    Left failure -> do
+      ResourcePool.destroyResource pool localPool resource
+      return (Left failure)

src/PostgREST/Workers.hs

@@ -12,8 +12,8 @@ import qualified Data.ByteString            as BS
 import qualified Data.ByteString.Lazy       as LBS
 import qualified Data.Text.Encoding         as T
 import qualified Hasql.Notifications        as SQL
-import qualified Hasql.Pool                 as SQL
 import qualified Hasql.Transaction.Sessions as SQL
+import qualified PostgREST.Pool             as SQL
 
 import Control.Retry    (RetryStatus, capDelay, exponentialBackoff,
                          retrying, rsPreviousDelay)