Companion.hs

{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE RankNTypes        #-}

-- | Companion threads, such as for printing messages saying we're still busy.
module Control.Concurrent.Companion
  ( withCompanion
  , onCompanionDone
  , Companion
  , Delay
  , StopCompanion
  ) where

import           RIO

-- | A companion thread which can perform arbitrary actions as well as delay.
type Companion m = Delay -> m ()

-- | Delay the given number of microseconds. If 'StopCompanion' is triggered
-- before the timer completes, a 'CompanionDone' exception will be thrown (which
-- is caught internally by 'withCompanion').
type Delay = forall mio. MonadIO mio => Int -> mio ()

-- | Tell the 'Companion' to stop. The next time 'Delay' is called, or if a
-- 'Delay' is currently blocking, the 'Companion' thread will exit with a
-- 'CompanionDone' exception.
type StopCompanion m = m ()

-- | When a delay was interrupted because we're told to stop, perform this
-- action.
onCompanionDone
  :: MonadUnliftIO m
  => m () -- ^ the delay
  -> m () -- ^ action to perform
  -> m ()
onCompanionDone theDelay theAction =
  theDelay `withException` \CompanionDone -> theAction

-- | Internal exception used by 'withCompanion' to allow short-circuiting of the
-- 'Companion'. Should not be used outside of this module.
data CompanionDone = CompanionDone
  deriving (Show, Typeable)

instance Exception CompanionDone

-- | Keep running the 'Companion' action until either the inner action completes
-- or calls the 'StopCompanion' action. This can be used to give the user status
-- information while running a long running operations.
withCompanion ::
     forall m a. MonadUnliftIO m
  => Companion m
  -> (StopCompanion m -> m a)
  -> m a
withCompanion companion inner = do
  -- Variable to indicate 'Delay'ing should result in a 'CompanionDone'
  -- exception.
  shouldStopVar <- newTVarIO False
  let -- Relatively simple: set shouldStopVar to True
      stopCompanion = atomically $ writeTVar shouldStopVar True

      delay :: Delay
      delay usec = do
        -- Register a delay with the runtime system
        delayDoneVar <- registerDelay usec
        join $ atomically $
          -- Delay has triggered, keep going
          (pure () <$ (readTVar delayDoneVar >>= checkSTM)) <|>
          -- Time to stop the companion, throw a 'CompanionDone' exception
          -- immediately
          (throwIO CompanionDone <$ (readTVar shouldStopVar >>= checkSTM))

  -- Run the 'Companion' and inner action together
  runConcurrently $
    -- Ignore a 'CompanionDone' exception from the companion, that's expected
    -- behavior
    Concurrently (companion delay `catch` \CompanionDone -> pure ()) *>
    -- Run the inner action, giving it the 'StopCompanion' action, and
    -- ensuring it is called regardless of exceptions.
    Concurrently (inner stopCompanion `finally` stopCompanion)