An attempt to model the masking and interrupt-safety of IO actions with phantom types on the IO monad.

withforkio.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Rank2Types #-}

import Control.Concurrent (ThreadId, forkIO, myThreadId, killThread)
import Control.Exception (Exception, SomeException, throwIO, throwTo, catch, mask, uninterruptibleMask)




-- * Annotated IO

data Unmasked
data Masked
data Interruptible
data Uninterruptible

-- | IO with phantom types for masking state and interruptibility.
newtype IO' msk int a = LiftIO { runIO :: IO a }
    deriving (Functor, Applicative, Monad)

-- | An interruptible context will run the same way in an uninterruptible
-- context.
liftI :: IO' msk Interruptible a -> IO' msk Uninterruptible a
liftI (LiftIO action) = LiftIO action




-- * Control.Exception


-- ** Throwing exceptions

-- | 'throwIO' is not interruptible [1] so it is safe to run in an
-- interruptible context.
--
-- [1]: https://hackage.haskell.org/package/base-4.18.0.0/docs/Control-Exception.html#g:13
throwIO' :: Exception e => e -> IO' msk Interruptible a
throwIO' = LiftIO . throwIO

-- | 'throwTo' is always interruptible even if it doesn't block [1] so it should
-- be run in an uninterruptible context.
--
-- [1]: https://hackage.haskell.org/package/base/docs/Control-Exception-Base.html#v:throwTo
throwTo' :: Exception e => ThreadId -> e -> IO' msk Uninterruptible ()
throwTo' tid = LiftIO . throwTo tid


-- ** Catching exceptions

-- | 'catch' catches an exception and runs its handler inside an implicit
-- masked [1,2] (against asynchronous exceptions) and interruptible context.
--
-- [1]: https://hackage.haskell.org/package/base-4.18.0.0/docs/Control-Exception.html#g:3
-- [2]: https://hackage.haskell.org/package/base-4.18.0.0/docs/Control-Exception.html#g:12
catch'
    :: Exception e
    => IO' msg int a
    -> (e -> IO' Masked Interruptible a)
    -> IO' msg int a
catch' action handler = LiftIO $ runIO action `catch` (runIO . handler)

-- Like 'catch' but strict in the action. Not for export.
catchException'
    :: Exception e
    => IO' msg int a
    -> (e -> IO' Masked Interruptible a)
    -> IO' msg int a
catchException' !action handler = catch' action handler

-- | Catches an exception and returns it in an either.
try' :: Exception e
    => IO' msk int a
    -> IO' msk int (Either e a)
try' action =
    catch'
        (action >>= \v -> return $ Right v)
        (\e -> return $ Left e)

-- | 'onException' runs a cleanup handler when there's an exception, in a
-- masked (from asynchronous exceptions) but interruptible context.
onException'
    :: IO' msk int a
    -> IO' Masked Interruptible b
    -> IO' msk int a
onException' action handler =
    action `catchException'` \e -> do
        _ <- handler
        throwIO' (e :: SomeException)


-- * Asynchronous exception control

-- | 'mask' runs a computation in a masked (from asynchronous exceptions) but
-- interruptible context and provides a function to restore the masking state
-- of the caller.
mask'
    :: ((forall a. IO' msk int a -> IO' Masked Interruptible a)
        -> IO' Masked Interruptible b)
    -> IO' msk int b
mask' action =
    LiftIO $ mask
        (\restore -> runIO $ action (LiftIO . restore . runIO))

-- | 'uninterruptibleMask' runs a computation in a masked (from asynchronous
-- exceptions) and uninterruptible context and provides a function to restore
-- the masking state of the caller.
--
-- This can cause programs to deadlock.
uninterruptibleMask'
    :: ((forall a. IO' msk int a -> IO' Masked Uninterruptible a)
        -> IO' Masked Uninterruptible b)
    -> IO' msk int b
uninterruptibleMask' action =
    LiftIO $ uninterruptibleMask
        (\restore -> runIO $ action (LiftIO . restore . runIO))

-- | 'mask_' runs a computation in a masked (from asynchronous exceptions) but
-- interruptible context.
mask_'
    :: IO' Masked Interruptible a
    -> IO' msk int a
mask_' action = mask' $ \_ -> action

-- | 'uninterruptibleMask' runs a computation in a masked (from asynchronous
-- exceptions) and uninterruptible context.
--
-- This can cause programs to deadlock.
uninterruptibleMask_'
    :: IO' Masked Uninterruptible a
    -> IO' msk int a
uninterruptibleMask_' action = uninterruptibleMask' $ \_ -> action


-- ** Utilities

-- | 'bracket' acquires and releases a resource in a masked (from asynchronous
-- exceptions) but interruptible context. The inner action is not masked.
bracket'
    :: IO' Masked Interruptible a
    -> (a -> IO' Masked Interruptible b)
    -> (a -> IO' msk int c)
    -> IO' msk int c
bracket' acquire release use =
    mask' $ \restore -> do
        a <- acquire
        r <- restore (use a) `onException'` release a
        _ <- release a
        return r

-- | 'finally' always runs a cleanup handler in a masked (from asynchronous
-- exceptions) but interruptible context.
finally'
    :: IO' msk int a
    -> IO' Masked Interruptible b
    -> IO' msk int a
action `finally'` after =
    mask' $ \restore -> do
        r <- restore action `onException'` after
        _ <- after
        return r




-- * Control.Concurrent

-- | A child of 'forkIO' inherits the masking state of the parent [1,2].
-- 'forkIO' is not interruptible [3] so it may run in an interruptible context.
--
-- [1]: https://hackage.haskell.org/package/base/docs/Control-Concurrent.html#v:forkIO
-- [2]: https://hackage.haskell.org/package/base/docs/Control-Exception.html#v:mask
-- [3]: https://hackage.haskell.org/package/base-4.18.0.0/docs/Control-Exception.html#g:13
forkIO'
    :: IO' msk int ()
    -> IO' msk Interruptible ThreadId
forkIO' = LiftIO . forkIO . runIO

myThreadId' :: IO' msk Interruptible ThreadId
myThreadId' = LiftIO myThreadId

-- | Useful to tell the parent when the child terminates.
--
-- ???: should the `int` type variable be fixed to 'Interruptible'? It seems
-- like it ought to be, since `forkFinally`, like `forkIO` is safe to run in an
-- interruptible context.
forkFinally'
    :: IO' msk int a
    -> (Either SomeException a -> IO' Masked Interruptible ())
    -> IO' msk int ThreadId
forkFinally' action handler =
    mask' $ \restore ->
        forkIO' $ try' (restore action) >>= handler

-- | 'killThread' is interruptible [1], so it requires an uninterruptible
-- context.
--
-- [1]: https://hackage.haskell.org/package/base/docs/Control-Exception.html#v:bracket
killThread'
    :: ThreadId
    -> IO' msk Uninterruptible ()
killThread' = LiftIO . killThread


-- ** Utilities

-- | Possibly like @async:Control.Async.withAsync@
withForkIO'
    :: IO' msk int ()
    -> IO' msk int a
    -> IO' msk int a
withForkIO' child parent =
    mask' $ \restore -> do
        t <- forkIO' $ restore child
        r <- restore parent
                `onException'` uninterruptibleMask_' (killThread' t)
        uninterruptibleMask_' (killThread' t)
        return r