Implement timeout variants

resource-pool.cabal

@@ -46,3 +46,28 @@ library
                     , LambdaCase
                     , RankNTypes
                     , TypeApplications
+
+test-suite tests
+  type:          exitcode-stdio-1.0
+  hs-source-dirs:  test
+  main-is: Main.hs
+  build-depends: 
+    base
+    , hedgehog
+    , hspec
+    , hspec-hedgehog
+    , resource-pool
+    , async
+  ghc-options: -Wall -Wcompat -threaded +RTS -N -RTS
+  build-tool-depends: hspec-discover:hspec-discover
+
+  default-language: Haskell2010
+
+  default-extensions: DeriveGeneric
+                    , LambdaCase
+                    , RankNTypes
+                    , TypeApplications
+
+  other-modules: 
+    Data.PoolSpec
+    Data.Pool.InternalSpec

src/Data/Pool.hs

@@ -9,7 +9,9 @@ module Data.Pool
 
     -- * Resource management
   , withResource
+  , withResourceTimeout
   , takeResource
+  , takeResourceTimeout
   , tryWithResource
   , tryTakeResource
   , putResource
@@ -23,6 +25,7 @@ module Data.Pool
 import Control.Concurrent
 import Control.Exception
 import Data.Time (NominalDiffTime)
+import System.Timeout
 
 import Data.Pool.Internal
 
@@ -50,6 +53,27 @@ withResource pool act = mask $ \unmask -> do
   putResource localPool res
   pure r
 
+-- | Attempt to acquire  a resource within given time in microseconds. If
+-- the resource is acquired at that time, provide @Just a@ to the action.
+-- If the timeout failed, then provide @Nothing@ to the callback.
+withResourceTimeout :: Int -> Pool a -> (Maybe a -> IO r) -> IO r
+withResourceTimeout i pool act = mask $ \unmask -> do
+  mres <- takeResourceTimeout i pool
+  case mres of
+    Nothing ->
+      unmask (act Nothing)
+    Just (res, localPool) -> do
+      r <- unmask (act (Just res)) `onException` destroyResource pool localPool res
+      putResource localPool res
+      pure r
+
+-- | Attempt to take a resource from the pool in the given number of
+-- microseconds. If the resource cannot be acquired in that time, return
+-- 'Nothing'.
+takeResourceTimeout :: Int -> Pool a -> IO (Maybe (a, LocalPool a))
+takeResourceTimeout i p =
+  timeout i $ takeResource p
+
 -- | Take a resource from the pool, following the same results as
 -- 'withResource'.
 --
@@ -109,6 +133,7 @@ createPool create free numStripes idleTime maxResources = newPool PoolConfig
 ----------------------------------------
 -- Helpers
 
+-- | The 'LocalPool' returned is the same as the input.
 takeAvailableResource
   :: Pool a
   -> LocalPool a

src/Data/Pool/Internal.hs

@@ -11,6 +11,7 @@ import Control.Monad
 import Data.Hashable (hash)
 import Data.IORef
 import Data.Primitive.SmallArray
+import Data.Traversable
 import GHC.Clock
 import qualified Data.List as L
 
@@ -282,3 +283,11 @@ reverseQueue = go Empty
     go acc = \case
       Empty      -> acc
       Queue x xs -> go (Queue x acc) xs
+
+-- | For diagnostic and test use.
+poolAvailableResources :: Pool a -> IO Int
+poolAvailableResources pool = do
+  avails <- for (localPools pool) $ \localPool -> do
+    stripe <- readMVar (stripeVar localPool)
+    pure (available stripe)
+  pure $! sum avails

test/Data/Pool/InternalSpec.hs

@@ -0,0 +1,68 @@
+{-# language RecordWildCards #-}
+{-# language OverloadedStrings #-}
+
+module Data.Pool.InternalSpec where
+
+import Control.Concurrent
+import Control.Exception
+import Control.Monad
+import Data.Maybe (fromMaybe)
+import Test.Hspec
+import Test.Hspec.Hedgehog
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Range as Range
+import Data.Pool.Internal
+import Data.Pool
+
+spec :: Spec
+spec = do
+  describe "newPool" $ do
+    it "throws an error if max resources is less than stripes" $ hedgehog $ do
+      maxResources <- forAll $ Gen.integral (Range.linear 1 100)
+      numStripes <- forAll $ Gen.maybe $ Gen.integral (Range.linear 1 100)
+      nothingStripes <- evalIO $ getNumCapabilities
+      eresult <- fmap void $ evalIO $ try $ newPool PoolConfig
+        { createResource = pure ()
+        , freeResource = \_ -> pure ()
+        , poolCacheTTL = 60.0
+        , poolMaxResources = maxResources
+        , poolNumStripes = numStripes
+        }
+      let actualStripes = fromMaybe nothingStripes numStripes
+      classify ("maxResources < actualStripes") (maxResources < actualStripes)
+      case eresult of
+        Left (ErrorCall msg) -> do
+          diff maxResources (<) actualStripes
+          msg === "poolMaxResources must not be smaller than numStripes"
+        Right _ ->
+          pure ()
+
+  describe "poolAvailableResources" $ do
+    let
+      mkPool =
+        newPool PoolConfig
+          { createResource = pure ()
+          , freeResource = \_ -> pure ()
+          , poolCacheTTL = 60.0
+          , poolMaxResources = 10
+          , poolNumStripes = Just 1
+          }
+    before mkPool $ do
+      it "works" $ \pool -> do
+        poolAvailableResources pool
+          `shouldReturn`
+            10
+      it "works when a resource has been taken" $ \pool -> do
+        withResource pool $ \_ ->
+          poolAvailableResources pool
+            `shouldReturn`
+              9
+      it "nested withResource works fine too" $ \pool -> do
+        withResource pool $ \_ -> do
+          poolAvailableResources pool
+            `shouldReturn`
+              9
+          withResource pool $ \_ -> do
+            poolAvailableResources pool
+              `shouldReturn`
+                8

test/Data/PoolSpec.hs

@@ -0,0 +1,144 @@
+{-# language OverloadedStrings #-}
+
+module Data.PoolSpec where
+
+import GHC.Clock
+import Control.Concurrent.Async
+import Control.Monad
+import Test.Hspec
+import Test.Hspec.Hedgehog
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Range as Range
+import Data.Pool
+import Data.Pool.Internal
+import System.Timeout
+import Control.Concurrent
+
+spec :: Spec
+spec = do
+  let
+    mkPool action =
+      newPool PoolConfig
+        { createResource = action
+        , freeResource = \() -> pure ()
+        , poolCacheTTL = 0.5
+        , poolMaxResources = 1
+        , poolNumStripes = Just 1
+        }
+    waitTime =
+      10000
+    waitAction =
+      threadDelay waitTime
+
+  describe "createPool" $ do
+    it "does not error for any legal set of inputs" $ hedgehog $ do
+      numStripes <- forAll $ Gen.integral $ Range.linear 1 100
+      idleTime <- forAll $ Gen.realFrac_ $ Range.linearFrac 0.5 100
+      maxResources <- forAll $ Gen.integral $ Range.linear 1 100
+      void $  evalIO $ createPool (pure ()) (\_ -> pure ()) numStripes idleTime maxResources
+
+  describe "takeResource" $ do
+    before (mkPool (pure ())) $ do
+      it "works when there is a resource available" $ \pool -> do
+        (a, _) <- takeResource pool
+        a `shouldBe` ()
+
+      it "blocks when there is not a resource available" $ \pool -> do
+        poolAvailableResources pool
+          `shouldReturn`
+            1
+        (a, _) <- takeResource pool
+        a `shouldBe` ()
+        poolAvailableResources pool
+          `shouldReturn`
+            0
+
+        mresult <- timeout 100 $ takeResource pool
+        void mresult `shouldBe` Nothing
+
+      it "returns when a resource is made available" $ \pool -> do
+        poolAvailableResources pool
+          `shouldReturn`
+            1
+        (a, lp) <- takeResource pool
+        a `shouldBe` ()
+        poolAvailableResources pool
+          `shouldReturn`
+            0
+
+        mresult <- timeout 100 $ takeResource pool
+        void mresult `shouldBe` Nothing
+
+        (mresult', ()) <- concurrently
+          (do
+            timeout waitTime $ takeResource pool)
+          (do
+            threadDelay (waitTime `div` 2)
+            putResource lp a)
+        void mresult' `shouldBe` Just ()
+
+    before (mkPool waitAction) $ do
+      it "returns a resource when it is timed out" $ \pool -> do
+        start <- getMonotonicTime
+        let timeoutTime = waitTime `div` 2
+        mresult <- timeout timeoutTime $ takeResource pool
+        end <- getMonotonicTime
+        let
+          -- end and start are both in seconds, so to get to microseconds, we
+          -- need to make them much bigger. at the same time, the RTS is
+          -- going to have some amount of drift, so we don't want ot be too
+          -- precise.
+          timeDiff =
+            100 * (end - start)
+          eps =
+            timeDiff - fromIntegral (timeoutTime `div` 10000)
+
+        abs eps `shouldSatisfy` (< 1)
+
+        void mresult `shouldBe` Nothing
+        poolAvailableResources pool
+          `shouldReturn`
+            1
+        ((), _) <- takeResource pool
+        poolAvailableResources pool
+          `shouldReturn`
+            0
+
+  describe "withResourceTimeout" $ do
+    before (mkPool waitAction) $ do
+      describe "when resource available" $ do
+        it "provides a resource if fast enough" $ \pool -> do
+          poolAvailableResources pool `shouldReturn` 1
+
+          withResourceTimeout (waitTime * 2) pool $ \munit -> do
+            munit `shouldBe` Just ()
+
+        it "provides Nothing when timeout happens" $ \pool -> do
+
+          withResourceTimeout (waitTime `div` 2) pool $ \munit -> do
+            munit `shouldBe` Nothing
+
+      describe "when resource not available" $ do
+        it "waits until one is available" $ \pool -> do
+          (a, lp) <- takeResource pool
+
+          var <- newEmptyMVar
+          _ <- forkIO $ do
+            takeMVar var
+            putResource lp a
+
+          -- definitely won't happen since the resource hasn't been returned
+          -- yet
+          withResourceTimeout (waitTime * 2) pool $ \munit -> do
+            munit `shouldBe` Nothing
+
+          _ <- forkIO $ do
+            -- fork the thread which will fill the signal to return the
+            -- resource
+            threadDelay $ waitTime * 2
+            putMVar var ()
+
+          withResourceTimeout (waitTime * 3) pool $ \munit -> do
+            -- since we're waiting longer than the forked thread will take,
+            -- this should succeed
+            munit `shouldBe` Just ()

test/Main.hs

@@ -0,0 +1 @@
+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}