-- | 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
  , tryWithResource
  , tryTakeResource
  , 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 Pool a -> LocalPool a -> Stripe a -> IO (a, LocalPool a)
forall a. Pool a -> LocalPool a -> Stripe a -> IO (a, LocalPool a)
takeAvailableResource Pool a
pool LocalPool a
lp Stripe a
stripe

-- | A variant of 'withResource' that doesn't execute the action and returns
-- 'Nothing' instead of blocking if the capability-local pool is exhausted.
tryWithResource :: Pool a -> (a -> IO r) -> IO (Maybe r)
tryWithResource :: Pool a -> (a -> IO r) -> IO (Maybe r)
tryWithResource Pool a
pool a -> IO r
act = ((forall a. IO a -> IO a) -> IO (Maybe r)) -> IO (Maybe r)
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO (Maybe r)) -> IO (Maybe r))
-> ((forall a. IO a -> IO a) -> IO (Maybe r)) -> IO (Maybe r)
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> Pool a -> IO (Maybe (a, LocalPool a))
forall a. Pool a -> IO (Maybe (a, LocalPool a))
tryTakeResource Pool a
pool IO (Maybe (a, LocalPool a))
-> (Maybe (a, LocalPool a) -> IO (Maybe r)) -> IO (Maybe r)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Just (a
res, LocalPool a
localPool) -> do
    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
    Maybe r -> IO (Maybe r)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (r -> Maybe r
forall a. a -> Maybe a
Just r
r)
  Maybe (a, LocalPool a)
Nothing -> Maybe r -> IO (Maybe r)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe r
forall a. Maybe a
Nothing

-- | A variant of 'takeResource' that returns 'Nothing' instead of blocking if
-- the capability-local pool is exhausted.
tryTakeResource :: Pool a -> IO (Maybe (a, LocalPool a))
tryTakeResource :: Pool a -> IO (Maybe (a, LocalPool a))
tryTakeResource Pool a
pool = IO (Maybe (a, LocalPool a)) -> IO (Maybe (a, LocalPool a))
forall a. IO a -> IO a
mask_ (IO (Maybe (a, LocalPool a)) -> IO (Maybe (a, LocalPool a)))
-> IO (Maybe (a, LocalPool a)) -> IO (Maybe (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 (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
stripe
      Maybe (a, LocalPool a) -> IO (Maybe (a, LocalPool a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (a, LocalPool a)
forall a. Maybe a
Nothing
    else (a, LocalPool a) -> Maybe (a, LocalPool a)
forall a. a -> Maybe a
Just ((a, LocalPool a) -> Maybe (a, LocalPool a))
-> IO (a, LocalPool a) -> IO (Maybe (a, LocalPool a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pool a -> LocalPool a -> Stripe a -> IO (a, LocalPool a)
forall a. Pool a -> LocalPool a -> Stripe a -> IO (a, LocalPool a)
takeAvailableResource Pool a
pool LocalPool a
lp Stripe a
stripe

{-# 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
  }

----------------------------------------
-- Helpers

takeAvailableResource
  :: Pool a
  -> LocalPool a
  -> Stripe a
  -> IO (a, LocalPool a)
takeAvailableResource :: Pool a -> LocalPool a -> Stripe a -> IO (a, LocalPool a)
takeAvailableResource Pool a
pool LocalPool a
lp Stripe a
stripe = 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)