feat(db): track source of database accesses

2021-02-26 11:00:44 +01:00 · 2021-02-26 11:00:44 +01:00 · 23ff9d9222
commit 23ff9d9222
parent fd57bbce3d
14 changed files with 217 additions and 102 deletions
--- a/src/Application.hs
+++ b/src/Application.hs
@ -207,7 +207,7 @@ makeFoundation appSettings''@AppSettings{..} = do
    -- temporary foundation without a real connection pool, get a log function
    -- from there, and then create the real foundation.
    let
-      mkFoundation :: _ -> (forall m. MonadIO m => Custom.Pool m SqlBackend) -> _
+      mkFoundation :: _ -> (forall m. MonadIO m => Custom.Pool' m DBConnLabel DBConnUseState SqlBackend) -> _
      mkFoundation appSettings' appConnPool appSmtpPool appLdapPool appCryptoIDKey appSessionStore appSecretBoxKey appWidgetMemcached appJSONWebKeySet appClusterID appMemcached appUploadCache appVerpSecret appAuthKey = UniWorX {..}
      -- The UniWorX {..} syntax is an example of record wild cards. For more
      -- information, see:
@ -257,8 +257,8 @@ makeFoundation appSettings''@AppSettings{..} = do
              liftIO $ connClose conn
              observeDatabaseConnectionClosed
              $logDebugS "SqlPool" "Closed connection"
-         in Custom.createPool (liftIO . flip runLoggingT logFunc) create destroy (Just . fromIntegral $ pgPoolIdleTimeout appDatabaseConf) (Just $ pgPoolSize appDatabaseConf)
-      let sqlPool :: forall m. MonadIO m => Custom.Pool m SqlBackend
+         in Custom.createPool' (liftIO . flip runLoggingT logFunc) create destroy ((flip runLoggingT logFunc .) . onUseDBConn) onReleaseDBConn (Just . fromIntegral $ pgPoolIdleTimeout appDatabaseConf) (Just $ pgPoolSize appDatabaseConf)
+      let sqlPool :: forall m. MonadIO m => Custom.Pool' m DBConnLabel DBConnUseState SqlBackend
          sqlPool = Custom.hoistPool (liftIO . flip runLoggingT logFunc) sqlPool'
      void . Prometheus.register . poolMetrics PoolDatabaseConnections $ sqlPool @IO

@ -332,7 +332,7 @@ mkSessionStore :: forall m.
                  , MonadResource m
                  )
               => AppSettings
-               -> (forall m'. MonadIO m' => Custom.Pool m' SqlBackend)
+               -> (forall m'. MonadIO m' => Custom.Pool' m' DBConnLabel DBConnUseState SqlBackend)
               -> ReaderT SqlBackend m SomeSessionStorage
 mkSessionStore AppSettings{..} mcdSqlConnPool
  | Just mcdConf@MemcachedConf{..} <- appSessionMemcachedConf = do
--- a/src/Foundation/Authorization.hs
+++ b/src/Foundation/Authorization.hs
@ -95,13 +95,13 @@ data AccessPredicate
  | APDB       (ByteString -> (forall m. MonadAP m => AuthTagsEval m) -> Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> ReaderT SqlReadBackend (HandlerFor UniWorX) AuthResult)

 class (MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => MonadAP m where
-  evalAccessPred :: AccessPredicate -> ByteString -> (forall m'. MonadAP m' => AuthTagsEval m') -> Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> m AuthResult
+  evalAccessPred :: HasCallStack => AccessPredicate -> ByteString -> (forall m'. MonadAP m' => AuthTagsEval m') -> Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> m AuthResult

 instance {-# INCOHERENT #-} (MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => MonadAP m where
  evalAccessPred aPred contCtx cont aid r w = liftHandler $ case aPred of
    (APPure    p) -> runReader (p aid r w) <$> getMsgRenderer
    (APHandler p) -> p aid r w
-    (APDB      p) -> runDBRead $ p contCtx cont aid r w
+    (APDB      p) -> runDBRead' callStack $ p contCtx cont aid r w

 instance (MonadHandler m, HandlerSite m ~ UniWorX, BackendCompatible SqlReadBackend backend, BearerAuthSite UniWorX) => MonadAP (ReaderT backend m) where
  evalAccessPred aPred contCtx cont aid r w = mapReaderT liftHandler . withReaderT (projectBackend @SqlReadBackend) $ case aPred of
@ -180,7 +180,8 @@ getAuthContext = liftHandler $ do
  return authCtx

 isDryRun :: forall m.
-            ( MonadHandler m, HandlerSite m ~ UniWorX
+            ( HasCallStack
+            , MonadHandler m, HandlerSite m ~ UniWorX
            , BearerAuthSite UniWorX
            )
         => m Bool
@ -1533,7 +1534,7 @@ routeAuthTags = fmap predDNFEntail . ofoldM parse defaultAuthDNF . routeAttrs
      Just t' -> Right . predDNFOr prev . PredDNF $ Set.singleton t'
      Nothing -> Left $ InvalidAuthTag t

-evalAuthTags :: forall ctx m. (Binary ctx, MonadAP m) => ctx -> AuthTagActive -> (forall m'. MonadAP m' => AuthTagsEval m') -> AuthTagsEval m
+evalAuthTags :: forall ctx m. (HasCallStack, Binary ctx, MonadAP m) => ctx -> AuthTagActive -> (forall m'. MonadAP m' => AuthTagsEval m') -> AuthTagsEval m
 -- ^ `tell`s disabled predicates, identified as pivots
 evalAuthTags ctx authActive@AuthTagActive{..} cont (map (Set.toList . toNullable) . Set.toList . dnfTerms -> authDNF') mAuthId route isWrite
  = do
@ -1579,7 +1580,7 @@ evalAuthTags ctx authActive@AuthTagActive{..} cont (map (Set.toList . toNullable

      return result

-evalAccessWithFor :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => [(AuthTag, Bool)] -> Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> m AuthResult
+evalAccessWithFor :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => [(AuthTag, Bool)] -> Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> m AuthResult
 evalAccessWithFor assumptions mAuthId route isWrite = do
  isSelf <- (== mAuthId) <$> liftHandler defaultMaybeAuthId
  tagActive <- if
@ -1597,42 +1598,42 @@ evalAccessWithFor assumptions mAuthId route isWrite = do
      tellSessionJson SessionInactiveAuthTags deactivated
    return result

-evalAccessFor :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> m AuthResult
+evalAccessFor :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> m AuthResult
 evalAccessFor = evalAccessWithFor []

-evalAccessForDB :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX, BackendCompatible SqlReadBackend backend) => Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> ReaderT backend m AuthResult
+evalAccessForDB :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX, BackendCompatible SqlReadBackend backend) => Maybe (AuthId UniWorX) -> Route UniWorX -> Bool -> ReaderT backend m AuthResult
 evalAccessForDB = evalAccessFor

-evalAccessWith :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => [(AuthTag, Bool)] -> Route UniWorX -> Bool -> m AuthResult
+evalAccessWith :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => [(AuthTag, Bool)] -> Route UniWorX -> Bool -> m AuthResult
 evalAccessWith assumptions route isWrite = do
  mAuthId <- liftHandler maybeAuthId
  evalAccessWithFor assumptions mAuthId route isWrite
  
-evalAccessWithDB :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX, BackendCompatible SqlReadBackend backend) => [(AuthTag, Bool)] -> Route UniWorX -> Bool -> ReaderT backend m AuthResult
+evalAccessWithDB :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX, BackendCompatible SqlReadBackend backend) => [(AuthTag, Bool)] -> Route UniWorX -> Bool -> ReaderT backend m AuthResult
 evalAccessWithDB = evalAccessWith

-evalAccess :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> Bool -> m AuthResult
+evalAccess :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> Bool -> m AuthResult
 evalAccess = evalAccessWith []

-evalAccessDB :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX, BackendCompatible SqlReadBackend backend) => Route UniWorX -> Bool -> ReaderT backend m AuthResult
+evalAccessDB :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX, BackendCompatible SqlReadBackend backend) => Route UniWorX -> Bool -> ReaderT backend m AuthResult
 evalAccessDB = evalAccess

 -- | Check whether the current user is authorized by `evalAccess` for the given route
 --   Convenience function for a commonly used code fragment
-hasAccessTo :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> Bool -> m Bool
+hasAccessTo :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> Bool -> m Bool
 hasAccessTo route isWrite = (== Authorized) <$> evalAccess route isWrite

 -- | Check whether the current user is authorized by `evalAccess` to read from the given route
 --   Convenience function for a commonly used code fragment
-hasReadAccessTo :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> m Bool
+hasReadAccessTo :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> m Bool
 hasReadAccessTo = flip hasAccessTo False

 -- | Check whether the current user is authorized by `evalAccess` to rwrite to the given route
 --   Convenience function for a commonly used code fragment
-hasWriteAccessTo :: (MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> m Bool
+hasWriteAccessTo :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX) => Route UniWorX -> m Bool
 hasWriteAccessTo = flip hasAccessTo True

-wouldHaveAccessTo :: ( MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX )
+wouldHaveAccessTo :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX)
                  => [(AuthTag, Bool)] -- ^ Assumptions
                  -> Route UniWorX
                  -> Bool
@ -1640,7 +1641,7 @@ wouldHaveAccessTo :: ( MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, Be
 wouldHaveAccessTo assumptions route isWrite = (== Authorized) <$> evalAccessWith assumptions route isWrite

 wouldHaveReadAccessTo, wouldHaveWriteAccessTo
-  :: ( MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX )
+  :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX)
  => [(AuthTag, Bool)] -- ^ Assumptions
  -> Route UniWorX
  -> m Bool
@ -1648,7 +1649,7 @@ wouldHaveReadAccessTo assumptions route = wouldHaveAccessTo assumptions route Fa
 wouldHaveWriteAccessTo assumptions route = wouldHaveAccessTo assumptions route True

 wouldHaveReadAccessToIff, wouldHaveWriteAccessToIff
-  :: ( MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX )
+  :: (HasCallStack, MonadThrow m, MonadHandler m, HandlerSite m ~ UniWorX, BearerAuthSite UniWorX)
  => [(AuthTag, Bool)] -- ^ Assumptions
  -> Route UniWorX
  -> m Bool
@ -1657,7 +1658,8 @@ wouldHaveWriteAccessToIff assumptions route = and2M (not <$> hasWriteAccessTo ro


 evalWorkflowRoleFor' :: forall m backend.
-                        ( MonadHandler m
+                        ( HasCallStack
+                        , MonadHandler m
                        , HandlerSite m ~ UniWorX
                        , BearerAuthSite UniWorX
                        , BackendCompatible SqlReadBackend backend
@ -1705,7 +1707,8 @@ evalWorkflowRoleFor' eval mAuthId mwwId wRole route isWrite = do
      return Authorized
    WorkflowRoleAuthorized{..} -> eval (predDNFEntail $ workflowRoleAuthorized `predDNFOr` defaultAuthDNF) mAuthId route isWrite

-evalWorkflowRoleFor :: ( MonadHandler m
+evalWorkflowRoleFor :: ( HasCallStack
+                       , MonadHandler m
                       , HandlerSite m ~ UniWorX
                       , BearerAuthSite UniWorX
                       , BackendCompatible SqlReadBackend backend
@ -1729,7 +1732,8 @@ evalWorkflowRoleFor mAuthId mwwId wRole route isWrite = do
    tellSessionJson SessionInactiveAuthTags deactivated
  return result
 
-hasWorkflowRole :: ( MonadHandler m
+hasWorkflowRole :: ( HasCallStack
+                   , MonadHandler m
                   , HandlerSite m ~ UniWorX
                   , BearerAuthSite UniWorX
                   , BackendCompatible SqlReadBackend backend
@ -1744,7 +1748,8 @@ hasWorkflowRole mwwId wRole route isWrite = do
  evalWorkflowRoleFor mAuthId mwwId wRole route isWrite

 mayViewWorkflowAction' :: forall backend m fileid.
-                          ( MonadHandler m
+                          ( HasCallStack
+                          , MonadHandler m
                          , HandlerSite m ~ UniWorX
                          , BearerAuthSite UniWorX
                          , BackendCompatible SqlReadBackend backend
@ -1774,7 +1779,8 @@ mayViewWorkflowAction' eval mAuthId wwId WorkflowAction{..} = hoist (withReaderT
  return True

 mayViewWorkflowAction :: forall backend m fileid.
-                         ( MonadHandler m
+                         ( HasCallStack
+                         , MonadHandler m
                         , HandlerSite m ~ UniWorX
                         , BearerAuthSite UniWorX
                         , BackendCompatible SqlReadBackend backend
--- a/src/Foundation/DB.hs
+++ b/src/Foundation/DB.hs
@ -1,6 +1,6 @@
 module Foundation.DB
-  ( runDBRead
-  , runSqlPoolRetry
+  ( runDBRead, runDBRead'
+  , runSqlPoolRetry, runSqlPoolRetry'
  ) where

 import Import.NoFoundation hiding (runDB, getDBRunner)
@ -16,14 +16,25 @@ import Database.Persist.Sql.Raw.QQ (executeQQ)
 import qualified Utils.Pool as Custom

  
-runSqlPoolRetry :: forall m a backend.
-                   ( MonadUnliftIO m, BackendCompatible SqlBackend backend
+runSqlPoolRetry :: forall m a backend c.
+                   ( HasCallStack
+                   , MonadUnliftIO m, BackendCompatible SqlBackend backend
                   , MonadLogger m, MonadMask m
                   )
                => ReaderT backend m a
-                -> Custom.Pool m backend
+                -> Custom.Pool' m DBConnLabel c backend
                -> m a
-runSqlPoolRetry action pool = do
+runSqlPoolRetry action pool = runSqlPoolRetry' action pool callStack
+  
+runSqlPoolRetry' :: forall m a backend c.
+                    ( MonadUnliftIO m, BackendCompatible SqlBackend backend
+                    , MonadLogger m, MonadMask m
+                    )
+                 => ReaderT backend m a
+                 -> Custom.Pool' m DBConnLabel c backend
+                 -> CallStack
+                 -> m a
+runSqlPoolRetry' action pool lbl = do
  let policy = Retry.fullJitterBackoff 1e3 & Retry.limitRetriesByCumulativeDelay 10e6
      handlers = Retry.skipAsyncExceptions `snoc` Retry.logRetries suggestRetry logRetry
        where suggestRetry :: IOException -> m Bool
@ -41,9 +52,12 @@ runSqlPoolRetry action pool = do

  Retry.recovering policy handlers $ \Retry.RetryStatus{..} -> do
    $logDebugS "runSqlPoolRetry" $ "rsIterNumber = " <> tshow rsIterNumber
-    customRunSqlPool action pool
+    customRunSqlPool' action pool lbl

-runDBRead :: ReaderT SqlReadBackend (HandlerFor UniWorX) a -> (HandlerFor UniWorX) a
-runDBRead action = do
+runDBRead :: HasCallStack => ReaderT SqlReadBackend (HandlerFor UniWorX) a -> (HandlerFor UniWorX) a
+runDBRead = runDBRead' callStack
+
+runDBRead' :: CallStack -> ReaderT SqlReadBackend (HandlerFor UniWorX) a -> (HandlerFor UniWorX) a
+runDBRead' lbl action = do
  $logDebugS "YesodPersist" "runDBRead"
-  runSqlPoolRetry (withReaderT SqlReadBackend $ [executeQQ|SET TRANSACTION READ ONLY|] *> action) . appConnPool =<< getYesod
+  flip (runSqlPoolRetry' . withReaderT SqlReadBackend $ [executeQQ|SET TRANSACTION READ ONLY|] *> action) lbl . appConnPool =<< getYesod
--- a/src/Foundation/Instances.hs
+++ b/src/Foundation/Instances.hs
@ -1,4 +1,5 @@
 {-# LANGUAGE UndecidableInstances #-} -- for `MonadCrypto` and `MonadSecretBox`
+{-# LANGUAGE InstanceSigs #-}
 {-# OPTIONS_GHC -fno-warn-orphans #-}

 module Foundation.Instances
@ -128,11 +129,13 @@ unsafeHandler f h = do

 -- How to run database actions.
 instance YesodPersist UniWorX where
-    type YesodPersistBackend UniWorX = SqlBackend
-    runDB = UniWorX.runDB
+  type YesodPersistBackend UniWorX = SqlBackend
+  runDB :: HasCallStack => YesodDB UniWorX a -> HandlerFor UniWorX a
+  runDB = UniWorX.runDB' callStack

 instance YesodPersistRunner UniWorX where
-    getDBRunner = UniWorX.getDBRunner
+  getDBRunner :: HasCallStack => HandlerFor UniWorX (DBRunner UniWorX, HandlerFor UniWorX ())
+  getDBRunner = UniWorX.getDBRunner' callStack

    
 instance YesodAuth UniWorX where
@ -187,7 +190,8 @@ instance YesodAuth UniWorX where
              hoistMaybe bearerImpersonate

 instance YesodAuthPersist UniWorX where
-  getAuthEntity = liftHandler . runDBRead . get
+  getAuthEntity :: (HasCallStack, MonadHandler m, HandlerSite m ~ UniWorX) => UserId -> m (Maybe User)
+  getAuthEntity = liftHandler . runDBRead' callStack . get

    
 instance YesodMail UniWorX where
--- a/src/Foundation/Navigation.hs
+++ b/src/Foundation/Navigation.hs
@ -16,7 +16,7 @@ module Foundation.Navigation
  , evalAccessCorrector
  ) where

-import Import.NoFoundation
+import Import.NoFoundation hiding (runDB)

 import Foundation.Type
 import Foundation.Routes
@ -525,7 +525,6 @@ defaultLinks :: ( MonadHandler m
                , HandlerSite m ~ UniWorX
                , MonadThrow m
                , BearerAuthSite UniWorX
-                , BackendCompatible SqlReadBackend (YesodPersistBackend UniWorX)
                ) => m [Nav]
 defaultLinks = fmap catMaybes . mapM runMaybeT $ -- Define the menu items of the header.
  [ return NavHeader
@ -710,7 +709,7 @@ defaultLinks = fmap catMaybes . mapM runMaybeT $ -- Define the menu items of the
      }
  , do
      authCtx <- getAuthContext
-      (haveInstances, haveWorkflows) <- $memcachedByHere (Just $ Right diffDay) authCtx . liftHandler . runDB $ (,) -- We don't expect haveTopWorkflowWorkflows to be relevant and haveTopWorkflowInstances shouldn't change often
+      (haveInstances, haveWorkflows) <- $memcachedByHere (Just $ Right diffDay) authCtx . liftHandler . runDBRead $ (,) -- We don't expect haveTopWorkflowWorkflows to be relevant and haveTopWorkflowInstances shouldn't change often
        <$> haveTopWorkflowInstances
        <*> haveTopWorkflowWorkflows

@ -2481,7 +2480,7 @@ pageActions route | Just (rScope, WorkflowInstanceListR) <- route ^? _WorkflowSc
    { navLink = NavLink
      { navLabel       = MsgMenuWorkflowWorkflowList
      , navRoute       = _WorkflowScopeRoute # (rScope, WorkflowWorkflowListR)
-      , navAccess'     = runDB $ haveWorkflowWorkflows rScope
+      , navAccess'     = runDBRead $ haveWorkflowWorkflows rScope
      , navType        = NavTypeLink { navModal = False }
      , navQuick'      = mempty
      , navForceActive = False
@ -2550,7 +2549,7 @@ pageActions TopWorkflowInstanceListR = return
    { navLink = NavLink
      { navLabel       = MsgMenuTopWorkflowWorkflowList
      , navRoute       = TopWorkflowWorkflowListR
-      , navAccess'     = runDB haveTopWorkflowWorkflows
+      , navAccess'     = runDBRead haveTopWorkflowWorkflows
      , navType        = NavTypeLink { navModal = False }
      , navQuick'      = mempty
      , navForceActive = False
--- a/src/Foundation/Type.hs
+++ b/src/Foundation/Type.hs
@ -28,6 +28,8 @@ import Data.IntervalMap.Strict (IntervalMap)

 import qualified Utils.Pool as Custom

+import Utils.Metrics (DBConnUseState)
+

 type SMTPPool = Pool SMTPConnection

@ -44,7 +46,7 @@ makePrisms ''SomeSessionStorage
 data UniWorX = UniWorX
    { appSettings'         :: AppSettings
    , appStatic            :: EmbeddedStatic -- ^ Settings for static file serving.
-    , appConnPool          :: forall m. MonadIO m => Custom.Pool m SqlBackend -- ^ Database connection pool.
+    , appConnPool          :: forall m. MonadIO m => Custom.Pool' m DBConnLabel DBConnUseState SqlBackend -- ^ Database connection pool.
    , appSmtpPool          :: Maybe SMTPPool
    , appLdapPool          :: Maybe (Failover (LdapConf, LdapPool))
    , appWidgetMemcached   :: Maybe Memcached.Connection -- ^ Actually a proper pool
--- a/src/Foundation/Yesod/Persist.hs
+++ b/src/Foundation/Yesod/Persist.hs
@ -1,5 +1,6 @@
 module Foundation.Yesod.Persist
  ( runDB, getDBRunner
+  , runDB', getDBRunner'
  , module Foundation.DB
  ) where

@ -19,30 +20,42 @@ import UnliftIO.Resource (allocate, unprotect)

 runDB :: ( YesodPersistBackend UniWorX ~ SqlBackend
         , BearerAuthSite UniWorX
+         , HasCallStack
         )
      => YesodDB UniWorX a -> HandlerFor UniWorX a
-runDB action = do
-  -- stack <- liftIO currentCallStack
-  -- $logDebugS "YesodPersist" . unlines $ "runDB" : map pack stack
+runDB = runDB' callStack
+
+runDB' :: ( YesodPersistBackend UniWorX ~ SqlBackend
+          , BearerAuthSite UniWorX
+          )
+       => CallStack -> YesodDB UniWorX a -> HandlerFor UniWorX a
+runDB' lbl action = do
  $logDebugS "YesodPersist" "runDB"
  dryRun <- isDryRun
  let action'
        | dryRun = action <* transactionUndo
        | otherwise = action

-  runSqlPoolRetry action' . appConnPool =<< getYesod
+  flip (runSqlPoolRetry' action') lbl . appConnPool =<< getYesod

 getDBRunner :: ( YesodPersistBackend UniWorX ~ SqlBackend
               , BearerAuthSite UniWorX
+               , HasCallStack
               )
            => HandlerFor UniWorX (DBRunner UniWorX, HandlerFor UniWorX ())
-getDBRunner = do
+getDBRunner = getDBRunner' callStack
+
+getDBRunner' :: ( YesodPersistBackend UniWorX ~ SqlBackend
+                , BearerAuthSite UniWorX
+                )
+             => CallStack -> HandlerFor UniWorX (DBRunner UniWorX, HandlerFor UniWorX ())
+getDBRunner' lbl = do
  pool <- getsYesod appConnPool
  UnliftIO{..} <- askUnliftIO
  let withPrep conn f = f (persistBackend conn) (SQL.getStmtConn $ persistBackend conn)
  (relKey, (conn, ident)) <- allocate
    (do
-        (conn, ident) <- unliftIO $ Custom.takeResource pool
+        (conn, ident) <- unliftIO $ Custom.takeResource' pool lbl
        withPrep conn (\c f -> SQL.connBegin c f Nothing)
        return (conn, ident)
    )
--- a/src/Handler/Utils/Form.hs
+++ b/src/Handler/Utils/Form.hs
@ -1452,38 +1452,34 @@ fsUniq :: (Text -> Text) -> Text -> FieldSettings site
 fsUniq mkUnique seed = "" { fsName = Just $ mkUnique seed }


-optionsPersistCryptoId :: forall site backend a msg.
-                          ( YesodPersist site
-                          , PersistQueryRead backend
-                          , HasCryptoUUID (Key a) (HandlerFor site)
+optionsPersistCryptoId :: forall backend a msg.
+                          ( PersistQueryRead backend
+                          , HasCryptoUUID (Key a) (HandlerFor UniWorX)
                          , KnownSymbol (CryptoIDNamespace UUID (Key a))
-                          , RenderMessage site msg
-                          , YesodPersistBackend site ~ backend
+                          , RenderMessage UniWorX msg
+                          , YesodPersistBackend UniWorX ~ backend
                          , PersistRecordBackend a backend
                          , PathPiece (Key a)
                          )
                       => [Filter a]
                       -> [SelectOpt a]
                       -> (a -> msg)
-                       -> HandlerFor site (OptionList (Entity a))
+                       -> HandlerFor UniWorX (OptionList (Entity a))
 optionsPersistCryptoId filts ords toDisplay = do
  ents <- runDB $ selectList filts ords
  optionsCryptoIdF ents (return . entityKey) (return . toDisplay . entityVal)

-optionsCryptoIdE :: forall site backend a msg.
-                    ( YesodPersist site
-                    , PersistQueryRead backend, PersistUniqueRead backend
-                    , HasCryptoUUID (Key a) (HandlerFor site)
+optionsCryptoIdE :: forall backend a msg.
+                    ( HasCryptoUUID (Key a) (HandlerFor UniWorX)
                    , KnownSymbol (CryptoIDNamespace UUID (Key a))
-                    , RenderMessage site msg
-                    , YesodPersistBackend site ~ backend
+                    , RenderMessage UniWorX msg
+                    , YesodPersistBackend UniWorX ~ backend
                    , PersistRecordBackend a backend
-                    , BackendCompatible SqlBackend backend
                    , PathPiece (Key a)
                    )
                 => E.SqlQuery (E.SqlExpr (Entity a))
                 -> (a -> msg)
-                 -> HandlerFor site (OptionList (Entity a))
+                 -> HandlerFor UniWorX (OptionList (Entity a))
 optionsCryptoIdE query toDisplay = do
  ents <- runDB $ E.select query
  optionsCryptoIdF ents (return . entityKey) (return . toDisplay . entityVal)
--- a/src/Import.hs
+++ b/src/Import.hs
@ -3,7 +3,8 @@ module Import
    ) where

 import Foundation            as Import
-import Import.NoFoundation   as Import
+import Foundation.Yesod.Persist as Import (runDB, getDBRunner)
+import Import.NoFoundation   as Import hiding (runDB, getDBRunner)
 import Model.Migration       as Import

 import Utils.SystemMessage   as Import
--- a/src/Import/NoModel.hs
+++ b/src/Import/NoModel.hs
@ -62,6 +62,7 @@ import Network.HaskellNet.SMTP as Import (SMTPConnection)
 import Data.Data              as Import (Data)
 import GHC.Exts               as Import (IsList)
 import Data.Ix                as Import (Ix)
+import GHC.Stack              as Import (CallStack, HasCallStack, callStack)

 import Data.Hashable          as Import
 import Data.List              as Import (elemIndex)
--- a/src/Utils/DB.hs
+++ b/src/Utils/DB.hs
@ -18,6 +18,8 @@ import qualified Utils.Pool as Custom

 import Database.Persist.Sql (runSqlConn)

+import GHC.Stack (HasCallStack, CallStack, callStack)
+

 emptyOrIn :: PersistField typ
          => E.SqlExpr (E.Value typ) -> Set typ -> E.SqlExpr (E.Value Bool)
@ -157,7 +159,19 @@ selectMaybe fltrs opts = listToMaybe <$> selectList fltrs (LimitTo 1 : opts')
        isLimit = \case
          LimitTo _ -> True
          _other    -> False
-  

-customRunSqlPool :: (MonadUnliftIO m, BackendCompatible SqlBackend backend) => ReaderT backend m a -> Custom.Pool m backend -> m a
-customRunSqlPool act p = Custom.withResource p $ runSqlConn act
+
+type DBConnLabel = CallStack
+
+customRunSqlPool :: (HasCallStack, MonadUnliftIO m, BackendCompatible SqlBackend backend)
+                 => ReaderT backend m a
+                 -> Custom.Pool' m DBConnLabel c backend
+                 -> m a
+customRunSqlPool act p = customRunSqlPool' act p callStack
+
+customRunSqlPool' :: (MonadUnliftIO m, BackendCompatible SqlBackend backend)
+                  => ReaderT backend m a
+                  -> Custom.Pool' m DBConnLabel c backend
+                  -> CallStack
+                  -> m a
+customRunSqlPool' act p label = Custom.withResource' p label $ runSqlConn act
--- a/src/Utils/Metrics.hs
+++ b/src/Utils/Metrics.hs
@ -23,6 +23,7 @@ module Utils.Metrics
  , PoolMetrics, PoolLabel(..)
  , poolMetrics
  , observeDatabaseConnectionOpened, observeDatabaseConnectionClosed
+  , onUseDBConn, onReleaseDBConn, DBConnUseState, DBConnLabel
  ) where

 import Import.NoModel hiding (Vector, Info)
@ -55,6 +56,8 @@ import qualified Data.Foldable as F

 import qualified Utils.Pool as Custom

+import GHC.Stack
+
 {-# ANN module ("HLint: ignore Use even" :: String) #-}


@ -355,7 +358,7 @@ data PoolLabel = PoolDatabaseConnections
 nullaryPathPiece ''PoolLabel $ camelToPathPiece' 1

 poolMetrics :: PoolLabel
-            -> Custom.Pool m a
+            -> Custom.Pool' m c' c a
            -> Metric PoolMetrics
 poolMetrics lbl pool = Metric $ return (PoolMetrics, collectPoolMetrics)
  where
@ -384,6 +387,35 @@ poolMetrics lbl pool = Metric $ return (PoolMetrics, collectPoolMetrics)
                     "Number of resources currently in use"
    usesInfo = Info "uni2work_pool_uses_count" 
                    "Number of takes executed against the pool"
+
+{-# NOINLINE databaseConnDuration #-}
+databaseConnDuration :: Vector Label1 Histogram
+databaseConnDuration = unsafeRegister . vector ("label") $ histogram info buckets
+  where
+    info = Info "uni2work_database_conn_duration_seconds"
+                "Duration of use of a database connection from the pool"
+    buckets = histogramBuckets 50e-6 5000
+  
+data DBConnUseState = DBConnUseState
+  { dbConnUseStart :: !TimeSpec
+  , dbConnUseLabel :: !CallStack
+  } deriving (Show, Typeable)
+
+onUseDBConn :: (MonadIO m, MonadLogger m) => CallStack -> a -> m DBConnUseState
+onUseDBConn dbConnUseLabel _ = do
+  $logDebugS "DB" $ case getCallStack dbConnUseLabel of
+    [] -> "no stack"
+    xs -> intercalate "; " $ map (\(f, loc) -> pack f <> " @(" <> pack (prettySrcLoc loc) <> ")") xs
+  dbConnUseStart <- liftIO $ getTime Monotonic
+  return DBConnUseState{..}
+
+onReleaseDBConn :: MonadIO m => DBConnUseState -> a -> m ()
+onReleaseDBConn DBConnUseState{..} _ = liftIO $ do
+  diff <- realToFrac . subtract dbConnUseStart <$> getTime Monotonic
+  let lbl = case reverse $ getCallStack dbConnUseLabel of
+        [] -> "unlabeled"
+        (_, SrcLoc{..}) : _ -> pack srcLocModule
+  withLabel databaseConnDuration lbl $ flip observe diff
    

 withHealthReportMetrics :: MonadIO m => m HealthReport -> m HealthReport
--- a/src/Utils/Pool.hs
+++ b/src/Utils/Pool.hs
@ -1,13 +1,16 @@
 {-# OPTIONS_GHC -Wno-error=unused-top-binds #-}

 module Utils.Pool
-  ( Pool, hoistPool
+  ( Pool', hoistPool
+  , PoolResourceIdent'
+  , Pool, PoolResourceIdent
  , getPoolAvailableCount, getPoolInUseCount, getPoolUsesCount
-  , createPool
+  , createPool, createPool'
  , purgePool
-  , withResource
+  , withResource, withResource'
  , destroyResources
  , takeResource, releaseResource
+  , takeResource'
  ) where

 import ClassyPrelude
@ -35,35 +38,42 @@ import Utils.NTop
 -- <https://hackage.haskell.org/package/ex-pool-0.2.1/docs/src/Data-Pool.html>


-newtype PoolResourceIdent = PoolResourceIdent Int
+data PoolResourceIdent' c = PoolResourceIdent Int c
  deriving (Eq, Ord, Show, Typeable)


-data Pool m a = Pool
+data Pool' m c' c a = Pool
  { create :: m a 
  , destroy :: a -> m ()
+  , onUse :: c' -> a -> m c
+  , onRelease :: c -> a -> m ()
  , idleTime :: !(Maybe Int)
  , maxAvailable :: !(Maybe Int)
-  , resources :: !(TVar (PoolResources a))
+  , resources :: !(TVar (PoolResources c a))
  , aliveRef :: !(IORef ())
  }

-data PoolResources a = PoolResources
+data PoolResources c a = PoolResources
  { inUseCount, availableCount :: !Int
  , inUse :: !(IntMap a)
  , available :: !(IntMap [a])
  , inUseTick :: !Int
  } deriving (Functor)

+type Pool m a = Pool' m () () a
+type PoolResourceIdent = PoolResourceIdent' ()

-hoistPool :: (forall b. m b -> n b) -> Pool m a -> Pool n a
+
+hoistPool :: (forall b. m b -> n b) -> Pool' m c' c a -> Pool' n c' c a
 hoistPool nat Pool{..} = Pool
  { create = nat create
  , destroy = nat . destroy
+  , onUse = (nat .) . onUse
+  , onRelease = (nat .) . onRelease
  , ..
  }

-getPoolAvailableCount, getPoolInUseCount, getPoolUsesCount :: Pool m a -> STM Int
+getPoolAvailableCount, getPoolInUseCount, getPoolUsesCount :: Pool' m c' c a -> STM Int
 getPoolAvailableCount Pool{..} = availableCount <$> readTVar resources
 getPoolInUseCount Pool{..} = inUseCount <$> readTVar resources
 getPoolUsesCount Pool{..} = inUseTick <$> readTVar resources
@ -83,7 +93,18 @@ createPool :: (MonadResource m, MonadUnliftIO m, MonadUnliftIO m')
           -> Maybe Int -- ^ Timeout in seconds
           -> Maybe Int -- ^ Max available
           -> m (Pool m' a)
-createPool nat create destroy (fmap $ max 0 -> idleTime) (fmap $ max 0 -> maxAvailable) = do
+createPool nat create destroy = createPool' nat create destroy (\() _ -> return ()) (\() _ -> return ())
+
+createPool' :: (MonadResource m, MonadUnliftIO m, MonadUnliftIO m')
+            => (forall b. m' b -> m b)
+            -> m' a -- ^ Create
+            -> (a -> m' ()) -- ^ Destroy
+            -> (c' -> a -> m' c) -- ^ onUse
+            -> (c -> a -> m' ()) -- ^ onRelease
+            -> Maybe Int -- ^ Timeout in seconds
+            -> Maybe Int -- ^ Max available
+            -> m (Pool' m' c' c a)
+createPool' nat create destroy onUse onRelease (fmap $ max 0 -> idleTime) (fmap $ max 0 -> maxAvailable) = do
  let
    inUseCount = 0
    availableCount = 0
@ -102,10 +123,10 @@ createPool nat create destroy (fmap $ max 0 -> idleTime) (fmap $ max 0 -> maxAva

  return pool

-purgePool :: MonadUnliftIO m => Pool m a -> m ()
+purgePool :: MonadUnliftIO m => Pool' m c' c a -> m ()
 purgePool = destroyResources $ const True

-reaper :: MonadUnliftIO m => (a -> m ()) -> TVar (PoolResources a) -> Int -> m ()
+reaper :: MonadUnliftIO m => (a -> m ()) -> TVar (PoolResources c a) -> Int -> m ()
 reaper destroy' resources' t = forever $ do
  atomically . waitDelay =<< liftIO (newDelay i)

@ -124,8 +145,11 @@ reaper destroy' resources' t = forever $ do
  where
    MkFixed (fromIntegral -> i) = 1 :: Micro

-takeResource :: MonadIO m => Pool m a -> m (a, PoolResourceIdent)
-takeResource Pool{..} = do
+takeResource :: MonadIO m => Pool m a -> m (a, PoolResourceIdent' ())
+takeResource p = takeResource' p ()
+
+takeResource' :: MonadIO m => Pool' m c' c a -> c' -> m (a, PoolResourceIdent' c)
+takeResource' Pool{..} stateInit = do
  takenAvailable <- atomically $ do
    PoolResources{..} <- readTVar resources
    case IntMap.maxViewWithKey available of
@ -147,7 +171,9 @@ takeResource Pool{..} = do
        return $ Just (av, inUseTick)
      _other -> return Nothing
  case takenAvailable of
-    Just (av, resTick) -> return (av, PoolResourceIdent resTick)
+    Just (av, resTick) -> do
+      hookData <- onUse stateInit av
+      return (av, PoolResourceIdent resTick hookData)
    Nothing -> do
      newResource <- create
      resTick <- atomically . stateTVar resources $ \res@PoolResources{..} ->
@ -157,12 +183,13 @@ takeResource Pool{..} = do
         in ( inUseTick
            , res{ inUseCount = inUseCount', inUse = inUse', inUseTick = inUseTick' }
            )
-      return (newResource, PoolResourceIdent resTick)
+      hookData <- onUse stateInit newResource
+      return (newResource, PoolResourceIdent resTick hookData)

 releaseResource :: MonadUnliftIO m
                => Bool -- ^ Destroy resource and don't return to pool?
-                -> Pool m a
-                -> (a, PoolResourceIdent)
+                -> Pool' m c' c a
+                -> (a, PoolResourceIdent' c)
                -> m ()
 releaseResource isLost p@Pool{..} (x, ident)
  | isLost = do
@ -171,27 +198,27 @@ releaseResource isLost p@Pool{..} (x, ident)
  | otherwise
  = markResourceAvailable p ident

-markResourceAvailable, markResourceLost :: MonadUnliftIO m => Pool m a -> PoolResourceIdent -> m ()
+markResourceAvailable, markResourceLost :: MonadUnliftIO m => Pool' m c' c a -> PoolResourceIdent' c -> m ()
 markResourceAvailable = returnResource True
 markResourceLost = returnResource False

 returnResource :: MonadUnliftIO m
               => Bool -- ^ return to available
-               -> Pool m a
-               -> PoolResourceIdent
+               -> Pool' m c' c a
+               -> PoolResourceIdent' c
               -> m ()
-returnResource toAvailable Pool{..} (PoolResourceIdent inUseKey) = do
+returnResource toAvailable Pool{..} (PoolResourceIdent inUseKey hookData) = do
  now <- if | toAvailable -> Just <$> currentSecond
            | otherwise   -> return Nothing
-  toDestroy <- atomically . stateTVar resources $ \res@PoolResources{..} -> case deleteView inUseKey inUse of
-    Nothing -> (Nothing, res)
+  (toDestroy, released) <- atomically . stateTVar resources $ \res@PoolResources{..} -> case deleteView inUseKey inUse of
+    Nothing -> ((Nothing, Nothing), res)
    Just (u, us) | NTop (Just availableCount) >= NTop maxAvailable
-      -> (Just u,) res
+      -> ((Just u, Just u),) res
           { inUse = us
           , inUseCount = pred inUseCount
           }
    Just (u, us)
-      -> (Nothing, ) PoolResources
+      -> ((Nothing, Just u), ) PoolResources
           { inUse = us
           , inUseCount = pred inUseCount
           , availableCount = bool id succ toAvailable availableCount
@ -199,6 +226,7 @@ returnResource toAvailable Pool{..} (PoolResourceIdent inUseKey) = do
           , inUseTick
           }

+  forM_ released $ \u -> onRelease hookData u
  forM_ toDestroy $ void . forkIO . destroy
  where
    deleteView :: Int -> IntMap a -> Maybe (a, IntMap a)
@ -207,9 +235,12 @@ returnResource toAvailable Pool{..} (PoolResourceIdent inUseKey) = do


 withResource :: forall b m a. MonadUnliftIO m => Pool m a -> (a -> m b) -> m b
-withResource p act = bracketOnError (takeResource p) (releaseResource True p) (\x'@(x, _) -> act x <* releaseResource False p x')
+withResource p = withResource' p ()
  
-destroyResources :: MonadUnliftIO m => (a -> Bool) -> Pool m a -> m ()
+withResource' :: forall b m c' c a. MonadUnliftIO m => Pool' m c' c a -> c' -> (a -> m b) -> m b
+withResource' p stateInit act = bracketOnError (takeResource' p stateInit) (releaseResource True p) (\x'@(x, _) -> act x <* releaseResource False p x')
+  
+destroyResources :: MonadUnliftIO m => (a -> Bool) -> Pool' m c' c a -> m ()
 destroyResources p Pool{..} = do
  toDestroy <- atomically . stateTVar resources $ \res@PoolResources{..}
    -> let partitioned = partition p <$> available
--- a/src/Web/ServerSession/Backend/Persistent/Memcached.hs
+++ b/src/Web/ServerSession/Backend/Persistent/Memcached.hs
@ -29,6 +29,8 @@ import qualified Data.ByteString.Lazy as LBS

 import qualified Data.ByteString.Base64.URL as Base64

+import Utils.Metrics (DBConnUseState)
+

 import Data.Text.Encoding (decodeUtf8')

@ -46,7 +48,7 @@ share [mkPersist sqlSettings, mkMigrate "migrateMemcachedSqlStorage"]


 data MemcachedSqlStorage sess = MemcachedSqlStorage
-  { mcdSqlConnPool :: forall m. MonadIO m => Custom.Pool m SqlBackend
+  { mcdSqlConnPool :: forall m. MonadIO m => Custom.Pool' m DBConnLabel DBConnUseState SqlBackend
  , mcdSqlMemcached :: Memcached.Connection
  , mcdSqlMemcachedKey :: AEAD.Key
  , mcdSqlMemcachedExpiration :: Maybe NominalDiffTime
@ -108,7 +110,7 @@ instance (IsSessionData sess, Binary (Decomposed sess)) => Storage (MemcachedSql
  type SessionData  (MemcachedSqlStorage sess) = sess
  type TransactionM (MemcachedSqlStorage sess) = SqlPersistT IO

-  runTransactionM MemcachedSqlStorage{..} = flip customRunSqlPool mcdSqlConnPool
+  runTransactionM MemcachedSqlStorage{..} act = customRunSqlPool act mcdSqlConnPool

  getSession MemcachedSqlStorage{..} sessId = exceptT (maybe (return Nothing) throwM) (return . Just) $ do
    encSession <- catchIfExceptT (const Nothing) Memcached.isKeyNotFound . liftIO . fmap LBS.toStrict $ Memcached.getAndTouch_ expiry (memcachedSqlSessionId # sessId) mcdSqlMemcached