{-# 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 :: forall (m :: * -> *). MonadUnliftIO m => m () -> m () -> m ()
onCompanionDone m ()
theDelay m ()
theAction =
  m ()
theDelay m () -> (CompanionDone -> m ()) -> m ()
forall (m :: * -> *) e a b.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m b) -> m a
`withException` \CompanionDone
CompanionDone -> m ()
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 (Int -> CompanionDone -> ShowS
[CompanionDone] -> ShowS
CompanionDone -> String
(Int -> CompanionDone -> ShowS)
-> (CompanionDone -> String)
-> ([CompanionDone] -> ShowS)
-> Show CompanionDone
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> CompanionDone -> ShowS
showsPrec :: Int -> CompanionDone -> ShowS
$cshow :: CompanionDone -> String
show :: CompanionDone -> String
$cshowList :: [CompanionDone] -> ShowS
showList :: [CompanionDone] -> ShowS
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 :: forall (m :: * -> *) a.
MonadUnliftIO m =>
Companion m -> (StopCompanion m -> m a) -> m a
withCompanion Companion m
companion StopCompanion m -> m a
inner = do
  -- Variable to indicate 'Delay'ing should result in a 'CompanionDone'

  -- exception.

  shouldStopVar <- Bool -> m (TVar Bool)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Bool
False
  let -- Relatively simple: set shouldStopVar to True

      stopCompanion = STM () -> StopCompanion m
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> StopCompanion m) -> STM () -> StopCompanion m
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
shouldStopVar Bool
True

      delay :: Delay
      delay Int
usec = do
        -- Register a delay with the runtime system

        delayDoneVar <- Int -> mio (TVar Bool)
forall (m :: * -> *). MonadIO m => Int -> m (TVar Bool)
registerDelay Int
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
CompanionDone -> () -> StopCompanion m
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) *>
    -- Run the inner action, giving it the 'StopCompanion' action, and

    -- ensuring it is called regardless of exceptions.

    Concurrently (inner stopCompanion `finally` stopCompanion)