Z-IO-1.0.0.0: Z/IO/Resource.hs
{-|
Module : Z.IO.Resource
Description : The Resource monad
Copyright : (c) Dong Han, 2017
License : BSD
Maintainer : winterland1989@gmail.com
Stability : experimental
Portability : non-portable
This module also implements Gabriel Gonzalez'd idea on 'Resource' applicative:
<http://www.haskellforall.com/2013/06/the-resource-applicative.html>. The 'Applicative' and 'Monad' instance is
especially useful when you want safely combine multiple resources.
A high performance resource pool is also provided.
-}
module Z.IO.Resource (
-- * Resource management
Resource(..)
, initResource
, initResource_
, withResource
, withResource'
-- * Resource pool
, Pool
, initPool
, withPool
, SimplePool
, initSimplePool
, withSimplePool
, statPool
-- * Re-export
, liftIO
) where
import Control.Concurrent
import Control.Monad
import qualified Control.Monad.Catch as MonadCatch
import Control.Monad.IO.Class
import qualified Data.Map.Strict as M
import Z.Data.PrimRef
import Z.Data.Array
import qualified Z.Data.Vector as V
import Data.IORef
import Z.IO.LowResTimer
import Z.IO.Exception
--------------------------------------------------------------------------------
-- | A 'Resource' is an 'IO' action which acquires some resource of type a and
-- also returns a finalizer of type IO () that releases the resource.
--
-- The only safe way to use a 'Resource' is 'withResource' and 'withResource'',
-- You should not use the 'acquire' field directly, unless you want to implement your own
-- resource management. In the later case, you should 'mask_' 'acquire' since
-- some resource initializations may assume async exceptions are masked.
--
-- 'MonadIO' instance is provided so that you can lift 'IO' computation inside
-- 'Resource', this is convenient for propagating 'Resource' around since many
-- 'IO' computations carry finalizers.
--
-- A convention in Z-IO is that functions returning a 'Resource' should be
-- named in @initXXX@ format, users are strongly recommended to follow this convention.
--
-- There're two additional guarantees we made in Z-IO:
--
-- * All resources in Z-IO can track its own liveness, throw 'ResourceVanished'
-- exception using 'throwECLOSED' or 'throwECLOSEDSTM' when used after resource
-- is closed.
--
-- * All resources' clean up action in Z-IO is idempotent.
--
-- Library authors providing 'initXXX' are also encouraged to provide these guarantees.
--
newtype Resource a = Resource { acquire :: IO (a, IO ()) }
-- | Create 'Resource' from create and release action.
--
-- Note, 'resource' doesn't open resource itself, resource is created when you use
-- 'with' \/ 'with''.
--
initResource :: IO a -> (a -> IO ()) -> Resource a
{-# INLINABLE initResource #-}
initResource create release = Resource $ do
r <- create
return $ (r, release r)
-- | Create 'Resource' from create and release action.
--
-- This function is useful when you want to add some initialization and clean up action
-- inside 'Resource' monad.
--
initResource_ :: IO a -> IO () -> Resource a
{-# INLINABLE initResource_ #-}
initResource_ create release = Resource $ do
r <- create
return $ (r, release)
instance Functor Resource where
{-# INLINE fmap #-}
fmap f resource = Resource $ do
(a, release) <- acquire resource
return (f a, release)
instance Applicative Resource where
{-# INLINE pure #-}
pure a = Resource (pure (a, pure ()))
{-# INLINE (<*>) #-}
resource1 <*> resource2 = Resource $ do
(f, release1) <- acquire resource1
(x, release2) <- acquire resource2 `onException` release1
return (f x, release2 >> release1)
instance Monad Resource where
{-# INLINE return #-}
return = pure
{-# INLINE (>>=) #-}
m >>= f = Resource $ do
(m', release1) <- acquire m
(x , release2) <- acquire (f m') `onException` release1
return (x, release2 >> release1)
instance MonadIO Resource where
{-# INLINE liftIO #-}
liftIO f = Resource $ fmap (\ a -> (a, dummyRelease)) f
where dummyRelease = return ()
-- | Create a new resource and run some computation, resource is guarantee to
-- be closed.
--
-- Be careful, don't leak the resource through the computation return value
-- because after the computation finishes, the resource is already closed.
--
withResource :: (MonadCatch.MonadMask m, MonadIO m, HasCallStack)
=> Resource a -> (a -> m b) -> m b
{-# INLINABLE withResource #-}
withResource resource k = MonadCatch.bracket
(liftIO (acquire resource))
(\(_, release) -> liftIO release)
(\(a, _) -> k a)
-- | Create a new resource and run some computation, resource is guarantee to
-- be closed.
--
-- The difference from 'with' is that the computation will receive an extra
-- close action, which can be used to close the resource early before the whole
-- computation finished, the close action can be called multiple times,
-- only the first call will clean up the resource.
--
withResource' :: (MonadCatch.MonadMask m, MonadIO m, HasCallStack)
=> Resource a -> (a -> m () -> m b) -> m b
{-# INLINABLE withResource' #-}
withResource' resource k = do
c <- liftIO (newCounter 0)
MonadCatch.bracket
(liftIO $ do
(a, release) <- (acquire resource)
let release' = do
c' <- atomicOrCounter c 1
when (c' == 0) release
return (a, release'))
(\(_, release) -> liftIO release)
(\(a, release) -> k a (liftIO release))
--------------------------------------------------------------------------------
-- | A entry linked-list annotated with size.
data Entry res
= EntryNil
| EntryCons
(res, IO ()) -- the resource and clean up action
{-# UNPACK #-} !Int -- size from this point on
{-# UNPACK #-} !Int -- the life remaining
(Entry res) -- next entry
-- | A high performance resource pool.
--
-- The Pool is first divided by GHC runtime capabilities, each capability maintains a map from key to living
-- resource list. Resource are fetched from living list first, create on demand if there's no living resource.
--
data Pool key res = Pool
{ _poolResource :: key -> Resource res -- ^ how to get a resource
, _poolLimitPerKey :: {-# UNPACK #-} !Int -- ^ max number for resource we keep alive after used
, _poolIdleTime :: {-# UNPACK #-} !Int -- ^ max idle time for resource we keep alive
, _poolArray :: {-# UNPACK #-} !(UnliftedArray (IORef (Maybe (M.Map key (Entry res)))))
}
-- | Dump the status of pool.
statPool :: Pool key res -> IO (SmallArray (M.Map key Int))
{-# INLINABLE statPool #-}
statPool (Pool _ _ _ arr) = (`V.traverse` arr) $ \ resMapRef -> do
mResMap <- readIORef resMapRef
case mResMap of
Just resMap -> return $ (`fmap` resMap) ( \ es ->
case es of EntryCons _ siz _ _ -> siz
_ -> 0)
_ -> throwECLOSED
-- | Initialize a resource pool with given 'Resource'
--
-- Like other initXXX functions, this function won't open a resource pool until you use 'withResource'.
initPool :: (key -> Resource res)
-> Int -- ^ maximum number of resources per local pool per key to be maintained.
-> Int -- ^ amount of time after which an unused resource can be released (in seconds).
-> Resource (Pool key res)
{-# INLINABLE initPool #-}
initPool resf limit itime = initResource createPool closePool
where
createPool = do
numCaps <- getNumCapabilities
marr <- newArr numCaps
forM_ [0..numCaps-1] $ \ i -> do
writeArr marr i =<< newIORef (Just M.empty)
arr <- unsafeFreezeArr marr
return (Pool resf limit itime arr)
closePool (Pool _ _ _ localPoolArr) = do
-- close all existed resource
(`V.traverse_` localPoolArr) $ \ resMapRef ->
atomicModifyIORef resMapRef $ \ mResMap ->
case mResMap of
Just resMap -> (Nothing, mapM_ closeEntry resMap)
_ -> (Nothing, return ())
closeEntry (EntryCons (_, close) _ _ _) = ignoreSync close
closeEntry EntryNil = return ()
-- | Open resource inside a given resource pool and do some computation.
--
-- This function is thread safe, concurrently usage will be guaranteed
-- to get different resource. If exception happens during computation,
-- resource will be closed(not return to pool).
withPool :: (MonadCatch.MonadMask m, MonadIO m, Ord key, HasCallStack)
=> Pool key res -> key -> (res -> m a) -> m a
{-# INLINABLE withPool #-}
withPool (Pool resf limitPerKey itime arr) key f = do
!resMapRef <- indexArr arr . fst <$> liftIO (threadCapability =<< myThreadId)
fst <$> MonadCatch.generalBracket
(liftIO $ takeFromPool resMapRef)
(\ r@(_, close) exit ->
case exit of
MonadCatch.ExitCaseSuccess _ -> liftIO (returnToPool resMapRef r)
_ -> liftIO close)
(\ (a, _) -> f a)
where
takeFromPool resMapRef =
join . atomicModifyIORef' resMapRef $ \ mResMap ->
case mResMap of
Just resMap ->
case M.lookup key resMap of
Just (EntryCons a _ _ es') ->
(Just $! M.adjust (const es') key resMap, return a)
_ -> (Just resMap, acquire (resf key))
_ -> (Nothing, throwECLOSED)
returnToPool resMapRef r = do
join . atomicModifyIORef' resMapRef $ \ mResMap ->
case mResMap of
Just resMap ->
case M.lookup key resMap of
Just (EntryCons _ siz _ _) ->
if siz < limitPerKey
-- if entries under given key do not exceed limit, we prepend res back to entries
then (Just $! M.adjust (EntryCons r (siz+1) itime) key resMap, return ())
-- otherwise we close it
else (Just resMap, snd r)
_ -> (Just $! M.insert key (EntryCons r 1 itime EntryNil) resMap,
scanLocalPool resMapRef)
_ -> (Nothing, snd r)
scanLocalPool resMapRef = do
registerLowResTimer_ 10 . join . atomicModifyIORef' resMapRef $ \ mResMap ->
case mResMap of
Just resMap ->
case M.lookup key resMap of
-- this is where we clean up empty keys
Just EntryNil -> (Just $! M.delete key resMap, return ())
Just es -> do
let (dead, living) = age es 0 [] EntryNil
case living of
-- no living resources any more, stop scanning
EntryNil -> (Just $! M.delete key resMap,
forM_ dead (ignoreSync . snd))
_ -> (Just $! M.adjust (const living) key resMap,
(do forM_ dead (ignoreSync . snd)
scanLocalPool resMapRef))
-- no living resources under given key, stop scanning
_ -> (Just resMap, return ())
_ -> (Nothing, return ())
age (EntryCons a _ life es) !livingNum dead living
| life > 1 = let !livingNum' = (livingNum+1)
in age es livingNum' dead (EntryCons a livingNum' (life-1) living)
| otherwise = age es livingNum (a:dead) living
age _ _ dead living = (dead, living)
-- | Simple resource pool where lookup via key is not needed.
type SimplePool res = Pool () res
-- | Initialize a 'SimplePool'.
initSimplePool :: Resource res
-> Int -- ^ maximum number of resources per local pool to be maintained.
-> Int -- ^ amount of time after which an unused resource can be released (in seconds).
-> Resource (SimplePool res)
{-# INLINABLE initSimplePool #-}
initSimplePool f = initPool (const f)
-- | Open resource with 'SimplePool', see 'withPool'
--
withSimplePool :: (MonadCatch.MonadMask m, MonadIO m, HasCallStack)
=> SimplePool res -> (res -> m a) -> m a
{-# INLINABLE withSimplePool #-}
withSimplePool pool = withPool pool ()