resource-pool-0.4.0.0: src/Data/Pool/Internal.hs
-- | Internal implementation details for "Data.Pool".
--
-- This module is intended for internal use only, and may change without warning
-- in subsequent releases.
{-# OPTIONS_HADDOCK not-home #-}
module Data.Pool.Internal where
import Control.Concurrent
import Control.Exception
import Control.Monad
import Data.Hashable (hash)
import Data.IORef
import Data.Primitive.SmallArray
import GHC.Clock
import qualified Data.List as L
-- | Striped resource pool based on "Control.Concurrent.QSem".
data Pool a = Pool
{ poolConfig :: !(PoolConfig a)
, localPools :: !(SmallArray (LocalPool a))
, reaperRef :: !(IORef ())
}
-- | A single, local pool.
data LocalPool a = LocalPool
{ stripeId :: !Int
, stripeVar :: !(MVar (Stripe a))
, cleanerRef :: !(IORef ())
}
-- | Stripe of a resource pool. If @available@ is 0, the list of threads waiting
-- for a resource (each with an associated 'MVar') is @queue ++ reverse queueR@.
data Stripe a = Stripe
{ available :: !Int
, cache :: ![Entry a]
, queue :: !(Queue a)
, queueR :: !(Queue a)
}
-- | An existing resource currently sitting in a pool.
data Entry a = Entry
{ entry :: a
, lastUsed :: !Double
}
-- | A queue of MVarS corresponding to threads waiting for resources.
--
-- Basically a monomorphic list to save two pointer indirections.
data Queue a = Queue !(MVar (Maybe a)) (Queue a) | Empty
-- | Configuration of a 'Pool'.
data PoolConfig a = PoolConfig
{ createResource :: !(IO a)
, freeResource :: !(a -> IO ())
, poolCacheTTL :: !Double
, poolMaxResources :: !Int
, poolNumStripes :: !(Maybe Int)
}
-- | Create a 'PoolConfig' with optional parameters having default values.
--
-- For setting optional parameters have a look at:
--
-- - 'setNumStripes'
--
-- @since 0.4.0.0
defaultPoolConfig
:: IO a
-- ^ The action that creates a new resource.
-> (a -> IO ())
-- ^ The action that destroys an existing resource.
-> Double
-- ^ The amount of seconds for which an unused resource is kept around. The
-- smallest acceptable value is @0.5@.
--
-- /Note:/ the elapsed time before destroying a resource may be a little
-- longer than requested, as the collector thread wakes at 1-second intervals.
-> Int
-- ^ The maximum number of resources to keep open __across all stripes__. The
-- smallest acceptable value is @1@.
--
-- /Note:/ for each stripe the number of resources is divided by the number of
-- stripes and rounded up, hence the pool might end up creating up to @N - 1@
-- resources more in total than specified, where @N@ is the number of stripes.
-> PoolConfig a
defaultPoolConfig create free cacheTTL maxResources = PoolConfig
{ createResource = create
, freeResource = free
, poolCacheTTL = cacheTTL
, poolMaxResources = maxResources
, poolNumStripes = Nothing
}
-- | Set the number of stripes in the pool.
--
-- If set to 'Nothing' (the default value), the pool will create the amount of
-- stripes equal to the number of capabilities. This ensures that threads never
-- compete over access to the same stripe and results in a very good performance
-- in a multi-threaded environment.
--
-- @since 0.4.0.0
setNumStripes :: Maybe Int -> PoolConfig a -> PoolConfig a
setNumStripes numStripes pc = pc { poolNumStripes = numStripes }
-- | Create a new striped resource pool.
--
-- /Note:/ although the runtime system will destroy all idle resources when the
-- pool is garbage collected, it's recommended to manually call
-- 'destroyAllResources' when you're done with the pool so that the resources
-- are freed up as soon as possible.
newPool :: PoolConfig a -> IO (Pool a)
newPool pc = do
when (poolCacheTTL pc < 0.5) $ do
error "poolCacheTTL must be at least 0.5"
when (poolMaxResources pc < 1) $ do
error "poolMaxResources must be at least 1"
numStripes <- maybe getNumCapabilities pure (poolNumStripes pc)
when (numStripes < 1) $ do
error "numStripes must be at least 1"
when (poolMaxResources pc < numStripes) $ do
error "poolMaxResources must not be smaller than numStripes"
pools <- fmap (smallArrayFromListN numStripes) . forM [1..numStripes] $ \n -> do
ref <- newIORef ()
stripe <- newMVar Stripe
{ available = poolMaxResources pc `quotCeil` numStripes
, cache = []
, queue = Empty
, queueR = Empty
}
-- When the local pool goes out of scope, free its resources.
void . mkWeakIORef ref $ cleanStripe (const True) (freeResource pc) stripe
pure LocalPool { stripeId = n
, stripeVar = stripe
, cleanerRef = ref
}
mask_ $ do
ref <- newIORef ()
collectorA <- forkIOWithUnmask $ \unmask -> unmask $ collector pools
void . mkWeakIORef ref $ do
-- When the pool goes out of scope, stop the collector. Resources existing
-- in stripes will be taken care by their cleaners.
killThread collectorA
pure Pool { poolConfig = pc
, localPools = pools
, reaperRef = ref
}
where
quotCeil :: Int -> Int -> Int
quotCeil x y =
-- Basically ceiling (x / y) without going through Double.
let (z, r) = x `quotRem` y in if r == 0 then z else z + 1
-- Collect stale resources from the pool once per second.
collector pools = forever $ do
threadDelay 1000000
now <- getMonotonicTime
let isStale e = now - lastUsed e > poolCacheTTL pc
mapM_ (cleanStripe isStale (freeResource pc) . stripeVar) pools
-- | Destroy a resource.
--
-- Note that this will ignore any exceptions in the destroy function.
destroyResource :: Pool a -> LocalPool a -> a -> IO ()
destroyResource pool lp a = do
uninterruptibleMask_ $ do -- Note [signal uninterruptible]
stripe <- takeMVar (stripeVar lp)
newStripe <- signal stripe Nothing
putMVar (stripeVar lp) newStripe
void . try @SomeException $ freeResource (poolConfig pool) a
-- | Return a resource to the given 'LocalPool'.
putResource :: LocalPool a -> a -> IO ()
putResource lp a = do
uninterruptibleMask_ $ do -- Note [signal uninterruptible]
stripe <- takeMVar (stripeVar lp)
newStripe <- signal stripe (Just a)
putMVar (stripeVar lp) newStripe
-- | Destroy all resources in all stripes in the pool.
--
-- Note that this will ignore any exceptions in the destroy function.
--
-- This function is useful when you detect that all resources in the pool are
-- broken. For example after a database has been restarted all connections
-- opened before the restart will be broken. In that case it's better to close
-- those connections so that 'takeResource' won't take a broken connection from
-- the pool but will open a new connection instead.
--
-- Another use-case for this function is that when you know you are done with
-- the pool you can destroy all idle resources immediately instead of waiting on
-- the garbage collector to destroy them, thus freeing up those resources
-- sooner.
destroyAllResources :: Pool a -> IO ()
destroyAllResources pool = forM_ (localPools pool) $ \lp -> do
cleanStripe (const True) (freeResource (poolConfig pool)) (stripeVar lp)
----------------------------------------
-- Helpers
-- | Get a local pool.
getLocalPool :: SmallArray (LocalPool a) -> IO (LocalPool a)
getLocalPool pools = do
sid <- if stripes == 1
-- If there is just one stripe, there is no choice.
then pure 0
else do
capabilities <- getNumCapabilities
-- If the number of stripes is smaller than the number of capabilities and
-- doesn't divide it, selecting a stripe by a capability the current
-- thread runs on wouldn't give equal load distribution across all stripes
-- (e.g. if there are 2 stripes and 3 capabilities, stripe 0 would be used
-- by capability 0 and 2, while stripe 1 would only be used by capability
-- 1, a 100% load difference). In such case we select based on the id of a
-- thread.
if stripes < capabilities && capabilities `rem` stripes /= 0
then hash <$> myThreadId
else fmap fst . threadCapability =<< myThreadId
pure $ pools `indexSmallArray` (sid `rem` stripes)
where
stripes = sizeofSmallArray pools
-- | Wait for the resource to be put into a given 'MVar'.
waitForResource :: MVar (Stripe a) -> MVar (Maybe a) -> IO (Maybe a)
waitForResource mstripe q = takeMVar q `onException` cleanup
where
cleanup = uninterruptibleMask_ $ do -- Note [signal uninterruptible]
stripe <- takeMVar mstripe
newStripe <- tryTakeMVar q >>= \case
Just ma -> do
-- Between entering the exception handler and taking ownership of
-- the stripe we got the resource we wanted. We don't need it
-- anymore though, so pass it to someone else.
signal stripe ma
Nothing -> do
-- If we're still waiting, fill up the MVar with an undefined value
-- so that 'signal' can discard our MVar from the queue.
putMVar q $ error "unreachable"
pure stripe
putMVar mstripe newStripe
-- | If an exception is received while a resource is being created, restore the
-- original size of the stripe.
restoreSize :: MVar (Stripe a) -> IO ()
restoreSize mstripe = uninterruptibleMask_ $ do
-- 'uninterruptibleMask_' is used since 'takeMVar' might block.
stripe <- takeMVar mstripe
putMVar mstripe $! stripe { available = available stripe + 1 }
-- | Free resource entries in the stripes that fulfil a given condition.
cleanStripe
:: (Entry a -> Bool)
-> (a -> IO ())
-> MVar (Stripe a)
-> IO ()
cleanStripe isStale free mstripe = mask $ \unmask -> do
-- Asynchronous exceptions need to be masked here to prevent leaking of
-- 'stale' resources before they're freed.
stale <- modifyMVar mstripe $ \stripe -> unmask $ do
let (stale, fresh) = L.partition isStale (cache stripe)
-- There's no need to update 'available' here because it only tracks
-- the number of resources taken from the pool.
newStripe = stripe { cache = fresh }
newStripe `seq` pure (newStripe, map entry stale)
-- We need to ignore exceptions in the 'free' function, otherwise if an
-- exception is thrown half-way, we leak the rest of the resources. Also,
-- asynchronous exceptions need to be hard masked here since freeing a
-- resource might in theory block.
uninterruptibleMask_ . forM_ stale $ try @SomeException . free
-- Note [signal uninterruptible]
--
-- If we have
--
-- bracket takeResource putResource (...)
--
-- and an exception arrives at the putResource, then we must not lose the
-- resource. The putResource is masked by bracket, but taking the MVar might
-- block, and so it would be interruptible. Hence we need an uninterruptible
-- variant of mask here.
signal :: Stripe a -> Maybe a -> IO (Stripe a)
signal stripe ma = if available stripe == 0
then loop (queue stripe) (queueR stripe)
else do
newCache <- case ma of
Just a -> do
now <- getMonotonicTime
pure $ Entry a now : cache stripe
Nothing -> pure $ cache stripe
pure $! stripe { available = available stripe + 1
, cache = newCache
}
where
loop Empty Empty = do
newCache <- case ma of
Just a -> do
now <- getMonotonicTime
pure [Entry a now]
Nothing -> pure []
pure $! Stripe { available = 1
, cache = newCache
, queue = Empty
, queueR = Empty
}
loop Empty qR = loop (reverseQueue qR) Empty
loop (Queue q qs) qR = tryPutMVar q ma >>= \case
-- This fails when 'waitForResource' went into the exception handler and
-- filled the MVar (with an undefined value) itself. In such case we
-- simply ignore it.
False -> loop qs qR
True -> pure $! stripe { available = 0
, queue = qs
, queueR = qR
}
reverseQueue :: Queue a -> Queue a
reverseQueue = go Empty
where
go acc = \case
Empty -> acc
Queue x xs -> go (Queue x acc) xs