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

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

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

type Action m a = Maybe Notification -> m a

-- | 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)
  -> Maybe (STM.STM ())
  -> Action m ()
  -> m ()
runNotificationConsumer :: forall (m :: * -> *).
MonadUnliftIO m =>
STM WorkerState
-> NominalDiffTime
-> TVar (Maybe Notification)
-> Maybe (STM ())
-> Action m ()
-> m ()
runNotificationConsumer STM WorkerState
readState NominalDiffTime
pollDelay TVar (Maybe Notification)
notifVar Maybe (STM ())
mWakeTrigger Action 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
forall b. Integral b => NominalDiffTime -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (NominalDiffTime
pollDelay NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
1_000_000)

    loop :: m ()
loop = do
      command <- m Command
nextCommand
      case command of
        Command
Halt -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Command
PauseCmd -> do
          next <- m Command
awaitUnpause
          case next of
            Command
Halt -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Command
_ -> Action 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
        NotificationRecv Notification
notification -> Action 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 -> Action 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
PauseCmd
          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` STM Command
waitWakeTrigger
              STM Command -> STM Command -> STM Command
forall a. STM a -> STM a -> STM a
`STM.orElse` STM Command
watchStateChange

    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
      mNotif <- TVar (Maybe Notification) -> STM (Maybe Notification)
forall a. TVar a -> STM a
STM.readTVar TVar (Maybe Notification)
notifVar
      case mNotif of
        Just Notification
notification -> do
          TVar (Maybe Notification) -> Maybe Notification -> STM ()
forall a. TVar a -> a -> STM ()
STM.writeTVar TVar (Maybe Notification)
notifVar Maybe Notification
forall a. Maybe a
Nothing
          Command -> STM Command
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Notification -> Command
NotificationRecv Notification
notification)
        Maybe Notification
Nothing -> STM Command
forall a. STM a
STM.retrySTM

    timerExpired :: TVar Bool -> STM Command
timerExpired TVar Bool
delayVar = do
      isExpired <- TVar Bool -> STM Bool
forall a. TVar a -> STM a
STM.readTVar TVar Bool
delayVar
      if isExpired then pure TimerExpired else STM.retrySTM

    waitWakeTrigger :: STM Command
waitWakeTrigger = case Maybe (STM ())
mWakeTrigger of
      Maybe (STM ())
Nothing -> STM Command
forall a. STM a
STM.retrySTM
      Just STM ()
trigger -> STM ()
trigger STM () -> STM Command -> STM Command
forall a b. STM a -> STM b -> STM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Command -> STM Command
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Command
TimerExpired

    watchStateChange :: STM Command
watchStateChange = do
      state <- STM WorkerState
readState
      case state 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
PauseCmd
        WorkerState
Running -> STM Command
forall a. STM a
STM.retrySTM

data Command
  = Halt
  | PauseCmd
  | NotificationRecv Notification
  | TimerExpired