{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE OverloadedStrings #-}

-- | The env, monad, pool state, and savepoint ladder shared by the pooled backends.
--
-- This is the API for backend authors.
module Arbiter.Core.Backend
  ( -- * Database monad
    Db (..)
  , Env (..)
  , Driver (..)
  , runDb
  , inTransaction

    -- * Environment creation
  , createEnvWithConfig
  , createEnvWithPool
  , destroyEnv
  , disableListener
  , useDedicatedListener
  , poolListener

    -- * Pool state
  , PoolState (..)
  , HasPoolState (..)
  , withConn
  , pinConnection
  , withSavepointTransaction
  ) where

import Control.Monad.Catch (MonadCatch, MonadMask, MonadThrow)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Reader (MonadReader, ReaderT (..), asks, local, runReaderT)
import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as BSC
import Data.Foldable (traverse_)
import Data.Pool (Pool, defaultPoolConfig, destroyAllResources, newPool, setNumStripes, withResource)
import Data.Text (Text)
import UnliftIO (MonadUnliftIO, mask, onException, withRunInIO)

import Arbiter.Core.Exceptions (throwInternal)
import Arbiter.Core.Job.Schema (SchemaName)
import Arbiter.Core.Listen (ListenConn, Listener, newListener)
import Arbiter.Core.PoolConfig (PoolConfig (..))
import Arbiter.Core.QueueRegistry (JobPayloadRegistry)

-- | Pool and the pinned connection with its transaction depth.
data PoolState conn = PoolState
  { forall conn. PoolState conn -> Maybe (Pool conn)
connectionPool :: Maybe (Pool conn)
  -- ^ The pool to check connections out of. 'Nothing' for a pinned-only env.
  , forall conn. PoolState conn -> Maybe (conn, Int)
pinned :: Maybe (conn, Int)
  -- ^ The pinned connection and its transaction depth. Depth 0 means no open transaction.
  }

-- | Ambient access to the pool state.
class (Monad m) => HasPoolState conn m | m -> conn where
  -- | The current pool state.
  getPoolState :: m (PoolState conn)

  -- | Run an action under an adjusted pool state.
  localPoolState :: (PoolState conn -> PoolState conn) -> m a -> m a

-- | What a connection type contributes to its env.
data Driver conn cfg = Driver
  { forall conn cfg.
Driver conn cfg -> conn -> (ListenConn -> IO ()) -> IO ()
withListenConn :: conn -> (ListenConn -> IO ()) -> IO ()
  -- ^ Run the listener loop on a connection's driver handle.
  , forall conn cfg. Driver conn cfg -> cfg
initialConfig :: cfg
  -- ^ The driver state a new env starts with.
  }

-- | Schema name, pool state, listener, and the driver's own state.
data Env conn cfg (registry :: JobPayloadRegistry) = Env
  { forall conn cfg (registry :: JobPayloadRegistry).
Env conn cfg registry -> Text
schema :: SchemaName
  -- ^ The schema the arbiter tables live in.
  , forall conn cfg (registry :: JobPayloadRegistry).
Env conn cfg registry -> PoolState conn
poolState :: PoolState conn
  -- ^ The pool and the pinned connection.
  , forall conn cfg (registry :: JobPayloadRegistry).
Env conn cfg registry -> Maybe Listener
listener :: Maybe Listener
  -- ^ Resolved LISTEN source. 'Nothing' runs poll-only.
  , forall conn cfg (registry :: JobPayloadRegistry).
Env conn cfg registry -> cfg
driverConfig :: cfg
  -- ^ The driver's own state.
  }

-- | A pooled backend's database monad.
newtype Db conn cfg (registry :: JobPayloadRegistry) m a = Db
  { forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Db conn cfg registry m a -> ReaderT (Env conn cfg registry) m a
unDb :: ReaderT (Env conn cfg registry) m a
  -- ^ The action as a reader over its env.
  }
  deriving newtype
    ( Functor (Db conn cfg registry m)
Functor (Db conn cfg registry m) =>
(forall a. a -> Db conn cfg registry m a)
-> (forall a b.
    Db conn cfg registry m (a -> b)
    -> Db conn cfg registry m a -> Db conn cfg registry m b)
-> (forall a b c.
    (a -> b -> c)
    -> Db conn cfg registry m a
    -> Db conn cfg registry m b
    -> Db conn cfg registry m c)
-> (forall a b.
    Db conn cfg registry m a
    -> Db conn cfg registry m b -> Db conn cfg registry m b)
-> (forall a b.
    Db conn cfg registry m a
    -> Db conn cfg registry m b -> Db conn cfg registry m a)
-> Applicative (Db conn cfg registry m)
forall a. a -> Db conn cfg registry m a
forall a b.
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m a
forall a b.
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
forall a b.
Db conn cfg registry m (a -> b)
-> Db conn cfg registry m a -> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
Applicative m =>
Functor (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Applicative m =>
a -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Applicative m =>
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Applicative m =>
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Applicative m =>
Db conn cfg registry m (a -> b)
-> Db conn cfg registry m a -> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b
       c.
Applicative m =>
(a -> b -> c)
-> Db conn cfg registry m a
-> Db conn cfg registry m b
-> Db conn cfg registry m c
forall a b c.
(a -> b -> c)
-> Db conn cfg registry m a
-> Db conn cfg registry m b
-> Db conn cfg registry m c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Applicative m =>
a -> Db conn cfg registry m a
pure :: forall a. a -> Db conn cfg registry m a
$c<*> :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Applicative m =>
Db conn cfg registry m (a -> b)
-> Db conn cfg registry m a -> Db conn cfg registry m b
<*> :: forall a b.
Db conn cfg registry m (a -> b)
-> Db conn cfg registry m a -> Db conn cfg registry m b
$cliftA2 :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b
       c.
Applicative m =>
(a -> b -> c)
-> Db conn cfg registry m a
-> Db conn cfg registry m b
-> Db conn cfg registry m c
liftA2 :: forall a b c.
(a -> b -> c)
-> Db conn cfg registry m a
-> Db conn cfg registry m b
-> Db conn cfg registry m c
$c*> :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Applicative m =>
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
*> :: forall a b.
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
$c<* :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Applicative m =>
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m a
<* :: forall a b.
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m a
Applicative
    , (forall a b.
 (a -> b) -> Db conn cfg registry m a -> Db conn cfg registry m b)
-> (forall a b.
    a -> Db conn cfg registry m b -> Db conn cfg registry m a)
-> Functor (Db conn cfg registry m)
forall a b.
a -> Db conn cfg registry m b -> Db conn cfg registry m a
forall a b.
(a -> b) -> Db conn cfg registry m a -> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Functor m =>
a -> Db conn cfg registry m b -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Functor m =>
(a -> b) -> Db conn cfg registry m a -> Db conn cfg registry m b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Functor m =>
(a -> b) -> Db conn cfg registry m a -> Db conn cfg registry m b
fmap :: forall a b.
(a -> b) -> Db conn cfg registry m a -> Db conn cfg registry m b
$c<$ :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Functor m =>
a -> Db conn cfg registry m b -> Db conn cfg registry m a
<$ :: forall a b.
a -> Db conn cfg registry m b -> Db conn cfg registry m a
Functor
    , Applicative (Db conn cfg registry m)
Applicative (Db conn cfg registry m) =>
(forall a b.
 Db conn cfg registry m a
 -> (a -> Db conn cfg registry m b) -> Db conn cfg registry m b)
-> (forall a b.
    Db conn cfg registry m a
    -> Db conn cfg registry m b -> Db conn cfg registry m b)
-> (forall a. a -> Db conn cfg registry m a)
-> Monad (Db conn cfg registry m)
forall a. a -> Db conn cfg registry m a
forall a b.
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
forall a b.
Db conn cfg registry m a
-> (a -> Db conn cfg registry m b) -> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
Monad m =>
Applicative (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Monad m =>
a -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Monad m =>
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Monad m =>
Db conn cfg registry m a
-> (a -> Db conn cfg registry m b) -> Db conn cfg registry m b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Monad m =>
Db conn cfg registry m a
-> (a -> Db conn cfg registry m b) -> Db conn cfg registry m b
>>= :: forall a b.
Db conn cfg registry m a
-> (a -> Db conn cfg registry m b) -> Db conn cfg registry m b
$c>> :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b.
Monad m =>
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
>> :: forall a b.
Db conn cfg registry m a
-> Db conn cfg registry m b -> Db conn cfg registry m b
$creturn :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Monad m =>
a -> Db conn cfg registry m a
return :: forall a. a -> Db conn cfg registry m a
Monad
    , MonadThrow (Db conn cfg registry m)
MonadThrow (Db conn cfg registry m) =>
(forall e a.
 (HasCallStack, Exception e) =>
 Db conn cfg registry m a
 -> (e -> Db conn cfg registry m a) -> Db conn cfg registry m a)
-> (forall e a.
    Exception e =>
    Db conn cfg registry m a
    -> (ExceptionWithContext e -> Db conn cfg registry m a)
    -> Db conn cfg registry m a)
-> MonadCatch (Db conn cfg registry m)
forall e a.
(HasCallStack, Exception e) =>
Db conn cfg registry m a
-> (e -> Db conn cfg registry m a) -> Db conn cfg registry m a
forall e a.
Exception e =>
Db conn cfg registry m a
-> (ExceptionWithContext e -> Db conn cfg registry m a)
-> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
MonadCatch m =>
MonadThrow (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadCatch m, HasCallStack, Exception e) =>
Db conn cfg registry m a
-> (e -> Db conn cfg registry m a) -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
Db conn cfg registry m a
-> (ExceptionWithContext e -> Db conn cfg registry m a)
-> Db conn cfg registry m a
forall (m :: * -> *).
MonadThrow m =>
(forall e a.
 (HasCallStack, Exception e) =>
 m a -> (e -> m a) -> m a)
-> (forall e a.
    Exception e =>
    m a -> (ExceptionWithContext e -> m a) -> m a)
-> MonadCatch m
$ccatch :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadCatch m, HasCallStack, Exception e) =>
Db conn cfg registry m a
-> (e -> Db conn cfg registry m a) -> Db conn cfg registry m a
catch :: forall e a.
(HasCallStack, Exception e) =>
Db conn cfg registry m a
-> (e -> Db conn cfg registry m a) -> Db conn cfg registry m a
$ccatchNoPropagate :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
Db conn cfg registry m a
-> (ExceptionWithContext e -> Db conn cfg registry m a)
-> Db conn cfg registry m a
catchNoPropagate :: forall e a.
Exception e =>
Db conn cfg registry m a
-> (ExceptionWithContext e -> Db conn cfg registry m a)
-> Db conn cfg registry m a
MonadCatch
    , Monad (Db conn cfg registry m)
Monad (Db conn cfg registry m) =>
(forall a. HasCallStack => String -> Db conn cfg registry m a)
-> MonadFail (Db conn cfg registry m)
forall a. HasCallStack => String -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
MonadFail m =>
Monad (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> Db conn cfg registry m a
forall (m :: * -> *).
Monad m =>
(forall a. HasCallStack => String -> m a) -> MonadFail m
$cfail :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> Db conn cfg registry m a
fail :: forall a. HasCallStack => String -> Db conn cfg registry m a
MonadFail
    , Monad (Db conn cfg registry m)
Monad (Db conn cfg registry m) =>
(forall a. IO a -> Db conn cfg registry m a)
-> MonadIO (Db conn cfg registry m)
forall a. IO a -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
MonadIO m =>
Monad (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
MonadIO m =>
IO a -> Db conn cfg registry m a
forall (m :: * -> *).
Monad m =>
(forall a. IO a -> m a) -> MonadIO m
$cliftIO :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
MonadIO m =>
IO a -> Db conn cfg registry m a
liftIO :: forall a. IO a -> Db conn cfg registry m a
MonadIO
    , MonadCatch (Db conn cfg registry m)
MonadCatch (Db conn cfg registry m) =>
(forall b.
 HasCallStack =>
 ((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
  -> Db conn cfg registry m b)
 -> Db conn cfg registry m b)
-> (forall b.
    HasCallStack =>
    ((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
     -> Db conn cfg registry m b)
    -> Db conn cfg registry m b)
-> (forall a b c.
    HasCallStack =>
    Db conn cfg registry m a
    -> (a -> ExitCase b -> Db conn cfg registry m c)
    -> (a -> Db conn cfg registry m b)
    -> Db conn cfg registry m (b, c))
-> MonadMask (Db conn cfg registry m)
forall b.
HasCallStack =>
((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
 -> Db conn cfg registry m b)
-> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
MonadMask m =>
MonadCatch (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) b.
(MonadMask m, HasCallStack) =>
((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
 -> Db conn cfg registry m b)
-> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b
       c.
(MonadMask m, HasCallStack) =>
Db conn cfg registry m a
-> (a -> ExitCase b -> Db conn cfg registry m c)
-> (a -> Db conn cfg registry m b)
-> Db conn cfg registry m (b, c)
forall a b c.
HasCallStack =>
Db conn cfg registry m a
-> (a -> ExitCase b -> Db conn cfg registry m c)
-> (a -> Db conn cfg registry m b)
-> Db conn cfg registry m (b, c)
forall (m :: * -> *).
MonadCatch m =>
(forall b. HasCallStack => ((forall a. m a -> m a) -> m b) -> m b)
-> (forall b.
    HasCallStack =>
    ((forall a. m a -> m a) -> m b) -> m b)
-> (forall a b c.
    HasCallStack =>
    m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c))
-> MonadMask m
$cmask :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) b.
(MonadMask m, HasCallStack) =>
((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
 -> Db conn cfg registry m b)
-> Db conn cfg registry m b
mask :: forall b.
HasCallStack =>
((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
 -> Db conn cfg registry m b)
-> Db conn cfg registry m b
$cuninterruptibleMask :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) b.
(MonadMask m, HasCallStack) =>
((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
 -> Db conn cfg registry m b)
-> Db conn cfg registry m b
uninterruptibleMask :: forall b.
HasCallStack =>
((forall a. Db conn cfg registry m a -> Db conn cfg registry m a)
 -> Db conn cfg registry m b)
-> Db conn cfg registry m b
$cgeneralBracket :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a b
       c.
(MonadMask m, HasCallStack) =>
Db conn cfg registry m a
-> (a -> ExitCase b -> Db conn cfg registry m c)
-> (a -> Db conn cfg registry m b)
-> Db conn cfg registry m (b, c)
generalBracket :: forall a b c.
HasCallStack =>
Db conn cfg registry m a
-> (a -> ExitCase b -> Db conn cfg registry m c)
-> (a -> Db conn cfg registry m b)
-> Db conn cfg registry m (b, c)
MonadMask
    , MonadReader (Env conn cfg registry)
    , Monad (Db conn cfg registry m)
Monad (Db conn cfg registry m) =>
(forall e a.
 (HasCallStack, Exception e) =>
 e -> Db conn cfg registry m a)
-> (forall e a.
    Exception e =>
    ExceptionWithContext e -> Db conn cfg registry m a)
-> MonadThrow (Db conn cfg registry m)
forall e a.
(HasCallStack, Exception e) =>
e -> Db conn cfg registry m a
forall e a.
Exception e =>
ExceptionWithContext e -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
MonadThrow m =>
Monad (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> Db conn cfg registry m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadThrow m, Exception e) =>
ExceptionWithContext e -> Db conn cfg registry m a
forall (m :: * -> *).
Monad m =>
(forall e a. (HasCallStack, Exception e) => e -> m a)
-> (forall e a. Exception e => ExceptionWithContext e -> m a)
-> MonadThrow m
$cthrowM :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> Db conn cfg registry m a
throwM :: forall e a.
(HasCallStack, Exception e) =>
e -> Db conn cfg registry m a
$crethrowM :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) e a.
(MonadThrow m, Exception e) =>
ExceptionWithContext e -> Db conn cfg registry m a
rethrowM :: forall e a.
Exception e =>
ExceptionWithContext e -> Db conn cfg registry m a
MonadThrow
    , MonadIO (Db conn cfg registry m)
MonadIO (Db conn cfg registry m) =>
(forall b.
 ((forall a. Db conn cfg registry m a -> IO a) -> IO b)
 -> Db conn cfg registry m b)
-> MonadUnliftIO (Db conn cfg registry m)
forall b.
((forall a. Db conn cfg registry m a -> IO a) -> IO b)
-> Db conn cfg registry m b
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *).
MonadUnliftIO m =>
MonadIO (Db conn cfg registry m)
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. Db conn cfg registry m a -> IO a) -> IO b)
-> Db conn cfg registry m b
forall (m :: * -> *).
MonadIO m =>
(forall b. ((forall a. m a -> IO a) -> IO b) -> m b)
-> MonadUnliftIO m
$cwithRunInIO :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. Db conn cfg registry m a -> IO a) -> IO b)
-> Db conn cfg registry m b
withRunInIO :: forall b.
((forall a. Db conn cfg registry m a -> IO a) -> IO b)
-> Db conn cfg registry m b
MonadUnliftIO
    )

instance (Monad m) => HasPoolState conn (Db conn cfg registry m) where
  getPoolState :: Db conn cfg registry m (PoolState conn)
getPoolState = (Env conn cfg registry -> PoolState conn)
-> Db conn cfg registry m (PoolState conn)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks Env conn cfg registry -> PoolState conn
forall conn cfg (registry :: JobPayloadRegistry).
Env conn cfg registry -> PoolState conn
poolState
  localPoolState :: forall a.
(PoolState conn -> PoolState conn)
-> Db conn cfg registry m a -> Db conn cfg registry m a
localPoolState PoolState conn -> PoolState conn
adjust = (Env conn cfg registry -> Env conn cfg registry)
-> Db conn cfg registry m a -> Db conn cfg registry m a
forall a.
(Env conn cfg registry -> Env conn cfg registry)
-> Db conn cfg registry m a -> Db conn cfg registry m a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (\Env conn cfg registry
env -> Env conn cfg registry
env {poolState = adjust (poolState env)})

-- | Run a 'Db' action in its env.
runDb :: Env conn cfg registry -> Db conn cfg registry m a -> m a
runDb :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Env conn cfg registry -> Db conn cfg registry m a -> m a
runDb Env conn cfg registry
env Db conn cfg registry m a
action = ReaderT (Env conn cfg registry) m a -> Env conn cfg registry -> m a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT (Db conn cfg registry m a -> ReaderT (Env conn cfg registry) m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Db conn cfg registry m a -> ReaderT (Env conn cfg registry) m a
unDb Db conn cfg registry m a
action) Env conn cfg registry
env

-- | Run a 'Db' action on one connection pinned as the caller's open transaction.
inTransaction :: Driver conn cfg -> conn -> SchemaName -> Db conn cfg registry m a -> m a
inTransaction :: forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Driver conn cfg -> conn -> Text -> Db conn cfg registry m a -> m a
inTransaction Driver conn cfg
drv conn
conn Text
schemaName =
  Env conn cfg registry -> Db conn cfg registry m a -> m a
forall conn cfg (registry :: JobPayloadRegistry) (m :: * -> *) a.
Env conn cfg registry -> Db conn cfg registry m a -> m a
runDb
    Env
      { schema :: Text
schema = Text
schemaName
      , poolState :: PoolState conn
poolState = PoolState {connectionPool :: Maybe (Pool conn)
connectionPool = Maybe (Pool conn)
forall a. Maybe a
Nothing, pinned :: Maybe (conn, Int)
pinned = (conn, Int) -> Maybe (conn, Int)
forall a. a -> Maybe a
Just (conn
conn, Int
outermostDepth)}
      , listener :: Maybe Listener
listener = Maybe Listener
forall a. Maybe a
Nothing
      , driverConfig :: cfg
driverConfig = Driver conn cfg -> cfg
forall conn cfg. Driver conn cfg -> cfg
initialConfig Driver conn cfg
drv
      }

-- | Close the idle connections in the env's pool. Connections in use stay open and go
-- back to the pool. The pool stays usable.
destroyEnv :: (MonadIO m) => Env conn cfg registry -> m ()
destroyEnv :: forall (m :: * -> *) conn cfg (registry :: JobPayloadRegistry).
MonadIO m =>
Env conn cfg registry -> m ()
destroyEnv Env conn cfg registry
env = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ (Pool conn -> IO ()) -> Maybe (Pool conn) -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Pool conn -> IO ()
forall a. Pool a -> IO ()
destroyAllResources (PoolState conn -> Maybe (Pool conn)
forall conn. PoolState conn -> Maybe (Pool conn)
connectionPool (Env conn cfg registry -> PoolState conn
forall conn cfg (registry :: JobPayloadRegistry).
Env conn cfg registry -> PoolState conn
poolState Env conn cfg registry
env))

-- | Turn off the shared LISTEN listener for an env, running poll-only.
disableListener :: Env conn cfg registry -> Env conn cfg registry
disableListener :: forall conn cfg (registry :: JobPayloadRegistry).
Env conn cfg registry -> Env conn cfg registry
disableListener Env conn cfg registry
env = Env conn cfg registry
env {listener = Nothing}

-- | Give the env a LISTEN connection of its own from a connection runner.
useDedicatedListener
  :: (MonadIO m) => ((ListenConn -> IO ()) -> IO ()) -> Env conn cfg registry -> m (Env conn cfg registry)
useDedicatedListener :: forall (m :: * -> *) conn cfg (registry :: JobPayloadRegistry).
MonadIO m =>
((ListenConn -> IO ()) -> IO ())
-> Env conn cfg registry -> m (Env conn cfg registry)
useDedicatedListener (ListenConn -> IO ()) -> IO ()
withDedicated Env conn cfg registry
env = IO (Env conn cfg registry) -> m (Env conn cfg registry)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Env conn cfg registry) -> m (Env conn cfg registry))
-> IO (Env conn cfg registry) -> m (Env conn cfg registry)
forall a b. (a -> b) -> a -> b
$ do
  lstn <- ((ListenConn -> IO ()) -> IO ()) -> IO Listener
newListener (ListenConn -> IO ()) -> IO ()
withDedicated
  pure env {listener = Just lstn}

-- | A listener that borrows one pool connection for the hub's lifetime.
poolListener :: Driver conn cfg -> Pool conn -> IO Listener
poolListener :: forall conn cfg. Driver conn cfg -> Pool conn -> IO Listener
poolListener Driver conn cfg
drv Pool conn
pool = ((ListenConn -> IO ()) -> IO ()) -> IO Listener
newListener (\ListenConn -> IO ()
action -> Pool conn -> (conn -> IO ()) -> IO ()
forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool conn
pool (\conn
conn -> Driver conn cfg -> conn -> (ListenConn -> IO ()) -> IO ()
forall conn cfg.
Driver conn cfg -> conn -> (ListenConn -> IO ()) -> IO ()
withListenConn Driver conn cfg
drv conn
conn ListenConn -> IO ()
action))

-- | Create an env over a new pool opened with the connect and release actions. While
-- anything listens, the listener holds one pool slot.
createEnvWithConfig
  :: (MonadIO m)
  => Driver conn cfg
  -> IO conn
  -> (conn -> IO ())
  -> SchemaName
  -> PoolConfig
  -> m (Env conn cfg registry)
createEnvWithConfig :: forall (m :: * -> *) conn cfg (registry :: JobPayloadRegistry).
MonadIO m =>
Driver conn cfg
-> IO conn
-> (conn -> IO ())
-> Text
-> PoolConfig
-> m (Env conn cfg registry)
createEnvWithConfig Driver conn cfg
drv IO conn
connect conn -> IO ()
release Text
schemaName PoolConfig
config = IO (Env conn cfg registry) -> m (Env conn cfg registry)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Env conn cfg registry) -> m (Env conn cfg registry))
-> IO (Env conn cfg registry) -> m (Env conn cfg registry)
forall a b. (a -> b) -> a -> b
$ do
  connPool <-
    PoolConfig conn -> IO (Pool conn)
forall a. PoolConfig a -> IO (Pool a)
newPool
      (PoolConfig conn -> IO (Pool conn))
-> PoolConfig conn -> IO (Pool conn)
forall a b. (a -> b) -> a -> b
$ Maybe Int -> PoolConfig conn -> PoolConfig conn
forall a. Maybe Int -> PoolConfig a -> PoolConfig a
setNumStripes (PoolConfig -> Maybe Int
poolStripes PoolConfig
config)
      (PoolConfig conn -> PoolConfig conn)
-> PoolConfig conn -> PoolConfig conn
forall a b. (a -> b) -> a -> b
$ IO conn -> (conn -> IO ()) -> Double -> Int -> PoolConfig conn
forall a. IO a -> (a -> IO ()) -> Double -> Int -> PoolConfig a
defaultPoolConfig IO conn
connect conn -> IO ()
release (Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Double) -> Int -> Double
forall a b. (a -> b) -> a -> b
$ PoolConfig -> Int
poolIdleTimeout PoolConfig
config) (PoolConfig -> Int
poolSize PoolConfig
config)
  createEnvWithPool drv connPool schemaName

-- | Create an env over a caller's own connection pool. While anything listens, the
-- listener holds one pool slot.
createEnvWithPool :: (MonadIO m) => Driver conn cfg -> Pool conn -> SchemaName -> m (Env conn cfg registry)
createEnvWithPool :: forall (m :: * -> *) conn cfg (registry :: JobPayloadRegistry).
MonadIO m =>
Driver conn cfg -> Pool conn -> Text -> m (Env conn cfg registry)
createEnvWithPool Driver conn cfg
drv Pool conn
connPool Text
schemaName = IO (Env conn cfg registry) -> m (Env conn cfg registry)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Env conn cfg registry) -> m (Env conn cfg registry))
-> IO (Env conn cfg registry) -> m (Env conn cfg registry)
forall a b. (a -> b) -> a -> b
$ do
  lstn <- Driver conn cfg -> Pool conn -> IO Listener
forall conn cfg. Driver conn cfg -> Pool conn -> IO Listener
poolListener Driver conn cfg
drv Pool conn
connPool
  pure
    Env
      { schema = schemaName
      , poolState = PoolState {connectionPool = Just connPool, pinned = Nothing}
      , listener = Just lstn
      , driverConfig = initialConfig drv
      }

-- | The pinned connection, or one checked out of the pool.
withConn :: (HasPoolState conn m, MonadUnliftIO m) => (conn -> m a) -> m a
withConn :: forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> m a) -> m a
withConn conn -> m a
action = do
  pool <- m (PoolState conn)
forall conn (m :: * -> *).
HasPoolState conn m =>
m (PoolState conn)
getPoolState
  case (fst <$> pinned pool, connectionPool pool) of
    (Just conn
conn, Maybe (Pool conn)
_) -> conn -> m a
action conn
conn
    (Maybe conn
Nothing, Just Pool conn
connPool) -> ((forall a. m a -> IO a) -> IO a) -> m a
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run -> Pool conn -> (conn -> IO a) -> IO a
forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool conn
connPool (m a -> IO a
forall a. m a -> IO a
run (m a -> IO a) -> (conn -> m a) -> conn -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. conn -> m a
action)
    (Maybe conn
Nothing, Maybe (Pool conn)
Nothing) -> Text -> m a
forall (m :: * -> *) a. MonadIO m => Text -> m a
throwInternal Text
noConnection

-- | Pin one pooled connection for the action.
pinConnection :: (HasPoolState conn m, MonadUnliftIO m) => m a -> m a
pinConnection :: forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
m a -> m a
pinConnection m a
action = (conn -> m a) -> m a
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> m a) -> m a
withConn ((conn -> m a) -> m a) -> (conn -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \conn
conn -> (PoolState conn -> PoolState conn) -> m a -> m a
forall a. (PoolState conn -> PoolState conn) -> m a -> m a
forall conn (m :: * -> *) a.
HasPoolState conn m =>
(PoolState conn -> PoolState conn) -> m a -> m a
localPoolState (\PoolState conn
st -> PoolState conn
st {pinned = Just (conn, maybe noTransactionDepth snd (pinned st))}) m a
action

-- | Transaction bracket over the backend's own statement runner and bracket. Nests via
-- savepoints.
withSavepointTransaction
  :: (HasPoolState conn m, MonadUnliftIO m)
  => (conn -> ByteString -> IO ())
  -> (conn -> IO a -> IO a)
  -> m a
  -> m a
withSavepointTransaction :: forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> ByteString -> IO ())
-> (conn -> IO a -> IO a) -> m a -> m a
withSavepointTransaction conn -> ByteString -> IO ()
runSql conn -> IO a -> IO a
transaction m a
action = do
  st <- m (PoolState conn)
forall conn (m :: * -> *).
HasPoolState conn m =>
m (PoolState conn)
getPoolState
  case pinned st of
    Just (conn
conn, Int
depth) | Int
depth Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
noTransactionDepth -> ((forall a. m a -> m a) -> m a) -> m a
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> m a) -> m b) -> m b
mask (((forall a. m a -> m a) -> m a) -> m a)
-> ((forall a. m a -> m a) -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore -> do
      let spName :: ByteString
spName = ByteString
"arbiter_sp_" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
BSC.pack (Int -> String
forall a. Show a => a -> String
show Int
depth)
      IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ conn -> ByteString -> IO ()
runSql conn
conn (ByteString
"SAVEPOINT " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
spName)
      result <-
        m a -> m a
forall a. m a -> m a
restore ((PoolState conn -> PoolState conn) -> m a -> m a
forall a. (PoolState conn -> PoolState conn) -> m a -> m a
forall conn (m :: * -> *) a.
HasPoolState conn m =>
(PoolState conn -> PoolState conn) -> m a -> m a
localPoolState (\PoolState conn
s -> PoolState conn
s {pinned = Just (conn, depth + 1)}) m a
action)
          m a -> m () -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`onException` IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (conn -> ByteString -> IO ()
runSql conn
conn (ByteString
"ROLLBACK TO SAVEPOINT " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
spName))
      liftIO $ runSql conn ("RELEASE SAVEPOINT " <> spName)
      pure result
    Maybe (conn, Int)
_ -> (conn -> m a) -> m a
forall conn (m :: * -> *) a.
(HasPoolState conn m, MonadUnliftIO m) =>
(conn -> m a) -> m a
withConn ((conn -> m a) -> m a) -> (conn -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \conn
conn -> ((forall a. m a -> IO a) -> IO a) -> m a
forall b. ((forall a. m a -> IO a) -> IO b) -> m b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. m a -> IO a) -> IO a) -> m a)
-> ((forall a. m a -> IO a) -> IO a) -> m a
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> IO a
run ->
      conn -> IO a -> IO a
transaction conn
conn
        (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ m a -> IO a
forall a. m a -> IO a
run
        (m a -> IO a) -> m a -> IO a
forall a b. (a -> b) -> a -> b
$ (PoolState conn -> PoolState conn) -> m a -> m a
forall a. (PoolState conn -> PoolState conn) -> m a -> m a
forall conn (m :: * -> *) a.
HasPoolState conn m =>
(PoolState conn -> PoolState conn) -> m a -> m a
localPoolState (\PoolState conn
s -> PoolState conn
s {pinned = Just (conn, outermostDepth)}) m a
action

-- | The transaction depth of a pinned connection outside any transaction.
noTransactionDepth :: Int
noTransactionDepth :: Int
noTransactionDepth = Int
0

-- | The transaction depth of a pinned connection inside its outermost transaction.
outermostDepth :: Int
outermostDepth :: Int
outermostDepth = Int
1

noConnection :: Text
noConnection :: Text
noConnection = Text
"No active connection and no connection pool available"