diff --git a/Data/Pool.hs b/Data/Pool.hs --- a/Data/Pool.hs +++ b/Data/Pool.hs @@ -1,4 +1,4 @@ -{-# LANGUAGE CPP, NamedFieldPuns, RecordWildCards, ScopedTypeVariables #-} +{-# LANGUAGE CPP, NamedFieldPuns, RecordWildCards, ScopedTypeVariables, RankNTypes #-} #if MIN_VERSION_monad_control(0,3,0) {-# LANGUAGE FlexibleContexts #-} @@ -12,7 +12,8 @@ -- Module: Data.Pool -- Copyright: (c) 2011 MailRank, Inc. -- License: BSD3 --- Maintainer: Bryan O'Sullivan +-- Maintainer: Bryan O'Sullivan , +-- Bas van Dijk -- Stability: experimental -- Portability: portable -- @@ -21,11 +22,11 @@ -- connections. -- -- \"Striped\" means that a single 'Pool' consists of several --- sub-pools, each managed independently. A stripe size of 1 is fine --- for many applications, and probably what you should choose by --- default. Larger stripe sizes will lead to reduced contention in --- high-performance multicore applications, at a trade-off of causing --- the maximum number of simultaneous resources in use to grow. +-- sub-pools, each managed independently. A single stripe is fine for +-- many applications, and probably what you should choose by default. +-- More stripes will lead to reduced contention in high-performance +-- multicore applications, at a trade-off of causing the maximum +-- number of simultaneous resources in use to grow. module Data.Pool ( Pool(idleTime, maxResources, numStripes) @@ -33,19 +34,23 @@ , createPool , withResource , takeResource + , tryWithResource + , tryTakeResource , destroyResource , putResource + , destroyAllResources ) where import Control.Applicative ((<$>)) -import Control.Concurrent (forkIO, killThread, myThreadId, threadDelay) +import Control.Concurrent (ThreadId, forkIOWithUnmask, killThread, myThreadId, threadDelay) import Control.Concurrent.STM -import Control.Exception (SomeException, onException) -import Control.Monad (forM_, forever, join, liftM2, unless, when) +import Control.Exception (SomeException, onException, mask_) +import Control.Monad (forM_, forever, join, liftM3, unless, when) import Data.Hashable (hash) +import Data.IORef (IORef, newIORef, mkWeakIORef) import Data.List (partition) import Data.Time.Clock (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime) -import System.Mem.Weak (addFinalizer) +import GHC.Conc.Sync (labelThread) import qualified Control.Exception as E import qualified Data.Vector as V @@ -80,6 +85,8 @@ -- ^ Count of open entries (both idle and in use). , entries :: TVar [Entry a] -- ^ Idle entries. + , lfin :: IORef () + -- ^ empty value used to attach a finalizer to (internal) } data Pool a = Pool { @@ -88,7 +95,7 @@ , destroy :: a -> IO () -- ^ Action for destroying an entry that is now done with. , numStripes :: Int - -- ^ Stripe count. The number of distinct sub-pools to maintain. + -- ^ The number of stripes (distinct sub-pools) to maintain. -- The smallest acceptable value is 1. , idleTime :: NominalDiffTime -- ^ Amount of time for which an unused resource is kept alive. @@ -99,12 +106,14 @@ , maxResources :: Int -- ^ Maximum number of resources to maintain per stripe. The -- smallest acceptable value is 1. - -- + -- -- Requests for resources will block if this limit is reached on a -- single stripe, even if other stripes have idle resources -- available. , localPools :: V.Vector (LocalPool a) -- ^ Per-capability resource pools. + , fin :: IORef () + -- ^ empty value used to attach a finalizer to (internal) } instance Show (Pool a) where @@ -112,13 +121,19 @@ "idleTime = " ++ show idleTime ++ ", " ++ "maxResources = " ++ show maxResources ++ "}" +-- | Create a striped resource pool. +-- +-- Although the garbage collector will destroy all idle resources when +-- the pool is garbage collected it's recommended to manually +-- 'destroyAllResources' when you're done with the pool so that the +-- resources are freed up as soon as possible. createPool :: IO a -- ^ Action that creates a new resource. -> (a -> IO ()) -- ^ Action that destroys an existing resource. -> Int - -- ^ Stripe count. The number of distinct sub-pools to maintain. + -- ^ The number of stripes (distinct sub-pools) to maintain. -- The smallest acceptable value is 1. -> NominalDiffTime -- ^ Amount of time for which an unused resource is kept open. @@ -130,7 +145,7 @@ -> Int -- ^ Maximum number of resources to keep open per stripe. The -- smallest acceptable value is 1. - -- + -- -- Requests for resources will block if this limit is reached on a -- single stripe, even if other stripes have idle resources -- available. @@ -142,9 +157,11 @@ modError "pool " $ "invalid idle time " ++ show idleTime when (maxResources < 1) $ modError "pool " $ "invalid maximum resource count " ++ show maxResources - localPools <- atomically . V.replicateM numStripes $ - liftM2 LocalPool (newTVar 0) (newTVar []) - reaperId <- forkIO $ reaper destroy idleTime localPools + localPools <- V.replicateM numStripes $ + liftM3 LocalPool (newTVarIO 0) (newTVarIO []) (newIORef ()) + reaperId <- forkIOLabeledWithUnmask "resource-pool: reaper" $ \unmask -> + unmask $ reaper destroy idleTime localPools + fin <- newIORef () let p = Pool { create , destroy @@ -152,10 +169,36 @@ , idleTime , maxResources , localPools + , fin } - addFinalizer p $ killThread reaperId + mkWeakIORef fin (killThread reaperId) >> + V.mapM_ (\lp -> mkWeakIORef (lfin lp) (purgeLocalPool destroy lp)) localPools return p +-- TODO: Propose 'forkIOLabeledWithUnmask' for the base library. + +-- | Sparks off a new thread using 'forkIOWithUnmask' to run the given +-- IO computation, but first labels the thread with the given label +-- (using 'labelThread'). +-- +-- The implementation makes sure that asynchronous exceptions are +-- masked until the given computation is executed. This ensures the +-- thread will always be labeled which guarantees you can always +-- easily find it in the GHC event log. +-- +-- Like 'forkIOWithUnmask', the given computation is given a function +-- to unmask asynchronous exceptions. See the documentation of that +-- function for the motivation of this. +-- +-- Returns the 'ThreadId' of the newly created thread. +forkIOLabeledWithUnmask :: String + -> ((forall a. IO a -> IO a) -> IO ()) + -> IO ThreadId +forkIOLabeledWithUnmask label m = mask_ $ forkIOWithUnmask $ \unmask -> do + tid <- myThreadId + labelThread tid label + m unmask + -- | Periodically go through all pools, closing any resources that -- have been left idle for too long. reaper :: (a -> IO ()) -> NominalDiffTime -> V.Vector (LocalPool a) -> IO () @@ -172,7 +215,18 @@ return (map entry stale) forM_ resources $ \resource -> do destroy resource `E.catch` \(_::SomeException) -> return () - + +-- | Destroy all idle resources of the given 'LocalPool' and remove them from +-- the pool. +purgeLocalPool :: (a -> IO ()) -> LocalPool a -> IO () +purgeLocalPool destroy LocalPool{..} = do + resources <- atomically $ do + idle <- swapTVar entries [] + modifyTVar_ inUse (subtract (length idle)) + return (map entry idle) + forM_ resources $ \resource -> + destroy resource `E.catch` \(_::SomeException) -> return () + -- | Temporarily take a resource from a 'Pool', perform an action with -- it, and return it to the pool afterwards. -- @@ -218,9 +272,8 @@ -- that it may either be destroyed (via 'destroyResource') or returned to the -- pool (via 'putResource'). takeResource :: Pool a -> IO (a, LocalPool a) -takeResource Pool{..} = do - i <- liftBase $ ((`mod` numStripes) . hash) <$> myThreadId - let pool@LocalPool{..} = localPools V.! i +takeResource pool@Pool{..} = do + local@LocalPool{..} <- getLocalPool pool resource <- liftBase . join . atomically $ do ents <- readTVar entries case ents of @@ -231,11 +284,70 @@ writeTVar inUse $! used + 1 return $ create `onException` atomically (modifyTVar_ inUse (subtract 1)) - return (resource, pool) + return (resource, local) #if __GLASGOW_HASKELL__ >= 700 {-# INLINABLE takeResource #-} #endif +-- | Similar to 'withResource', but only performs the action if a resource could +-- be taken from the pool /without blocking/. Otherwise, 'tryWithResource' +-- returns immediately with 'Nothing' (ie. the action function is /not/ called). +-- Conversely, if a resource can be borrowed from the pool without blocking, the +-- action is performed and it's result is returned, wrapped in a 'Just'. +tryWithResource :: +#if MIN_VERSION_monad_control(0,3,0) + (MonadBaseControl IO m) +#else + (MonadControlIO m) +#endif + => Pool a -> (a -> m b) -> m (Maybe b) +tryWithResource pool act = control $ \runInIO -> mask $ \restore -> do + res <- tryTakeResource pool + case res of + Just (resource, local) -> do + ret <- restore (runInIO (Just <$> act resource)) `onException` + destroyResource pool local resource + putResource local resource + return ret + Nothing -> restore . runInIO $ return Nothing +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE tryWithResource #-} +#endif + +-- | A non-blocking version of 'takeResource'. The 'tryTakeResource' function +-- returns immediately, with 'Nothing' if the pool is exhausted, or @'Just' (a, +-- 'LocalPool' a)@ if a resource could be borrowed from the pool successfully. +tryTakeResource :: Pool a -> IO (Maybe (a, LocalPool a)) +tryTakeResource pool@Pool{..} = do + local@LocalPool{..} <- getLocalPool pool + resource <- liftBase . join . atomically $ do + ents <- readTVar entries + case ents of + (Entry{..}:es) -> writeTVar entries es >> return (return . Just $ entry) + [] -> do + used <- readTVar inUse + if used == maxResources + then return (return Nothing) + else do + writeTVar inUse $! used + 1 + return $ Just <$> + create `onException` atomically (modifyTVar_ inUse (subtract 1)) + return $ (flip (,) local) <$> resource +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE tryTakeResource #-} +#endif + +-- | Get a (Thread-)'LocalPool' +-- +-- Internal, just to not repeat code for 'takeResource' and 'tryTakeResource' +getLocalPool :: Pool a -> IO (LocalPool a) +getLocalPool Pool{..} = do + i <- liftBase $ ((`mod` numStripes) . hash) <$> myThreadId + return $ localPools V.! i +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE getLocalPool #-} +#endif + -- | Destroy a resource. Note that this will ignore any exceptions in the -- destroy function. destroyResource :: Pool a -> LocalPool a -> a -> IO () @@ -254,6 +366,23 @@ #if __GLASGOW_HASKELL__ >= 700 {-# INLINABLE putResource #-} #endif + +-- | Destroy all resources in all stripes in the pool. Note that this +-- will ignore any exceptions in the destroy function. +-- +-- This function is useful when you detect that all resources in the +-- pool are broken. For example after a database has been restarted +-- all connections opened before the restart will be broken. In that +-- case it's better to close those connections so that 'takeResource' +-- won't take a broken connection from the pool but will open a new +-- connection instead. +-- +-- Another use-case for this function is that when you know you are +-- done with the pool you can destroy all idle resources immediately +-- instead of waiting on the garbage collector to destroy them, thus +-- freeing up those resources sooner. +destroyAllResources :: Pool a -> IO () +destroyAllResources Pool{..} = V.forM_ localPools $ purgeLocalPool destroy modifyTVar_ :: TVar a -> (a -> a) -> STM () modifyTVar_ v f = readTVar v >>= \a -> writeTVar v $! f a diff --git a/resource-pool.cabal b/resource-pool.cabal --- a/resource-pool.cabal +++ b/resource-pool.cabal @@ -1,5 +1,5 @@ name: resource-pool -version: 0.2.1.1 +version: 0.2.2.0 synopsis: A high-performance striped resource pooling implementation description: A high-performance striped pooling abstraction for managing @@ -10,7 +10,8 @@ license: BSD3 license-file: LICENSE author: Bryan O'Sullivan -maintainer: Bryan O'Sullivan +maintainer: Bryan O'Sullivan , + Bas van Dijk copyright: Copyright 2011 MailRank, Inc. category: Data, Database, Network build-type: Simple @@ -22,13 +23,14 @@ flag developer description: operate in developer mode default: False + manual: True library - exposed-modules: + exposed-modules: Data.Pool - - build-depends: - base == 4.*, + + build-depends: + base >= 4.4 && < 5, hashable, monad-control >= 0.2.0.1, transformers,