Add support for assigning a label to the pool

Author: arybczak

CHANGELOG.md

@@ -9,6 +9,7 @@
 * 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.

resource-pool.cabal

@@ -36,6 +36,7 @@ library
                , hashable >= 1.1.0.0
                , primitive >= 0.7
                , stm
+               , text
                , time
 
   ghc-options:      -Wall

src/Data/Pool.hs

@@ -10,6 +10,7 @@ module Data.Pool
   , PoolConfig
   , defaultPoolConfig
   , setNumStripes
+  , setPoolLabel
 
     -- * Resource management
   , withResource
@@ -27,6 +28,7 @@ module Data.Pool
 import Control.Concurrent.STM
 import Control.Exception
 import Control.Monad
+import Data.Text qualified as T
 import Data.Time (NominalDiffTime)
 
 import Data.Pool.Internal
@@ -116,6 +118,7 @@ createPool create free numStripes idleTime maxResources =
       , poolCacheTTL = realToFrac idleTime
       , poolMaxResources = numStripes * maxResources
       , poolNumStripes = Just numStripes
+      , pcLabel = T.empty
       }
 
 ----------------------------------------

src/Data/Pool/Internal.hs

@@ -14,9 +14,10 @@ import Data.Either
 import Data.Hashable (hash)
 import Data.IORef
 import Data.List qualified as L
+import Data.Text qualified as T
 import Data.Primitive.SmallArray
 import GHC.Clock (getMonotonicTime)
-import GHC.Conc (unsafeIOToSTM)
+import GHC.Conc (labelThread, unsafeIOToSTM)
 
 -- | Striped resource pool based on "Control.Concurrent.QSem".
 data Pool a = Pool
@@ -60,6 +61,7 @@ data PoolConfig a = PoolConfig
   , poolCacheTTL :: !Double
   , poolMaxResources :: !Int
   , poolNumStripes :: !(Maybe Int)
+  , pcLabel :: !T.Text
   }
 
 -- | Create a 'PoolConfig' with optional parameters having default values.
@@ -94,6 +96,7 @@ defaultPoolConfig create free cacheTTL maxResources =
     , poolCacheTTL = cacheTTL
     , poolMaxResources = maxResources
     , poolNumStripes = Just 1
+    , pcLabel = T.empty
     }
 
 -- | Set the number of stripes (sub-pools) in the pool.
@@ -111,6 +114,15 @@ defaultPoolConfig create free cacheTTL maxResources =
 setNumStripes :: Maybe Int -> PoolConfig a -> PoolConfig a
 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
@@ -149,7 +161,10 @@ newPool pc = do
         }
   mask_ $ do
     ref <- newIORef ()
-    collectorA <- forkIOWithUnmask $ \unmask -> unmask $ collector pools
+    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.

src/Data/Pool/Introspection.hs

@@ -9,6 +9,7 @@ module Data.Pool.Introspection
   , PoolConfig
   , defaultPoolConfig
   , setNumStripes
+  , setPoolLabel
 
     -- * Resource management
   , Resource (..)
@@ -25,6 +26,7 @@ module Data.Pool.Introspection
 import Control.Concurrent.STM
 import Control.Exception
 import Control.Monad
+import Data.Text qualified as T
 import GHC.Clock (getMonotonicTime)
 import GHC.Generics (Generic)
 
@@ -33,6 +35,7 @@ import Data.Pool.Internal
 -- | A resource taken from the pool along with additional information.
 data Resource a = Resource
   { resource :: a
+  , poolLabel :: !T.Text
   , stripeNumber :: !Int
   , availableResources :: !Int
   , acquisition :: !Acquisition
@@ -75,6 +78,7 @@ takeResource pool = mask_ $ do
               let res =
                     Resource
                       { resource = a
+                      , poolLabel = pcLabel $ poolConfig pool
                       , stripeNumber = stripeId lp
                       , availableResources = 0
                       , acquisition = Delayed
@@ -89,6 +93,7 @@ takeResource pool = mask_ $ do
               let res =
                     Resource
                       { resource = a
+                      , poolLabel = pcLabel $ poolConfig pool
                       , stripeNumber = stripeId lp
                       , availableResources = 0
                       , acquisition = Delayed
@@ -143,6 +148,7 @@ takeAvailableResource pool t1 lp stripe = case cache stripe of
       let res =
             Resource
               { resource = a
+              , poolLabel = pcLabel $ poolConfig pool
               , stripeNumber = stripeId lp
               , availableResources = newAvailable
               , acquisition = Immediate
@@ -158,6 +164,7 @@ takeAvailableResource pool t1 lp stripe = case cache stripe of
       let res =
             Resource
               { resource = a
+              , poolLabel = pcLabel $ poolConfig pool
               , stripeNumber = stripeId lp
               , availableResources = newAvailable
               , acquisition = Immediate