diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,12 @@ +# 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 diff --git a/resource-pool.cabal b/resource-pool.cabal --- a/resource-pool.cabal +++ b/resource-pool.cabal @@ -1,7 +1,7 @@ cabal-version: 3.0 build-type: Simple name: resource-pool -version: 0.5.0.0 +version: 0.5.0.1 license: BSD-3-Clause license-file: LICENSE category: Data, Database, Network @@ -14,7 +14,7 @@ flexibly-sized collections of resources such as database connections. -tested-with: GHC == { 8.10.7, 9.0.2, 9.2.8, 9.4.8, 9.6.7, 9.8.4, 9.10.2, 9.12.2 } +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 @@ -51,5 +51,35 @@ , 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 @@ -99,9 +99,7 @@ join . atomically $ do stripe <- readTVar (stripeVar lp) if available stripe == 0 - then do - writeTVar (stripeVar lp) stripe - pure $ pure Nothing + then pure $ pure Nothing else fmap Just <$> takeAvailableResource pool lp stripe {-# DEPRECATED createPool "Use newPool instead" #-} 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 @@ -195,7 +195,8 @@ -- | 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 = mask_ $ do atomically $ do @@ -241,17 +242,27 @@ 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 + -- 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 - pure $ pools `indexSmallArray` (sid `rem` stripes) + -- '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 @@ -279,7 +290,11 @@ -- original size of the stripe. restoreSize :: TVar (Stripe a) -> IO () restoreSize mstripe = atomically $ do - modifyTVar' mstripe $ \stripe -> stripe {available = available stripe + 1} + 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 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 @@ -123,9 +123,7 @@ join . atomically $ do stripe <- readTVar (stripeVar lp) if available stripe == 0 - then do - writeTVar (stripeVar lp) stripe - pure $ pure Nothing + then pure $ pure Nothing else fmap Just <$> takeAvailableResource pool t1 lp stripe ---------------------------------------- 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)