1  {- Resource pools.
 2   -
 3   - Copyright 2020 Joey Hess <id@joeyh.name>
 4   -
 5   - License: BSD-2-clause
 6   -}
 7  
 8  {-# LANGUAGE BangPatterns #-}
 9  
10  module Utility.ResourcePool (
11  	ResourcePool(..),
12  	mkResourcePool,
13  	mkResourcePoolNonConcurrent,
14  	withResourcePool,
15  	freeResourcePool,
16  ) where
17  
18  import Common
19  
20  import Control.Concurrent.STM
21  import Control.Monad.IO.Class
22  import Data.Either
23  
24  data ResourcePool r
25  	= ResourcePool Int (TVar Int) (TVar [r])
26  	| ResourcePoolNonConcurrent r
27  
28  {- Make a new resource pool, that can grow to contain the specified number
29   - of resources. -}
30  mkResourcePool :: MonadIO m => Int -> m (ResourcePool r)
31  mkResourcePool maxsz = liftIO $
32  	ResourcePool maxsz
33  		<$> newTVarIO 0
34  		<*> newTVarIO []
35  
36  {- When there will not be multiple threads that may 
37   - concurrently try to use it, using this is more
38   - efficient than mkResourcePool.
39   -}
40  mkResourcePoolNonConcurrent :: (MonadMask m, MonadIO m) => m r -> m (ResourcePool r)
41  mkResourcePoolNonConcurrent allocresource =
42  	ResourcePoolNonConcurrent <$> allocresource
43  
44  {- Runs an action with a resource.
45   -
46   - If no free resource is available in the pool,
47   - will run the action the allocate a new resource if the pool's size
48   - allows. Or will block a resource becomes available to use.
49   -
50   - The resource is returned to the pool at the end.
51   -}
52  withResourcePool :: (MonadMask m, MonadIO m) => ResourcePool r -> m r -> (r -> m a) -> m a
53  withResourcePool (ResourcePoolNonConcurrent r) _ a = a r
54  withResourcePool (ResourcePool maxsz currsz p) allocresource a =
55  	bracket setup cleanup a
56    where
57  	setup = do
58  		mr <- liftIO $ atomically $ do
59  			l <- readTVar p
60  			case l of
61  				(r:rs) -> do
62  					writeTVar p rs
63  					return (Just r)
64  				[] -> do
65  					n <- readTVar currsz
66  					if n < maxsz
67  						then do
68  							let !n' = succ n
69  							writeTVar currsz n'
70  							return Nothing
71  						else retry
72  		case mr of
73  			Just r -> return r
74  			Nothing -> allocresource
75  	cleanup r = liftIO $ atomically $ modifyTVar' p (r:)
76  
77  {- Frees all resources in use in the pool, running the supplied action on
78   - each. (If any of the actions throw an exception, it will be rethrown
79   - after all the actions have completed.)
80   -
81   - The pool should not have any resources in use when this is called,
82   - and the pool should not be used again after calling this.
83   -}
84  freeResourcePool :: (MonadMask m, MonadIO m) => ResourcePool r -> (r -> m ()) -> m ()
85  freeResourcePool (ResourcePoolNonConcurrent r) freeresource = freeresource r
86  freeResourcePool (ResourcePool _ currsz p) freeresource = do
87  	rs <- liftIO $ atomically $ do
88  		writeTVar currsz 0
89  		swapTVar p []
90  	res <- forM rs $ tryNonAsync . freeresource
91  	case lefts res of
92  		[] -> return ()
93  		(e:_) -> throwM e
94