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

-- | 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 :: 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.

  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 :: Delay
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 a. a -> mio a
forall (f :: * -> *) a. Applicative f => a -> f a
pure () mio () -> STM () -> STM (mio ())
forall a b. a -> STM b -> STM a
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 a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> STM ()
checkSTM)) STM (mio ()) -> STM (mio ()) -> STM (mio ())
forall a. STM a -> STM a -> STM a
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 a b. a -> STM b -> STM a
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 a b. STM a -> (a -> STM b) -> STM b
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 Int -> mio ()
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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) Concurrently m () -> Concurrently m a -> Concurrently m a
forall a b.
Concurrently m a -> Concurrently m b -> Concurrently m b
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)