Pool.hs

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

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


-- |
-- A pool of connections to DB.
data Pool =
  Pool
    Connection.Settings
    {-^ Connection settings. -}
    (TQueue ActiveConnection)
    {-^ Queue of established connections. -}
    (TVar Int)
    {-^ Slots available for establishing new connections. -}
    (TVar Bool)
    {-^ Flag signaling whether pool's alive. -}

data ActiveConnection =
  ActiveConnection {
    activeConnectionLastUseTimestamp :: Int,
    activeConnectionConnection :: Connection
  }

loopCollectingGarbage :: Int -> TQueue ActiveConnection -> TVar Int -> TVar Bool -> IO ()
loopCollectingGarbage timeout establishedQueue slotsAvailVar aliveVar =
  decide
  where
    decide =
      do
        ts <- getMillisecondsSinceEpoch
        join $ atomically $ do
          alive <- readTVar aliveVar
          if alive
            then let
              tryToRelease =
                tryReadTQueue establishedQueue >>= \ case
                  -- The queue is empty. Just wait for changes in the state.
                  Nothing ->
                    retry
                  Just entry@(ActiveConnection lastUseTs connection) ->
                    let
                      outdatingTs =
                        lastUseTs + timeout
                      in
                        -- Check whether it's outdated.
                        if outdatingTs < ts
                          -- Fetch the current value of available slots and
                          -- release this one and other connections.
                          then do
                            slotsAvail <- readTVar slotsAvailVar
                            collectAndRelease slotsAvail [connection] outdatingTs
                          -- Return it to the front of the queue and
                          -- wait until it's outdating time.
                          else do
                            unGetTQueue establishedQueue entry
                            return (sleep outdatingTs *> decide)
              collectAndRelease !slotsAvail !outdatedList outdatingTs =
                tryReadTQueue establishedQueue >>= \ case
                  Nothing ->
                    finalizeAndRelease slotsAvail outdatedList outdatingTs
                  Just entry@(ActiveConnection lastUseTs connection) ->
                    let
                      outdatingTs =
                        lastUseTs + timeout
                      in
                        if outdatingTs < ts
                          then do
                            unGetTQueue establishedQueue entry
                            finalizeAndRelease slotsAvail outdatedList outdatingTs
                          else
                            collectAndRelease (succ slotsAvail) (connection : outdatedList) outdatingTs
              finalizeAndRelease slotsAvail outdatedList outdatingTs =
                do
                  writeTVar slotsAvailVar slotsAvail
                  return (release outdatedList *> sleep outdatingTs *> decide)
              in tryToRelease
            else do
              list <- flushTQueue establishedQueue
              return (release (fmap activeConnectionConnection list))
    sleep untilTs =
      do
        ts <- getMillisecondsSinceEpoch
        let
          diff =
            untilTs - ts
          in if diff > 0
            then threadDelay (diff * 1000)
            else return ()
    release =
      traverse_ Connection.release

-- |
-- Settings of the connection pool. Consist of:
-- 
-- * Pool-size.
-- 
-- * Timeout.   
-- An amount of time in milliseconds for which the unused connections are kept open.
-- 
-- * Connection settings.
-- 
type Settings =
  (Int, Int, Connection.Settings)

-- |
-- Given the pool-size, timeout and connection settings
-- create a connection-pool.
acquire :: Settings -> IO Pool
acquire (size, timeout, connectionSettings) =
  do
    establishedQueue <- newTQueueIO
    slotsAvailVar <- newTVarIO size
    aliveVar <- newTVarIO (size > 0)
    forkIO $ loopCollectingGarbage timeout establishedQueue slotsAvailVar aliveVar
    return (Pool connectionSettings establishedQueue slotsAvailVar aliveVar)

-- |
-- Release the connection-pool.
release :: Pool -> IO ()
release (Pool _ _ _ aliveVar) =
  atomically (writeTVar aliveVar False)

-- |
-- A union over the connection establishment error and the session error.
data UsageError =
  ConnectionUsageError Connection.ConnectionError |
  SessionUsageError Session.QueryError |
  PoolIsReleasedUsageError
  deriving (Show, 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 connectionSettings establishedQueue slotsAvailVar aliveVar) session =
  join $ atomically $ do
    alive <- readTVar aliveVar
    if alive
      then
        tryReadTQueue establishedQueue >>= \ case
          -- No established connection avail at the moment.
          Nothing -> do
            slotsAvail <- readTVar slotsAvailVar
            -- Do we have any slots left for establishing new connections?
            if slotsAvail > 0
              -- Reduce the available slots var and instruct to
              -- establish and use a new connection.
              then do
                writeTVar slotsAvailVar $! pred slotsAvail
                return acquireConnectionThenUseThenPutItToQueue
              -- Wait until the state changes and retry.
              else
                retry
          Just (ActiveConnection _ connection) ->
            return (useConnectionThenPutItToQueue connection)
      else
        return (return (Left PoolIsReleasedUsageError))
  where
    acquireConnectionThenUseThenPutItToQueue =
      do
        res <- Connection.acquire connectionSettings
        case res of
          -- Failed to acquire, so release an availability slot,
          -- returning the error details.
          Left acquisitionError -> do
            atomically $ modifyTVar' slotsAvailVar succ
            return (Left (ConnectionUsageError acquisitionError))
          Right connection ->
            useConnectionThenPutItToQueue connection
    useConnectionThenPutItToQueue connection =
      do
        res <- Session.run session connection
        case res of
          Left queryError -> do
            -- Check whether the error is on client-side,
            -- and in that case release the connection.
            case queryError of
              Session.QueryError _ _ (Session.ClientError _) ->
                releaseConnection connection
              _ ->
                putConnectionToPool connection
            return (Left (SessionUsageError queryError))
          Right res -> do
            putConnectionToPool connection
            return (Right res)
    putConnectionToPool connection =
      do
        ts <- getMillisecondsSinceEpoch
        atomically $ writeTQueue establishedQueue (ActiveConnection ts connection)
    releaseConnection connection =
      do
        atomically $ modifyTVar' slotsAvailVar succ
        Connection.release connection