module Hasql.Pool
( Pool
, Settings(..)
, UsageError(..)
, ConnectionGetter
, Stats(..)
, stats
, getPoolUsageStat
, acquire
, acquireWith
, release
, use
, useWithObserver
, withResourceOnEither
)
where
import qualified Data.Pool as ResourcePool
import qualified Data.Pool.Internal as Unstable
import System.Clock (Clock(Monotonic), diffTimeSpec, getTime, toNanoSecs)
import Hasql.Pool.Prelude
import qualified Hasql.Connection
import qualified Hasql.Session
import Hasql.Pool.Observer (Observed(..), ObserverAction)
-- |
-- A pool of open DB connections.
newtype Pool =
Pool (ResourcePool.Pool (Either Hasql.Connection.ConnectionError Hasql.Connection.Connection))
type PoolSize = Int
type PoolStripes = Int
type ResidenceTimeout = NominalDiffTime
-- |
-- Connection getter action that allows for obtaining Postgres connection settings
-- via external resources such as AWS tokens etc.
type ConnectionGetter = IO (Either Hasql.Connection.ConnectionError Hasql.Connection.Connection)
-- |
-- 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 =
(PoolSize, ResidenceTimeout, Hasql.Connection.Settings)
-- |
-- Given the pool-size, timeout and connection settings
-- create a connection-pool.
acquire :: Settings -> IO Pool
acquire settings@(_size, _timeout, connectionSettings) =
acquireWith (Hasql.Connection.acquire connectionSettings) settings
-- |
-- Similar to 'acquire', allows for finer configuration.
acquireWith :: ConnectionGetter
-> Settings
-> IO Pool
acquireWith connGetter (maxSize, timeout, connectionSettings) =
fmap Pool $ createPool connGetter release timeout maxSize
where
release = either (const (pure ())) Hasql.Connection.release
createPool :: IO a
-> (a -> IO ())
-> NominalDiffTime
-> Int
-> IO (ResourcePool.Pool a)
createPool create free idleTime maxResources = ResourcePool.newPool cfg where
-- defaultPoolConfig create free cacheTTL maxResources = PoolConfig
cfg = ResourcePool.defaultPoolConfig create free (realToFrac idleTime) maxResources
-- |
-- Release the connection-pool by closing and removing all connections.
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 = useWithObserver Nothing
-- |
-- Same as 'use' but allows for a custom observer action. You can use it for gathering latency metrics.
useWithObserver :: Maybe ObserverAction
-> Pool
-> Hasql.Session.Session a
-> IO (Either UsageError a)
useWithObserver observer (Pool pool) session =
fmap (either (Left . ConnectionError) (either (Left . SessionError) Right)) $
withResourceOnEither pool $
traverse runQuery
where
runQuery dbConn = maybe action (runWithObserver action) observer
where
action = Hasql.Session.run session dbConn
runWithObserver action doObserve = do
let measure = getTime Monotonic
start <- measure
result <- action
end <- measure
let nsRatio = 1000000000
observed = Observed { latency = toRational (toNanoSecs (end `diffTimeSpec` start) % nsRatio)
}
doObserve observed >> pure result
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
pure $ Right success
Left failure -> do
ResourcePool.destroyResource pool localPool resource
pure $ Left failure
data Stats = Stats
{ currentUsage :: !Int
-- ^ Current number of items.
, available :: !Int
-- ^ Total items available for consumption.
} deriving Show
stats :: Pool -> IO Stats
stats (Pool pool) = currentlyAvailablePerStripe >>= collect where
-- attributes extraction and counting
collect xs = pure $ Stats inUse avail where
inUse = maxResources - avail
avail = sum xs
currentlyAvailablePerStripe = traverse id peekAvailable
peekAvailable = (fmap stripeAvailability) <$> allStripes -- array of IO Int
stripeAvailability ms = maybe quotaPerStripe Unstable.available ms -- if the stripe ref is uninitialised, count the default availability
allStripes = peekStripe <$> Unstable.localPools pool -- array of IO Maybe
peekStripe = tryReadMVar . Unstable.stripeVar
-- data from the pool
quotaPerStripe = maxResources `quotCeil` numStripes
numStripes = length $ Unstable.localPools pool -- can be 'sizeofSmallArray' but requires 'primitive' as dependency
maxResources = Unstable.poolMaxResources . Unstable.poolConfig $ pool
quotCeil x y = let (z, r) = x `quotRem` y in if r == 0 then z else z + 1 -- copied from 'Data.Pool.Internal'
getPoolUsageStat :: Pool -> IO PoolSize
getPoolUsageStat pool = stats pool >>= (pure . currentUsage)