-- | A high-performance pooling abstraction for managing flexibly-sized
-- collections of resources such as database connections.
module Data.Pool
  ( -- * Pool
    PoolConfig(..)
  , Pool
  , LocalPool
  , newPool

    -- * Resource management
  , withResource
  , takeResource
  , putResource
  , destroyResource
  , destroyAllResources

    -- * Compatibility with 0.2
  , createPool
  ) where

import Control.Concurrent
import Control.Exception
import Data.Time (NominalDiffTime)

import Data.Pool.Internal

-- | Take a resource from the pool, perform an action with it and return it to
-- the pool afterwards.
--
-- * If the pool has an idle resource available, it is used immediately.
--
-- * Otherwise, if the maximum number of resources has not yet been reached, a
--   new resource is created and used.
--
-- * If the maximum number of resources has been reached, this function blocks
--   until a resource becomes available.
--
-- If the action throws an exception of any type, the resource is destroyed and
-- not returned to the pool.
--
-- It probably goes without saying that you should never manually destroy a
-- pooled resource, as doing so will almost certainly cause a subsequent user
-- (who expects the resource to be valid) to throw an exception.
withResource :: Pool a -> (a -> IO r) -> IO r
withResource :: Pool a -> (a -> IO r) -> IO r
withResource Pool a
pool a -> IO r
act = ((forall a. IO a -> IO a) -> IO r) -> IO r
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO r) -> IO r)
-> ((forall a. IO a -> IO a) -> IO r) -> IO r
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
  (a
res, LocalPool a
localPool) <- Pool a -> IO (a, LocalPool a)
forall a. Pool a -> IO (a, LocalPool a)
takeResource Pool a
pool
  r
r                <- IO r -> IO r
forall a. IO a -> IO a
unmask (a -> IO r
act a
res) IO r -> IO () -> IO r
forall a b. IO a -> IO b -> IO a
`onException` Pool a -> LocalPool a -> a -> IO ()
forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
localPool a
res
  LocalPool a -> a -> IO ()
forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
localPool a
res
  r -> IO r
forall (f :: * -> *) a. Applicative f => a -> f a
pure r
r

-- | Take a resource from the pool, following the same results as
-- 'withResource'.
--
-- /Note:/ this function returns both a resource and the 'LocalPool' it came
-- from so that it may either be destroyed (via 'destroyResource') or returned
-- to the pool (via 'putResource').
takeResource :: Pool a -> IO (a, LocalPool a)
takeResource :: Pool a -> IO (a, LocalPool a)
takeResource Pool a
pool = IO (a, LocalPool a) -> IO (a, LocalPool a)
forall a. IO a -> IO a
mask_ (IO (a, LocalPool a) -> IO (a, LocalPool a))
-> IO (a, LocalPool a) -> IO (a, LocalPool a)
forall a b. (a -> b) -> a -> b
$ do
  LocalPool a
lp <- SmallArray (LocalPool a) -> IO (LocalPool a)
forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool (Pool a -> SmallArray (LocalPool a)
forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool)
  Stripe a
stripe <- MVar (Stripe a) -> IO (Stripe a)
forall a. MVar a -> IO a
takeMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
  if Stripe a -> Int
forall a. Stripe a -> Int
available Stripe a
stripe Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      MVar (Maybe a)
q <- IO (MVar (Maybe a))
forall a. IO (MVar a)
newEmptyMVar
      MVar (Stripe a) -> Stripe a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> IO ()) -> Stripe a -> IO ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { queueR :: Queue a
queueR = MVar (Maybe a) -> Queue a -> Queue a
forall a. MVar (Maybe a) -> Queue a -> Queue a
Queue MVar (Maybe a)
q (Stripe a -> Queue a
forall a. Stripe a -> Queue a
queueR Stripe a
stripe) }
      MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
forall a. MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
waitForResource (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) MVar (Maybe a)
q IO (Maybe a)
-> (Maybe a -> IO (a, LocalPool a)) -> IO (a, LocalPool a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just a
a  -> (a, LocalPool a) -> IO (a, LocalPool a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
        Maybe a
Nothing -> do
          a
a <- PoolConfig a -> IO a
forall a. PoolConfig a -> IO a
createResource (Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`onException` MVar (Stripe a) -> IO ()
forall a. MVar (Stripe a) -> IO ()
restoreSize (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
          (a, LocalPool a) -> IO (a, LocalPool a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
    else case Stripe a -> [Entry a]
forall a. Stripe a -> [Entry a]
cache Stripe a
stripe of
      [] -> do
        MVar (Stripe a) -> Stripe a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> IO ()) -> Stripe a -> IO ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = Stripe a -> Int
forall a. Stripe a -> Int
available Stripe a
stripe Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 }
        a
a <- PoolConfig a -> IO a
forall a. PoolConfig a -> IO a
createResource (Pool a -> PoolConfig a
forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) IO a -> IO () -> IO a
forall a b. IO a -> IO b -> IO a
`onException` MVar (Stripe a) -> IO ()
forall a. MVar (Stripe a) -> IO ()
restoreSize (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
        (a, LocalPool a) -> IO (a, LocalPool a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
      Entry a
a Double
_ : [Entry a]
as -> do
        MVar (Stripe a) -> Stripe a -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (LocalPool a -> MVar (Stripe a)
forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> IO ()) -> Stripe a -> IO ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe
         { available :: Int
available = Stripe a -> Int
forall a. Stripe a -> Int
available Stripe a
stripe Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
         , cache :: [Entry a]
cache = [Entry a]
as
         }
        (a, LocalPool a) -> IO (a, LocalPool a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)

{-# DEPRECATED createPool "Use newPool instead" #-}
-- | Provided for compatibility with @resource-pool < 0.3@.
--
-- Use 'newPool' instead.
createPool :: IO a -> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
createPool :: IO a
-> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
createPool IO a
create a -> IO ()
free Int
numStripes NominalDiffTime
idleTime Int
maxResources = PoolConfig a -> IO (Pool a)
forall a. PoolConfig a -> IO (Pool a)
newPool PoolConfig :: forall a. IO a -> (a -> IO ()) -> Double -> Int -> PoolConfig a
PoolConfig
  { createResource :: IO a
createResource   = IO a
create
  , freeResource :: a -> IO ()
freeResource     = a -> IO ()
free
  , poolCacheTTL :: Double
poolCacheTTL     = NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
idleTime
  , poolMaxResources :: Int
poolMaxResources = Int
numStripes Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
maxResources
  }