resource-pool 0.4.0.0 → 0.5.0.1
raw patch · 7 files changed
Files
- CHANGELOG.md +22/−0
- README.md +1/−2
- resource-pool.cabal +43/−6
- src/Data/Pool.hs +52/−41
- src/Data/Pool/Internal.hs +221/−161
- src/Data/Pool/Introspection.hs +101/−84
- test/Main.hs +320/−0
CHANGELOG.md view
@@ -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.
README.md view
@@ -1,8 +1,7 @@ # resource-pool -[](https://github.com/scrive/pool/actions?query=branch%3Amaster)+[](https://github.com/scrive/pool/actions/workflows/haskell-ci.yml) [](https://hackage.haskell.org/package/resource-pool)-[](https://packdeps.haskellers.com/feed?needle=andrzej@rybczak.net) [](https://www.stackage.org/lts/package/resource-pool) [](https://www.stackage.org/nightly/package/resource-pool)
resource-pool.cabal view
@@ -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
src/Data/Pool.hs view
@@ -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)
src/Data/Pool/Internal.hs view
@@ -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
src/Data/Pool/Introspection.hs view
@@ -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)
+ test/Main.hs view
@@ -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)