{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE NoImplicitPrelude #-}
-- | Companion threads, such as for printing messages saying we're
-- still busy. Ultimately this could be put into its own package. This
-- is a non-standard API for use by Pantry and Stack, please /DO NOT
-- DEPEND ON IT/.
module Pantry.Internal.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 :: 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
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 :: 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 <- Bool -> m (TVar Bool)
forall (m :: * -> *) a. MonadIO m => a -> m (TVar a)
newTVarIO Bool
False
  let -- Relatively simple: set shouldStopVar to True
      stopCompanion :: StopCompanion m
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 -> mio ()
delay Int
usec = do
        -- Register a delay with the runtime system
        TVar Bool
delayDoneVar <- Int -> mio (TVar Bool)
forall (m :: * -> *). MonadIO m => Int -> m (TVar Bool)
registerDelay Int
usec
        mio (mio ()) -> mio ()
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (mio (mio ()) -> mio ()) -> mio (mio ()) -> mio ()
forall a b. (a -> b) -> a -> b
$ STM (mio ()) -> mio (mio ())
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM (mio ()) -> mio (mio ())) -> STM (mio ()) -> mio (mio ())
forall a b. (a -> b) -> a -> b
$
          -- Delay has triggered, keep going
          (() -> mio ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure () mio () -> STM () -> STM (mio ())
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
delayDoneVar STM Bool -> (Bool -> STM ()) -> STM ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
checkSTM)) STM (mio ()) -> STM (mio ()) -> STM (mio ())
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|>
          -- Time to stop the companion, throw a 'CompanionDone' exception immediately
          (CompanionDone -> mio ()
forall (m :: * -> *) e a. (MonadIO m, Exception e) => e -> m a
throwIO CompanionDone
CompanionDone mio () -> STM () -> STM (mio ())
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ (TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar TVar Bool
shouldStopVar STM Bool -> (Bool -> STM ()) -> STM ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
checkSTM))

  -- Run the 'Companion' and inner action together
  Concurrently m a -> m a
forall (m :: * -> *) a. Concurrently m a -> m a
runConcurrently (Concurrently m a -> m a) -> Concurrently m a -> m a
forall a b. (a -> b) -> a -> b
$
    -- Ignore a 'CompanionDone' exception from the companion, that's expected behavior
    StopCompanion m -> Concurrently m ()
forall (m :: * -> *) a. m a -> Concurrently m a
Concurrently (Companion m
companion Delay
delay StopCompanion m
-> (CompanionDone -> StopCompanion m) -> StopCompanion m
forall (m :: * -> *) e a.
(MonadUnliftIO m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` \CompanionDone
CompanionDone -> () -> StopCompanion m
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) Concurrently m () -> Concurrently m a -> Concurrently m a
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.
    m a -> Concurrently m a
forall (m :: * -> *) a. m a -> Concurrently m a
Concurrently (StopCompanion m -> m a
inner StopCompanion m
stopCompanion m a -> StopCompanion m -> m a
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` StopCompanion m
stopCompanion)