{-# LANGUAGE OverloadedRecordDot #-}
module Hasql.Pool
(   Pool
,   PoolSize
,   Settings
,   ConnectionSettings(..)
,   UsageError(..)
,   ConnectionGetter
,   Stats(..)
,   TimeoutSetting(..)
,   TimeUnit(..)
,   errorToDetailedMsg
,   errorIsTransient
,   stats
,   getPoolUsageStat
,   acquire
,   acquireWith
,   release
,   use
,   useWithObserver
,   useWithPoolAcquisitionTimeout
,   useWithObserverAndPoolAcquisitionTimeout
,   withConnectionWithPoolAcquisitionTimeout
,   withResourceOnEither
,   extendedConnectionSettings
)
where

import qualified    Data.Pool                                   as ResourcePool
import qualified    Data.Text                                   as T
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.Connection.Settings
import qualified    Hasql.Errors
import qualified    Hasql.Session
import              Hasql.Pool.Observer                         (Observed(..), ObserverAction)


-- |
-- A pool of open DB connections.
newtype Pool =
    Pool (ResourcePool.Pool (Either Hasql.Errors.ConnectionError Hasql.Connection.Connection))



type PoolSize         = 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.Errors.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, ConnectionSettings)


data TimeUnit
    =  Microseconds
    |  Milliseconds
    |  Seconds
    |  Minutes
    |  Hours
    |  Days

-- https://www.postgresql.org/docs/18/config-setting.html#CONFIG-SETTING-NAMES-VALUES
instance Show TimeUnit where
    show :: TimeUnit -> String
show TimeUnit
Microseconds   = String
"us"
    show TimeUnit
Milliseconds   = String
"ms"
    show TimeUnit
Seconds        = String
"s"
    show TimeUnit
Minutes        = String
"min"
    show TimeUnit
Hours          = String
"h"
    show TimeUnit
Days           = String
"d"


data TimeoutSetting = TimeoutSetting Word16 TimeUnit

instance Show TimeoutSetting where
    show :: TimeoutSetting -> String
show (TimeoutSetting Word16
v TimeUnit
u) = Word16 -> String
forall a. Show a => a -> String
show Word16
v String -> ShowS
forall a. Semigroup a => a -> a -> a
<> TimeUnit -> String
forall a. Show a => a -> String
show TimeUnit
u


-- | Extended connection settings
data ConnectionSettings = ConnectionSettings
    {   ConnectionSettings -> Text
host                :: T.Text
    ,   ConnectionSettings -> Word16
port                :: Word16
    ,   ConnectionSettings -> Text
user                :: T.Text
    ,   ConnectionSettings -> Text
password            :: T.Text
    ,   ConnectionSettings -> Text
dbName              :: T.Text
    ,   ConnectionSettings -> Word16
connAcqTimeout      :: Word16               -- ^ In seconds: zero, negative, or not specified means wait indefinitely. Doesn't support unit suffixes.
    ,   ConnectionSettings -> TimeoutSetting
txIdleTimeout       :: TimeoutSetting       -- ^ Sets explicit `idle_in_transaction_session_timeout`: zero, negative, or not specified means wait indefinitely.
    ,   ConnectionSettings -> TimeoutSetting
stmtTimeout         :: TimeoutSetting       -- ^ Sets explicit `statement_timeout`: zero, negative, or not specified means wait indefinitely.
    ,   ConnectionSettings -> Text
sslMode             :: T.Text               -- ^ See https://www.postgresql.org/docs/17/libpq-connect.html#LIBPQ-CONNECT-SSLMODE
    ,   ConnectionSettings -> Text
sslRootCert         :: T.Text               -- ^ See https://www.postgresql.org/docs/17/libpq-connect.html#LIBPQ-CONNECT-SSLROOTCERT
    }

-- | https://www.postgresql.org/docs/18/libpq-connect.html#LIBPQ-CONNECT-CONNECT-TIMEOUT
connectTimeout :: Text -> Settings
connectTimeout      = Text -> Text -> Settings
Hasql.Connection.Settings.other Text
"connect_timeout"

-- | https://www.postgresql.org/docs/18/runtime-config-client.html#GUC-TRANSACTION-TIMEOUT
serverOptions :: Text -> Settings
serverOptions = Text -> Text -> Settings
Hasql.Connection.Settings.other Text
"options"

-- | https://www.postgresql.org/docs/18/libpq-connect.html#LIBPQ-CONNECT-SSLMODE
sslmode :: Text -> Settings
sslmode             = Text -> Text -> Settings
Hasql.Connection.Settings.other Text
"sslmode"

-- | https://www.postgresql.org/docs/18/libpq-connect.html#LIBPQ-CONNECT-SSLROOTCERT
sslrootcert :: Text -> Settings
sslrootcert         = Text -> Text -> Settings
Hasql.Connection.Settings.other Text
"sslrootcert"


-- |
-- Given the pool-size, timeout and connection settings
-- create a connection-pool.
acquire :: Settings -> IO Pool
acquire :: Settings -> IO Pool
acquire settings :: Settings
settings@(Int
_, ResidenceTimeout
_, ConnectionSettings
cset) =
    ConnectionGetter -> Settings -> IO Pool
acquireWith
        (Settings -> ConnectionGetter
Hasql.Connection.acquire (Settings -> ConnectionGetter)
-> (ConnectionSettings -> Settings)
-> ConnectionSettings
-> ConnectionGetter
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ConnectionSettings -> Settings
extendedConnectionSettings (ConnectionSettings -> ConnectionGetter)
-> ConnectionSettings -> ConnectionGetter
forall a b. (a -> b) -> a -> b
$ ConnectionSettings
cset)
        Settings
settings


-- | Produce connection settings suitable for acquiring a connection, from an extended set of parameters covering ssl options.
extendedConnectionSettings :: ConnectionSettings -> Hasql.Connection.Settings.Settings
extendedConnectionSettings :: ConnectionSettings -> Settings
extendedConnectionSettings ConnectionSettings
cset =
    (Settings -> Settings -> Settings)
-> Settings -> [Settings] -> Settings
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Settings -> Settings -> Settings
forall a. Semigroup a => a -> a -> a
(<>) Settings
forall a. Monoid a => a
mempty
        [   Text -> Word16 -> Settings
Hasql.Connection.Settings.hostAndPort               ConnectionSettings
cset.host ConnectionSettings
cset.port
        ,   Text -> Settings
Hasql.Connection.Settings.user                      ConnectionSettings
cset.user
        ,   Text -> Settings
Hasql.Connection.Settings.password                  ConnectionSettings
cset.password
        ,   Text -> Settings
Hasql.Connection.Settings.dbname                    ConnectionSettings
cset.dbName
        ,   (Text -> Settings
connectTimeout (Text -> Settings) -> (Word16 -> Text) -> Word16 -> Settings
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. String -> Text
T.pack (String -> Text) -> (Word16 -> String) -> Word16 -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Word16 -> String
forall a. Show a => a -> String
show)                    ConnectionSettings
cset.connAcqTimeout
        ,   Text -> Settings
sslmode                                             ConnectionSettings
cset.sslMode
        ,   Text -> Settings
sslrootcert                                         ConnectionSettings
cset.sslRootCert
        ,   Text -> Settings
serverOptions
                ([(Text, Text)] -> Text
serverOpts
                        -- https://www.postgresql.org/docs/18/runtime-config-client.html#GUC-IDLE-IN-TRANSACTION-SESSION-TIMEOUT
                    [   (Text
"idle_in_transaction_session_timeout", String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ TimeoutSetting -> String
forall a. Show a => a -> String
show ConnectionSettings
cset.txIdleTimeout)
                        -- https://www.postgresql.org/docs/18/runtime-config-client.html#GUC-STATEMENT-TIMEOUT
                    ,   (Text
"statement_timeout",                   String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ TimeoutSetting -> String
forall a. Show a => a -> String
show ConnectionSettings
cset.stmtTimeout)
                    ]
                )
        ]


serverOpts :: [(T.Text,  T.Text)] -> T.Text
serverOpts :: [(Text, Text)] -> Text
serverOpts = (Text -> (Text, Text) -> Text) -> Text -> [(Text, Text)] -> Text
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Text
acc (Text
k ,Text
v) -> Text
acc Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" -c " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
k Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"=" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
v) Text
forall a. Monoid a => a
mempty


-- |
-- Similar to 'acquire', allows for finer configuration.
acquireWith :: ConnectionGetter
            -> Settings
            -> IO Pool
acquireWith :: ConnectionGetter -> Settings -> IO Pool
acquireWith ConnectionGetter
connGetter (Int
maxSize, ResidenceTimeout
sTimeout, ConnectionSettings
_connectionSettings) =
    Pool (Either ConnectionError Connection) -> Pool
Pool (Pool (Either ConnectionError Connection) -> Pool)
-> IO (Pool (Either ConnectionError Connection)) -> IO Pool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ConnectionGetter
-> (Either ConnectionError Connection -> IO ())
-> ResidenceTimeout
-> Int
-> IO (Pool (Either ConnectionError Connection))
forall a.
IO a -> (a -> IO ()) -> ResidenceTimeout -> Int -> IO (Pool a)
createPool ConnectionGetter
connGetter Either ConnectionError Connection -> IO ()
forall {a}. Either a Connection -> IO ()
releaseConn ResidenceTimeout
sTimeout Int
maxSize
    where
        releaseConn :: Either a Connection -> IO ()
releaseConn = (a -> IO ())
-> (Connection -> IO ()) -> Either a Connection -> IO ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (IO () -> a -> IO ()
forall a b. a -> b -> a
const (() -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())) Connection -> IO ()
Hasql.Connection.release


acquisitionTimeoutMicros :: Int -> Maybe Int
acquisitionTimeoutMicros :: Int -> Maybe Int
acquisitionTimeoutMicros Int
seconds
    | Int
seconds Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 =
        Maybe Int
forall a. Maybe a
Nothing
    | Bool
otherwise =
        Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int
forall a. Bounded a => a
maxBound :: Int) (Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
seconds Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Integer
1000000))


createPool :: IO a
           -> (a -> IO ())
           -> NominalDiffTime
           -> PoolSize
           -> IO (ResourcePool.Pool a)
createPool :: forall a.
IO a -> (a -> IO ()) -> ResidenceTimeout -> Int -> IO (Pool a)
createPool IO a
create a -> IO ()
free ResidenceTimeout
idleTime Int
maxResources = PoolConfig a -> IO (Pool a)
forall a. PoolConfig a -> IO (Pool a)
ResourcePool.newPool PoolConfig a
cfg where
    -- defaultPoolConfig create free cacheTTL maxResources = PoolConfig
    cfg :: PoolConfig a
cfg = IO a -> (a -> IO ()) -> Double -> Int -> PoolConfig a
forall a. IO a -> (a -> IO ()) -> Double -> Int -> PoolConfig a
ResourcePool.defaultPoolConfig IO a
create a -> IO ()
free (ResidenceTimeout -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac ResidenceTimeout
idleTime) Int
maxResources


-- |
-- Release the connection-pool by closing and removing all connections.
release :: Pool -> IO ()
release :: Pool -> IO ()
release (Pool Pool (Either ConnectionError Connection)
pool) =
    Pool (Either ConnectionError Connection) -> IO ()
forall a. Pool a -> IO ()
ResourcePool.destroyAllResources Pool (Either ConnectionError Connection)
pool


-- |
-- A union over the connection establishment error and the session error.
data UsageError
    =   AcquisitionTimeoutUsageError
    |   ConnectionError Hasql.Errors.ConnectionError
    |   SessionError    Hasql.Errors.SessionError
    deriving (Int -> UsageError -> ShowS
[UsageError] -> ShowS
UsageError -> String
(Int -> UsageError -> ShowS)
-> (UsageError -> String)
-> ([UsageError] -> ShowS)
-> Show UsageError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> UsageError -> ShowS
showsPrec :: Int -> UsageError -> ShowS
$cshow :: UsageError -> String
show :: UsageError -> String
$cshowList :: [UsageError] -> ShowS
showList :: [UsageError] -> ShowS
Show, UsageError -> UsageError -> Bool
(UsageError -> UsageError -> Bool)
-> (UsageError -> UsageError -> Bool) -> Eq UsageError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: UsageError -> UsageError -> Bool
== :: UsageError -> UsageError -> Bool
$c/= :: UsageError -> UsageError -> Bool
/= :: UsageError -> UsageError -> Bool
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 :: forall a. Pool -> Session a -> IO (Either UsageError a)
use = Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
forall a.
Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
useWithObserver Maybe ObserverAction
forall a. Maybe a
Nothing

-- |
-- Same as 'use' but bounds the time spent waiting for an available pool slot.
-- The timeout is in seconds; zero means wait indefinitely.
useWithPoolAcquisitionTimeout :: Int
                              -> Pool
                              -> Hasql.Session.Session a
                              -> IO (Either UsageError a)
useWithPoolAcquisitionTimeout :: forall a. Int -> Pool -> Session a -> IO (Either UsageError a)
useWithPoolAcquisitionTimeout =
    Maybe ObserverAction
-> Int -> Pool -> Session a -> IO (Either UsageError a)
forall a.
Maybe ObserverAction
-> Int -> Pool -> Session a -> IO (Either UsageError a)
useWithObserverAndPoolAcquisitionTimeout Maybe ObserverAction
forall a. Maybe a
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 :: forall a.
Maybe ObserverAction
-> Pool -> Session a -> IO (Either UsageError a)
useWithObserver Maybe ObserverAction
observer =
    Maybe ObserverAction
-> Int -> Pool -> Session a -> IO (Either UsageError a)
forall a.
Maybe ObserverAction
-> Int -> Pool -> Session a -> IO (Either UsageError a)
useWithObserverAndPoolAcquisitionTimeout Maybe ObserverAction
observer Int
0

-- |
-- Same as 'useWithObserver' but bounds the time spent waiting for an available pool slot.
-- The timeout is in seconds; zero means wait indefinitely.
useWithObserverAndPoolAcquisitionTimeout :: Maybe ObserverAction
                                         -> Int
                                         -> Pool
                                         -> Hasql.Session.Session a
                                         -> IO (Either UsageError a)
useWithObserverAndPoolAcquisitionTimeout :: forall a.
Maybe ObserverAction
-> Int -> Pool -> Session a -> IO (Either UsageError a)
useWithObserverAndPoolAcquisitionTimeout Maybe ObserverAction
observer Int
poolAcquisitionTimeout (Pool Pool (Either ConnectionError Connection)
pool) Session a
session =
    (Either UsageError (Either SessionError a) -> Either UsageError a)
-> IO (Either UsageError (Either SessionError a))
-> IO (Either UsageError a)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((UsageError -> Either UsageError a)
-> (Either SessionError a -> Either UsageError a)
-> Either UsageError (Either SessionError a)
-> Either UsageError a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either UsageError -> Either UsageError a
forall a b. a -> Either a b
Left ((SessionError -> Either UsageError a)
-> (a -> Either UsageError a)
-> Either SessionError a
-> Either UsageError a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (UsageError -> Either UsageError a
forall a b. a -> Either a b
Left (UsageError -> Either UsageError a)
-> (SessionError -> UsageError)
-> SessionError
-> Either UsageError a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. SessionError -> UsageError
SessionError) a -> Either UsageError a
forall a b. b -> Either a b
Right)) (IO (Either UsageError (Either SessionError a))
 -> IO (Either UsageError a))
-> IO (Either UsageError (Either SessionError a))
-> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$
    Maybe Int
-> UsageError
-> Pool (Either ConnectionError Connection)
-> (Either ConnectionError Connection
    -> IO (Either UsageError (Either SessionError a)))
-> IO (Either UsageError (Either SessionError a))
forall failure resource success.
Maybe Int
-> failure
-> Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
withResourceOnEitherTimeout (Int -> Maybe Int
acquisitionTimeoutMicros Int
poolAcquisitionTimeout) UsageError
AcquisitionTimeoutUsageError Pool (Either ConnectionError Connection)
pool ((Either ConnectionError Connection
  -> IO (Either UsageError (Either SessionError a)))
 -> IO (Either UsageError (Either SessionError a)))
-> (Either ConnectionError Connection
    -> IO (Either UsageError (Either SessionError a)))
-> IO (Either UsageError (Either SessionError a))
forall a b. (a -> b) -> a -> b
$
    (ConnectionError -> IO (Either UsageError (Either SessionError a)))
-> (Connection -> IO (Either UsageError (Either SessionError a)))
-> Either ConnectionError Connection
-> IO (Either UsageError (Either SessionError a))
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Either UsageError (Either SessionError a)
-> IO (Either UsageError (Either SessionError a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either UsageError (Either SessionError a)
 -> IO (Either UsageError (Either SessionError a)))
-> (ConnectionError -> Either UsageError (Either SessionError a))
-> ConnectionError
-> IO (Either UsageError (Either SessionError a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. UsageError -> Either UsageError (Either SessionError a)
forall a b. a -> Either a b
Left (UsageError -> Either UsageError (Either SessionError a))
-> (ConnectionError -> UsageError)
-> ConnectionError
-> Either UsageError (Either SessionError a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ConnectionError -> UsageError
ConnectionError) ((Either SessionError a
 -> Either UsageError (Either SessionError a))
-> IO (Either SessionError a)
-> IO (Either UsageError (Either SessionError a))
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Either SessionError a -> Either UsageError (Either SessionError a)
forall a b. b -> Either a b
Right (IO (Either SessionError a)
 -> IO (Either UsageError (Either SessionError a)))
-> (Connection -> IO (Either SessionError a))
-> Connection
-> IO (Either UsageError (Either SessionError a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Connection -> IO (Either SessionError a)
runQuery)
    where
        runQuery :: Connection -> IO (Either SessionError a)
runQuery Connection
dbConn = IO (Either SessionError a)
-> (ObserverAction -> IO (Either SessionError a))
-> Maybe ObserverAction
-> IO (Either SessionError a)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO (Either SessionError a)
action (IO (Either SessionError a)
-> ObserverAction -> IO (Either SessionError a)
forall {b} {a}. IO b -> (Observed -> IO a) -> IO b
runWithObserver IO (Either SessionError a)
action) Maybe ObserverAction
observer
            where
                action :: IO (Either SessionError a)
action = Connection -> Session a -> IO (Either SessionError a)
forall a. Connection -> Session a -> IO (Either SessionError a)
Hasql.Connection.use Connection
dbConn Session a
session

        runWithObserver :: IO b -> (Observed -> IO a) -> IO b
runWithObserver IO b
action Observed -> IO a
doObserve = do
            let measure :: IO TimeSpec
measure = Clock -> IO TimeSpec
getTime Clock
Monotonic
            start  <- IO TimeSpec
measure
            result <- action
            end    <- measure
            let nsRatio  = a
1000000000
                observed = Observed {   latency :: Ratio Integer
latency = Ratio Integer -> Ratio Integer
forall a. Real a => a -> Ratio Integer
toRational (TimeSpec -> Integer
toNanoSecs (TimeSpec
end TimeSpec -> TimeSpec -> TimeSpec
`diffTimeSpec` TimeSpec
start) Integer -> Integer -> Ratio Integer
forall a. Integral a => a -> a -> Ratio a
% Integer
forall {a}. Num a => a
nsRatio)
                                    }
            doObserve observed >> pure result


-- |
-- Borrow a live connection from the pool and run a callback with it.
--
-- This is useful for libraries that need to pin one connection across multiple
-- operations, for example while managing their own transaction state.
--
-- The timeout is in seconds; zero means wait indefinitely.
withConnectionWithPoolAcquisitionTimeout :: Int
                                         -> Pool
                                         -> (Hasql.Connection.Connection -> IO (Either UsageError a))
                                         -> IO (Either UsageError a)
withConnectionWithPoolAcquisitionTimeout :: forall a.
Int
-> Pool
-> (Connection -> IO (Either UsageError a))
-> IO (Either UsageError a)
withConnectionWithPoolAcquisitionTimeout Int
poolAcquisitionTimeout (Pool Pool (Either ConnectionError Connection)
pool) Connection -> IO (Either UsageError a)
act =
    Maybe Int
-> UsageError
-> Pool (Either ConnectionError Connection)
-> (Either ConnectionError Connection -> IO (Either UsageError a))
-> IO (Either UsageError a)
forall failure resource success.
Maybe Int
-> failure
-> Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
withResourceOnEitherTimeout (Int -> Maybe Int
acquisitionTimeoutMicros Int
poolAcquisitionTimeout) UsageError
AcquisitionTimeoutUsageError Pool (Either ConnectionError Connection)
pool ((Either ConnectionError Connection -> IO (Either UsageError a))
 -> IO (Either UsageError a))
-> (Either ConnectionError Connection -> IO (Either UsageError a))
-> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$
    (ConnectionError -> IO (Either UsageError a))
-> (Connection -> IO (Either UsageError a))
-> Either ConnectionError Connection
-> IO (Either UsageError a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (Either UsageError a -> IO (Either UsageError a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either UsageError a -> IO (Either UsageError a))
-> (ConnectionError -> Either UsageError a)
-> ConnectionError
-> IO (Either UsageError a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. UsageError -> Either UsageError a
forall a b. a -> Either a b
Left (UsageError -> Either UsageError a)
-> (ConnectionError -> UsageError)
-> ConnectionError
-> Either UsageError a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ConnectionError -> UsageError
ConnectionError) Connection -> IO (Either UsageError a)
act


withResourceOnEither :: ResourcePool.Pool resource
                     -> (resource -> IO (Either failure success))
                     -> IO (Either failure success)
withResourceOnEither :: forall resource failure success.
Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
withResourceOnEither Pool resource
pool resource -> IO (Either failure success)
act = IO (Either failure success) -> IO (Either failure success)
forall a. IO a -> IO a
mask_ (IO (Either failure success) -> IO (Either failure success))
-> IO (Either failure success) -> IO (Either failure success)
forall a b. (a -> b) -> a -> b
$ do
    (resource, localPool) <- Pool resource -> IO (resource, LocalPool resource)
forall a. Pool a -> IO (a, LocalPool a)
ResourcePool.takeResource Pool resource
pool
    failureOrSuccess      <- act resource `onException` ResourcePool.destroyResource pool localPool resource
    case failureOrSuccess of
        Right success
success -> do
            LocalPool resource -> resource -> IO ()
forall a. LocalPool a -> a -> IO ()
ResourcePool.putResource LocalPool resource
localPool resource
resource
            Either failure success -> IO (Either failure success)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either failure success -> IO (Either failure success))
-> Either failure success -> IO (Either failure success)
forall a b. (a -> b) -> a -> b
$ success -> Either failure success
forall a b. b -> Either a b
Right success
success
        Left failure
failure -> do
            Pool resource -> LocalPool resource -> resource -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
ResourcePool.destroyResource Pool resource
pool LocalPool resource
localPool resource
resource
            Either failure success -> IO (Either failure success)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either failure success -> IO (Either failure success))
-> Either failure success -> IO (Either failure success)
forall a b. (a -> b) -> a -> b
$ failure -> Either failure success
forall a b. a -> Either a b
Left failure
failure


withResourceOnEitherTimeout :: Maybe Int
                            -> failure
                            -> ResourcePool.Pool resource
                            -> (resource -> IO (Either failure success))
                            -> IO (Either failure success)
withResourceOnEitherTimeout :: forall failure resource success.
Maybe Int
-> failure
-> Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
withResourceOnEitherTimeout Maybe Int
Nothing failure
_ Pool resource
pool resource -> IO (Either failure success)
act =
    Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
forall resource failure success.
Pool resource
-> (resource -> IO (Either failure success))
-> IO (Either failure success)
withResourceOnEither Pool resource
pool resource -> IO (Either failure success)
act
withResourceOnEitherTimeout (Just Int
acquisitionTimeout) failure
timeoutFailure Pool resource
pool resource -> IO (Either failure success)
act = ((forall a. IO a -> IO a) -> IO (Either failure success))
-> IO (Either failure success)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO (Either failure success))
 -> IO (Either failure success))
-> ((forall a. IO a -> IO a) -> IO (Either failure success))
-> IO (Either failure success)
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
    resourceOrTimeout <- Int -> Pool resource -> IO (Maybe (resource, LocalPool resource))
forall resource.
Int -> Pool resource -> IO (Maybe (resource, LocalPool resource))
takeResourceWithin Int
acquisitionTimeout Pool resource
pool
    case resourceOrTimeout of
        Maybe (resource, LocalPool resource)
Nothing ->
            Either failure success -> IO (Either failure success)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either failure success -> IO (Either failure success))
-> Either failure success -> IO (Either failure success)
forall a b. (a -> b) -> a -> b
$ failure -> Either failure success
forall a b. a -> Either a b
Left failure
timeoutFailure
        Just (resource
resource, LocalPool resource
localPool) -> do
            failureOrSuccess <- IO (Either failure success) -> IO (Either failure success)
forall a. IO a -> IO a
restore (resource -> IO (Either failure success)
act resource
resource) IO (Either failure success) -> IO () -> IO (Either failure success)
forall a b. IO a -> IO b -> IO a
`onException` Pool resource -> LocalPool resource -> resource -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
ResourcePool.destroyResource Pool resource
pool LocalPool resource
localPool resource
resource
            case failureOrSuccess of
                Right success
success -> do
                    LocalPool resource -> resource -> IO ()
forall a. LocalPool a -> a -> IO ()
ResourcePool.putResource LocalPool resource
localPool resource
resource
                    Either failure success -> IO (Either failure success)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either failure success -> IO (Either failure success))
-> Either failure success -> IO (Either failure success)
forall a b. (a -> b) -> a -> b
$ success -> Either failure success
forall a b. b -> Either a b
Right success
success
                Left failure
failure -> do
                    Pool resource -> LocalPool resource -> resource -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
ResourcePool.destroyResource Pool resource
pool LocalPool resource
localPool resource
resource
                    Either failure success -> IO (Either failure success)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either failure success -> IO (Either failure success))
-> Either failure success -> IO (Either failure success)
forall a b. (a -> b) -> a -> b
$ failure -> Either failure success
forall a b. a -> Either a b
Left failure
failure


-- We cannot implement this as `timeout acquisitionTimeout (ResourcePool.takeResource pool)`.
-- In practice that waits until resource-pool eventually returns a resource in exhausted-pool
-- scenarios. Instead, mirror resource-pool's checkout logic and race availability against
-- our timeout in STM, like hasql-pool does.
takeResourceWithin :: Int
                   -> ResourcePool.Pool resource
                   -> IO (Maybe (resource, ResourcePool.LocalPool resource))
takeResourceWithin :: forall resource.
Int -> Pool resource -> IO (Maybe (resource, LocalPool resource))
takeResourceWithin Int
acquisitionTimeout Pool resource
pool = IO (Maybe (resource, LocalPool resource))
-> IO (Maybe (resource, LocalPool resource))
forall a. IO a -> IO a
mask_ (IO (Maybe (resource, LocalPool resource))
 -> IO (Maybe (resource, LocalPool resource)))
-> IO (Maybe (resource, LocalPool resource))
-> IO (Maybe (resource, LocalPool resource))
forall a b. (a -> b) -> a -> b
$ do
    delay <- Int -> IO (TVar Bool)
newDelay Int
acquisitionTimeout
    localPool <- Unstable.getLocalPool (Unstable.localPools pool)
    join . atomically $
        asum
            [ do
                stripe <- readTVar (Unstable.stripeVar localPool)
                case stripe of
                    Unstable.Stripe Int
0 [Entry resource]
_ Queue resource
_ Queue resource
_ ->
                        STM (IO (Maybe (resource, LocalPool resource)))
forall a. STM a
retry
                    Stripe resource
_ ->
                        ((resource, LocalPool resource)
 -> Maybe (resource, LocalPool resource))
-> IO (resource, LocalPool resource)
-> IO (Maybe (resource, LocalPool resource))
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (resource, LocalPool resource)
-> Maybe (resource, LocalPool resource)
forall a. a -> Maybe a
Just (IO (resource, LocalPool resource)
 -> IO (Maybe (resource, LocalPool resource)))
-> STM (IO (resource, LocalPool resource))
-> STM (IO (Maybe (resource, LocalPool resource)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pool resource
-> LocalPool resource
-> Stripe resource
-> STM (IO (resource, LocalPool resource))
forall resource.
Pool resource
-> LocalPool resource
-> Stripe resource
-> STM (IO (resource, LocalPool resource))
takeAvailableResource Pool resource
pool LocalPool resource
localPool Stripe resource
stripe
            , do
                timedOut <- readTVar delay
                if timedOut
                    then pure $ pure Nothing
                    else retry
            ]


newDelay :: Int -> IO (TVar Bool)
newDelay :: Int -> IO (TVar Bool)
newDelay Int
delayMicros = do
    delay <- Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False
    void . forkIO $ do
        threadDelay delayMicros
        atomically $ writeTVar delay True
    pure delay


takeAvailableResource :: ResourcePool.Pool resource
                      -> ResourcePool.LocalPool resource
                      -> Unstable.Stripe resource
                      -> STM (IO (resource, ResourcePool.LocalPool resource))
takeAvailableResource :: forall resource.
Pool resource
-> LocalPool resource
-> Stripe resource
-> STM (IO (resource, LocalPool resource))
takeAvailableResource Pool resource
pool LocalPool resource
localPool (Unstable.Stripe Int
available [Entry resource]
cached Queue resource
queue Queue resource
queueR) =
    case [Entry resource]
cached of
        [] -> do
            TVar (Stripe resource) -> Stripe resource -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool resource -> TVar (Stripe resource)
forall a. LocalPool a -> TVar (Stripe a)
Unstable.stripeVar LocalPool resource
localPool) (Stripe resource -> STM ()) -> Stripe resource -> STM ()
forall a b. (a -> b) -> a -> b
$! Int
-> [Entry resource]
-> Queue resource
-> Queue resource
-> Stripe resource
forall a. Int -> [Entry a] -> Queue a -> Queue a -> Stripe a
Unstable.Stripe (Int
available Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Entry resource]
cached Queue resource
queue Queue resource
queueR
            IO (resource, LocalPool resource)
-> STM (IO (resource, LocalPool resource))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (resource, LocalPool resource)
 -> STM (IO (resource, LocalPool resource)))
-> IO (resource, LocalPool resource)
-> STM (IO (resource, LocalPool resource))
forall a b. (a -> b) -> a -> b
$ do
                resource <-
                    PoolConfig resource -> IO resource
forall a. PoolConfig a -> IO a
Unstable.createResource (Pool resource -> PoolConfig resource
forall a. Pool a -> PoolConfig a
Unstable.poolConfig Pool resource
pool)
                        IO resource -> IO () -> IO resource
forall a b. IO a -> IO b -> IO a
`onException` TVar (Stripe resource) -> IO ()
forall a. TVar (Stripe a) -> IO ()
Unstable.restoreSize (LocalPool resource -> TVar (Stripe resource)
forall a. LocalPool a -> TVar (Stripe a)
Unstable.stripeVar LocalPool resource
localPool)
                pure (resource, localPool)
        Unstable.Entry resource
resource Double
_ : [Entry resource]
remainingCached -> do
            TVar (Stripe resource) -> Stripe resource -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool resource -> TVar (Stripe resource)
forall a. LocalPool a -> TVar (Stripe a)
Unstable.stripeVar LocalPool resource
localPool) (Stripe resource -> STM ()) -> Stripe resource -> STM ()
forall a b. (a -> b) -> a -> b
$! Int
-> [Entry resource]
-> Queue resource
-> Queue resource
-> Stripe resource
forall a. Int -> [Entry a] -> Queue a -> Queue a -> Stripe a
Unstable.Stripe (Int
available Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Entry resource]
remainingCached Queue resource
queue Queue resource
queueR
            IO (resource, LocalPool resource)
-> STM (IO (resource, LocalPool resource))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (resource, LocalPool resource)
 -> STM (IO (resource, LocalPool resource)))
-> IO (resource, LocalPool resource)
-> STM (IO (resource, LocalPool resource))
forall a b. (a -> b) -> a -> b
$ (resource, LocalPool resource) -> IO (resource, LocalPool resource)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (resource
resource, LocalPool resource
localPool)


data Stats = Stats
    {   Stats -> Int
currentUsage  :: !Int
        -- ^ Current number of items.
    ,   Stats -> Int
available     :: !Int
        -- ^ Total items available for consumption.
    } deriving Int -> Stats -> ShowS
[Stats] -> ShowS
Stats -> String
(Int -> Stats -> ShowS)
-> (Stats -> String) -> ([Stats] -> ShowS) -> Show Stats
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Stats -> ShowS
showsPrec :: Int -> Stats -> ShowS
$cshow :: Stats -> String
show :: Stats -> String
$cshowList :: [Stats] -> ShowS
showList :: [Stats] -> ShowS
Show


stats :: Pool -> IO Stats
stats :: Pool -> IO Stats
stats (Pool Pool (Either ConnectionError Connection)
pool) = IO (SmallArray Int)
currentlyAvailablePerStripe IO (SmallArray Int) -> (SmallArray Int -> IO Stats) -> IO Stats
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SmallArray Int -> IO Stats
collect where
    -- attributes extraction and counting
    collect :: SmallArray Int -> IO Stats
collect SmallArray Int
xs = Stats -> IO Stats
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Stats -> IO Stats) -> Stats -> IO Stats
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Stats
Stats Int
inUse Int
avail where
        inUse :: Int
inUse = Int
maxResources Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
avail
        avail :: Int
avail = SmallArray Int -> Int
forall a. Num a => SmallArray a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum SmallArray Int
xs

    currentlyAvailablePerStripe :: IO (SmallArray Int)
currentlyAvailablePerStripe = (IO Int -> IO Int) -> SmallArray (IO Int) -> IO (SmallArray Int)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> SmallArray a -> f (SmallArray b)
traverse IO Int -> IO Int
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id SmallArray (IO Int)
peekAvailable
    peekAvailable :: SmallArray (IO Int)
peekAvailable               = ((Stripe (Either ConnectionError Connection) -> Int)
-> IO (Stripe (Either ConnectionError Connection)) -> IO Int
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Stripe (Either ConnectionError Connection) -> Int
forall {a}. Stripe a -> Int
stripeAvailability) (IO (Stripe (Either ConnectionError Connection)) -> IO Int)
-> SmallArray (IO (Stripe (Either ConnectionError Connection)))
-> SmallArray (IO Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallArray (IO (Stripe (Either ConnectionError Connection)))
allStripes    -- array of IO Int
    stripeAvailability :: Stripe a -> Int
stripeAvailability Stripe a
ms       = Stripe a -> Int
forall {a}. Stripe a -> Int
Unstable.available Stripe a
ms                       -- stripe is always initialized with value, as it's from TVar
    allStripes :: SmallArray (IO (Stripe (Either ConnectionError Connection)))
allStripes                  = LocalPool (Either ConnectionError Connection)
-> IO (Stripe (Either ConnectionError Connection))
forall {a}. LocalPool a -> IO (Stripe a)
peekStripe (LocalPool (Either ConnectionError Connection)
 -> IO (Stripe (Either ConnectionError Connection)))
-> SmallArray (LocalPool (Either ConnectionError Connection))
-> SmallArray (IO (Stripe (Either ConnectionError Connection)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pool (Either ConnectionError Connection)
-> SmallArray (LocalPool (Either ConnectionError Connection))
forall a. Pool a -> SmallArray (LocalPool a)
Unstable.localPools Pool (Either ConnectionError Connection)
pool     -- array of IO vals
    peekStripe :: LocalPool a -> IO (Stripe a)
peekStripe                  = TVar (Stripe a) -> IO (Stripe a)
forall a. TVar a -> IO a
readTVarIO (TVar (Stripe a) -> IO (Stripe a))
-> (LocalPool a -> TVar (Stripe a)) -> LocalPool a -> IO (Stripe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
Unstable.stripeVar

    -- data from the pool
    maxResources :: Int
maxResources                = PoolConfig (Either ConnectionError Connection) -> Int
forall a. PoolConfig a -> Int
Unstable.poolMaxResources (PoolConfig (Either ConnectionError Connection) -> Int)
-> (Pool (Either ConnectionError Connection)
    -> PoolConfig (Either ConnectionError Connection))
-> Pool (Either ConnectionError Connection)
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Pool (Either ConnectionError Connection)
-> PoolConfig (Either ConnectionError Connection)
forall a. Pool a -> PoolConfig a
Unstable.poolConfig (Pool (Either ConnectionError Connection) -> Int)
-> Pool (Either ConnectionError Connection) -> Int
forall a b. (a -> b) -> a -> b
$ Pool (Either ConnectionError Connection)
pool
    _quotaPerStripe :: Int
_quotaPerStripe             = Int
maxResources Int -> Int -> Int
forall {a}. Integral a => a -> a -> a
`quotCeil` Int
_numStripes
    _numStripes :: Int
_numStripes                 = SmallArray (LocalPool (Either ConnectionError Connection)) -> Int
forall a. SmallArray a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (SmallArray (LocalPool (Either ConnectionError Connection)) -> Int)
-> SmallArray (LocalPool (Either ConnectionError Connection))
-> Int
forall a b. (a -> b) -> a -> b
$ Pool (Either ConnectionError Connection)
-> SmallArray (LocalPool (Either ConnectionError Connection))
forall a. Pool a -> SmallArray (LocalPool a)
Unstable.localPools Pool (Either ConnectionError Connection)
pool  -- can be 'sizeofSmallArray' but requires 'primitive' as dependency
    quotCeil :: a -> a -> a
quotCeil a
x a
y                = let (a
z, a
r) = a
x a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
`quotRem` a
y in if a
r a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
0 then a
z else a
z a -> a -> a
forall a. Num a => a -> a -> a
+ a
1  -- copied from 'Data.Pool.Internal'


getPoolUsageStat :: Pool -> IO PoolSize
getPoolUsageStat :: Pool -> IO Int
getPoolUsageStat Pool
pool = Stats -> Int
currentUsage (Stats -> Int) -> IO Stats -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pool -> IO Stats
stats Pool
pool


errorToDetailedMsg :: UsageError -> T.Text
errorToDetailedMsg :: UsageError -> Text
errorToDetailedMsg = \case
    UsageError
AcquisitionTimeoutUsageError ->
        Text
"Connection acquisition timeout"
    ConnectionError ConnectionError
err ->
        ConnectionError -> Text
forall e. IsError e => e -> Text
Hasql.Errors.toDetailedText ConnectionError
err
    SessionError SessionError
err ->
        SessionError -> Text
forall e. IsError e => e -> Text
Hasql.Errors.toDetailedText SessionError
err


errorIsTransient :: UsageError -> Bool
errorIsTransient :: UsageError -> Bool
errorIsTransient = \case
    UsageError
AcquisitionTimeoutUsageError ->
        Bool
True
    ConnectionError ConnectionError
err ->
        ConnectionError -> Bool
forall a. IsError a => a -> Bool
Hasql.Errors.isTransient ConnectionError
err
    SessionError SessionError
err ->
        SessionError -> Bool
forall a. IsError a => a -> Bool
Hasql.Errors.isTransient SessionError
err