|
| 1 | +-- | Allocate resources from a pool, guaranteeing resource handling via the |
| 2 | +-- ResourceT transformer. |
| 3 | +module Data.Conduit.Pool |
| 4 | + ( ManagedResource (..) |
| 5 | + , takeResource |
| 6 | + , takeResourceCheck |
| 7 | + , P.Pool |
| 8 | + , P.createPool |
| 9 | + , P.withResource |
| 10 | + ) where |
| 11 | + |
| 12 | +import qualified Data.Pool as P |
| 13 | +import Control.Monad.Trans.Resource |
| 14 | +import Control.Monad.IO.Class (liftIO) |
| 15 | +import qualified Data.IORef as I |
| 16 | + |
| 17 | +-- | The result of taking a resource. |
| 18 | +data ManagedResource m a = ManagedResource |
| 19 | + { mrValue :: a -- ^ The actual resource. |
| 20 | + , mrReuse :: Bool -> ResourceT m () |
| 21 | + -- ^ Let's you specify whether the resource should be returned to the pool |
| 22 | + -- (via 'P.putResource') or destroyed (via 'P.destroyResource') on release. |
| 23 | + -- This defaults to destruction, in case of exceptions. |
| 24 | + , mrRelease :: ResourceT m () |
| 25 | + -- ^ Release this resource, either destroying it or returning it to the |
| 26 | + -- pool. |
| 27 | + } |
| 28 | + |
| 29 | +-- | Take a resource from the pool and register a release action. |
| 30 | +takeResource :: ResourceIO m => P.Pool a -> ResourceT m (ManagedResource m a) |
| 31 | +takeResource pool = do |
| 32 | + onRelRef <- liftIO $ I.newIORef False |
| 33 | + (relKey, (a, _)) <- withIO |
| 34 | + (P.takeResource pool) |
| 35 | + (\(a, local) -> do |
| 36 | + onRel <- I.readIORef onRelRef |
| 37 | + if onRel |
| 38 | + then P.putResource local a |
| 39 | + else P.destroyResource pool local a) |
| 40 | + return ManagedResource |
| 41 | + { mrValue = a |
| 42 | + , mrReuse = liftIO . I.writeIORef onRelRef |
| 43 | + , mrRelease = release relKey |
| 44 | + } |
| 45 | + |
| 46 | +-- | Same as 'takeResource', but apply some action to check if a resource is |
| 47 | +-- still valid. |
| 48 | +takeResourceCheck :: ResourceIO m |
| 49 | + => P.Pool a |
| 50 | + -> (a -> ResourceT m Bool) |
| 51 | + -> ResourceT m (ManagedResource m a) |
| 52 | +takeResourceCheck pool check = do |
| 53 | + mr <- takeResource pool |
| 54 | + isValid <- check $ mrValue mr |
| 55 | + if isValid |
| 56 | + then return mr |
| 57 | + else do |
| 58 | + mrRelease mr |
| 59 | + takeResourceCheck pool check |
0 commit comments