EasyPool.hs

{-# LANGUAGE PackageImports #-}

module EasyPool
    ( EasyPool
    , withEasyPool
    , newEasyPool
    ) where

import Pool
import Control.Monad.STM
import Control.Concurrent.STM.TVar
import Control.Monad.IO.Class
import "MonadCatchIO-transformers" Control.Monad.CatchIO

data EasyPool r m = EasyPool
    { epPool :: Pool r
    , epMake :: m (Maybe r)
    }

withEasyPool :: MonadCatchIO m => EasyPool r m -> (r -> m a) -> m (Maybe a)
withEasyPool (EasyPool pool mk) = withPool pool mk

newEasyPool :: MonadCatchIO m => Int -> m r -> m (EasyPool r m)
newEasyPool count mk = do
    pool <- new
    texist <- liftIO $ newTVarIO 0
    return $ EasyPool pool $ mk' texist
  where
    mk' texist = do
        exist <- liftIO $ atomically $ readTVar texist
        if exist >= count
            then return Nothing
            else do
                r <- mk
                liftIO $ atomically $ do
                    exist <- readTVar texist
                    if exist >= count
                        then return Nothing
                        else return $ Just r

Pool.hs

{-# LANGUAGE NoImplicitPrelude #-} -- (I like to be explicit)
{-# LANGUAGE PackageImports #-}

module Pool (Pool, new, withPool) where

import Data.Function               ( ($), (.) )
import Data.Maybe                  ( Maybe(Nothing,Just), maybe )
import Data.Functor                ( (<$>) )
import Control.Monad               ( return, (>>=), (>>), (=<<), fail, join, liftM )
import Control.Monad.STM           ( atomically )
import Control.Concurrent.STM.TVar ( TVar, newTVarIO, readTVar, writeTVar )
import "MonadCatchIO-transformers"
       Control.Monad.CatchIO       ( MonadCatchIO, block, finally )
import Control.Monad.IO.Class      ( liftIO )

newtype Pool r = Pool (TVar [r])

new :: MonadCatchIO m => m (Pool r)
new = liftIO $ Pool <$> newTVarIO []

withPool :: MonadCatchIO m => Pool r -> m (Maybe r) -> (r -> m a) -> m (Maybe a)
withPool (Pool tv) mk f = block $ join $ liftIO $ atomically $ do
 rrs <- readTVar tv
 case rrs of
   [] -> return $ mk >>= maybe (return Nothing) with
   r:rs -> writeTVar tv rs >> return (with r)
   where
     with r = liftM Just (f r)
               `finally`
                 liftIO (atomically $ writeTVar tv . (r:) =<< readTVar tv)