{-# LANGUAGE OverloadedStrings #-}

-- | The pool's dispatcher: claims jobs and hands them to worker threads.
module Arbiter.Worker.Dispatcher
  ( runDispatcher
  ) where

import Arbiter.Core.HighLevel (QueueOperation)
import Arbiter.Core.Job.Types (JobRead)
import Arbiter.Core.Listen (Notification)
import Arbiter.Core.Operations qualified as Ops
import Control.Monad (when)
import Data.Bifunctor (first)
import Data.Foldable (for_, traverse_)
import Data.List.NonEmpty (NonEmpty (..))
import UnliftIO.STM qualified as STM

import Arbiter.Worker.Config
  ( HandlerMode (..)
  , WorkerConfig (..)
  , WorkerState (..)
  , pulseHeartbeat
  , readEffectiveState
  )
import Arbiter.Worker.Logger (LogLevel (..), newFailureGate, tryLog, tryReported)
import Arbiter.Worker.Logger.Internal (withJobContext)
import Arbiter.Worker.NotificationListener (runNotificationConsumer)
import Arbiter.Worker.WorkQueue (WorkQueue, awaitFinished, inFlight, pushWork)

-- | Wake on NOTIFY, poll timer, or worker-finished, then claim up to capacity.
-- @notifVar@ is filled from the shared hub in "Arbiter.Core.Listen".
runDispatcher
  :: forall payload m
   . (QueueOperation m payload)
  => WorkerConfig m payload
  -> Int
  -> Ops.JobStatements
  -> WorkQueue (NonEmpty (JobRead payload))
  -> STM.TVar (Maybe Notification)
  -> m ()
runDispatcher :: forall payload (m :: * -> *).
QueueOperation m payload =>
WorkerConfig m payload
-> Int
-> JobStatements
-> WorkQueue (NonEmpty (JobRead payload))
-> TVar (Maybe Notification)
-> m ()
runDispatcher WorkerConfig m payload
config Int
workerCapacity JobStatements
statements WorkQueue (NonEmpty (JobRead payload))
workQueue TVar (Maybe Notification)
notifVar = do
  claimGate <- m FailureGate
forall (m :: * -> *). MonadIO m => m FailureGate
newFailureGate
  deadLetterGate <- newFailureGate
  let
    getFreeWorkers :: STM.STM (Maybe Int)
    getFreeWorkers = do
      free <- (Int
workerCapacity Int -> Int -> Int
forall a. Num a => a -> a -> a
-) (Int -> Int) -> STM Int -> STM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkQueue (NonEmpty (JobRead payload)) -> STM Int
forall a. WorkQueue a -> STM Int
inFlight WorkQueue (NonEmpty (JobRead payload))
workQueue
      pure $ if free > 0 then Just free else Nothing

    claimAndEnqueue :: Int -> m ()
    claimAndEnqueue Int
freeWorkers = do
      eJobs <- LogConfig
-> LogLevel
-> FailureGate
-> Text
-> m ([NonEmpty (JobRead payload)], [RejectedRow payload])
-> m (Either
        SomeException
        ([NonEmpty (JobRead payload)], [RejectedRow payload]))
forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig
-> LogLevel
-> FailureGate
-> Text
-> m a
-> m (Either SomeException a)
tryReported (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) LogLevel
Error FailureGate
claimGate Text
"Dispatcher claim" (m ([NonEmpty (JobRead payload)], [RejectedRow payload])
 -> m (Either
         SomeException
         ([NonEmpty (JobRead payload)], [RejectedRow payload])))
-> m ([NonEmpty (JobRead payload)], [RejectedRow payload])
-> m (Either
        SomeException
        ([NonEmpty (JobRead payload)], [RejectedRow payload]))
forall a b. (a -> b) -> a -> b
$
        case WorkerConfig m payload -> HandlerMode m payload
forall (m :: * -> *) payload.
WorkerConfig m payload -> HandlerMode m payload
handlerMode WorkerConfig m payload
config of
          SingleJobMode JobHandler m payload (ResultOf m payload)
_ ->
            ([JobRead payload] -> [NonEmpty (JobRead payload)])
-> ([JobRead payload], [RejectedRow payload])
-> ([NonEmpty (JobRead payload)], [RejectedRow payload])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ((JobRead payload -> NonEmpty (JobRead payload))
-> [JobRead payload] -> [NonEmpty (JobRead payload)]
forall a b. (a -> b) -> [a] -> [b]
map (JobRead payload -> [JobRead payload] -> NonEmpty (JobRead payload)
forall a. a -> [a] -> NonEmpty a
:| [])) (([JobRead payload], [RejectedRow payload])
 -> ([NonEmpty (JobRead payload)], [RejectedRow payload]))
-> m ([JobRead payload], [RejectedRow payload])
-> m ([NonEmpty (JobRead payload)], [RejectedRow payload])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> JobStatements
-> Int -> m ([JobRead payload], [RejectedRow payload])
forall (m :: * -> *) payload.
(JobPayload payload, MonadArbiter m) =>
JobStatements
-> Int -> m ([JobRead payload], [RejectedRow payload])
Ops.claimJobsCached JobStatements
statements Int
freeWorkers
          BatchedJobsMode Int
_ NonEmpty (JobRead payload)
-> BatchCallbacks m payload (ResultOf m payload) -> m ()
_ ->
            JobStatements
-> Int -> m ([NonEmpty (JobRead payload)], [RejectedRow payload])
forall (m :: * -> *) payload.
(JobPayload payload, MonadArbiter m) =>
JobStatements
-> Int -> m ([NonEmpty (JobRead payload)], [RejectedRow payload])
Ops.claimJobsBatchedCached JobStatements
statements Int
freeWorkers
      for_ eJobs $ \([NonEmpty (JobRead payload)]
jobs, [RejectedRow payload]
rejected) -> do
        WorkQueue (NonEmpty (JobRead payload))
-> [NonEmpty (JobRead payload)] -> m ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
WorkQueue a -> [a] -> m ()
pushWork WorkQueue (NonEmpty (JobRead payload))
workQueue [NonEmpty (JobRead payload)]
jobs
        (RejectedRow payload -> m ()) -> [RejectedRow payload] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ RejectedRow payload -> m ()
deadLetter [RejectedRow payload]
rejected
        -- An all-poison claim proved the queue has rows, so claim again while still running.
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([NonEmpty (JobRead payload)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [NonEmpty (JobRead payload)]
jobs Bool -> Bool -> Bool
&& Bool -> Bool
not ([RejectedRow payload] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [RejectedRow payload]
rejected)) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
          state <- STM WorkerState -> m WorkerState
forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically (WorkerConfig m payload -> STM WorkerState
forall (n :: * -> *) payload.
WorkerConfig n payload -> STM WorkerState
readEffectiveState WorkerConfig m payload
config)
          when (state == Running) claimOnWakeup
      -- Pulse on every attempt, including a failed claim.
      STM.atomically (pulseHeartbeat config)

    deadLetter :: Ops.RejectedRow payload -> m ()
    deadLetter rejected :: RejectedRow payload
rejected@(JobRead (Stored payload)
row, Text
err) = do
      moved <-
        LogConfig
-> LogLevel
-> FailureGate
-> Text
-> m Int64
-> m (Either SomeException Int64)
forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig
-> LogLevel
-> FailureGate
-> Text
-> m a
-> m (Either SomeException a)
tryReported (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) LogLevel
Error FailureGate
deadLetterGate Text
"Dead-letter undecodable job" (RejectedRow payload -> m Int64
forall (m :: * -> *) payload.
MonadArbiter m =>
RejectedRow payload -> m Int64
Ops.deadLetterRejected RejectedRow payload
rejected)
      for_ moved $ \Int64
n ->
        Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int64
n Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
> Int64
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog (LogConfig -> NonEmpty (JobRead (Stored payload)) -> LogConfig
forall payload.
LogConfig -> NonEmpty (JobRead payload) -> LogConfig
withJobContext (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) (JobRead (Stored payload)
row JobRead (Stored payload)
-> [JobRead (Stored payload)]
-> NonEmpty (JobRead (Stored payload))
forall a. a -> [a] -> NonEmpty a
:| [])) LogLevel
Error (Text
"Job moved to the DLQ, " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
err)

    claimOnWakeup :: m ()
    claimOnWakeup = STM (Maybe Int) -> m (Maybe Int)
forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically STM (Maybe Int)
getFreeWorkers m (Maybe Int) -> (Maybe Int -> 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
>>= (Int -> m ()) -> Maybe Int -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Int -> m ()
claimAndEnqueue

  runNotificationConsumer
    (readEffectiveState config)
    (pollInterval config)
    notifVar
    (awaitFinished workQueue)
    (const claimOnWakeup)