{-|
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.PrimIORef
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 { Resource a -> IO (a, IO ())
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
{-# INLINE initResource #-}
initResource :: IO a -> (a -> IO ()) -> Resource a
initResource IO a
create a -> IO ()
release = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource (IO (a, IO ()) -> Resource a) -> IO (a, IO ()) -> Resource a
forall a b. (a -> b) -> a -> b
$ do
    a
r <- IO a
create
    (a, IO ()) -> IO (a, IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return ((a, IO ()) -> IO (a, IO ())) -> (a, IO ()) -> IO (a, IO ())
forall a b. (a -> b) -> a -> b
$ (a
r, a -> IO ()
release a
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
{-# INLINE initResource_ #-}
initResource_ :: IO a -> IO () -> Resource a
initResource_ IO a
create IO ()
release = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource (IO (a, IO ()) -> Resource a) -> IO (a, IO ()) -> Resource a
forall a b. (a -> b) -> a -> b
$ do
    a
r <- IO a
create
    (a, IO ()) -> IO (a, IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return ((a, IO ()) -> IO (a, IO ())) -> (a, IO ()) -> IO (a, IO ())
forall a b. (a -> b) -> a -> b
$ (a
r, IO ()
release)

instance Functor Resource where
    {-# INLINE fmap #-}
    fmap :: (a -> b) -> Resource a -> Resource b
fmap a -> b
f Resource a
resource = IO (b, IO ()) -> Resource b
forall a. IO (a, IO ()) -> Resource a
Resource (IO (b, IO ()) -> Resource b) -> IO (b, IO ()) -> Resource b
forall a b. (a -> b) -> a -> b
$ do
        (a
a, IO ()
release) <- Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource
        (b, IO ()) -> IO (b, IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b
f a
a, IO ()
release)

instance Applicative Resource where
    {-# INLINE pure #-}
    pure :: a -> Resource a
pure a
a = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource ((a, IO ()) -> IO (a, IO ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()))
    {-# INLINE (<*>) #-}
    Resource (a -> b)
resource1 <*> :: Resource (a -> b) -> Resource a -> Resource b
<*> Resource a
resource2 = IO (b, IO ()) -> Resource b
forall a. IO (a, IO ()) -> Resource a
Resource (IO (b, IO ()) -> Resource b) -> IO (b, IO ()) -> Resource b
forall a b. (a -> b) -> a -> b
$ do
        (a -> b
f, IO ()
release1) <- Resource (a -> b) -> IO (a -> b, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource (a -> b)
resource1
        (a
x, IO ()
release2) <- Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource2 IO (a, IO ()) -> IO () -> IO (a, IO ())
forall a b. IO a -> IO b -> IO a
`onException` IO ()
release1
        (b, IO ()) -> IO (b, IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b
f a
x, IO ()
release2 IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
release1)

instance Monad Resource where
    {-# INLINE return #-}
    return :: a -> Resource a
return = a -> Resource a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    {-# INLINE (>>=) #-}
    Resource a
m >>= :: Resource a -> (a -> Resource b) -> Resource b
>>= a -> Resource b
f = IO (b, IO ()) -> Resource b
forall a. IO (a, IO ()) -> Resource a
Resource (IO (b, IO ()) -> Resource b) -> IO (b, IO ()) -> Resource b
forall a b. (a -> b) -> a -> b
$ do
        (a
m', IO ()
release1) <- Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
m
        (b
x , IO ()
release2) <- Resource b -> IO (b, IO ())
forall a. Resource a -> IO (a, IO ())
acquire (a -> Resource b
f a
m') IO (b, IO ()) -> IO () -> IO (b, IO ())
forall a b. IO a -> IO b -> IO a
`onException` IO ()
release1
        (b, IO ()) -> IO (b, IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return (b
x, IO ()
release2 IO () -> IO () -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
release1)

instance MonadIO Resource where
    {-# INLINE liftIO #-}
    liftIO :: IO a -> Resource a
liftIO IO a
f = IO (a, IO ()) -> Resource a
forall a. IO (a, IO ()) -> Resource a
Resource (IO (a, IO ()) -> Resource a) -> IO (a, IO ()) -> Resource a
forall a b. (a -> b) -> a -> b
$ (a -> (a, IO ())) -> IO a -> IO (a, IO ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ a
a -> (a
a, IO ()
dummyRelease)) IO a
f
        where dummyRelease :: IO ()
dummyRelease = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
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 a -> (a -> m b) -> m b
withResource Resource a
resource a -> m b
k = m (a, IO ()) -> ((a, IO ()) -> m ()) -> ((a, IO ()) -> m b) -> m b
forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
MonadCatch.bracket
    (IO (a, IO ()) -> m (a, IO ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource))
    (\(a
_, IO ()
release) -> IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release)
    (\(a
a, IO ()
_) -> a -> m b
k a
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 a -> (a -> m () -> m b) -> m b
withResource' Resource a
resource a -> m () -> m b
k = do
    Counter
c <- IO Counter -> m Counter
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Int -> IO Counter
newCounter Int
0)
    m (a, IO ()) -> ((a, IO ()) -> m ()) -> ((a, IO ()) -> m b) -> m b
forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
MonadCatch.bracket
        (IO (a, IO ()) -> m (a, IO ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (a, IO ()) -> m (a, IO ())) -> IO (a, IO ()) -> m (a, IO ())
forall a b. (a -> b) -> a -> b
$ do
            (a
a, IO ()
release) <- (Resource a -> IO (a, IO ())
forall a. Resource a -> IO (a, IO ())
acquire Resource a
resource)
            let release' :: IO ()
release' = do
                    Int
c' <- Counter -> Int -> IO Int
atomicOrCounter Counter
c Int
1
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
c' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) IO ()
release
            (a, IO ()) -> IO (a, IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, IO ()
release'))
        (\(a
_, IO ()
release) -> IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release)
        (\(a
a, IO ()
release) -> a -> m () -> m b
k a
a (IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
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
    { Pool key res -> key -> Resource res
_poolResource     :: key -> Resource res      -- ^ how to get a resource
    , Pool key res -> Int
_poolLimitPerKey  :: {-# UNPACK #-} !Int      -- ^ max number for resource we keep alive after used
    , Pool key res -> Int
_poolIdleTime     :: {-# UNPACK #-} !Int      -- ^ max idle time for resource we keep alive
    , Pool key res -> UnliftedArray (IORef (Maybe (Map key (Entry res))))
_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))
statPool :: Pool key res -> IO (SmallArray (Map key Int))
statPool (Pool key -> Resource res
_ Int
_ Int
_ UnliftedArray (IORef (Maybe (Map key (Entry res))))
arr) = ((IORef (Maybe (Map key (Entry res))) -> IO (Map key Int))
-> UnliftedArray (IORef (Maybe (Map key (Entry res))))
-> IO (SmallArray (Map key Int))
forall (v :: * -> *) a (u :: * -> *) b (f :: * -> *).
(Vec v a, Vec u b, Applicative f) =>
(a -> f b) -> v a -> f (u b)
`V.traverseVec` UnliftedArray (IORef (Maybe (Map key (Entry res))))
arr) ((IORef (Maybe (Map key (Entry res))) -> IO (Map key Int))
 -> IO (SmallArray (Map key Int)))
-> (IORef (Maybe (Map key (Entry res))) -> IO (Map key Int))
-> IO (SmallArray (Map key Int))
forall a b. (a -> b) -> a -> b
$ \ IORef (Maybe (Map key (Entry res)))
resMapRef -> do
    Maybe (Map key (Entry res))
mResMap <- IORef (Maybe (Map key (Entry res)))
-> IO (Maybe (Map key (Entry res)))
forall a. IORef a -> IO a
readIORef IORef (Maybe (Map key (Entry res)))
resMapRef
    case Maybe (Map key (Entry res))
mResMap of
        Just Map key (Entry res)
resMap -> Map key Int -> IO (Map key Int)
forall (m :: * -> *) a. Monad m => a -> m a
return (Map key Int -> IO (Map key Int))
-> Map key Int -> IO (Map key Int)
forall a b. (a -> b) -> a -> b
$ ((Entry res -> Int) -> Map key (Entry res) -> Map key Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Map key (Entry res)
resMap) ( \ Entry res
es ->
                case Entry res
es of EntryCons (res, IO ())
_ Int
siz Int
_ Entry res
_ -> Int
siz
                           Entry res
_                   -> Int
0)
        Maybe (Map key (Entry res))
_ -> IO (Map key Int)
forall a. HasCallStack => IO a
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)
initPool :: (key -> Resource res) -> Int -> Int -> Resource (Pool key res)
initPool key -> Resource res
resf Int
limit Int
itime = IO (Pool key res)
-> (Pool key res -> IO ()) -> Resource (Pool key res)
forall a. IO a -> (a -> IO ()) -> Resource a
initResource IO (Pool key res)
createPool Pool key res -> IO ()
forall key res. Pool key res -> IO ()
closePool
  where
    createPool :: IO (Pool key res)
createPool = do
        Int
numCaps <- IO Int
getNumCapabilities
        MutableUnliftedArray
  RealWorld (IORef (Maybe (Map key (Entry res))))
marr <- Int
-> IO
     (MArr
        UnliftedArray RealWorld (IORef (Maybe (Map key (Entry res)))))
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
Int -> m (MArr arr s a)
newArr Int
numCaps
        [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
numCapsInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Int
i -> do
            MArr UnliftedArray RealWorld (IORef (Maybe (Map key (Entry res))))
-> Int -> IORef (Maybe (Map key (Entry res))) -> IO ()
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> Int -> a -> m ()
writeArr MArr UnliftedArray RealWorld (IORef (Maybe (Map key (Entry res))))
MutableUnliftedArray
  RealWorld (IORef (Maybe (Map key (Entry res))))
marr Int
i (IORef (Maybe (Map key (Entry res))) -> IO ())
-> IO (IORef (Maybe (Map key (Entry res)))) -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe (Map key (Entry res))
-> IO (IORef (Maybe (Map key (Entry res))))
forall a. a -> IO (IORef a)
newIORef (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just Map key (Entry res)
forall k a. Map k a
M.empty)
        UnliftedArray (IORef (Maybe (Map key (Entry res))))
arr <- MArr UnliftedArray RealWorld (IORef (Maybe (Map key (Entry res))))
-> IO (UnliftedArray (IORef (Maybe (Map key (Entry res)))))
forall (arr :: * -> *) a (m :: * -> *) s.
(Arr arr a, PrimMonad m, PrimState m ~ s) =>
MArr arr s a -> m (arr a)
unsafeFreezeArr MArr UnliftedArray RealWorld (IORef (Maybe (Map key (Entry res))))
MutableUnliftedArray
  RealWorld (IORef (Maybe (Map key (Entry res))))
marr
        Pool key res -> IO (Pool key res)
forall (m :: * -> *) a. Monad m => a -> m a
return ((key -> Resource res)
-> Int
-> Int
-> UnliftedArray (IORef (Maybe (Map key (Entry res))))
-> Pool key res
forall key res.
(key -> Resource res)
-> Int
-> Int
-> UnliftedArray (IORef (Maybe (Map key (Entry res))))
-> Pool key res
Pool key -> Resource res
resf Int
limit Int
itime UnliftedArray (IORef (Maybe (Map key (Entry res))))
arr)

    closePool :: Pool key res -> IO ()
closePool (Pool key -> Resource res
_ Int
_ Int
_ UnliftedArray (IORef (Maybe (Map key (Entry res))))
localPoolArr) = do
        -- close all existed resource
        ((IORef (Maybe (Map key (Entry res))) -> IO (IO ()))
-> UnliftedArray (IORef (Maybe (Map key (Entry res)))) -> IO ()
forall (v :: * -> *) a (f :: * -> *) b.
(Vec v a, Applicative f) =>
(a -> f b) -> v a -> f ()
`V.traverseVec_` UnliftedArray (IORef (Maybe (Map key (Entry res))))
localPoolArr) ((IORef (Maybe (Map key (Entry res))) -> IO (IO ())) -> IO ())
-> (IORef (Maybe (Map key (Entry res))) -> IO (IO ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ IORef (Maybe (Map key (Entry res)))
resMapRef ->
            IORef (Maybe (Map key (Entry res)))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef (Maybe (Map key (Entry res)))
resMapRef ((Maybe (Map key (Entry res))
  -> (Maybe (Map key (Entry res)), IO ()))
 -> IO (IO ()))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO (IO ())
forall a b. (a -> b) -> a -> b
$ \ Maybe (Map key (Entry res))
mResMap ->
                case Maybe (Map key (Entry res))
mResMap of
                    Just Map key (Entry res)
resMap -> (Maybe (Map key (Entry res))
forall a. Maybe a
Nothing, (Entry res -> IO ()) -> Map key (Entry res) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Entry res -> IO ()
forall res. Entry res -> IO ()
closeEntry Map key (Entry res)
resMap)
                    Maybe (Map key (Entry res))
_ -> (Maybe (Map key (Entry res))
forall a. Maybe a
Nothing, () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())

    closeEntry :: Entry res -> IO ()
closeEntry (EntryCons (res
_, IO ()
close) Int
_ Int
_ Entry res
_) = IO () -> IO ()
forall a. IO a -> IO ()
ignoreSync IO ()
close
    closeEntry Entry res
EntryNil = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
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
withPool :: Pool key res -> key -> (res -> m a) -> m a
withPool (Pool key -> Resource res
resf Int
limitPerKey Int
itime UnliftedArray (IORef (Maybe (Map key (Entry res))))
arr) key
key res -> m a
f = do
    !IORef (Maybe (Map key (Entry res)))
resMapRef <- UnliftedArray (IORef (Maybe (Map key (Entry res))))
-> Int -> IORef (Maybe (Map key (Entry res)))
forall (arr :: * -> *) a. Arr arr a => arr a -> Int -> a
indexArr UnliftedArray (IORef (Maybe (Map key (Entry res))))
arr (Int -> IORef (Maybe (Map key (Entry res))))
-> ((Int, Bool) -> Int)
-> (Int, Bool)
-> IORef (Maybe (Map key (Entry res)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Bool) -> Int
forall a b. (a, b) -> a
fst ((Int, Bool) -> IORef (Maybe (Map key (Entry res))))
-> m (Int, Bool) -> m (IORef (Maybe (Map key (Entry res))))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Int, Bool) -> m (Int, Bool)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ThreadId -> IO (Int, Bool)
threadCapability (ThreadId -> IO (Int, Bool)) -> IO ThreadId -> IO (Int, Bool)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO ThreadId
myThreadId)
    (a, ()) -> a
forall a b. (a, b) -> a
fst ((a, ()) -> a) -> m (a, ()) -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (res, IO ())
-> ((res, IO ()) -> ExitCase a -> m ())
-> ((res, IO ()) -> m a)
-> m (a, ())
forall (m :: * -> *) a b c.
MonadMask m =>
m a -> (a -> ExitCase b -> m c) -> (a -> m b) -> m (b, c)
MonadCatch.generalBracket
        (IO (res, IO ()) -> m (res, IO ())
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (res, IO ()) -> m (res, IO ()))
-> IO (res, IO ()) -> m (res, IO ())
forall a b. (a -> b) -> a -> b
$ IORef (Maybe (Map key (Entry res))) -> IO (res, IO ())
takeFromPool IORef (Maybe (Map key (Entry res)))
resMapRef)
        (\ r :: (res, IO ())
r@(res
_, IO ()
close) ExitCase a
exit ->
            case ExitCase a
exit of
                MonadCatch.ExitCaseSuccess a
_ -> IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef (Maybe (Map key (Entry res))) -> (res, IO ()) -> IO ()
returnToPool IORef (Maybe (Map key (Entry res)))
resMapRef (res, IO ())
r)
                ExitCase a
_ -> IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
close)
        (\ (res
a, IO ()
_) -> res -> m a
f res
a)
  where
    takeFromPool :: IORef (Maybe (Map key (Entry res))) -> IO (res, IO ())
takeFromPool IORef (Maybe (Map key (Entry res)))
resMapRef =
        IO (IO (res, IO ())) -> IO (res, IO ())
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO (res, IO ())) -> IO (res, IO ()))
-> ((Maybe (Map key (Entry res))
     -> (Maybe (Map key (Entry res)), IO (res, IO ())))
    -> IO (IO (res, IO ())))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO (res, IO ())))
-> IO (res, IO ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef (Maybe (Map key (Entry res)))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO (res, IO ())))
-> IO (IO (res, IO ()))
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Maybe (Map key (Entry res)))
resMapRef ((Maybe (Map key (Entry res))
  -> (Maybe (Map key (Entry res)), IO (res, IO ())))
 -> IO (res, IO ()))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO (res, IO ())))
-> IO (res, IO ())
forall a b. (a -> b) -> a -> b
$ \ Maybe (Map key (Entry res))
mResMap ->
            case Maybe (Map key (Entry res))
mResMap of
                Just Map key (Entry res)
resMap ->
                    case key -> Map key (Entry res) -> Maybe (Entry res)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup key
key Map key (Entry res)
resMap of
                        Just (EntryCons (res, IO ())
a Int
_ Int
_ Entry res
es') ->
                            (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just (Map key (Entry res) -> Maybe (Map key (Entry res)))
-> Map key (Entry res) -> Maybe (Map key (Entry res))
forall a b. (a -> b) -> a -> b
$! (Entry res -> Entry res)
-> key -> Map key (Entry res) -> Map key (Entry res)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (Entry res -> Entry res -> Entry res
forall a b. a -> b -> a
const Entry res
es') key
key Map key (Entry res)
resMap, (res, IO ()) -> IO (res, IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return (res, IO ())
a)
                        Maybe (Entry res)
_ ->  (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just Map key (Entry res)
resMap, Resource res -> IO (res, IO ())
forall a. Resource a -> IO (a, IO ())
acquire (key -> Resource res
resf key
key))
                Maybe (Map key (Entry res))
_ -> (Maybe (Map key (Entry res))
forall a. Maybe a
Nothing, IO (res, IO ())
forall a. HasCallStack => IO a
throwECLOSED)

    returnToPool :: IORef (Maybe (Map key (Entry res))) -> (res, IO ()) -> IO ()
returnToPool IORef (Maybe (Map key (Entry res)))
resMapRef (res, IO ())
r = do
        IO (IO ()) -> IO ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO ()) -> IO ())
-> ((Maybe (Map key (Entry res))
     -> (Maybe (Map key (Entry res)), IO ()))
    -> IO (IO ()))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef (Maybe (Map key (Entry res)))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Maybe (Map key (Entry res)))
resMapRef ((Maybe (Map key (Entry res))
  -> (Maybe (Map key (Entry res)), IO ()))
 -> IO ())
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO ()
forall a b. (a -> b) -> a -> b
$ \ Maybe (Map key (Entry res))
mResMap ->
            case Maybe (Map key (Entry res))
mResMap of
                Just Map key (Entry res)
resMap ->
                    case key -> Map key (Entry res) -> Maybe (Entry res)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup key
key Map key (Entry res)
resMap of
                        Just (EntryCons (res, IO ())
_ Int
siz Int
_ Entry res
_) ->
                            if Int
siz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
limitPerKey
                            -- if entries under given key do not exceed limit, we prepend res back to entries
                            then (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just (Map key (Entry res) -> Maybe (Map key (Entry res)))
-> Map key (Entry res) -> Maybe (Map key (Entry res))
forall a b. (a -> b) -> a -> b
$! (Entry res -> Entry res)
-> key -> Map key (Entry res) -> Map key (Entry res)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust ((res, IO ()) -> Int -> Int -> Entry res -> Entry res
forall res. (res, IO ()) -> Int -> Int -> Entry res -> Entry res
EntryCons (res, IO ())
r (Int
sizInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
itime) key
key Map key (Entry res)
resMap, () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
                            -- otherwise we close it
                            else (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just Map key (Entry res)
resMap, (res, IO ()) -> IO ()
forall a b. (a, b) -> b
snd (res, IO ())
r)
                        Maybe (Entry res)
_ -> (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just (Map key (Entry res) -> Maybe (Map key (Entry res)))
-> Map key (Entry res) -> Maybe (Map key (Entry res))
forall a b. (a -> b) -> a -> b
$! key -> Entry res -> Map key (Entry res) -> Map key (Entry res)
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert key
key ((res, IO ()) -> Int -> Int -> Entry res -> Entry res
forall res. (res, IO ()) -> Int -> Int -> Entry res -> Entry res
EntryCons (res, IO ())
r Int
1 Int
itime Entry res
forall res. Entry res
EntryNil) Map key (Entry res)
resMap,
                                IORef (Maybe (Map key (Entry res))) -> IO ()
scanLocalPool IORef (Maybe (Map key (Entry res)))
resMapRef)
                Maybe (Map key (Entry res))
_ -> (Maybe (Map key (Entry res))
forall a. Maybe a
Nothing, (res, IO ()) -> IO ()
forall a b. (a, b) -> b
snd (res, IO ())
r)

    scanLocalPool :: IORef (Maybe (Map key (Entry res))) -> IO ()
scanLocalPool IORef (Maybe (Map key (Entry res)))
resMapRef = do
        Int -> IO () -> IO ()
registerLowResTimer_ Int
10 (IO () -> IO ())
-> ((Maybe (Map key (Entry res))
     -> (Maybe (Map key (Entry res)), IO ()))
    -> IO ())
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (IO ()) -> IO ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO ()) -> IO ())
-> ((Maybe (Map key (Entry res))
     -> (Maybe (Map key (Entry res)), IO ()))
    -> IO (IO ()))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef (Maybe (Map key (Entry res)))
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO (IO ())
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Maybe (Map key (Entry res)))
resMapRef ((Maybe (Map key (Entry res))
  -> (Maybe (Map key (Entry res)), IO ()))
 -> IO ())
-> (Maybe (Map key (Entry res))
    -> (Maybe (Map key (Entry res)), IO ()))
-> IO ()
forall a b. (a -> b) -> a -> b
$ \ Maybe (Map key (Entry res))
mResMap ->
            case Maybe (Map key (Entry res))
mResMap of
                Just Map key (Entry res)
resMap ->
                    case key -> Map key (Entry res) -> Maybe (Entry res)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup key
key Map key (Entry res)
resMap of
                        -- this is where we clean up empty keys
                        Just Entry res
EntryNil -> (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just (Map key (Entry res) -> Maybe (Map key (Entry res)))
-> Map key (Entry res) -> Maybe (Map key (Entry res))
forall a b. (a -> b) -> a -> b
$! key -> Map key (Entry res) -> Map key (Entry res)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete key
key Map key (Entry res)
resMap, () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        Just Entry res
es -> do
                            let ([(res, IO ())]
dead, Entry res
living) = Entry res
-> Int
-> [(res, IO ())]
-> Entry res
-> ([(res, IO ())], Entry res)
forall res.
Entry res
-> Int
-> [(res, IO ())]
-> Entry res
-> ([(res, IO ())], Entry res)
age Entry res
es Int
0 [] Entry res
forall res. Entry res
EntryNil
                            case Entry res
living of
                                -- no living resources any more, stop scanning
                                Entry res
EntryNil -> (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just (Map key (Entry res) -> Maybe (Map key (Entry res)))
-> Map key (Entry res) -> Maybe (Map key (Entry res))
forall a b. (a -> b) -> a -> b
$! key -> Map key (Entry res) -> Map key (Entry res)
forall k a. Ord k => k -> Map k a -> Map k a
M.delete key
key Map key (Entry res)
resMap,
                                    [(res, IO ())] -> ((res, IO ()) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(res, IO ())]
dead (IO () -> IO ()
forall a. IO a -> IO ()
ignoreSync (IO () -> IO ())
-> ((res, IO ()) -> IO ()) -> (res, IO ()) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (res, IO ()) -> IO ()
forall a b. (a, b) -> b
snd))
                                Entry res
_ ->  (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just (Map key (Entry res) -> Maybe (Map key (Entry res)))
-> Map key (Entry res) -> Maybe (Map key (Entry res))
forall a b. (a -> b) -> a -> b
$! (Entry res -> Entry res)
-> key -> Map key (Entry res) -> Map key (Entry res)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (Entry res -> Entry res -> Entry res
forall a b. a -> b -> a
const Entry res
living) key
key Map key (Entry res)
resMap,
                                    (do [(res, IO ())] -> ((res, IO ()) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(res, IO ())]
dead (IO () -> IO ()
forall a. IO a -> IO ()
ignoreSync (IO () -> IO ())
-> ((res, IO ()) -> IO ()) -> (res, IO ()) -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (res, IO ()) -> IO ()
forall a b. (a, b) -> b
snd)
                                        IORef (Maybe (Map key (Entry res))) -> IO ()
scanLocalPool IORef (Maybe (Map key (Entry res)))
resMapRef))
                        -- no living resources under given key, stop scanning
                        Maybe (Entry res)
_ -> (Map key (Entry res) -> Maybe (Map key (Entry res))
forall a. a -> Maybe a
Just Map key (Entry res)
resMap, () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
                Maybe (Map key (Entry res))
_ -> (Maybe (Map key (Entry res))
forall a. Maybe a
Nothing, () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())

    age :: Entry res
-> Int
-> [(res, IO ())]
-> Entry res
-> ([(res, IO ())], Entry res)
age (EntryCons (res, IO ())
a Int
_ Int
life Entry res
es) !Int
livingNum [(res, IO ())]
dead Entry res
living
        | Int
life Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1  = let !livingNum' :: Int
livingNum' = (Int
livingNumInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                      in Entry res
-> Int
-> [(res, IO ())]
-> Entry res
-> ([(res, IO ())], Entry res)
age Entry res
es Int
livingNum' [(res, IO ())]
dead ((res, IO ()) -> Int -> Int -> Entry res -> Entry res
forall res. (res, IO ()) -> Int -> Int -> Entry res -> Entry res
EntryCons (res, IO ())
a Int
livingNum' (Int
lifeInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Entry res
living)
        | Bool
otherwise = Entry res
-> Int
-> [(res, IO ())]
-> Entry res
-> ([(res, IO ())], Entry res)
age Entry res
es Int
livingNum ((res, IO ())
a(res, IO ()) -> [(res, IO ())] -> [(res, IO ())]
forall a. a -> [a] -> [a]
:[(res, IO ())]
dead) Entry res
living
    age Entry res
_ Int
_ [(res, IO ())]
dead Entry res
living = ([(res, IO ())]
dead, Entry res
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)
initSimplePool :: Resource res -> Int -> Int -> Resource (SimplePool res)
initSimplePool Resource res
f = (() -> Resource res) -> Int -> Int -> Resource (SimplePool res)
forall key res.
(key -> Resource res) -> Int -> Int -> Resource (Pool key res)
initPool (Resource res -> () -> Resource res
forall a b. a -> b -> a
const Resource res
f)

-- | Open resource with 'SimplePool', see 'withPool'
--
withSimplePool :: (MonadCatch.MonadMask m, MonadIO m, HasCallStack)
               => SimplePool res -> (res -> m a) -> m a
withSimplePool :: SimplePool res -> (res -> m a) -> m a
withSimplePool SimplePool res
pool = SimplePool res -> () -> (res -> m a) -> m a
forall (m :: * -> *) key res a.
(MonadMask m, MonadIO m, Ord key, HasCallStack) =>
Pool key res -> key -> (res -> m a) -> m a
withPool SimplePool res
pool ()