Roman Cheplyaka

Decompose ContT

Published on

Can we decompose a ContT action into separate acquire and release functions?

Motivation

Consider some resource, such as a file handle, socket, database connection etc. The two actions common to all resources are acquiring and releasing.

There are two ways to represent these actions: as a pair of functions

acquireResource :: IO Resource
releaseResource :: Resource -> IO ()

or as a single function

withResource :: (Resource -> IO a) -> IO a

As I explained in Why would you use ContT?, the latter can be nicely wrapped into a continuation monad:

getResource :: ContT r IO Resource
getResource = ContT withResource

We can go from an (acquireResource, releaseResource) pair to the ContT version by using bracket from Control.Exception:

getResource :: ContT r IO Resource
getResource = ContT $ bracket acquireResource releaseResource

But can we go in the opposite direction, i.e. decompose a value of type ContT r IO Resource into separate acquire and release functions?

To give an example of why we may want to do that, say a library only provides the withResource or getResource form, but we want to cache the allocated resource instances, deciding dynamically which ones to free. Another example would be combining resources and coroutines, see Understanding ResourceT by Michael Snoyman.

Note that not every ContT action follows the pattern “acquire; call the continuation; release”. Instead, the continuation may be called multiple times like so:

getInt :: ContT r IO Int
getInt = ContT $ \k -> last <$> mapM k [1..5]

(The parametricity ensures, however, that an action of type forall r . ContT r m a will call its continuation at least once or throw an exception.)

So what we are really asking is, assuming the ContT action has a certain form such as bracket acquire release, can we recover the acquire and release1 actions?

A beautiful but useless solution

If our ContT action has the form

withResource :: MonadIO m => ContT r m Resource
withResource = ContT $ \k -> do
  x <- acquire
  r <- k x
  release x
  return r

then we can recover the acquire and release x actions by… adding another ContT layer! This may sound bizarre, but the intuition behind this is simple: the release x is itself a continuation of k x—it is what happens when k x returns—and the ContT monad allows us to capture that continuation. Here’s the code:

{-# LANGUAGE RankNTypes, ScopedTypeVariables #-}
import Control.Monad.Cont
import Control.Monad.IO.Class (liftIO)

runC :: Monad m => ContT a m a -> m a
runC ca = runContT ca return

shift :: Monad m => ((a -> m w) -> ContT w m w) -> ContT w m a
shift f = ContT $ runC . f

decomposeContT
  :: forall a w .
     (forall r m . MonadIO m => ContT r m a)
  -> IO (a, IO ()) -- the resource and its release action
decomposeContT ca = runC $ runContT ca $ \a ->
  shift $ \k -> do
    let
      r :: (a, IO ())
      r = (a, void $ k r)
    return r

Here, shift captures the current continuation delimited by runC; see Oleg’s delimited continuations tutorial.

Here is an example of how we can use decomposeContT. We generate two random numbers and then “release” them in their increasing order:

import System.Random
import Text.Printf (printf)

withRandomInt :: MonadIO m => ContT r m Int
withRandomInt = ContT $ \k -> do
  n :: Int <- liftIO $ randomRIO (1,100)
  liftIO $ printf "Acquired number %d\n" n
  r <- k n
  liftIO $ printf "Released number %d\n" n
  return r

main = do
  (num1, release1) <- decomposeContT withRandomInt
  (num2, release2) <- decomposeContT withRandomInt
  if num1 < num2
    then release1 >> release2
    else release2 >> release1

However, this solution has a major drawback. It relies on the ability to run our ContT action in a non-standard base monad (another ContT). Therefore, it cannot be used to decompose ContT $ bracket acquire release, because bracket cannot run in the ContT monad.

A more practical solution

A less elegant solution, but one that actually solves real-world problems, uses threads. Every resource is acquired in its own thread. The thread then blocks on an MVar until it receives a signal that the resource is no longer needed, in which case it proceeds to release it.

Here is the code.

{-# LANGUAGE RankNTypes, ScopedTypeVariables #-}
import Control.Monad.Cont
import Control.Concurrent
import Control.Exception

data ContCall a
  = ContException !SomeException
    -- ^ an exception occurred before the continuation was called
  | ContCall a
    -- ^ the continuation was called with this argument
  | ContNoCall
    -- ^ the ContT action returned wihtout ever calling the continuation

-- | Retrieve a resource from a 'ContT' action and return an action to
-- release it.
decomposeContT :: forall a . (forall r . ContT r IO a) -> IO (a, IO ())
decomposeContT ca = mask $ \restore -> do
  -- mvar_a is used to pass the 'a' result (or an exception) from the ContT thread to the
  -- calling thread
  mvar_a :: MVar (ContCall a) <- newEmptyMVar
  -- mvar_r is used to signal the ContT thread that its resources can be
  -- freed
  mvar_r :: MVar () <- newEmptyMVar
  -- mvar_e is used to communicate a possible final exception to the
  -- calling thread
  mvar_e :: MVar (Maybe SomeException) <- newEmptyMVar
  let
    -- tell the ContT thread to release its resources
    freeResources :: IO ()
    freeResources = void $ tryPutMVar mvar_r ()
    -- like freeResources, but also check and propagate any exception that
    -- arose while trying to free the resource
    freeResourcesCheckException :: IO ()
    freeResourcesCheckException = do
      void $ tryPutMVar mvar_r ()
      mb_e <- readMVar mvar_e
      maybe (return ()) throwIO mb_e
  pid <- forkIO $ do
    ----------------------------------------------------------------------
    --             The code below runs in a new thread
    ----------------------------------------------------------------------
    r <- try . restore $ runContT ca $ \a -> do
      -- Try recording the argument of the continuation call in the MVar.
      -- Might fail (return False) if this is not the first continuation
      -- call, but we don't care.
      _ :: Bool <- tryPutMVar mvar_a (ContCall a)
      -- Then wait until it's ok to free the resources.
      () <- readMVar mvar_r
      return ()
    case r of
      Right () -> do
        -- The ContT action returned successfully. We don't know how many
        -- times the continuation was called. Try putting ContNoCall in
        -- case it wasn't called at all. If it was called, then tryPutMVar
        -- will return False; we don't care.
        _ :: Bool <- tryPutMVar mvar_a ContNoCall
        _ :: Bool <- tryPutMVar mvar_e Nothing
        return ()
      Left e -> do
        -- An exception was raised. We don't know whether it was before or
        -- after the continuation was called. Try putting the exception in
        -- both mvar_r and mvar_e so that it surfaces exactly once
        -- (assuming the cleanup function eventually gets called) in the
        -- calling thread.
        _ :: Bool <- tryPutMVar mvar_a $ ContException e
        _ :: Bool <- tryPutMVar mvar_e $ Just e
        return ()
    ----------------------------------------------------------------------
    --             The code above runs in a new thread
    ----------------------------------------------------------------------

  -- Now, in the calling thread, we wait until mvar_a is filled
  -- readMVar is blocking, so it may receive exceptions. If there is
  -- an exception, we also send it to the ContT thread.
  contCall <- readMVar mvar_a `catch` (\(e :: SomeException) -> do throwTo pid e; throwIO e)
  case contCall of
    ContCall a -> return (a, freeResourcesCheckException)
    ContException e -> do
      freeResources
      throwIO e
    ContNoCall -> do
      freeResources
      throwIO $ ErrorCall "decomposeContT: the continuation was never called"