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

    -- ** Configuration
  , PoolConfig
  , defaultPoolConfig
  , setNumStripes

    -- * 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 :: forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool a
pool a -> IO r
act = forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
  (a
res, LocalPool a
localPool) <- forall a. Pool a -> IO (a, LocalPool a)
takeResource Pool a
pool
  r
r                <- forall a. IO a -> IO a
unmask (a -> IO r
act a
res) forall a b. IO a -> IO b -> IO a
`onException` forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
localPool a
res
  forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
localPool a
res
  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 :: forall a. Pool a -> IO (a, LocalPool a)
takeResource Pool a
pool = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
  LocalPool a
lp <- forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool (forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool)
  Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
  if forall a. Stripe a -> Int
available Stripe a
stripe forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      MVar (Maybe a)
q <- forall a. IO (MVar a)
newEmptyMVar
      forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { queueR :: Queue a
queueR = forall a. MVar (Maybe a) -> Queue a -> Queue a
Queue MVar (Maybe a)
q (forall a. Stripe a -> Queue a
queueR Stripe a
stripe) }
      forall a. MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
waitForResource (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) MVar (Maybe a)
q forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just a
a  -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
        Maybe a
Nothing -> do
          a
a <- forall a. PoolConfig a -> IO a
createResource (forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) forall a b. IO a -> IO b -> IO a
`onException` forall a. MVar (Stripe a) -> IO ()
restoreSize (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
    else 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 local pool is exhausted.
tryWithResource :: Pool a -> (a -> IO r) -> IO (Maybe r)
tryWithResource :: forall a r. Pool a -> (a -> IO r) -> IO (Maybe r)
tryWithResource Pool a
pool a -> IO r
act = forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> forall a. Pool a -> IO (Maybe (a, LocalPool a))
tryTakeResource Pool a
pool forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  Just (a
res, LocalPool a
localPool) -> do
    r
r <- forall a. IO a -> IO a
unmask (a -> IO r
act a
res) forall a b. IO a -> IO b -> IO a
`onException` forall a. Pool a -> LocalPool a -> a -> IO ()
destroyResource Pool a
pool LocalPool a
localPool a
res
    forall a. LocalPool a -> a -> IO ()
putResource LocalPool a
localPool a
res
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Maybe a
Just r
r)
  Maybe (a, LocalPool a)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing

-- | A variant of 'takeResource' that returns 'Nothing' instead of blocking if
-- the local pool is exhausted.
tryTakeResource :: Pool a -> IO (Maybe (a, LocalPool a))
tryTakeResource :: forall a. Pool a -> IO (Maybe (a, LocalPool a))
tryTakeResource Pool a
pool = forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
  LocalPool a
lp <- forall a. SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool (forall a. Pool a -> SmallArray (LocalPool a)
localPools Pool a
pool)
  Stripe a
stripe <- forall a. MVar a -> IO a
takeMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
  if forall a. Stripe a -> Int
available Stripe a
stripe forall a. Eq a => a -> a -> Bool
== Int
0
    then do
      forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) Stripe a
stripe
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
    else forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 :: forall a.
IO a
-> (a -> IO ()) -> Int -> NominalDiffTime -> Int -> IO (Pool a)
createPool IO a
create a -> IO ()
free Int
numStripes NominalDiffTime
idleTime Int
maxResources = forall a. PoolConfig a -> IO (Pool a)
newPool PoolConfig
  { createResource :: IO a
createResource   = IO a
create
  , freeResource :: a -> IO ()
freeResource     = a -> IO ()
free
  , poolCacheTTL :: Double
poolCacheTTL     = forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
idleTime
  , poolMaxResources :: Int
poolMaxResources = Int
numStripes forall a. Num a => a -> a -> a
* Int
maxResources
  , poolNumStripes :: Maybe Int
poolNumStripes   = forall a. a -> Maybe a
Just Int
numStripes
  }

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

takeAvailableResource
  :: Pool a
  -> LocalPool a
  -> Stripe a
  -> IO (a, LocalPool a)
takeAvailableResource :: forall a. Pool a -> LocalPool a -> Stripe a -> IO (a, LocalPool a)
takeAvailableResource Pool a
pool LocalPool a
lp Stripe a
stripe = case forall a. Stripe a -> [Entry a]
cache Stripe a
stripe of
  [] -> do
    forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) forall a b. (a -> b) -> a -> b
$! Stripe a
stripe { available :: Int
available = forall a. Stripe a -> Int
available Stripe a
stripe forall a. Num a => a -> a -> a
- Int
1 }
    a
a <- forall a. PoolConfig a -> IO a
createResource (forall a. Pool a -> PoolConfig a
poolConfig Pool a
pool) forall a b. IO a -> IO b -> IO a
`onException` forall a. MVar (Stripe a) -> IO ()
restoreSize (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp)
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
  Entry a
a Double
_ : [Entry a]
as -> do
    forall a. MVar a -> a -> IO ()
putMVar (forall a. LocalPool a -> MVar (Stripe a)
stripeVar LocalPool a
lp) forall a b. (a -> b) -> a -> b
$! Stripe a
stripe
     { available :: Int
available = forall a. Stripe a -> Int
available Stripe a
stripe forall a. Num a => a -> a -> a
- Int
1
     , cache :: [Entry a]
cache = [Entry a]
as
     }
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)