module ZooKeeper.Internal.Utils
  ( -- * Resource management
    Resource (..)
  , initResource
  , initResource_
  , withResource
  , withResource'
  ) where

import           Control.Exception      (onException)
import           Control.Monad          (when)
import qualified Control.Monad.Catch    as MonadCatch
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           GHC.Stack              (HasCallStack)
import           Z.Data.PrimRef         (atomicOrCounter, newCounter)

--------------------------------------------------------------------------------

-- | 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.
newtype Resource a = Resource { forall a. 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
{-# INLINABLE initResource #-}
initResource :: forall a. 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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (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
{-# INLINABLE initResource_ #-}
initResource_ :: forall a. 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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
r, IO ()
release)

instance Functor Resource where
    {-# INLINE fmap #-}
    fmap :: forall a b. (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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b
f a
a, IO ()
release)

instance Applicative Resource where
    {-# INLINE pure #-}
    pure :: forall a. 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 a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
a, () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()))
    {-# INLINE (<*>) #-}
    Resource (a -> b)
resource1 <*> :: forall a b. 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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b
f a
x, IO ()
release2 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
release1)

instance Monad Resource where
    {-# INLINE return #-}
    return :: forall a. a -> Resource a
return = a -> Resource a
forall a. a -> Resource a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
    {-# INLINE (>>=) #-}
    Resource a
m >>= :: forall a b. 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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (b
x, IO ()
release2 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
release1)

instance MonadIO Resource where
    {-# INLINE liftIO #-}
    liftIO :: forall a. 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 a b. (a -> b) -> IO a -> IO b
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 a. a -> IO a
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 -> (HasCallStack => a -> m b) -> m b
{-# INLINABLE withResource #-}
withResource :: forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (HasCallStack => a -> m b) -> m b
withResource Resource a
resource HasCallStack => 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 a. IO a -> m a
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 a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release)
    (\(a
a, IO ()
_) -> a -> m b
HasCallStack => 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 -> (HasCallStack => a -> m () -> m b) -> m b
{-# INLINABLE withResource' #-}
withResource' :: forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (HasCallStack => a -> m () -> m b) -> m b
withResource' Resource a
resource HasCallStack => a -> m () -> m b
k = do
    Counter
c <- IO Counter -> m Counter
forall a. IO a -> m a
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 a. IO a -> m a
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 a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, IO ()
release'))
        (\(a
_, IO ()
release) -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release)
        (\(a
a, IO ()
release) -> a -> m () -> m b
HasCallStack => a -> m () -> m b
k a
a (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ()
release))