module Hasql.Pool
  ( Pool,
    acquire,
    release,
    UsageError (..),
    use,
  )
where

import Hasql.Connection (Connection)
import qualified Hasql.Connection as Connection
import Hasql.Pool.Prelude
import Hasql.Session (Session)
import qualified Hasql.Session as Session

-- |
-- A pool of connections to DB.
data Pool = Pool
  { -- | Connection settings.
    Pool -> Settings
poolConnectionSettings :: Connection.Settings,
    -- | Avail connections.
    Pool -> TQueue Connection
poolConnectionQueue :: TQueue Connection,
    -- | Capacity.
    Pool -> TVar Int
poolCapacity :: TVar Int,
    -- | Alive.
    Pool -> TVar Bool
poolAlive :: TVar Bool
  }

-- | Given the pool-size and connection settings create a connection-pool.
acquire :: Int -> Connection.Settings -> IO Pool
acquire :: Int -> Settings -> IO Pool
acquire Int
poolSize Settings
connectionSettings = do
  Settings -> TQueue Connection -> TVar Int -> TVar Bool -> Pool
Pool Settings
connectionSettings
    (TQueue Connection -> TVar Int -> TVar Bool -> Pool)
-> IO (TQueue Connection) -> IO (TVar Int -> TVar Bool -> Pool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (TQueue Connection)
forall a. IO (TQueue a)
newTQueueIO
    IO (TVar Int -> TVar Bool -> Pool)
-> IO (TVar Int) -> IO (TVar Bool -> Pool)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> IO (TVar Int)
forall a. a -> IO (TVar a)
newTVarIO Int
poolSize
    IO (TVar Bool -> Pool) -> IO (TVar Bool) -> IO Pool
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
True

-- |
-- Release all the connections in the pool.
release :: Pool -> IO ()
release :: Pool -> IO ()
release Pool {Settings
TVar Bool
TVar Int
TQueue Connection
poolAlive :: TVar Bool
poolCapacity :: TVar Int
poolConnectionQueue :: TQueue Connection
poolConnectionSettings :: Settings
poolAlive :: Pool -> TVar Bool
poolCapacity :: Pool -> TVar Int
poolConnectionQueue :: Pool -> TQueue Connection
poolConnectionSettings :: Pool -> Settings
..} = do
  [Connection]
connections <- STM [Connection] -> IO [Connection]
forall a. STM a -> IO a
atomically (STM [Connection] -> IO [Connection])
-> STM [Connection] -> IO [Connection]
forall a b. (a -> b) -> a -> b
$ do
    TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
poolAlive Bool
False
    TQueue Connection -> STM [Connection]
forall a. TQueue a -> STM [a]
flushTQueue TQueue Connection
poolConnectionQueue
  [Connection] -> (Connection -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Connection]
connections Connection -> IO ()
Connection.release

-- |
-- A union over the connection establishment error and the session error.
data UsageError
  = ConnectionUsageError Connection.ConnectionError
  | SessionUsageError Session.QueryError
  | PoolIsReleasedUsageError
  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
showList :: [UsageError] -> ShowS
$cshowList :: [UsageError] -> ShowS
show :: UsageError -> String
$cshow :: UsageError -> String
showsPrec :: Int -> UsageError -> ShowS
$cshowsPrec :: Int -> UsageError -> ShowS
Show, UsageError -> UsageError -> Bool
(UsageError -> UsageError -> Bool)
-> (UsageError -> UsageError -> Bool) -> Eq UsageError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: UsageError -> UsageError -> Bool
$c/= :: UsageError -> UsageError -> Bool
== :: UsageError -> UsageError -> Bool
$c== :: 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 -> Session.Session a -> IO (Either UsageError a)
use :: Pool -> Session a -> IO (Either UsageError a)
use Pool {Settings
TVar Bool
TVar Int
TQueue Connection
poolAlive :: TVar Bool
poolCapacity :: TVar Int
poolConnectionQueue :: TQueue Connection
poolConnectionSettings :: Settings
poolAlive :: Pool -> TVar Bool
poolCapacity :: Pool -> TVar Int
poolConnectionQueue :: Pool -> TQueue Connection
poolConnectionSettings :: Pool -> Settings
..} Session a
sess =
  IO (IO (Either UsageError a)) -> IO (Either UsageError a)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO (Either UsageError a)) -> IO (Either UsageError a))
-> (STM (IO (Either UsageError a))
    -> IO (IO (Either UsageError a)))
-> STM (IO (Either UsageError a))
-> IO (Either UsageError a)
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. STM (IO (Either UsageError a)) -> IO (IO (Either UsageError a))
forall a. STM a -> IO a
atomically (STM (IO (Either UsageError a)) -> IO (Either UsageError a))
-> STM (IO (Either UsageError a)) -> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$ do
    Bool
alive <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
poolAlive
    if Bool
alive
      then
        [STM (IO (Either UsageError a))] -> STM (IO (Either UsageError a))
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum
          [ TQueue Connection -> STM Connection
forall a. TQueue a -> STM a
readTQueue TQueue Connection
poolConnectionQueue STM Connection
-> (Connection -> IO (Either UsageError a))
-> STM (IO (Either UsageError a))
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Connection -> IO (Either UsageError a)
onConn,
            do
              Int
capVal <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
poolCapacity
              if Int
capVal Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                then do
                  TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Int
poolCapacity (Int -> STM ()) -> Int -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> Int
forall a. Enum a => a -> a
pred Int
capVal
                  IO (Either UsageError a) -> STM (IO (Either UsageError a))
forall (m :: * -> *) a. Monad m => a -> m a
return IO (Either UsageError a)
onNewConn
                else STM (IO (Either UsageError a))
forall a. STM a
retry
          ]
      else IO (Either UsageError a) -> STM (IO (Either UsageError a))
forall (m :: * -> *) a. Monad m => a -> m a
return (IO (Either UsageError a) -> STM (IO (Either UsageError a)))
-> (UsageError -> IO (Either UsageError a))
-> UsageError
-> STM (IO (Either UsageError a))
forall k (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Either UsageError a -> IO (Either UsageError a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either UsageError a -> IO (Either UsageError a))
-> (UsageError -> Either UsageError a)
-> UsageError
-> IO (Either UsageError a)
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 -> STM (IO (Either UsageError a)))
-> UsageError -> STM (IO (Either UsageError a))
forall a b. (a -> b) -> a -> b
$ UsageError
PoolIsReleasedUsageError
  where
    onNewConn :: IO (Either UsageError a)
onNewConn = do
      Either ConnectionError Connection
connRes <- Settings -> IO (Either ConnectionError Connection)
Connection.acquire Settings
poolConnectionSettings
      case Either ConnectionError Connection
connRes of
        Left ConnectionError
connErr -> do
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
poolCapacity Int -> Int
forall a. Enum a => a -> a
succ
          Either UsageError a -> IO (Either UsageError a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either UsageError a -> IO (Either UsageError a))
-> Either UsageError a -> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$ UsageError -> Either UsageError a
forall a b. a -> Either a b
Left (UsageError -> Either UsageError a)
-> UsageError -> Either UsageError a
forall a b. (a -> b) -> a -> b
$ ConnectionError -> UsageError
ConnectionUsageError ConnectionError
connErr
        Right Connection
conn -> Connection -> IO (Either UsageError a)
onConn Connection
conn
    onConn :: Connection -> IO (Either UsageError a)
onConn Connection
conn = do
      Either QueryError a
sessRes <- Session a -> Connection -> IO (Either QueryError a)
forall a. Session a -> Connection -> IO (Either QueryError a)
Session.run Session a
sess Connection
conn
      case Either QueryError a
sessRes of
        Left QueryError
err -> case QueryError
err of
          Session.QueryError Settings
_ [Text]
_ (Session.ClientError ConnectionError
_) -> do
            STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
poolCapacity Int -> Int
forall a. Enum a => a -> a
succ
            Either UsageError a -> IO (Either UsageError a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either UsageError a -> IO (Either UsageError a))
-> Either UsageError a -> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$ UsageError -> Either UsageError a
forall a b. a -> Either a b
Left (UsageError -> Either UsageError a)
-> UsageError -> Either UsageError a
forall a b. (a -> b) -> a -> b
$ QueryError -> UsageError
SessionUsageError QueryError
err
          QueryError
_ -> do
            IO ()
returnConn
            Either UsageError a -> IO (Either UsageError a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either UsageError a -> IO (Either UsageError a))
-> Either UsageError a -> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$ UsageError -> Either UsageError a
forall a b. a -> Either a b
Left (UsageError -> Either UsageError a)
-> UsageError -> Either UsageError a
forall a b. (a -> b) -> a -> b
$ QueryError -> UsageError
SessionUsageError QueryError
err
        Right a
res -> do
          IO ()
returnConn
          Either UsageError a -> IO (Either UsageError a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either UsageError a -> IO (Either UsageError a))
-> Either UsageError a -> IO (Either UsageError a)
forall a b. (a -> b) -> a -> b
$ a -> Either UsageError a
forall a b. b -> Either a b
Right a
res
      where
        returnConn :: IO ()
returnConn = do
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            Bool
alive <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
poolAlive
            Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
alive (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$ TQueue Connection -> Connection -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Connection
poolConnectionQueue Connection
conn