-- | Pause-aware consumer loop the dispatcher drives off the shared listener.
module Arbiter.Worker.NotificationListener
  ( runNotificationConsumer
  ) where

import Arbiter.Core.Listen (Notification)
import Arbiter.Core.Operations qualified as Ops
import Control.Monad (when)
import Data.Time (NominalDiffTime)
import UnliftIO (MonadUnliftIO)
import UnliftIO.STM qualified as STM

import Arbiter.Worker.WorkerState (WorkerState (..))

-- | Loop until 'ShuttingDown'. Per iteration wait on notification, poll timer,
-- wake trigger, or state change. Fires @action Nothing@ once at startup if the
-- state is 'Running'.
runNotificationConsumer
  :: (MonadUnliftIO m)
  => STM.STM WorkerState
  -> NominalDiffTime
  -> STM.TVar (Maybe Notification)
  -> STM.STM ()
  -> (Maybe Notification -> m ())
  -> m ()
runNotificationConsumer :: forall (m :: * -> *).
MonadUnliftIO m =>
STM WorkerState
-> NominalDiffTime
-> TVar (Maybe Notification)
-> STM ()
-> (Maybe Notification -> m ())
-> m ()
runNotificationConsumer STM WorkerState
readState NominalDiffTime
pollDelay TVar (Maybe Notification)
notifVar STM ()
wakeTrigger Maybe Notification -> m ()
action = do
  state <- STM WorkerState -> m WorkerState
forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically STM WorkerState
readState
  when (state == Running) (action Nothing)
  loop
  where
    pollMicros :: Int
pollMicros = NominalDiffTime -> Int
Ops.micros NominalDiffTime
pollDelay

    loop :: m ()
loop = m Command
nextCommand m Command -> (Command -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Command -> m ()
step

    step :: Command -> m ()
step = \case
      Command
Halt -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Command
PauseRecv -> m Command
awaitUnpause m Command -> (Command -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Command -> m ()
step
      NotificationRecv Notification
notification -> Maybe Notification -> m ()
action (Notification -> Maybe Notification
forall a. a -> Maybe a
Just Notification
notification) m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m ()
loop
      Command
TimerExpired -> Maybe Notification -> m ()
action Maybe Notification
forall a. Maybe a
Nothing m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m ()
loop

    nextCommand :: m Command
nextCommand = do
      delayVar <- Int -> m (TVar Bool)
forall (m :: * -> *). MonadIO m => Int -> m (TVar Bool)
STM.registerDelay Int
pollMicros
      STM.atomically $ do
        status <- readState
        case status of
          WorkerState
ShuttingDown -> Command -> STM Command
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Command
Halt
          WorkerState
Paused -> Command -> STM Command
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Command
PauseRecv
          WorkerState
Running ->
            STM Command
consumeNotification
              STM Command -> STM Command -> STM Command
forall a. STM a -> STM a -> STM a
`STM.orElse` TVar Bool -> STM Command
timerExpired TVar Bool
delayVar
              STM Command -> STM Command -> STM Command
forall a. STM a -> STM a -> STM a
`STM.orElse` (Command
TimerExpired Command -> STM () -> STM Command
forall a b. a -> STM b -> STM a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ STM ()
wakeTrigger)

    awaitUnpause :: m Command
awaitUnpause =
      STM Command -> m Command
forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically (STM Command -> m Command) -> STM Command -> m Command
forall a b. (a -> b) -> a -> b
$ do
        state <- STM WorkerState
readState
        case state of
          WorkerState
Paused -> STM Command
forall a. STM a
STM.retrySTM
          WorkerState
ShuttingDown -> Command -> STM Command
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Command
Halt
          WorkerState
Running -> Command -> STM Command
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Command
TimerExpired

    consumeNotification :: STM Command
consumeNotification = do
      notification <- TVar (Maybe Notification) -> STM (Maybe Notification)
forall a. TVar a -> STM a
STM.readTVar TVar (Maybe Notification)
notifVar STM (Maybe Notification)
-> (Maybe Notification -> STM Notification) -> STM Notification
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STM Notification
-> (Notification -> STM Notification)
-> Maybe Notification
-> STM Notification
forall b a. b -> (a -> b) -> Maybe a -> b
maybe STM Notification
forall a. STM a
STM.retrySTM Notification -> STM Notification
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
      STM.writeTVar notifVar Nothing
      pure (NotificationRecv notification)

    timerExpired :: TVar Bool -> STM Command
timerExpired TVar Bool
delayVar =
      Command
TimerExpired Command -> STM () -> STM Command
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
STM.readTVar TVar Bool
delayVar 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 ()
STM.checkSTM)

data Command
  = Halt
  | PauseRecv
  | NotificationRecv Notification
  | TimerExpired