diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,25 @@ +# resource-pool-0.5.0.1 (2026-07-08) +* Fix a bug where a thread waiting for a resource would get stuck in the queue + indefinitely if resource creation failed in another thread. +* Fix a potential out-of-bounds array access in stripe selection when the hash + of a thread id is negative. +* Fix a bug where some of the stripes would never be used if the number of + stripes was larger than the number of capabilities. +* Add a test suite. + +# resource-pool-0.5.0.0 (2025-06-13) +* Drop support for GHC < 8.10. +* Use STM based lockless implementation as it results in much better throughput + in a multi-threaded environment when number of stripes is not equal to the + number of capabilities (in particular with a single stripe). +* Stop running resource freeing functions within `uninterruptibleMask`. +* `destroyResource` no longer ignores exceptions thrown from resource releasing + functions. +* Change the default number of stripes to 1. +* Do not exceed the maximum number of resources if the number of stripes does + not divide it. +* Add support for assigning a label to the pool. + # resource-pool-0.4.0.0 (2023-01-16) * Require `poolMaxResources` to be not smaller than the number of stripes. * Add support for setting the number of stripes. diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -1,8 +1,7 @@ # resource-pool -[![Build Status](https://github.com/scrive/pool/workflows/Haskell-CI/badge.svg?branch=master)](https://github.com/scrive/pool/actions?query=branch%3Amaster) +[![CI](https://github.com/scrive/pool/actions/workflows/haskell-ci.yml/badge.svg?branch=master)](https://github.com/scrive/pool/actions/workflows/haskell-ci.yml) [![Hackage](https://img.shields.io/hackage/v/resource-pool.svg)](https://hackage.haskell.org/package/resource-pool) -[![Dependencies](https://img.shields.io/hackage-deps/v/resource-pool.svg)](https://packdeps.haskellers.com/feed?needle=andrzej@rybczak.net) [![Stackage LTS](https://www.stackage.org/package/resource-pool/badge/lts)](https://www.stackage.org/lts/package/resource-pool) [![Stackage Nightly](https://www.stackage.org/package/resource-pool/badge/nightly)](https://www.stackage.org/nightly/package/resource-pool) diff --git a/resource-pool.cabal b/resource-pool.cabal --- a/resource-pool.cabal +++ b/resource-pool.cabal @@ -1,7 +1,7 @@ -cabal-version: 2.4 +cabal-version: 3.0 build-type: Simple name: resource-pool -version: 0.4.0.0 +version: 0.5.0.1 license: BSD-3-Clause license-file: LICENSE category: Data, Database, Network @@ -14,8 +14,7 @@ flexibly-sized collections of resources such as database connections. -tested-with: GHC ==8.4.4 || ==8.6.5 || ==8.8.4 || ==8.10.7 || ==9.0.2 || ==9.2.5 - || ==9.4.3 +tested-with: GHC == { 8.10.7, 9.0.2, 9.2.8, 9.4.8, 9.6.7, 9.8.4, 9.10.3, 9.12.4, 9.14.1 } extra-doc-files: CHANGELOG.md @@ -33,16 +32,54 @@ Data.Pool.Internal Data.Pool.Introspection - build-depends: base >= 4.11 && < 5 + build-depends: base >= 4.14 && < 5 , hashable >= 1.1.0.0 , primitive >= 0.7 + , stm + , text , time - ghc-options: -Wall -Wcompat + ghc-options: -Wall + -Wcompat + -Wmissing-deriving-strategies + -Werror=prepositive-qualified-module default-language: Haskell2010 default-extensions: DeriveGeneric + , DerivingStrategies + , ImportQualifiedPost , LambdaCase , RankNTypes + , ScopedTypeVariables + , TypeApplications + +test-suite test + type: exitcode-stdio-1.0 + hs-source-dirs: test + + main-is: Main.hs + + build-depends: base + , async + , resource-pool + , stm + , tasty + , tasty-hunit + , text + + ghc-options: -Wall + -Wcompat + -Wmissing-deriving-strategies + -Werror=prepositive-qualified-module + -threaded + -rtsopts + "-with-rtsopts=-N4" + + default-language: Haskell2010 + + default-extensions: DerivingStrategies + , ImportQualifiedPost + , LambdaCase + , ScopedTypeVariables , TypeApplications diff --git a/src/Data/Pool.hs b/src/Data/Pool.hs --- a/src/Data/Pool.hs +++ b/src/Data/Pool.hs @@ -10,6 +10,7 @@ , PoolConfig , defaultPoolConfig , setNumStripes + , setPoolLabel -- * Resource management , withResource @@ -24,8 +25,10 @@ , createPool ) where -import Control.Concurrent +import Control.Concurrent.STM import Control.Exception +import Control.Monad +import Data.Text qualified as T import Data.Time (NominalDiffTime) import Data.Pool.Internal @@ -50,7 +53,7 @@ withResource :: Pool a -> (a -> IO r) -> IO r withResource pool act = mask $ \unmask -> do (res, localPool) <- takeResource pool - r <- unmask (act res) `onException` destroyResource pool localPool res + r <- unmask (act res) `onException` destroyResource pool localPool res putResource localPool res pure r @@ -63,52 +66,58 @@ takeResource :: Pool a -> IO (a, LocalPool a) takeResource pool = mask_ $ do lp <- getLocalPool (localPools pool) - stripe <- takeMVar (stripeVar lp) - if available stripe == 0 - then do - q <- newEmptyMVar - putMVar (stripeVar lp) $! stripe { queueR = Queue q (queueR stripe) } - waitForResource (stripeVar lp) q >>= \case - Just a -> pure (a, lp) - Nothing -> do - a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) - pure (a, lp) - else takeAvailableResource pool lp stripe + join . atomically $ do + stripe <- readTVar (stripeVar lp) + if available stripe == 0 + then do + q <- newEmptyTMVar + writeTVar (stripeVar lp) $! stripe {queueR = Queue q (queueR stripe)} + pure + $ waitForResource (stripeVar lp) q >>= \case + Just a -> pure (a, lp) + Nothing -> do + a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) + pure (a, lp) + else takeAvailableResource pool lp 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 pool act = mask $ \unmask -> tryTakeResource pool >>= \case - Just (res, localPool) -> do - r <- unmask (act res) `onException` destroyResource pool localPool res - putResource localPool res - pure (Just r) - Nothing -> pure Nothing +tryWithResource pool act = mask $ \unmask -> + tryTakeResource pool >>= \case + Just (res, localPool) -> do + r <- unmask (act res) `onException` destroyResource pool localPool res + putResource localPool res + pure (Just r) + Nothing -> pure 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 pool = mask_ $ do lp <- getLocalPool (localPools pool) - stripe <- takeMVar (stripeVar lp) - if available stripe == 0 - then do - putMVar (stripeVar lp) stripe - pure Nothing - else Just <$> takeAvailableResource pool lp stripe + join . atomically $ do + stripe <- readTVar (stripeVar lp) + if available stripe == 0 + then pure $ pure Nothing + else fmap Just <$> takeAvailableResource pool lp 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 create free numStripes idleTime maxResources = newPool PoolConfig - { createResource = create - , freeResource = free - , poolCacheTTL = realToFrac idleTime - , poolMaxResources = numStripes * maxResources - , poolNumStripes = Just numStripes - } +createPool create free numStripes idleTime maxResources = + newPool + PoolConfig + { createResource = create + , freeResource = free + , poolCacheTTL = realToFrac idleTime + , poolMaxResources = numStripes * maxResources + , poolNumStripes = Just numStripes + , pcLabel = T.empty + } ---------------------------------------- -- Helpers @@ -117,15 +126,17 @@ :: Pool a -> LocalPool a -> Stripe a - -> IO (a, LocalPool a) + -> STM (IO (a, LocalPool a)) takeAvailableResource pool lp stripe = case cache stripe of [] -> do - putMVar (stripeVar lp) $! stripe { available = available stripe - 1 } - a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) - pure (a, lp) + writeTVar (stripeVar lp) $! stripe {available = available stripe - 1} + pure $ do + a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) + pure (a, lp) Entry a _ : as -> do - putMVar (stripeVar lp) $! stripe - { available = available stripe - 1 - , cache = as - } - pure (a, lp) + writeTVar (stripeVar lp) + $! stripe + { available = available stripe - 1 + , cache = as + } + pure $ pure (a, lp) diff --git a/src/Data/Pool/Internal.hs b/src/Data/Pool/Internal.hs --- a/src/Data/Pool/Internal.hs +++ b/src/Data/Pool/Internal.hs @@ -1,60 +1,67 @@ +{-# OPTIONS_HADDOCK not-home #-} + -- | 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.Concurrent.STM import Control.Exception import Control.Monad +import Data.Either import Data.Hashable (hash) import Data.IORef +import Data.List qualified as L import Data.Primitive.SmallArray -import GHC.Clock -import qualified Data.List as L +import Data.Text qualified as T +import GHC.Clock (getMonotonicTime) +import GHC.Conc (labelThread, unsafeIOToSTM) -- | Striped resource pool based on "Control.Concurrent.QSem". data Pool a = Pool - { poolConfig :: !(PoolConfig a) - , localPools :: !(SmallArray (LocalPool a)) - , reaperRef :: !(IORef ()) + { poolConfig :: !(PoolConfig a) + , localPools :: !(SmallArray (LocalPool a)) + , reaperRef :: !(IORef ()) } -- | A single, local pool. data LocalPool a = LocalPool - { stripeId :: !Int - , stripeVar :: !(MVar (Stripe a)) + { stripeId :: !Int + , stripeVar :: !(TVar (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@. +-- for a resource (each with an associated 'TMVar') is @queue ++ reverse queueR@ +-- to ensure fairness. data Stripe a = Stripe { available :: !Int - , cache :: ![Entry a] - , queue :: !(Queue a) - , queueR :: !(Queue a) + , cache :: ![Entry a] + , queue :: !(Queue a) + , queueR :: !(Queue a) } -- | An existing resource currently sitting in a pool. data Entry a = Entry - { entry :: a + { entry :: a , lastUsed :: !Double } --- | A queue of MVarS corresponding to threads waiting for resources. +-- | A queue of TMVarS 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 +data Queue a = Queue !(TMVar (Maybe a)) (Queue a) | Empty -- | Configuration of a 'Pool'. data PoolConfig a = PoolConfig - { createResource :: !(IO a) - , freeResource :: !(a -> IO ()) - , poolCacheTTL :: !Double + { createResource :: !(IO a) + , freeResource :: !(a -> IO ()) + , poolCacheTTL :: !Double , poolMaxResources :: !Int - , poolNumStripes :: !(Maybe Int) + , poolNumStripes :: !(Maybe Int) + , pcLabel :: !T.Text } -- | Create a 'PoolConfig' with optional parameters having default values. @@ -70,38 +77,52 @@ -> (a -> IO ()) -- ^ The action that destroys an existing resource. -> Double - -- ^ The amount of seconds for which an unused resource is kept around. The + -- ^ The number 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@. + -- smallest acceptable value is @1@ per stripe. -- - -- /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. + -- /Note:/ if the number of stripes does not divide the number of resources, + -- some of the stripes will have 1 more resource available than the others. -> PoolConfig a -defaultPoolConfig create free cacheTTL maxResources = PoolConfig - { createResource = create - , freeResource = free - , poolCacheTTL = cacheTTL - , poolMaxResources = maxResources - , poolNumStripes = Nothing - } +defaultPoolConfig create free cacheTTL maxResources = + PoolConfig + { createResource = create + , freeResource = free + , poolCacheTTL = cacheTTL + , poolMaxResources = maxResources + , poolNumStripes = Just 1 + , pcLabel = T.empty + } --- | Set the number of stripes in the pool. +-- | Set the number of stripes (sub-pools) 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. +-- If not explicitly set, the default number of stripes is 1, which should be +-- good for typical use (when in doubt, profile your application first). -- +-- If set to 'Nothing', the pool will create the number of stripes equal to the +-- number of capabilities. +-- +-- /Note:/ usage of multiple stripes reduces contention, but can also result in +-- suboptimal use of resources since stripes are separated from each other. +-- -- @since 0.4.0.0 setNumStripes :: Maybe Int -> PoolConfig a -> PoolConfig a -setNumStripes numStripes pc = pc { poolNumStripes = numStripes } +setNumStripes numStripes pc = pc {poolNumStripes = numStripes} +-- | Assign a label to the pool. +-- +-- The label will appear in a label of the collector thread as well as +-- t'Data.Pool.Introspection.Resource'. +-- +-- @since 0.5.0.0 +setPoolLabel :: T.Text -> PoolConfig a -> PoolConfig a +setPoolLabel label pc = pc {pcLabel = label} + -- | Create a new striped resource pool. -- -- /Note:/ although the runtime system will destroy all idle resources when the @@ -119,36 +140,51 @@ 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 + let mkArray = fmap (smallArrayFromListN numStripes) + pools <- mkArray . forM (stripeResources numStripes) $ \(n, resources) -> do ref <- newIORef () - stripe <- newMVar Stripe - { available = poolMaxResources pc `quotCeil` numStripes - , cache = [] - , queue = Empty - , queueR = Empty - } + stripe <- + newTVarIO + Stripe + { available = resources + , 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 - } + pure + LocalPool + { stripeId = n + , stripeVar = stripe + , cleanerRef = ref + } mask_ $ do - ref <- newIORef () - collectorA <- forkIOWithUnmask $ \unmask -> unmask $ collector pools + ref <- newIORef () + collectorA <- forkIOWithUnmask $ \unmask -> unmask $ do + tid <- myThreadId + labelThread tid $ "resource-pool: collector (" ++ T.unpack (pcLabel pc) ++ ")" + 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 - } + 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 + stripeResources :: Int -> [(Int, Int)] + stripeResources numStripes = + let (base, rest) = quotRem (poolMaxResources pc) numStripes + in zip [1 .. numStripes] $ addRest (replicate numStripes base) rest + where + addRest [] = error "unreachable" + addRest acc@(r : rs) = \case + 0 -> acc + rest -> r + 1 : addRest rs (rest - 1) -- Collect stale resources from the pool once per second. collector pools = forever $ do @@ -159,22 +195,22 @@ -- | Destroy a resource. -- --- Note that this will ignore any exceptions in the destroy function. +-- /Note:/ since version 0.5.0.0 exceptions thrown by the destroy function are +-- no longer ignored. destroyResource :: Pool a -> LocalPool a -> a -> IO () -destroyResource pool lp a = do - uninterruptibleMask_ $ do -- Note [signal uninterruptible] - stripe <- takeMVar (stripeVar lp) +destroyResource pool lp a = mask_ $ do + atomically $ do + stripe <- readTVar (stripeVar lp) newStripe <- signal stripe Nothing - putMVar (stripeVar lp) newStripe - void . try @SomeException $ freeResource (poolConfig pool) a + writeTVar (stripeVar lp) $! newStripe + 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 +putResource lp a = atomically $ do + stripe <- readTVar (stripeVar lp) + newStripe <- signal stripe (Just a) + writeTVar (stripeVar lp) $! newStripe -- | Destroy all resources in all stripes in the pool. -- @@ -200,121 +236,145 @@ -- | 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) + sid <- + if stripes == 1 + then -- If there is just one stripe, there is no choice. + pure 0 + else do + capabilities <- getNumCapabilities + -- Selecting a stripe by a capability the current thread runs on gives + -- equal load distribution only if the number of stripes divides the + -- number of capabilities. Otherwise: + -- + -- - If the number of stripes is smaller than the number of + -- capabilities, load distribution across all stripes wouldn't be + -- equal (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). + -- + -- - If the number of stripes is larger than the number of capabilities, + -- stripes with an id larger than the highest capability would never + -- be selected. + -- + -- In such cases we select based on the id of a thread. + if capabilities `rem` stripes /= 0 + then hash <$> myThreadId + else fmap fst . threadCapability =<< myThreadId + -- 'mod' is used instead of 'rem' since the hash of a thread id might be + -- negative, which would result in an out-of-bounds array access. + pure $ pools `indexSmallArray` (sid `mod` 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 +-- | Wait for the resource to be put into a given 'TMVar'. +waitForResource :: TVar (Stripe a) -> TMVar (Maybe a) -> IO (Maybe a) +waitForResource mstripe q = atomically (takeTMVar 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 + cleanup = atomically $ do + stripe <- readTVar mstripe + newStripe <- + tryTakeTMVar 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 TMVar with an undefined value + -- so that 'signal' can discard our TMVar from the queue. + putTMVar q $ error "unreachable" + pure stripe + writeTVar 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 } +restoreSize :: TVar (Stripe a) -> IO () +restoreSize mstripe = atomically $ do + stripe <- readTVar mstripe + -- Signal needs to be called so that if there are threads waiting for a + -- resource, one of them wakes up and attempts the creation itself. + newStripe <- signal stripe Nothing + writeTVar mstripe $! newStripe -- | Free resource entries in the stripes that fulfil a given condition. cleanStripe :: (Entry a -> Bool) -> (a -> IO ()) - -> MVar (Stripe a) + -> TVar (Stripe a) -> IO () -cleanStripe isStale free mstripe = mask $ \unmask -> do +cleanStripe isStale free mstripe = mask_ $ do -- Asynchronous exceptions need to be masked here to prevent leaking of -- 'stale' resources before they're freed. - stale <- modifyMVar mstripe $ \stripe -> unmask $ do + stale <- atomically $ do + stripe <- readTVar mstripe 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) + -- There's no need to update 'available' here because it only tracks + -- the number of resources taken from the pool. + writeTVar mstripe $! stripe {cache = fresh} + pure $ 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 + -- asynchronous exceptions need to be hard masked here we need to run 'free' + -- for all resources. + uninterruptibleMask $ \release -> do + rs <- forM stale $ try @SomeException . release . free + -- If any async exception arrived in between, propagate it. + rethrowFirstAsyncException $ lefts rs + where + rethrowFirstAsyncException = \case + [] -> pure () + e : es + | Just SomeAsyncException {} <- fromException e -> throwIO e + | otherwise -> rethrowFirstAsyncException es --- 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 - } +signal :: forall a. Stripe a -> Maybe a -> STM (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 <- unsafeIOToSTM getMonotonicTime + pure $ Entry a now : cache stripe + Nothing -> pure $ cache stripe + pure + stripe + { available = available stripe + 1 + , cache = newCache + } where + loop :: Queue a -> Queue a -> STM (Stripe a) loop Empty Empty = do newCache <- case ma of Just a -> do - now <- getMonotonicTime + now <- unsafeIOToSTM 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 - } + pure + Stripe + { available = 1 + , cache = newCache + , queue = Empty + , queueR = Empty + } + loop Empty qR = loop (reverseQueue qR) Empty + loop (Queue q qs) qR = + tryPutTMVar q ma >>= \case + -- This fails when 'waitForResource' went into the exception handler and + -- filled the TMVar (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 + reverseQueue :: Queue a -> Queue a + reverseQueue = go Empty + where + go acc = \case + Empty -> acc + Queue x xs -> go (Queue x acc) xs diff --git a/src/Data/Pool/Introspection.hs b/src/Data/Pool/Introspection.hs --- a/src/Data/Pool/Introspection.hs +++ b/src/Data/Pool/Introspection.hs @@ -5,14 +5,15 @@ , LocalPool , newPool - -- ** Configuration + -- ** Configuration , PoolConfig , defaultPoolConfig , setNumStripes + , setPoolLabel -- * Resource management - , Resource(..) - , Acquisition(..) + , Resource (..) + , Acquisition (..) , withResource , takeResource , tryWithResource @@ -22,30 +23,34 @@ , destroyAllResources ) where -import Control.Concurrent +import Control.Concurrent.STM import Control.Exception -import GHC.Clock +import Control.Monad +import Data.Text qualified as T +import GHC.Clock (getMonotonicTime) import GHC.Generics (Generic) import Data.Pool.Internal -- | A resource taken from the pool along with additional information. data Resource a = Resource - { resource :: a - , stripeNumber :: !Int + { resource :: a + , poolLabel :: !T.Text + , stripeNumber :: !Int , availableResources :: !Int - , acquisition :: !Acquisition - , acquisitionTime :: !Double - , creationTime :: !(Maybe Double) - } deriving (Eq, Show, Generic) + , acquisition :: !Acquisition + , acquisitionTime :: !Double + , creationTime :: !(Maybe Double) + } + deriving stock (Eq, Generic, Show) -- | Describes how a resource was acquired from the pool. data Acquisition - = Immediate - -- ^ A resource was taken from the pool immediately. - | Delayed - -- ^ The thread had to wait until a resource was released. - deriving (Eq, Show, Generic) + = -- | A resource was taken from the pool immediately. + Immediate + | -- | The thread had to wait until a resource was released. + Delayed + deriving stock (Eq, Generic, Show) -- | 'Data.Pool.withResource' with introspection capabilities. withResource :: Pool a -> (Resource a -> IO r) -> IO r @@ -60,47 +65,54 @@ takeResource pool = mask_ $ do t1 <- getMonotonicTime lp <- getLocalPool (localPools pool) - stripe <- takeMVar (stripeVar lp) - if available stripe == 0 - then do - q <- newEmptyMVar - putMVar (stripeVar lp) $! stripe { queueR = Queue q (queueR stripe) } - waitForResource (stripeVar lp) q >>= \case - Just a -> do - t2 <- getMonotonicTime - let res = Resource - { resource = a - , stripeNumber = stripeId lp - , availableResources = 0 - , acquisition = Delayed - , acquisitionTime = t2 - t1 - , creationTime = Nothing - } - pure (res, lp) - Nothing -> do - t2 <- getMonotonicTime - a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) - t3 <- getMonotonicTime - let res = Resource - { resource = a - , stripeNumber = stripeId lp - , availableResources = 0 - , acquisition = Delayed - , acquisitionTime = t2 - t1 - , creationTime = Just $! t3 - t2 - } - pure (res, lp) - else takeAvailableResource pool t1 lp stripe + join . atomically $ do + stripe <- readTVar (stripeVar lp) + if available stripe == 0 + then do + q <- newEmptyTMVar + writeTVar (stripeVar lp) $! stripe {queueR = Queue q (queueR stripe)} + pure + $ waitForResource (stripeVar lp) q >>= \case + Just a -> do + t2 <- getMonotonicTime + let res = + Resource + { resource = a + , poolLabel = pcLabel $ poolConfig pool + , stripeNumber = stripeId lp + , availableResources = 0 + , acquisition = Delayed + , acquisitionTime = t2 - t1 + , creationTime = Nothing + } + pure (res, lp) + Nothing -> do + t2 <- getMonotonicTime + a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) + t3 <- getMonotonicTime + let res = + Resource + { resource = a + , poolLabel = pcLabel $ poolConfig pool + , stripeNumber = stripeId lp + , availableResources = 0 + , acquisition = Delayed + , acquisitionTime = t2 - t1 + , creationTime = Just $! t3 - t2 + } + pure (res, lp) + else takeAvailableResource pool t1 lp 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 -> (Resource a -> IO r) -> IO (Maybe r) -tryWithResource pool act = mask $ \unmask -> tryTakeResource pool >>= \case - Just (res, localPool) -> do - r <- unmask (act res) `onException` destroyResource pool localPool (resource res) - putResource localPool (resource res) - pure (Just r) - Nothing -> pure Nothing +tryWithResource pool act = mask $ \unmask -> + tryTakeResource pool >>= \case + Just (res, localPool) -> do + r <- unmask (act res) `onException` destroyResource pool localPool (resource res) + putResource localPool (resource res) + pure (Just r) + Nothing -> pure Nothing -- | A variant of 'takeResource' that returns 'Nothing' instead of blocking if -- the local pool is exhausted. @@ -108,12 +120,11 @@ tryTakeResource pool = mask_ $ do t1 <- getMonotonicTime lp <- getLocalPool (localPools pool) - stripe <- takeMVar (stripeVar lp) - if available stripe == 0 - then do - putMVar (stripeVar lp) stripe - pure Nothing - else Just <$> takeAvailableResource pool t1 lp stripe + join . atomically $ do + stripe <- readTVar (stripeVar lp) + if available stripe == 0 + then pure $ pure Nothing + else fmap Just <$> takeAvailableResource pool t1 lp stripe ---------------------------------------- -- Helpers @@ -123,33 +134,39 @@ -> Double -> LocalPool a -> Stripe a - -> IO (Resource a, LocalPool a) + -> STM (IO (Resource a, LocalPool a)) takeAvailableResource pool t1 lp stripe = case cache stripe of [] -> do let newAvailable = available stripe - 1 - putMVar (stripeVar lp) $! stripe { available = newAvailable } - t2 <- getMonotonicTime - a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) - t3 <- getMonotonicTime - let res = Resource - { resource = a - , stripeNumber = stripeId lp - , availableResources = newAvailable - , acquisition = Immediate - , acquisitionTime = t2 - t1 - , creationTime = Just $! t3 - t2 - } - pure (res, lp) + writeTVar (stripeVar lp) $! stripe {available = newAvailable} + pure $ do + t2 <- getMonotonicTime + a <- createResource (poolConfig pool) `onException` restoreSize (stripeVar lp) + t3 <- getMonotonicTime + let res = + Resource + { resource = a + , poolLabel = pcLabel $ poolConfig pool + , stripeNumber = stripeId lp + , availableResources = newAvailable + , acquisition = Immediate + , acquisitionTime = t2 - t1 + , creationTime = Just $! t3 - t2 + } + pure (res, lp) Entry a _ : as -> do let newAvailable = available stripe - 1 - putMVar (stripeVar lp) $! stripe { available = newAvailable, cache = as } - t2 <- getMonotonicTime - let res = Resource - { resource = a - , stripeNumber = stripeId lp - , availableResources = newAvailable - , acquisition = Immediate - , acquisitionTime = t2 - t1 - , creationTime = Nothing - } - pure (res, lp) + writeTVar (stripeVar lp) $! stripe {available = newAvailable, cache = as} + pure $ do + t2 <- getMonotonicTime + let res = + Resource + { resource = a + , poolLabel = pcLabel $ poolConfig pool + , stripeNumber = stripeId lp + , availableResources = newAvailable + , acquisition = Immediate + , acquisitionTime = t2 - t1 + , creationTime = Nothing + } + pure (res, lp) diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,320 @@ +module Main (main) where + +import Control.Concurrent +import Control.Concurrent.Async +import Control.Concurrent.STM +import Control.Exception +import Control.Monad +import Data.Foldable qualified as F +import Data.IORef +import Data.List qualified as L +import Data.Maybe +import Data.Text qualified as T +import System.Timeout (timeout) +import Test.Tasty hiding (withResource) +import Test.Tasty.HUnit + +import Data.Pool +import Data.Pool.Internal (LocalPool (..), Pool (..), Queue (..), Stripe (..)) +import Data.Pool.Introspection qualified as I + +main :: IO () +main = + defaultMain . localOption (mkTimeout 60000000) + $ testGroup + "resource-pool" + [ configValidationTests + , basicTests + , concurrencyTests + , introspectionTests + , stressTests + ] + +---------------------------------------- +-- Configuration validation + +configValidationTests :: TestTree +configValidationTests = + testGroup + "config validation" + [ testCase "rejects too small poolCacheTTL" $ do + expectError . newPool $ poolConfig_ 0.4 1 + , testCase "rejects non-positive poolMaxResources" $ do + expectError . newPool $ poolConfig_ 100 0 + , testCase "rejects non-positive number of stripes" $ do + expectError . newPool . setNumStripes (Just 0) $ poolConfig_ 100 1 + , testCase "rejects poolMaxResources smaller than the number of stripes" $ do + expectError . newPool . setNumStripes (Just 4) $ poolConfig_ 100 2 + ] + where + poolConfig_ = defaultPoolConfig (pure ()) (\_ -> pure ()) + + expectError :: IO a -> Assertion + expectError io = + try @ErrorCall (void io) >>= \case + Left _ -> pure () + Right _ -> assertFailure "expected newPool to error out" + +---------------------------------------- +-- Basic tests + +basicTests :: TestTree +basicTests = + testGroup + "basic" + [ testCase "an idle resource is reused" $ do + tp <- mkPool 5 + r1 <- withResource (testPool tp) pure + r2 <- withResource (testPool tp) pure + assertEqual "the same resource is taken" r1 r2 + readIORef (createdCount tp) >>= assertEqual "created resources" 1 + , testCase "concurrently taken resources are distinct" $ do + tp <- mkPool 5 + (r1, lp1) <- takeResource (testPool tp) + (r2, lp2) <- takeResource (testPool tp) + assertBool "resources are distinct" $ r1 /= r2 + putResource lp1 r1 + putResource lp2 r2 + , testCase "tryTakeResource returns Nothing iff the pool is exhausted" $ do + tp <- mkPool 1 + Just (r, lp) <- tryTakeResource (testPool tp) + mr <- tryTakeResource (testPool tp) + assertBool "Nothing when exhausted" $ isNothing mr + putResource lp r + mr2 <- tryTakeResource (testPool tp) + assertBool "Just when a resource is available" $ isJust mr2 + , testCase "tryWithResource returns Nothing iff the pool is exhausted" $ do + tp <- mkPool 1 + (r, lp) <- takeResource (testPool tp) + mr <- tryWithResource (testPool tp) pure + assertEqual "Nothing when exhausted" Nothing mr + putResource lp r + mr2 <- tryWithResource (testPool tp) pure + assertEqual "Just when a resource is available" (Just r) mr2 + , testCase "withResource destroys the resource on exception" $ do + tp <- mkPool 5 + r <- try @Boom $ withResource (testPool tp) $ \_ -> throwIO Boom :: IO () + assertEqual "the exception is propagated" (Left Boom) r + readIORef (freedCount tp) >>= assertEqual "freed resources" 1 + -- The pool is still usable and creates a fresh resource. + _ <- withResource (testPool tp) pure + readIORef (createdCount tp) >>= assertEqual "created resources" 2 + , testCase "destroyResource frees the resource and its slot" $ do + tp <- mkPool 1 + (r1, lp) <- takeResource (testPool tp) + destroyResource (testPool tp) lp r1 + readIORef (freedCount tp) >>= assertEqual "freed resources" 1 + r2 <- withResource (testPool tp) pure + assertBool "resource is fresh" $ r2 /= r1 + , testCase "destroyAllResources frees all idle resources" $ do + tp <- mkPool 5 + rs <- replicateM 3 $ takeResource (testPool tp) + forM_ rs $ \(r, lp) -> putResource lp r + destroyAllResources (testPool tp) + readIORef (freedCount tp) >>= assertEqual "freed resources" 3 + -- The pool is still usable and creates a fresh resource. + _ <- withResource (testPool tp) pure + readIORef (createdCount tp) >>= assertEqual "created resources" 4 + , testCase "resources are distributed evenly across stripes" $ do + tp <- mkPoolWith (setNumStripes $ Just 2) 5 + ss <- stripeStates (testPool tp) + assertEqual "available resources per stripe" [2, 3] + $ L.sort (map available ss) + , testCase "idle resources exceeding the TTL are collected" $ do + createdC <- newIORef (0 :: Int) + freedC <- newIORef (0 :: Int) + pool <- + newPool + $ defaultPoolConfig + (atomicModifyIORef' createdC $ \n -> (n + 1, n + 1)) + (\_ -> atomicModifyIORef' freedC $ \n -> (n + 1, ())) + 0.5 + 5 + _ <- withResource pool pure + -- The collector thread wakes up every second. + waitUntil "the resource is collected" $ (== 1) <$> readIORef freedC + readIORef createdC >>= assertEqual "created resources" 1 + ] + +---------------------------------------- +-- Concurrency tests + +concurrencyTests :: TestTree +concurrencyTests = + testGroup + "concurrency" + [ testCase "the number of created resources never exceeds the maximum" $ do + tp <- mkPool 3 + mapConcurrently_ + (\_ -> replicateM_ 10 . withResource (testPool tp) $ \_ -> threadDelay 1000) + [1 .. 20 :: Int] + created <- readIORef (createdCount tp) + assertBool ("created " ++ show created ++ " resources") $ created <= 3 + ss <- stripeStates (testPool tp) + assertEqual "all permits are available" 3 $ sum (map available ss) + , testCase "a blocked thread is woken up by putResource" $ do + tp <- mkPool 1 + (r, lp) <- takeResource (testPool tp) + done <- newEmptyMVar + _ <- forkIO $ withResource (testPool tp) pure >>= putMVar done + waitUntil "the thread is queued" + $ any ((> 0) . queueLength) <$> stripeStates (testPool tp) + putResource lp r + mr <- timeout 5000000 $ takeMVar done + assertEqual "the returned resource is received" (Just r) mr + , testCase "a blocked thread is woken up when resource creation fails" $ do + entered <- newEmptyMVar + gate <- newEmptyMVar + let create = do + putMVar entered () + takeMVar gate + throwIO Boom :: IO Int + pool <- newPool $ defaultPoolConfig create (\_ -> pure ()) 100 1 + t1 <- newEmptyMVar + _ <- forkIO $ try @Boom (withResource pool $ \_ -> pure ()) >>= putMVar t1 + takeMVar entered + t2 <- newEmptyMVar + _ <- forkIO $ try @Boom (withResource pool $ \_ -> pure ()) >>= putMVar t2 + waitUntil "the second thread is queued" + $ any ((> 0) . queueLength) <$> stripeStates pool + -- Fail creation in the first thread. The second one needs to be woken + -- up and attempt creation itself. + putMVar gate () + takeMVar t1 >>= assertEqual "the first thread failed" (Left Boom) + mr <- timeout 5000000 $ do + takeMVar entered + putMVar gate () + takeMVar t2 + assertEqual "the second thread attempted creation" (Just (Left Boom)) mr + ] + +---------------------------------------- +-- Introspection tests + +introspectionTests :: TestTree +introspectionTests = + testGroup + "introspection" + [ testCase "resource metadata is filled in correctly" $ do + pool <- + newPool . setPoolLabel (T.pack "label") + $ defaultPoolConfig (pure (0 :: Int)) (\_ -> pure ()) 100 1 + (res, lp) <- I.takeResource pool + assertEqual "poolLabel" (T.pack "label") (I.poolLabel res) + assertEqual "stripeNumber" 1 (I.stripeNumber res) + assertEqual "acquisition" I.Immediate (I.acquisition res) + assertBool "creationTime is set for a fresh resource" + $ isJust (I.creationTime res) + putResource lp $ I.resource res + (res2, lp2) <- I.takeResource pool + assertEqual "acquisition" I.Immediate (I.acquisition res2) + assertEqual "creationTime is not set for a cached resource" Nothing + $ I.creationTime res2 + putResource lp2 $ I.resource res2 + , testCase "acquisition of a blocked thread is Delayed" $ do + tp <- mkPool 1 + (res, lp) <- I.takeResource (testPool tp) + done <- newEmptyMVar + _ <- forkIO $ I.takeResource (testPool tp) >>= putMVar done + waitUntil "the thread is queued" + $ any ((> 0) . queueLength) <$> stripeStates (testPool tp) + putResource lp $ I.resource res + mr <- timeout 5000000 $ takeMVar done + case mr of + Nothing -> assertFailure "the blocked thread wasn't woken up" + Just (res2, lp2) -> do + assertEqual "acquisition" I.Delayed (I.acquisition res2) + putResource lp2 $ I.resource res2 + ] + +---------------------------------------- +-- Stress tests + +stressTests :: TestTree +stressTests = + testGroup + "stress" + [ testCase "pool invariants hold under load (1 stripe)" $ stressTest 1 + , testCase "pool invariants hold under load (4 stripes)" $ stressTest 4 + ] + where + stressTest numStripes = do + let maxResources = 8 + tp <- mkPoolWith (setNumStripes $ Just numStripes) maxResources + mapConcurrently_ + ( \i -> replicateM_ 50 $ do + r <- try @Boom . withResource (testPool tp) $ \_ -> do + threadDelay 100 + -- Some of the threads always fail, destroying the resource. + when (i `rem` 5 == 0) $ throwIO Boom + case r of + Left Boom -> when (i `rem` 5 /= 0) $ throwIO Boom + Right () -> pure () + ) + [1 .. 30 :: Int] + -- Once the pool is idle again, all permits are available, all live + -- resources sit in stripe caches and no thread is queued. + ss <- stripeStates (testPool tp) + assertEqual "all permits are available" maxResources + $ sum (map available ss) + created <- readIORef (createdCount tp) + freed <- readIORef (freedCount tp) + assertEqual "all live resources are cached" (created - freed) + $ sum (map (length . cache) ss) + assertEqual "no thread is queued" 0 $ sum (map queueLength ss) + +---------------------------------------- +-- Helpers + +data Boom = Boom + deriving stock (Eq, Show) + +instance Exception Boom + +data TestPool = TestPool + { testPool :: Pool Int + , createdCount :: IORef Int + , freedCount :: IORef Int + } + +mkPool :: Int -> IO TestPool +mkPool = mkPoolWith id + +mkPoolWith :: (PoolConfig Int -> PoolConfig Int) -> Int -> IO TestPool +mkPoolWith adjust maxResources = do + createdC <- newIORef 0 + freedC <- newIORef 0 + pool <- + newPool . adjust + $ defaultPoolConfig + (atomicModifyIORef' createdC $ \n -> (n + 1, n + 1)) + (\_ -> atomicModifyIORef' freedC $ \n -> (n + 1, ())) + 100 + maxResources + pure + TestPool + { testPool = pool + , createdCount = createdC + , freedCount = freedC + } + +stripeStates :: Pool a -> IO [Stripe a] +stripeStates pool = mapM (readTVarIO . stripeVar) . F.toList $ localPools pool + +queueLength :: Stripe a -> Int +queueLength stripe = go (queue stripe) + go (queueR stripe) + where + go = \case + Empty -> 0 + Queue _ q -> 1 + go q + +-- | Wait (up to 10 seconds) until a condition is met. +waitUntil :: String -> IO Bool -> Assertion +waitUntil label cond = go (100 :: Int) + where + go n = do + ok <- cond + unless ok + $ if n == 0 + then assertFailure $ "timed out waiting until " ++ label + else threadDelay 100000 >> go (n - 1)