-- |
-- Module      : Streamly.Internal.Data.IOFinalizer
-- Copyright   : (c) 2020 Composewell Technologies and Contributors
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- A value associated with an IO action that is automatically called whenever
-- the value is garbage collected.

module Streamly.Internal.Data.IOFinalizer
    (
      IOFinalizer
    , newIOFinalizer
    , runIOFinalizer
    , clearingIOFinalizer
    )
where

import Control.Exception (mask_)
import Control.Monad (void)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.Control (MonadBaseControl, control)
import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
import Streamly.Internal.Control.Concurrent (captureMonadState, runInIO)

-- | An 'IOFinalizer' has an associated IO action that is automatically called
-- whenever the finalizer is garbage collected. The action can be run and
-- cleared prematurely.
--
-- You can hold a reference to the finalizer in your data structure, if the
-- data structure gets garbage collected the finalizer will be called.
--
-- It is implemented using 'mkWeakIORef'.
--
-- /Pre-release/
newtype IOFinalizer = IOFinalizer (IORef (Maybe (IO ())))

-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
mkIOFinalizer :: m b -> m (IO ())
mkIOFinalizer m b
f = do
    RunInIO m
mrun <- m (RunInIO m)
forall (m :: * -> *). MonadBaseControl IO m => m (RunInIO m)
captureMonadState
    IO () -> m (IO ())
forall (m :: * -> *) a. Monad m => a -> m a
return (IO () -> m (IO ())) -> IO () -> m (IO ())
forall a b. (a -> b) -> a -> b
$
        IO () -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            StM m b
_ <- RunInIO m -> m b -> IO (StM m b)
forall (m :: * -> *). RunInIO m -> forall b. m b -> IO (StM m b)
runInIO RunInIO m
mrun m b
f
            () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | GC hook to run an IO action stored in a finalized IORef.
runFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
runFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
runFinalizerGC IORef (Maybe (IO ()))
ref = do
    Maybe (IO ())
res <- IORef (Maybe (IO ())) -> IO (Maybe (IO ()))
forall a. IORef a -> IO a
readIORef IORef (Maybe (IO ()))
ref
    case Maybe (IO ())
res of
        Maybe (IO ())
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Just IO ()
f -> IO ()
f

-- | Create a finalizer that calls the supplied function automatically when the
-- it is garbage collected.
--
-- /The finalizer is always run using the state of the monad that is captured
-- at the time of calling 'newFinalizer'./
--
-- Note: To run it on garbage collection we have no option but to use the monad
-- state captured at some earlier point of time.  For the case when the
-- finalizer is run manually before GC we could run it with the current state
-- of the monad but we want to keep both the cases consistent.
--
-- /Pre-release/
newIOFinalizer :: (MonadIO m, MonadBaseControl IO m) => m a -> m IOFinalizer
newIOFinalizer :: m a -> m IOFinalizer
newIOFinalizer m a
finalizer = do
    IO ()
f <- m a -> m (IO ())
forall (m :: * -> *) b. MonadBaseControl IO m => m b -> m (IO ())
mkIOFinalizer m a
finalizer
    IORef (Maybe (IO ()))
ref <- IO (IORef (Maybe (IO ()))) -> m (IORef (Maybe (IO ())))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (IORef (Maybe (IO ()))) -> m (IORef (Maybe (IO ()))))
-> IO (IORef (Maybe (IO ()))) -> m (IORef (Maybe (IO ())))
forall a b. (a -> b) -> a -> b
$ Maybe (IO ()) -> IO (IORef (Maybe (IO ())))
forall a. a -> IO (IORef a)
newIORef (Maybe (IO ()) -> IO (IORef (Maybe (IO ()))))
-> Maybe (IO ()) -> IO (IORef (Maybe (IO ())))
forall a b. (a -> b) -> a -> b
$ IO () -> Maybe (IO ())
forall a. a -> Maybe a
Just IO ()
f
    Weak (IORef (Maybe (IO ())))
_ <- IO (Weak (IORef (Maybe (IO ()))))
-> m (Weak (IORef (Maybe (IO ()))))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Weak (IORef (Maybe (IO ()))))
 -> m (Weak (IORef (Maybe (IO ())))))
-> IO (Weak (IORef (Maybe (IO ()))))
-> m (Weak (IORef (Maybe (IO ()))))
forall a b. (a -> b) -> a -> b
$ IORef (Maybe (IO ())) -> IO () -> IO (Weak (IORef (Maybe (IO ()))))
forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef (Maybe (IO ()))
ref (IORef (Maybe (IO ())) -> IO ()
runFinalizerGC IORef (Maybe (IO ()))
ref)
    IOFinalizer -> m IOFinalizer
forall (m :: * -> *) a. Monad m => a -> m a
return (IOFinalizer -> m IOFinalizer) -> IOFinalizer -> m IOFinalizer
forall a b. (a -> b) -> a -> b
$ IORef (Maybe (IO ())) -> IOFinalizer
IOFinalizer IORef (Maybe (IO ()))
ref

-- | Run the action associated with the finalizer and deactivate it so that it
-- never runs again.  Note, the finalizing action runs with async exceptions
-- masked.
--
-- /Pre-release/
runIOFinalizer :: MonadIO m => IOFinalizer -> m ()
runIOFinalizer :: IOFinalizer -> m ()
runIOFinalizer (IOFinalizer IORef (Maybe (IO ()))
ref) = IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    Maybe (IO ())
res <- IORef (Maybe (IO ())) -> IO (Maybe (IO ()))
forall a. IORef a -> IO a
readIORef IORef (Maybe (IO ()))
ref
    case Maybe (IO ())
res of
        Maybe (IO ())
Nothing -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Just IO ()
action -> do
            -- if an async exception comes after writing 'Nothing' then the
            -- finalizing action will never be run. We need to do this
            -- atomically wrt async exceptions.
            IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                IORef (Maybe (IO ())) -> Maybe (IO ()) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IO ()))
ref Maybe (IO ())
forall a. Maybe a
Nothing
                IO ()
action

-- | Run an action clearing the finalizer atomically wrt async exceptions. The
-- action is run with async exceptions masked.
--
-- /Pre-release/
clearingIOFinalizer :: MonadBaseControl IO m => IOFinalizer -> m a -> m a
clearingIOFinalizer :: IOFinalizer -> m a -> m a
clearingIOFinalizer (IOFinalizer IORef (Maybe (IO ()))
ref) m a
action = do
    (RunInBase m IO -> IO (StM m a)) -> m a
forall (b :: * -> *) (m :: * -> *) a.
MonadBaseControl b m =>
(RunInBase m b -> b (StM m a)) -> m a
control ((RunInBase m IO -> IO (StM m a)) -> m a)
-> (RunInBase m IO -> IO (StM m a)) -> m a
forall a b. (a -> b) -> a -> b
$ \RunInBase m IO
runinio ->
        IO (StM m a) -> IO (StM m a)
forall a. IO a -> IO a
mask_ (IO (StM m a) -> IO (StM m a)) -> IO (StM m a) -> IO (StM m a)
forall a b. (a -> b) -> a -> b
$ do
            IORef (Maybe (IO ())) -> Maybe (IO ()) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IO ()))
ref Maybe (IO ())
forall a. Maybe a
Nothing
            m a -> IO (StM m a)
RunInBase m IO
runinio m a
action