Exception.hs

{-# LANGUAGE CPP #-}
-- |
-- Module      : Streamly.Internal.Control.Exception
-- Copyright   : (c) 2019 Composewell Technologies
--
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Additional "Control.Exception" utilities.

module Streamly.Internal.Control.Exception
    (
    -- * Verify
      verify
    , verifyM

    -- * Resource Management
    -- | Exception safe, thread safe resource managment operations, similar to
    -- but more powerful than the @bracket@ and @finally@ operations available
    -- in the base package.
    --
    -- These operations support allocation and free only in the IO monad,
    -- hence the IO suffix.
    --
    , AcquireIO(..)
    , Priority(..)
    , allocator
    , releaser
    , withAcquireIO
    , acquireWith
    , acquire
    , acquire_
    , registerWith
    , register
    , hook
    )
where

-- import Control.Concurrent (myThreadId)
import Control.Monad (void)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Exception (mask_)
import Control.Monad.Catch (MonadMask)
import Data.IntMap.Strict (IntMap)
import Data.IORef (IORef, newIORef, atomicModifyIORef')

import qualified Control.Monad.Catch as MC
import qualified Data.IntMap.Strict as Map

#include "DocTestControlException.hs"

-------------------------------------------------------------------------------
-- Asserts
-------------------------------------------------------------------------------

-- | Like 'assert' but is not removed by the compiler, it is always present in
-- production code.
--
-- /Pre-release/
--
{-# INLINE verify #-}
verify :: Bool -> a -> a
verify predicate val =
    if predicate
    -- XXX it would be nice if we can print the predicate expr.
    then error "verify failed"
    else val

-- Like 'verify' but returns @()@ in an 'Applicative' context so that it can be
-- used as an independent statement in a @do@ block.
--
-- /Pre-release/
--
{-# INLINE verifyM #-}
verifyM :: Applicative f => Bool -> f ()
verifyM predicate = verify predicate (pure ())

-------------------------------------------------------------------------------
-- Resource management
-------------------------------------------------------------------------------

-- XXX In a manual release mechanism of resources we always have the risk of
-- using the resource by some persisting thread even after it has been freed.
-- Ideally, we should use the GC to clean up resources because that way we do
-- not need to worry about references, we can pass around resources to other
-- threads and we get an automatic reference counting. Is it possible to use
-- compact regions to confine resource to smaller areas so that we can perform
-- a limited GC to free them? We can then just put gc sync barriers at points
-- where we want to ensure that resources are freed.

-- | Resources with 'Priority1' are freed before 'Priority2'. Priority is
-- especially introduced to take care of the case where we need to free
-- concurrency channels, so that all the workers of the channel are cleaned up
-- before we free the resources allocated by the workers of the channel.
-- Otherwise we might free the resources and workers may be trying to use them
-- and start misbehaving.
--
data Priority = Priority1 | Priority2 deriving Show

-- To keep the type signatures simple and to avoid inference problems we should
-- use this newtype. We cannot pass around a foralled type without wrapping
-- it in a newtype.

-- | @AcquireIO@ is used to acquire a resource safely such that it is
-- automatically released if not released manually.
--
-- See 'withAcquireIO'.
--
newtype AcquireIO = AcquireIO
    (forall b c. Priority -> IO b -> (b -> IO c) -> IO (b, IO ()))

-- | /Internal/.
allocator :: MonadIO m =>
       IORef (Int, IntMap (IO ()), IntMap (IO ()))
    -> Priority
    -> IO a
    -> (a -> IO b)
    -> m (a, m ())
allocator ref pri alloc free = do
    let insertResource r (i, mp1, mp2) =
            case pri of
                Priority1 ->
                    ((i + 1, Map.insert i (void $ free r) mp1, mp2), i)
                Priority2 ->
                    ((i + 1, mp1, Map.insert i (void $ free r) mp2), i)

    (r, index) <-
        liftIO $ mask_ $ do
            -- tid <- myThreadId
            r <- alloc
            idx <- atomicModifyIORef' ref (insertResource r)
            -- liftIO $ putStrLn $ "insert: " ++ show pri
            --      ++ " " ++ show idx ++ " " ++ show tid
            return (r, idx)

    let deleteResource (i, mp1, mp2) =
            case pri of
                Priority1 ->
                    let res = Map.lookup index mp1
                     in ((i, Map.delete index mp1, mp2), res)
                Priority2 ->
                    let res = Map.lookup index mp2
                     in ((i, mp1, Map.delete index mp2), res)

        release =
            -- IMPORTANT: do not use interruptible operations in this
            -- critical section. Even putStrLn can make tests fail.
            liftIO $ mask_ $ do
                -- tid <- myThreadId
                -- liftIO $ putStrLn $ "releasing index: " ++ show index
                --      ++ " " ++ show tid
                f <- atomicModifyIORef' ref deleteResource
                -- restoring exceptions makes it non-atomic, tests fail.
                -- Can use allowInterrupt in "free" if desired.
                sequence_ f
    return (r, release)

-- XXX can we ensure via GC that the resources that we are freeing are all
-- dead, there are no other references to them?

-- | We ensure that all async workers for concurrent streams are stopped
-- before we release the resources so that nobody could be using the
-- resource after they are freed.
--
-- The only other possibility, could be user issued forkIO not being
-- tracked by us, however, that would be a programming error and any such
-- threads could misbehave if we freed the resources from under them.
--
-- We use GC based hooks in 'Stream.bracketIO\'' so there could be async threads
-- spawned by GC, releasing resources concurrently with us. For that reason we
-- need to make sure that the "release" in the bracket end action is executed
-- only once in that case.
--
-- /Internal/.
releaser :: MonadIO m => IORef (a, IntMap (IO b), IntMap (IO b)) -> m ()
releaser ref =
    liftIO $ mask_ $ do
        -- Delete the map from the ref first so that anyone else (GC)
        -- releasing concurrently cannot find the map.
        -- liftIO $ putStrLn "cleaning up priority 1"
        mp1 <- atomicModifyIORef' ref
            (\(i, mp1,mp2) -> ((i, Map.empty, mp2), mp1))
        -- Note that the channel cleanup function is interruptible because
        -- it has blocking points.
        sequence_ mp1
        -- Now nobody would be changing mp2, we can read it safely
        -- liftIO $ putStrLn "cleaning up priority 2"
        mp2 <- atomicModifyIORef' ref
            (\(i, mp,mp2) -> ((i, mp, Map.empty), mp2))
        sequence_ mp2
        -- XXX We can now assert that the IORef has both maps empty.

-- | @withAcquireIO action@ runs the given @action@, providing it with a
-- an 'AcquireIO' reference called @ref@ as argument. @ref@ is used for resource
-- acquisition or hook registeration within the scope of @action@. An @acquire
-- ref alloc free@ call can be used within @action@ any number of times to
-- acquire resources that are automatically freed when the scope of @action@
-- ends or if an exception occurs at any time. @alloc@ is a function supplied
-- by the user to allocate a resource and @free@ is supplied to free the
-- allocated resource. @acquire@ returns @(resource, release)@ -- the acquired
-- @resource@ and a @release@ action to release it.
--
-- @acquire@ allocates a resource in an exception safe manner and sets up its
-- automatic release on exception or when the scope of @action@ ends. The
-- @release@ function returned by @acquire@ can be used to free the resource
-- manually at any time. @release@ is guaranteed to free the resource once and
-- only once even if it is called concurrently or multiple times.
--
-- Here is an example to allocate resources that are guaranteed to be released
-- automatically, and can be released manually as well:
--
-- >>> :{
-- close x h = do
--  putStrLn $ "closing: " ++ x
--  hClose h
-- :}
--
-- >>> :{
-- action ref =
--      Stream.fromList ["file1", "file2"]
--    & Stream.mapM
--        (\x -> do
--            (h, release) <- Exception.acquire ref (openFile x ReadMode) (close x)
--            -- use h here
--            threadDelay 1000000
--            when (x == "file1") $ do
--                putStrLn $ "Manually releasing: " ++ x
--                release
--            return x
--        )
--    & Stream.trace print
--    & Stream.fold Fold.drain
-- :}
--
-- >>> run = Exception.withAcquireIO action
--
-- In the above code, you should see the \"closing:\" message for both the
-- files, and only once for each file. Even if you interrupt the program with
-- CTRL-C you should still see the \"closing:\" message for the files opened
-- before the interrupt. Make sure you create "file1" and "file2" before
-- running this code snippet.
--
-- Cleanup is guaranteed to happen as soon as the scope of 'action'
-- finishes or if an exception occurs.
--
-- Here is an example for just registering hooks to be called eventually:
--
-- >>> :{
-- action ref =
--      Stream.fromList ["file1", "file2"]
--    & Stream.mapM
--        (\x -> do
--            Exception.register ref $ putStrLn $ "saw: " ++ x
--            threadDelay 1000000
--            return x
--        )
--    & Stream.trace print
--    & Stream.fold Fold.drain
-- :}
--
-- >>> run = Exception.withAcquireIO action
--
-- In the above code, even if you interrupt the program with CTRL-C you should
-- still see the "saw:" message for the elements seen before the interrupt.
--
-- The registered hooks are guaranteed to be invoked as soon as the scope of
-- 'action' finishes or if an exception occurs.
--
-- This function provides functionality similar to the @bracket@ function
-- available in the base library. However, it is more powerful as any number of
-- resources can be allocated and released within the scope of 'action'.
--
-- Exception safe, thread safe.
{-# INLINE withAcquireIO #-}
withAcquireIO :: (MonadIO m, MonadMask m) => (AcquireIO -> m a) -> m a
withAcquireIO action = do
    -- Assuming 64-bit int counter will never overflow
    ref <- liftIO $ newIORef (0 :: Int, Map.empty, Map.empty)
    action (AcquireIO (allocator ref)) `MC.finally` releaser ref

-- | Like 'acquire' but allows specifying a priority for releasing the
-- resource. 'Priority1' resources are released before 'Priority2'. This allows
-- us to specify a dependency between resource release.
{-# INLINE acquireWith #-}
acquireWith :: Priority -> AcquireIO -> IO b -> (b -> IO c) -> IO (b, IO ())
acquireWith pri (AcquireIO f) = f pri

-- | @acquire ref alloc free@ is used in bracket-style safe resource allocation
-- functions, where @alloc@ is a function supplied by the user to allocate a
-- resource and @free@ is supplied to free it. @acquire@ returns a tuple
-- @(resource, release)@ where @resource@ is the allocated resource and
-- @release@ is an action that can be called later to release the resource.
-- Both @alloc@ and @free@ are invoked with async signals masked. You can use
-- @allowInterrupt@ from base package for allowing interrupts if required.
--
-- The @release@ action can be called multiple times or even concurrently from
-- multiple threads,  but it will release the resource only once. If @release@
-- is never called by the programmer it will be automatically called at the end
-- of the bracket scope.
--
acquire :: AcquireIO -> IO b -> (b -> IO c) -> IO (b, IO ())
acquire = acquireWith Priority2

-- | Like 'acquire' but does not return a release action. The resource is freed
-- automatically only.
acquire_ :: AcquireIO -> IO b -> (b -> IO c) -> IO b
acquire_ a b c = fmap fst $ acquire a b c

-- | Like 'register' but specifies a 'Priority' for calling the hook.
{-# INLINE registerWith #-}
registerWith :: Priority -> AcquireIO -> IO () -> IO ()
registerWith pri (AcquireIO f) g = void $ f pri (return ()) (\() -> g)

-- | Register a hook to be executed at the end of a bracket.
register :: AcquireIO -> IO () -> IO ()
register = registerWith Priority2

-- | Like 'register' but returns a hook release function as well. When the
-- returned hook release function is called, the hook is invoked and removed.
-- If the returned function is never called by the programmer then it is
-- automatically invoked at the end of the bracket. The hook is invoked once
-- and only once.
--
hook :: AcquireIO -> IO () -> IO (IO())
hook (AcquireIO f) g = fmap snd $ f Priority2 (return ()) (\() -> g)