{-# 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 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
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CompanionDone] -> ShowS
$cshowList :: [CompanionDone] -> ShowS
show :: CompanionDone -> String
$cshow :: CompanionDone -> String
showsPrec :: Int -> CompanionDone -> ShowS
$cshowsPrec :: Int -> 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.

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

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

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

        TVar Bool
delayDoneVar <- forall (m :: * -> *). MonadIO m => Int -> m (TVar Bool)
registerDelay Int
usec
        forall (m :: * -> *) a. Monad m => m (m a) -> m a
join forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically forall a b. (a -> b) -> a -> b
$
          -- Delay has triggered, keep going

          (forall (f :: * -> *) a. Applicative f => a -> f a
pure () forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (forall a. TVar a -> STM a
readTVar TVar Bool
delayDoneVar forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
checkSTM)) forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|>
          -- Time to stop the companion, throw a 'CompanionDone' exception

          -- immediately

          (forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO CompanionDone
CompanionDone forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (forall a. TVar a -> STM a
readTVar TVar Bool
shouldStopVar forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
checkSTM))

  -- Run the 'Companion' and inner action together

  forall (m :: * -> *) a. Concurrently m a -> m a
runConcurrently forall a b. (a -> b) -> a -> b
$
    -- Ignore a 'CompanionDone' exception from the companion, that's expected

    -- behavior

    forall (m :: * -> *) a. m a -> Concurrently m a
Concurrently (Companion m
companion Delay
delay forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` \CompanionDone
CompanionDone -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*>
    -- Run the inner action, giving it the 'StopCompanion' action, and

    -- ensuring it is called regardless of exceptions.

    forall (m :: * -> *) a. m a -> Concurrently m a
Concurrently (StopCompanion m -> m a
inner StopCompanion m
stopCompanion forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` StopCompanion m
stopCompanion)