Don't exceed maximum resources

CHANGELOG.md

@@ -1,6 +1,7 @@
 # resource-pool-0.4.0.0 (????-??-??)
 * Require `poolMaxResources` to be not smaller than the number of stripes.
 * Add support for setting the number of stripes.
+* The pool will not not exceed maximum resources. [#16](https://github.com/scrive/pool/pull/16)
 
 # resource-pool-0.3.1.0 (2022-06-15)
 * Add `tryWithResource` and `tryTakeResource`.

resource-pool.cabal

@@ -46,3 +46,26 @@ 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
+  ghc-options: -Wall -Wcompat
+  build-tool-depends: hspec-discover:hspec-discover
+
+  default-language: Haskell2010
+
+  default-extensions: DeriveGeneric
+                    , LambdaCase
+                    , RankNTypes
+                    , TypeApplications
+
+  other-modules: 
+    Data.Pool.InternalSpec

src/Data/Pool.hs

@@ -102,7 +102,7 @@ createPool create free numStripes idleTime maxResources = newPool PoolConfig
   { createResource   = create
   , freeResource     = free
   , poolCacheTTL     = realToFrac idleTime
-  , poolMaxResources = numStripes * maxResources
+  , poolMaxResources = maxResources
   , poolNumStripes   = Just numStripes
   }
 

src/Data/Pool/Internal.hs

@@ -3,6 +3,7 @@
 -- 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
@@ -64,15 +65,16 @@ data PoolConfig a = PoolConfig
   -- ^ The maximum number of resources to keep open across all stripes. The
   -- smallest acceptable value is @1@.
   --
-  -- /Note:/ for each stripe the number of resources is divided by the number of
-  -- stripes and rounded up, hence the pool might end up creating up to @N - 1@
-  -- resources more in total than specified, where @N@ is the number of stripes.
+  -- /Note/: Each stripe will try to have an equal amount of resources per
+  -- stripe. If the number of resources does not evenly divide into the
+  -- stripes, then some stripes will have one fewer resource than others.
   , poolNumStripes :: !(Maybe Int)
-  -- ^ The number of stripes in the pool.
+  -- ^ The number of stripes in the pool. If 'Nothing' is provided, then
+  -- this will use 'getNumCapabilities'.
   --
-  -- /Note:/ if set to 'Nothing', it defaults to the number of capabilities,
-  -- which ensures that threads never compete over access to the same stripe and
-  -- results in a very good performance in a multi-threaded environment.
+  -- The resource count in 'poolMaxResources' will be split evenly among
+  -- the available stripes. If there are fewer stripes than resources, then
+  -- this will only create as many stripes as resources.
   }
 
 -- | Create a new striped resource pool.
@@ -87,22 +89,34 @@ newPool pc = do
     error "poolCacheTTL must be at least 0.5"
   when (poolMaxResources pc < 1) $ do
     error "poolMaxResources must be at least 1"
-  numStripes <- maybe getNumCapabilities pure (poolNumStripes pc)
-  when (numStripes < 1) $ do
+  numStripesRequested <- maybe getNumCapabilities pure (poolNumStripes pc)
+  when (numStripesRequested < 1) $ do
     error "numStripes must be at least 1"
+  let stripeResourceAllocation =
+        howManyStripes Input
+          { inputMaxResources = poolMaxResources pc
+          , inputStripes = numStripesRequested
+          }
+      stripeAllocations =
+        robin stripeResourceAllocation
+      indexedAllocations =
+        zip [1..] stripeAllocations
+      numStripes =
+        allowedStripes stripeResourceAllocation
+
   when (poolMaxResources pc < numStripes) $ do
     error "poolMaxResources must not be smaller than numStripes"
-  pools <- fmap (smallArrayFromListN numStripes) . forM [1..numStripes] $ \n -> do
+  pools <- fmap (smallArrayFromListN numStripes) . forM indexedAllocations $ \(index, allocation) -> do
     ref <- newIORef ()
     stripe <- newMVar Stripe
-      { available = poolMaxResources pc `quotCeil` numStripes
+      { available = allocation
       , 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
+    pure LocalPool { stripeId   = index
                    , stripeVar  = stripe
                    , cleanerRef = ref
                    }
@@ -118,18 +132,45 @@ newPool pc = do
               , reaperRef  = ref
               }
   where
-    quotCeil :: Int -> Int -> Int
-    quotCeil x y =
-      -- Basically ceiling (x / y) without going through Double.
-      let (z, r) = x `quotRem` y in if r == 0 then z else z + 1
-
     -- Collect stale resources from the pool once per second.
     collector pools = forever $ do
       threadDelay 1000000
       now <- getMonotonicTime
       let isStale e = now - lastUsed e > poolCacheTTL pc
       mapM_ (cleanStripe isStale (freeResource pc) . stripeVar) pools
 
+-- | A datatype representing the requested maximum resources and count of
+-- stripes. We don't use these figurs directly, but instead calculate  a
+data Input = Input { inputMaxResources :: !Int, inputStripes :: !Int }
+  deriving Show
+
+data StripeResourceAllocation = StripeResourceAllocation { poolInput :: !Input, allowedStripes :: !Int }
+  deriving Show
+
+-- | Determine how many resources should be allocated to each stripe.
+--
+-- The output list contains a single `Int` per stripe, with the 'Int'
+-- representing the amount of resources available to that stripe.
+robin :: StripeResourceAllocation -> [Int]
+robin stripeResourceAllocation =
+  let
+    numStripes =
+      allowedStripes stripeResourceAllocation
+    (baseCount, remainder) =
+      inputMaxResources (poolInput stripeResourceAllocation)
+        `divMod` numStripes
+  in
+    replicate remainder (baseCount + 1) ++ replicate (numStripes - remainder) baseCount
+
+howManyStripes :: Input -> StripeResourceAllocation
+howManyStripes inp = StripeResourceAllocation
+  { allowedStripes =
+      if inputStripes inp > inputMaxResources inp
+      then inputMaxResources inp
+      else inputStripes inp
+  , poolInput = inp
+  }
+
 -- | Destroy a resource.
 --
 -- Note that this will ignore any exceptions in the destroy function.

test/Data/Pool/InternalSpec.hs

@@ -0,0 +1,46 @@
+{-# language RecordWildCards #-}
+{-# language OverloadedStrings #-}
+
+module Data.Pool.InternalSpec where
+
+import Test.Hspec
+import Test.Hspec.Hedgehog
+import qualified Hedgehog.Gen as Gen
+import qualified Hedgehog.Range as Range
+import Data.Pool.Internal
+
+spec :: Spec
+spec = do
+  describe "howManyStripes" $ do
+    it "has one stripe per resource" $ hedgehog $ do
+      input@Input {..} <- forAll inputs
+      let allocation = howManyStripes input
+      -- actual stripes may be at most requested stripes
+      if inputStripes >= inputMaxResources
+        then do
+          label "More stripes than resources"
+        else do
+          label "Fewer stripes than resources"
+          inputStripes === allowedStripes allocation
+      diff (allowedStripes allocation) (<=) inputStripes
+
+  describe "robin" $ do
+    it "sum of resources is always equal to input" $ hedgehog $ do
+      input <- forAll inputs
+      let resourceAllocations = robin (howManyStripes input)
+      sum resourceAllocations === inputMaxResources input
+    it "the difference between smallest and largest stripe is at most 1" $ hedgehog $ do
+      input <- forAll inputs
+      let resourceAllocations = robin (howManyStripes input)
+      diff (minimum resourceAllocations) (\a b -> abs (a - b) <= 1) (maximum resourceAllocations)
+    it "there is a resource allocation for each stripe" $ hedgehog $ do
+      input <- forAll inputs
+      let stripeAndResource = howManyStripes input
+      let resourceAllocations = robin stripeAndResource
+      length resourceAllocations === allowedStripes stripeAndResource
+
+inputs :: Gen Input
+inputs = do
+  resources <- Gen.int (Range.exponentialFrom 8 1 1000)
+  stripes <- Gen.int (Range.exponentialFrom 20 1 100)
+  pure $ Input resources stripes

test/Main.hs

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