-- | 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
  , setPoolLabel

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

    -- * Compatibility with 0.2
  , createPool
  ) where

import Control.Concurrent.STM
import Control.Exception
import Control.Monad
import Data.Text qualified as T
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 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 a. a -> IO a
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 = 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)
  IO (IO (a, LocalPool a)) -> IO (a, LocalPool a)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO (a, LocalPool a)) -> IO (a, LocalPool a))
-> (STM (IO (a, LocalPool a)) -> IO (IO (a, LocalPool a)))
-> STM (IO (a, LocalPool a))
-> IO (a, LocalPool a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM (IO (a, LocalPool a)) -> IO (IO (a, LocalPool a))
forall a. STM a -> IO a
atomically (STM (IO (a, LocalPool a)) -> IO (a, LocalPool a))
-> STM (IO (a, LocalPool a)) -> IO (a, LocalPool a)
forall a b. (a -> b) -> a -> b
$ do
    Stripe a
stripe <- TVar (Stripe a) -> STM (Stripe a)
forall a. TVar a -> STM a
readTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (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
        TMVar (Maybe a)
q <- STM (TMVar (Maybe a))
forall a. STM (TMVar a)
newEmptyTMVar
        TVar (Stripe a) -> Stripe a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> STM ()) -> Stripe a -> STM ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe {queueR = Queue q (queueR stripe)}
        IO (a, LocalPool a) -> STM (IO (a, LocalPool a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
          (IO (a, LocalPool a) -> STM (IO (a, LocalPool a)))
-> IO (a, LocalPool a) -> STM (IO (a, LocalPool a))
forall a b. (a -> b) -> a -> b
$ TVar (Stripe a) -> TMVar (Maybe a) -> IO (Maybe a)
forall a. TVar (Stripe a) -> TMVar (Maybe a) -> IO (Maybe a)
waitForResource (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp) TMVar (Maybe a)
q IO (Maybe a)
-> (Maybe a -> IO (a, LocalPool a)) -> IO (a, LocalPool a)
forall a b. IO a -> (a -> IO b) -> IO b
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 a. a -> IO 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` TVar (Stripe a) -> IO ()
forall a. TVar (Stripe a) -> IO ()
restoreSize (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp)
              (a, LocalPool a) -> IO (a, LocalPool a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
      else Pool a -> LocalPool a -> Stripe a -> STM (IO (a, LocalPool a))
forall a.
Pool a -> LocalPool a -> Stripe a -> STM (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 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 a b. IO a -> (a -> IO b) -> IO b
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 a. a -> IO a
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 a. a -> IO a
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 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 = 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)
  IO (IO (Maybe (a, LocalPool a))) -> IO (Maybe (a, LocalPool a))
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO (Maybe (a, LocalPool a))) -> IO (Maybe (a, LocalPool a)))
-> (STM (IO (Maybe (a, LocalPool a)))
    -> IO (IO (Maybe (a, LocalPool a))))
-> STM (IO (Maybe (a, LocalPool a)))
-> IO (Maybe (a, LocalPool a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM (IO (Maybe (a, LocalPool a)))
-> IO (IO (Maybe (a, LocalPool a)))
forall a. STM a -> IO a
atomically (STM (IO (Maybe (a, LocalPool a))) -> IO (Maybe (a, LocalPool a)))
-> STM (IO (Maybe (a, LocalPool a))) -> IO (Maybe (a, LocalPool a))
forall a b. (a -> b) -> a -> b
$ do
    Stripe a
stripe <- TVar (Stripe a) -> STM (Stripe a)
forall a. TVar a -> STM a
readTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (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
        TVar (Stripe a) -> Stripe a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp) Stripe a
stripe
        IO (Maybe (a, LocalPool a)) -> STM (IO (Maybe (a, LocalPool a)))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (Maybe (a, LocalPool a)) -> STM (IO (Maybe (a, LocalPool a))))
-> IO (Maybe (a, LocalPool a)) -> STM (IO (Maybe (a, LocalPool a)))
forall a b. (a -> b) -> a -> b
$ Maybe (a, LocalPool a) -> IO (Maybe (a, LocalPool a))
forall a. a -> IO 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))
-> IO (a, LocalPool a) -> IO (Maybe (a, LocalPool a))
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, LocalPool a) -> Maybe (a, LocalPool a)
forall a. a -> Maybe a
Just (IO (a, LocalPool a) -> IO (Maybe (a, LocalPool a)))
-> STM (IO (a, LocalPool a)) -> STM (IO (Maybe (a, LocalPool a)))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pool a -> LocalPool a -> Stripe a -> STM (IO (a, LocalPool a))
forall a.
Pool a -> LocalPool a -> Stripe a -> STM (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 =
  PoolConfig a -> IO (Pool a)
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 = 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
      , poolNumStripes :: Maybe Int
poolNumStripes = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
numStripes
      , pcLabel :: Text
pcLabel = Text
T.empty
      }

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

takeAvailableResource
  :: Pool a
  -> LocalPool a
  -> Stripe a
  -> STM (IO (a, LocalPool a))
takeAvailableResource :: forall a.
Pool a -> LocalPool a -> Stripe a -> STM (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
    TVar (Stripe a) -> Stripe a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp) (Stripe a -> STM ()) -> Stripe a -> STM ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe {available = available stripe - 1}
    IO (a, LocalPool a) -> STM (IO (a, LocalPool a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (a, LocalPool a) -> STM (IO (a, LocalPool a)))
-> IO (a, LocalPool a) -> STM (IO (a, LocalPool a))
forall a b. (a -> b) -> a -> b
$ 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` TVar (Stripe a) -> IO ()
forall a. TVar (Stripe a) -> IO ()
restoreSize (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp)
      (a, LocalPool a) -> IO (a, LocalPool a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)
  Entry a
a Double
_ : [Entry a]
as -> do
    TVar (Stripe a) -> Stripe a -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalPool a -> TVar (Stripe a)
forall a. LocalPool a -> TVar (Stripe a)
stripeVar LocalPool a
lp)
      (Stripe a -> STM ()) -> Stripe a -> STM ()
forall a b. (a -> b) -> a -> b
$! Stripe a
stripe
        { available = available stripe - 1
        , cache = as
        }
    IO (a, LocalPool a) -> STM (IO (a, LocalPool a))
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (IO (a, LocalPool a) -> STM (IO (a, LocalPool a)))
-> IO (a, LocalPool a) -> STM (IO (a, LocalPool a))
forall a b. (a -> b) -> a -> b
$ (a, LocalPool a) -> IO (a, LocalPool a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, LocalPool a
lp)