<commit_before>{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE PackageImports #-}
module Data.Pool
    ( -- * Using pools
      createPool
    , createPoolCheckAlive
    , withPool
    , withPool'
    , withPoolAllocate
    , Pool
      -- * Diagnostics
    , PoolStats (..)
    , poolStats
    ) where

import Data.IORef (IORef, newIORef, atomicModifyIORef, readIORef)
import Control.Exception (throwIO, Exception)
import qualified Control.Exception as E
import Data.Typeable
import qualified Control.Monad.IO.Control as I
import qualified Control.Exception.Control as I
import Control.Monad.IO.Class
import Control.Monad

data PoolData a = PoolData
    { poolAvail :: ![a]
    , poolCreated :: !Int
    }

data Pool a = Pool
    { poolMax :: Int
    , poolData :: IORef (PoolData a)
    , poolMake :: IO a
    , poolFree :: a -> IO ()
    , poolCheckAlive :: a -> IO Bool
    }

data PoolStats = PoolStats
    { poolStatsMax :: Int
    , poolStatsAvailable :: Int
    , poolStatsCreated :: Int
    }

poolStats :: Pool a -> IO PoolStats
poolStats p = do
    d <- readIORef $ poolData p
    return $ PoolStats (poolMax p) (length $ poolAvail d) (poolCreated d)

createPool :: I.MonadControlIO m
           => IO a -> (a -> IO ()) -> Int -> (Pool a -> m b) -> m b
createPool mk fr mx f = createPoolCheckAlive mk fr mx f $ const $ return True

createPoolCheckAlive
    :: I.MonadControlIO m
    => IO a -> (a -> IO ()) -> Int -> (Pool a -> m b)
    -> (a -> IO Bool) -- ^ is the resource alive?
    -> m b
createPoolCheckAlive mk fr mx f ca = do
    pd <- liftIO $ newIORef $ PoolData [] 0
    I.finally (f $ Pool mx pd mk fr ca) $ liftIO $ do
        PoolData ress _ <- readIORef pd
        mapM_ fr ress

data PoolExhaustedException = PoolExhaustedException
    deriving (Show, Typeable)
instance Exception PoolExhaustedException

-- | This function throws a 'PoolExhaustedException' when no resources are
-- available. See 'withPoolAllocate' to avoid this.
withPool' :: I.MonadControlIO m => Pool a -> (a -> m b) -> m b
withPool' p f = do
    x <- withPool p f
    case x of
        Nothing -> liftIO $ throwIO PoolExhaustedException
        Just x' -> return x'

-- | Same as @withPool'@, but instead of throwing a 'PoolExhaustedException'
-- when there the maximum number of resources are created and allocated, it
-- allocates a new resource, passes it to the subprocess and then frees it.
withPoolAllocate :: I.MonadControlIO m => Pool a -> (a -> m b) -> m b
withPoolAllocate p f = do
    x <- withPool p f
    case x of
        Just x' -> return x'
        Nothing ->
            I.bracket
                (liftIO $ poolMake p)
                (liftIO . poolFree p)
                f

withPool :: I.MonadControlIO m => Pool a -> (a -> m b) -> m (Maybe b)
withPool p f = I.mask $ \unmask -> do
    eres <- liftIO $ atomicModifyIORef (poolData p) $ \pd ->
        case poolAvail pd of
            x:xs -> (pd { poolAvail = xs }, Right x)
            [] -> (pd, Left $ poolCreated pd)
    case eres of
        Left pc ->
            if pc >= poolMax p
                then return Nothing
                else I.bracket
                    (liftIO $ poolMake p)
                    (insertResource 1)
                    (liftM Just . unmask . f)
        Right res -> do
            isAlive <- I.try $ unmask $ liftIO $ poolCheckAlive p res
            case isAlive :: Either E.SomeException Bool of
                Right True ->
                    I.finally
                        (liftM Just $ unmask $ f res)
                        (insertResource 0 res)
                _ -> do
                    -- decrement the poolCreated count and then start over
                    liftIO $ atomicModifyIORef (poolData p) $ \pd ->
                        (pd { poolCreated = poolCreated pd - 1}, ())
                    unmask $ withPool p f
  where
    insertResource i x = liftIO $ atomicModifyIORef (poolData p) $ \pd ->
        (pd { poolAvail = x : poolAvail pd
                , poolCreated = i + poolCreated pd
                }, ())<commit_msg>Use liftIOOp in combination with the regular bracket and finally Since the acquire and release computations are in IO it's more efficient to use the regular bracket and finally and lift them using liftIOOp.<commit_after>{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE PackageImports #-}
module Data.Pool
    ( -- * Using pools
      createPool
    , createPoolCheckAlive
    , withPool
    , withPool'
    , withPoolAllocate
    , Pool
      -- * Diagnostics
    , PoolStats (..)
    , poolStats
    ) where

import Data.IORef (IORef, newIORef, atomicModifyIORef, readIORef)
import Control.Exception (throwIO, Exception, bracket, finally)
import qualified Control.Exception as E
import Data.Typeable
import qualified Control.Monad.IO.Control as I
import qualified Control.Exception.Control as I
import Control.Monad.IO.Class
import Control.Monad

data PoolData a = PoolData
    { poolAvail :: ![a]
    , poolCreated :: !Int
    }

data Pool a = Pool
    { poolMax :: Int
    , poolData :: IORef (PoolData a)
    , poolMake :: IO a
    , poolFree :: a -> IO ()
    , poolCheckAlive :: a -> IO Bool
    }

data PoolStats = PoolStats
    { poolStatsMax :: Int
    , poolStatsAvailable :: Int
    , poolStatsCreated :: Int
    }

poolStats :: Pool a -> IO PoolStats
poolStats p = do
    d <- readIORef $ poolData p
    return $ PoolStats (poolMax p) (length $ poolAvail d) (poolCreated d)

createPool :: I.MonadControlIO m
           => IO a -> (a -> IO ()) -> Int -> (Pool a -> m b) -> m b
createPool mk fr mx f = createPoolCheckAlive mk fr mx f $ const $ return True

createPoolCheckAlive
    :: I.MonadControlIO m
    => IO a -> (a -> IO ()) -> Int -> (Pool a -> m b)
    -> (a -> IO Bool) -- ^ is the resource alive?
    -> m b
createPoolCheckAlive mk fr mx f ca = do
    pd <- liftIO $ newIORef $ PoolData [] 0
    finallyIO (f $ Pool mx pd mk fr ca) $ do
        PoolData ress _ <- readIORef pd
        mapM_ fr ress

finallyIO :: I.MonadControlIO m => m a -> IO b -> m a
finallyIO a io = I.controlIO $ \runInIO -> finally (runInIO a) io
                           
data PoolExhaustedException = PoolExhaustedException
    deriving (Show, Typeable)
instance Exception PoolExhaustedException

-- | This function throws a 'PoolExhaustedException' when no resources are
-- available. See 'withPoolAllocate' to avoid this.
withPool' :: I.MonadControlIO m => Pool a -> (a -> m b) -> m b
withPool' p f = do
    x <- withPool p f
    case x of
        Nothing -> liftIO $ throwIO PoolExhaustedException
        Just x' -> return x'

-- | Same as @withPool'@, but instead of throwing a 'PoolExhaustedException'
-- when there the maximum number of resources are created and allocated, it
-- allocates a new resource, passes it to the subprocess and then frees it.
withPoolAllocate :: I.MonadControlIO m => Pool a -> (a -> m b) -> m b
withPoolAllocate p f = do
    x <- withPool p f
    case x of
        Just x' -> return x'
        Nothing -> I.liftIOOp (bracket (poolMake p) (poolFree p)) f

withPool :: I.MonadControlIO m => Pool a -> (a -> m b) -> m (Maybe b)
withPool p f = I.mask $ \unmask -> do
    eres <- liftIO $ atomicModifyIORef (poolData p) $ \pd ->
        case poolAvail pd of
            x:xs -> (pd { poolAvail = xs }, Right x)
            [] -> (pd, Left $ poolCreated pd)
    case eres of
        Left pc ->
            if pc >= poolMax p
                then return Nothing
                else I.liftIOOp (bracket (poolMake p) (insertResource 1))
                                (liftM Just . unmask . f)
        Right res -> do
            isAlive <- I.try $ unmask $ liftIO $ poolCheckAlive p res
            case isAlive :: Either E.SomeException Bool of
                Right True -> finallyIO (liftM Just $ unmask $ f res)
                                        (insertResource 0 res)
                _ -> do
                    -- decrement the poolCreated count and then start over
                    liftIO $ atomicModifyIORef (poolData p) $ \pd ->
                        (pd { poolCreated = poolCreated pd - 1}, ())
                    unmask $ withPool p f
  where
    insertResource i x = atomicModifyIORef (poolData p) $ \pd ->
        (pd { poolAvail = x : poolAvail pd
                , poolCreated = i + poolCreated pd
                }, ())
{'fraction_non_alphanumeric': 0.09870317002881844, 'fraction_numerical': 0.0009606147934678194, 'mean_word_length': 1.9863750448189315, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 99, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3455638', 'n_tokens_mistral': 2565, 'n_tokens_neox': 2475, 'n_words': 918}