{-# LANGUAGE OverloadedStrings #-}

-- | The worker thread: takes batches off the queue and runs each through
-- "Arbiter.Worker.Batch".
module Arbiter.Worker.Processing
  ( workerLoop
  ) where

import Arbiter.Core.Exceptions (JobForceCancelled (..), displayEx)
import Arbiter.Core.Job.Types qualified as Job
import Arbiter.Core.Threads (micros)
import Control.Exception qualified as E
import Control.Monad (forever)
import Control.Monad.IO.Class (liftIO)
import Data.List.NonEmpty (NonEmpty)
import Data.Time (NominalDiffTime)
import UnliftIO
  ( MonadUnliftIO
  , atomically
  , finally
  , mask
  )
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.Exception (isAsyncException)

import Arbiter.Worker.Batch (newHandoff, runBatch)
import Arbiter.Worker.Config
import Arbiter.Worker.Heartbeat (PoolGuard)
import Arbiter.Worker.Logger
import Arbiter.Worker.Settlement (PoolEffects, PoolMode, batchLog)
import Arbiter.Worker.WorkQueue (WorkQueue, finishWork, popWork)

-- | The pause after a batch escapes with an unexpected exception.
exceptionPause :: NominalDiffTime
exceptionPause :: NominalDiffTime
exceptionPause = NominalDiffTime
2

-- | Main loop for a single worker thread.
workerLoop
  :: forall payload m
   . (MonadUnliftIO m)
  => WorkerConfig m payload
  -> PoolGuard payload
  -> PoolMode m payload
  -> PoolEffects m payload
  -> WorkQueue (NonEmpty (Job.JobRead payload))
  -> m ()
workerLoop :: forall payload (m :: * -> *).
MonadUnliftIO m =>
WorkerConfig m payload
-> PoolGuard payload
-> PoolMode m payload
-> PoolEffects m payload
-> WorkQueue (NonEmpty (JobRead payload))
-> m ()
workerLoop WorkerConfig m payload
config PoolGuard payload
guard PoolMode m payload
mode PoolEffects m payload
effectsFor WorkQueue (NonEmpty (JobRead payload))
workQueue =
  m () -> m ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ ((forall a. m a -> m a) -> m ()) -> m ()
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> m a) -> m b) -> m b
mask (((forall a. m a -> m a) -> m ()) -> m ())
-> ((forall a. m a -> m a) -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \forall a. m a -> m a
restore -> do
    -- Mask covers the window between taking a batch (which moves it from queued
    -- to busy) and entering the finally block that frees the busy slot.
    jobBatch <- WorkQueue (NonEmpty (JobRead payload))
-> m (NonEmpty (JobRead payload))
forall (m :: * -> *) a. MonadUnliftIO m => WorkQueue a -> m a
popWork WorkQueue (NonEmpty (JobRead payload))
workQueue

    let process = do
          handoff <- IO (Handoff IO (JobRead payload))
-> m (Handoff IO (JobRead payload))
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (PoolGuard payload -> IO (Handoff IO (JobRead payload))
forall (n :: * -> *) job.
(MonadSTM n, MonadTime n) =>
HeartbeatGuard n job -> n (Handoff n job)
newHandoff PoolGuard payload
guard)
          result <- restore (liftIO (E.try (runBatch (effectsFor jobBatch) guard mode handoff jobBatch)))
          case result of
            Right () -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Left SomeException
exception
              | Just JobForceCancelled {} <- SomeException -> Maybe JobForceCancelled
forall e. Exception e => SomeException -> Maybe e
E.fromException SomeException
exception -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
              | SomeException -> Bool
forall e. Exception e => e -> Bool
isAsyncException SomeException
exception -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (SomeException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO SomeException
exception)
              | Bool
otherwise -> do
                  LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog (WorkerConfig m payload -> NonEmpty (JobRead payload) -> LogConfig
forall (m :: * -> *) payload.
WorkerConfig m payload -> NonEmpty (JobRead payload) -> LogConfig
batchLog WorkerConfig m payload
config NonEmpty (JobRead payload)
jobBatch) LogLevel
Error (Text -> m ()) -> Text -> m ()
forall a b. (a -> b) -> a -> b
$ Text
"Worker exception: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SomeException -> Text
displayEx SomeException
exception
                  Int -> m ()
forall (m :: * -> *). MonadIO m => Int -> m ()
threadDelay (NominalDiffTime -> Int
micros NominalDiffTime
exceptionPause)
    process `finally` atomically (finishWork workQueue)