Permalink
Browse files

Merge pull request #11 from kim/master

Non-blocking versions of 'withResource', 'takeResource'
  • Loading branch information...
2 parents 12bfbb7 + 95f04b1 commit b0197df65564aa6544e8a19474a401f779c3d520 @bos committed Feb 12, 2013
Showing with 67 additions and 7 deletions.
  1. +67 −7 Data/Pool.hs
View
74 Data/Pool.hs
@@ -33,6 +33,8 @@ module Data.Pool
, createPool
, withResource
, takeResource
+ , tryWithResource
+ , tryTakeResource
, destroyResource
, putResource
) where
@@ -99,7 +101,7 @@ data Pool a = Pool {
, 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.
@@ -130,7 +132,7 @@ createPool
-> 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.
@@ -172,7 +174,7 @@ reaper destroy idleTime pools = forever $ do
return (map entry stale)
forM_ resources $ \resource -> do
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 +220,8 @@ withResource pool act = control $ \runInIO -> mask $ \restore -> do
-- 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 +232,70 @@ takeResource Pool{..} = do
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 ()

0 comments on commit b0197df

Please sign in to comment.