Pool.hs

-- | Connection pooling with health-check validation and retry configuration.
module Database.Bolty.Pool
  ( BoltPool(..)
  , PoolConfig(..)
  , defaultPoolConfig
  , ValidationStrategy(..)
  , RetryConfig(..)
  , defaultRetryConfig
  , PoolCounters(..)
  , createPool
  , destroyPool
  , withConnection
  , CheckedOutConnection(..)
  , acquireConnection
  , releaseConnection
  , releaseConnectionOnError
  , poolCounters
  ) where

import           Control.Exception                (SomeException, mask, try, fromException, throwIO)
import           Data.IORef                       (IORef, newIORef, readIORef, atomicModifyIORef')
import           Data.Pool                        (Pool)
import qualified Data.Pool                        as Pool
import           Data.Word                        (Word64)
import           GHC.Clock                        (getMonotonicTimeNSec)
import           GHC.Stack                        (HasCallStack)

import           Data.Kind                        (Type)
import           Database.Bolty.Connection.Type   (Connection, ValidatedConfig, Error(..))
import qualified Database.Bolty.Connection.Pipe   as P
import           Database.Bolty.Connection.Pipe   (touchConnection, connectionLastActivity,
                                                   connectionServerIdleTimeout)


-- | Configuration for retry logic on transient failures.
type RetryConfig :: Type
data RetryConfig = RetryConfig
  { maxRetries   :: !Int
  -- ^ Maximum number of retry attempts. Default: 5.
  , initialDelay :: !Int
  -- ^ Initial delay in microseconds before first retry. Default: 200000 (200ms).
  , maxDelay     :: !Int
  -- ^ Maximum delay in microseconds between retries. Default: 5000000 (5s).
  }
  deriving stock (Show, Eq)

-- | Default retry configuration: 5 retries, 200ms initial delay, 5s max delay.
defaultRetryConfig :: RetryConfig
defaultRetryConfig = RetryConfig
  { maxRetries   = 5
  , initialDelay = 200_000
  , maxDelay     = 5_000_000
  }


-- | Strategy for validating connections when they are checked out of the pool.
type ValidationStrategy :: Type
data ValidationStrategy
  = AlwaysPing
  -- ^ Always send RESET before use (current default, safest).
  | PingIfIdle !Int
  -- ^ Only ping if the connection has been idle for more than N seconds.
  -- Connections used within the last N seconds are assumed healthy.
  | NeverPing
  -- ^ Skip health check entirely (fastest, use only in trusted environments).
  deriving stock (Show, Eq)

-- | Configuration for the connection pool.
type PoolConfig :: Type
data PoolConfig = PoolConfig
  { maxConnections     :: Int
  -- ^ Maximum number of open connections. Default: 10.
  , idleTimeout        :: Double
  -- ^ Seconds before an idle connection is closed. Default: 60.
  , maxPingRetries     :: Int
  -- ^ Number of additional attempts after a ping failure. Default: 1.
  -- Total attempts = 1 + maxPingRetries.
  , retryConfig        :: RetryConfig
  -- ^ Retry configuration for transient failures. Default: 'defaultRetryConfig'.
  , validationStrategy :: ValidationStrategy
  -- ^ How to validate connections on checkout. Default: 'AlwaysPing'.
  }

-- | Default pool configuration: 10 max connections, 60 second idle timeout, 1 ping retry.
defaultPoolConfig :: PoolConfig
defaultPoolConfig = PoolConfig
  { maxConnections     = 10
  , idleTimeout        = 60
  , maxPingRetries     = 1
  , retryConfig        = defaultRetryConfig
  , validationStrategy = PingIfIdle 30
  }


-- | A pool of Neo4j connections with health-check and retry configuration.
type BoltPool :: Type
data BoltPool = BoltPool
  { bpPool        :: !(Pool Connection)
  , bpMaxRetries  :: !Int
  , bpRetryConfig :: !RetryConfig
  , bpValidation  :: !ValidationStrategy
  , bpActiveConns :: !(IORef Int)
  , bpTotalAcqs   :: !(IORef Int)
  , bpTotalWaitNs :: !(IORef Word64)
  }


-- | Snapshot of pool-level counters.
type PoolCounters :: Type
data PoolCounters = PoolCounters
  { pcActive      :: !Int
  -- ^ Number of connections currently checked out.
  , pcTotalAcqs   :: !Int
  -- ^ Lifetime count of connection acquisitions.
  , pcTotalWaitNs :: !Word64
  -- ^ Cumulative nanoseconds spent waiting for connections.
  } deriving stock (Show, Eq)


-- | Create a new connection pool.
-- If the server advertises @connection.recv_timeout_seconds@, the pool's
-- idle timeout is capped to @min(configured, hint - 5)@ so connections
-- are recycled before the server closes them.
createPool :: HasCallStack => ValidatedConfig -> PoolConfig -> IO BoltPool
createPool cfg PoolConfig{..} = do
  -- Probe a connection to discover the server's idle timeout hint.
  probe <- P.connect cfg
  let effectiveIdle = case connectionServerIdleTimeout probe of
        Just serverSecs | serverSecs > 5 ->
          min idleTimeout (fromIntegral (serverSecs - 5))
        _ -> idleTimeout
  P.close probe
  let poolCfg = Pool.setNumStripes (Just 1)
              $ Pool.defaultPoolConfig
                  (P.connect cfg)           -- create
                  P.close                   -- destroy
                  effectiveIdle             -- idle timeout (seconds)
                  maxConnections            -- max resources
  pool <- Pool.newPool poolCfg
  activeRef <- newIORef 0
  acqRef    <- newIORef 0
  waitRef   <- newIORef 0
  pure BoltPool { bpPool = pool, bpMaxRetries = maxPingRetries
                , bpRetryConfig = retryConfig, bpValidation = validationStrategy
                , bpActiveConns = activeRef, bpTotalAcqs = acqRef, bpTotalWaitNs = waitRef }


-- | Destroy the pool, closing all connections.
destroyPool :: BoltPool -> IO ()
destroyPool BoltPool{bpPool} = Pool.destroyAllResources bpPool


-- | Acquire a healthy connection from the pool, run an action, then release it.
-- Validates the connection using the configured 'ValidationStrategy'.
-- On connection failure during the action, discards the dead connection and
-- retries once with a fresh one.
withConnection :: HasCallStack => BoltPool -> (Connection -> IO a) -> IO a
withConnection BoltPool{bpPool, bpMaxRetries, bpValidation, bpActiveConns, bpTotalAcqs, bpTotalWaitNs} action =
    go (bpMaxRetries + 1) True
  where
    go 0 _ = fail "withConnection: no healthy connection available"
    go n canRetryDead = mask $ \restore -> do
      t0 <- getMonotonicTimeNSec
      (conn, localPool) <- Pool.takeResource bpPool
      t1 <- getMonotonicTimeNSec
      atomicModifyIORef' bpActiveConns $ \x -> (x + 1, ())
      atomicModifyIORef' bpTotalAcqs $ \x -> (x + 1, ())
      atomicModifyIORef' bpTotalWaitNs $ \x -> (x + (t1 - t0), ())
      let decActive = atomicModifyIORef' bpActiveConns $ \x -> (x - 1, ())
      shouldPing <- needsPing bpValidation conn
      healthy <- if shouldPing then P.ping conn else pure True
      if healthy
        then do
          result <- try $ restore (action conn)
          case result of
            Right x -> do
              touchConnection conn
              Pool.putResource localPool conn
              decActive
              pure x
            Left (e :: SomeException)
              | canRetryDead, isConnectionError e -> do
                  Pool.destroyResource bpPool localPool conn
                  decActive
                  restore $ go n False   -- retry with fresh connection, no more retries
              | otherwise -> do
                  Pool.destroyResource bpPool localPool conn
                  decActive
                  restore $ throwIOSome e
        else do
          Pool.destroyResource bpPool localPool conn
          decActive
          restore $ go (n - 1) canRetryDead

    throwIOSome :: SomeException -> IO a
    throwIOSome = throwIO' where
      throwIO' :: SomeException -> IO a
      throwIO' = Control.Exception.throwIO

    isConnectionError :: SomeException -> Bool
    isConnectionError e = case fromException e :: Maybe Error of
      Just (NonboltyError _) -> True
      _                      -> False


-- | Check if a connection needs a ping based on the validation strategy.
needsPing :: ValidationStrategy -> Connection -> IO Bool
needsPing AlwaysPing _    = pure True
needsPing NeverPing _     = pure False
needsPing (PingIfIdle secs) conn = do
  now <- getMonotonicTimeNSec
  lastAct <- connectionLastActivity conn
  let idleNs = now - lastAct
  let thresholdNs = fromIntegral secs * 1_000_000_000
  pure (idleNs >= thresholdNs)


-- | A checked-out connection handle, bundling the connection with its
-- local pool reference for proper release.
type CheckedOutConnection :: Type
data CheckedOutConnection = CheckedOutConnection
  { cocConnection :: !Connection
  , cocLocalPool  :: !(Pool.LocalPool Connection)
  , cocBoltPool   :: !BoltPool
  }


-- | Acquire a validated connection from the pool. The caller is responsible
-- for releasing it with 'releaseConnection' or 'releaseConnectionOnError'.
--
-- This is the low-level primitive behind 'withConnection'. Prefer
-- 'withConnection' for simple request-response patterns.
-- Use 'acquireConnection' when you need the connection to outlive a callback
-- (e.g. for streaming with @bracketIO@).
acquireConnection :: BoltPool -> IO CheckedOutConnection
acquireConnection bp@BoltPool{bpPool, bpMaxRetries, bpValidation, bpActiveConns, bpTotalAcqs, bpTotalWaitNs} =
    go (bpMaxRetries + 1)
  where
    go 0 = fail "acquireConnection: no healthy connection available"
    go n = do
      t0 <- getMonotonicTimeNSec
      (conn, localPool) <- Pool.takeResource bpPool
      t1 <- getMonotonicTimeNSec
      atomicModifyIORef' bpActiveConns $ \x -> (x + 1, ())
      atomicModifyIORef' bpTotalAcqs $ \x -> (x + 1, ())
      atomicModifyIORef' bpTotalWaitNs $ \x -> (x + (t1 - t0), ())
      shouldPing <- needsPing bpValidation conn
      healthy <- if shouldPing then P.ping conn else pure True
      if healthy
        then pure CheckedOutConnection
              { cocConnection = conn, cocLocalPool = localPool, cocBoltPool = bp }
        else do
          Pool.destroyResource bpPool localPool conn
          atomicModifyIORef' bpActiveConns $ \x -> (x - 1, ())
          go (n - 1)


-- | Release a connection back to the pool after successful use.
releaseConnection :: CheckedOutConnection -> IO ()
releaseConnection CheckedOutConnection{cocConnection, cocLocalPool, cocBoltPool} = do
  touchConnection cocConnection
  Pool.putResource cocLocalPool cocConnection
  atomicModifyIORef' (bpActiveConns cocBoltPool) $ \x -> (x - 1, ())


-- | Release a connection after an error, destroying it instead of returning
-- it to the pool (since it may be in a bad state).
releaseConnectionOnError :: CheckedOutConnection -> IO ()
releaseConnectionOnError CheckedOutConnection{cocConnection, cocLocalPool, cocBoltPool} = do
  Pool.destroyResource (bpPool cocBoltPool) cocLocalPool cocConnection
  atomicModifyIORef' (bpActiveConns cocBoltPool) $ \x -> (x - 1, ())


-- | Read a snapshot of the pool's lifetime counters.
poolCounters :: BoltPool -> IO PoolCounters
poolCounters BoltPool{bpActiveConns, bpTotalAcqs, bpTotalWaitNs} = do
  active <- readIORef bpActiveConns
  acqs   <- readIORef bpTotalAcqs
  waitNs <- readIORef bpTotalWaitNs
  pure PoolCounters{pcActive = active, pcTotalAcqs = acqs, pcTotalWaitNs = waitNs}