ConnectionPool.hs

{-# LANGUAGE BangPatterns                #-}

module Network.ConnectionPool
  ( withConnection
  , newConnectionPool
  , closeConnectionPool
  , ConnectionPool
  ) where

------------------------------------------------------------------------
import           Control.Concurrent.MVar.Strict
import           Control.Exception
import           Control.Parallel.Strategies
------------------------------------------------------------------------

data NFData conn => ConnectionPool conn =
    ConnectionPool
    { poolCreateConn  :: IO conn
    , poolDestroyConn :: conn -> IO ()
    , poolConnections :: MVar [conn]
    , poolMax         :: Int }


-- | given a creation function and a teardown function, make a new connection
-- | pool
newConnectionPool :: NFData conn =>
                     IO conn
                  -> (conn -> IO ())
                  -> Int
                  -> IO (ConnectionPool conn)
newConnectionPool create destroy maxconns = do
    mvar <- newMVar []
    return $ ConnectionPool create destroy mvar maxconns


-- | Execute an action with a connection from the pool.
withConnection :: NFData conn =>
                  ConnectionPool conn
               -> (conn -> IO a)
               -> IO a
withConnection !pool !action =
    bracketOnError (getConnection pool)
                   (poolDestroyConn pool)
                   (\c -> do
                      r <- action c
                      releaseConnection pool c
                      return r)


-- | Close all open connection in a connection pool
closeConnectionPool :: NFData conn =>
                       ConnectionPool conn -> IO ()
closeConnectionPool !pool =
    modifyMVar_ (poolConnections pool) $ \cs -> do
      mapM_ (poolDestroyConn pool) cs
      return []


------------------------------------------------------------------------


getConnection :: NFData conn =>
                 ConnectionPool conn -> IO conn
getConnection !pool =
    modifyMVar (poolConnections pool) $ \cs ->
        if null cs then do
            newconn <- (poolCreateConn pool)
            return ([], newconn)
          else
            return (tail cs, head cs)


releaseConnection :: NFData conn =>
                     ConnectionPool conn -> conn -> IO ()
releaseConnection pool c =
    modifyMVar_ (poolConnections pool) $ \cs ->
        if length cs >= poolMax pool - 1 then do
            poolDestroyConn pool c
            return cs
        else
            return (cs ++ [c])