{-# LANGUAGE OverloadedStrings #-}

-- | Execution of one worker pool and its job lifecycle. Multi-queue
-- orchestration lives in "Arbiter.Worker.MultiQueue".
module Arbiter.Worker.Pool
  ( runWorkerPool
  , throwIfInvalid
  ) where

import Arbiter.Core.Enum (showText)
import Arbiter.Core.Exceptions (displayEx)
import Arbiter.Core.HighLevel (QueueOperation)
import Arbiter.Core.HighLevel qualified as Arb
import Arbiter.Core.Job.Schema (SchemaName)
import Arbiter.Core.Job.Schema qualified as Schema
import Arbiter.Core.JobResult
import Arbiter.Core.Listen qualified as Listen
import Arbiter.Core.MonadArbiter (MonadArbiter (..))
import Arbiter.Core.Operations qualified as Ops
import Arbiter.Core.QueueRegistry (RegistryTables (..))
import Arbiter.Core.Trace (ConsumeShape, consumeSpanFor, toConsumeShape)
import Control.Exception qualified as E
import Control.Monad (forever, replicateM, unless, void)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Cont (ContT (..), evalContT)
import Data.ByteString (ByteString)
import Data.Foldable (traverse_)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text.Encoding qualified as TE
import Data.Time (NominalDiffTime)
import System.Directory (removeFile)
import UnliftIO
  ( MonadUnliftIO
  , STM
  , TVar
  , atomically
  , checkSTM
  , finally
  , newTVarIO
  , readTVar
  , readTVarIO
  , registerDelay
  , takeTMVar
  , tryAny
  , waitAnyCatch
  , writeTVar
  )
import UnliftIO.Async qualified as Async
import UnliftIO.Concurrent (threadDelay)

import Arbiter.Worker.ChannelHandlers
  ( handleCancelNotif
  , handleCronRunNotif
  , handlePauseNotif
  )
import Arbiter.Worker.Config
import Arbiter.Worker.Cron.Scheduler (runCronScheduler)
import Arbiter.Worker.Cron.Types (CronJob (..))
import Arbiter.Worker.Dispatcher
import Arbiter.Worker.Heartbeat (PoolGuard, newHeartbeatGuard, runHeartbeatGuard)
import Arbiter.Worker.Logger
import Arbiter.Worker.Logger.Internal (tryWarn, tryWarnWith)
import Arbiter.Worker.Processing (workerLoop)
import Arbiter.Worker.Reaper (MaintenancePace (..), reaperLoop)
import Arbiter.Worker.Retry (spawnRetried)
import Arbiter.Worker.Settlement (poolEffects, poolMode)
import Arbiter.Worker.WorkQueue (WorkQueue, busyCount, inFlight, newWorkQueue, queuedCount)
import Arbiter.Worker.WorkerState (isShuttingDown)

-- | Throw 'WorkerConfigException' for an invalid config.
throwIfInvalid :: (MonadIO m) => WorkerConfig n payload -> m ()
throwIfInvalid :: forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
throwIfInvalid = (Text -> m ()) -> (() -> m ()) -> Either Text () -> m ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (Text -> IO ()) -> Text -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WorkerConfigException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO (WorkerConfigException -> IO ())
-> (Text -> WorkerConfigException) -> Text -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> WorkerConfigException
WorkerConfigException) () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text () -> m ())
-> (WorkerConfig n payload -> Either Text ())
-> WorkerConfig n payload
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WorkerConfig n payload -> Either Text ()
forall (m :: * -> *) payload.
WorkerConfig m payload -> Either Text ()
validateWorkerConfig

-- | The span shape for claims made by a pool.
poolSpanShape :: WorkerConfig m payload -> ConsumeShape
poolSpanShape :: forall (m :: * -> *) payload.
WorkerConfig m payload -> ConsumeShape
poolSpanShape = Int -> ConsumeShape
toConsumeShape (Int -> ConsumeShape)
-> (WorkerConfig m payload -> Int)
-> WorkerConfig m payload
-> ConsumeShape
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WorkerConfig m payload -> Int
forall (m :: * -> *) payload. WorkerConfig m payload -> Int
handlerBatchSize

-- | Reaper timing for a worker pool.
reaperPace :: WorkerConfig m payload -> MaintenancePace
reaperPace :: forall (m :: * -> *) payload.
WorkerConfig m payload -> MaintenancePace
reaperPace WorkerConfig m payload
config =
  MaintenancePace
    { paceWindow :: NominalDiffTime
paceWindow = WorkerConfig m payload -> NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> NominalDiffTime
reaperInterval WorkerConfig m payload
config
    , paceSparseWindow :: NominalDiffTime
paceSparseWindow = WorkerConfig m payload -> NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> NominalDiffTime
reaperSparseInterval WorkerConfig m payload
config
    , paceBucketIdle :: NominalDiffTime
paceBucketIdle = WorkerConfig m payload -> NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> NominalDiffTime
reaperBucketIdle WorkerConfig m payload
config
    }

-- | Run a worker pool until 'shutdownWorkerPool', then drain within
-- 'gracefulShutdownTimeout'. Throws 'WorkerConfigException' for an invalid config.
runWorkerPool
  :: forall payload m
   . ( Arb.RegistryAdmissionPolicies (RegistryOf m)
     , EncodeJobResult (ResultOf m payload)
     , QueueOperation m payload
     , RegistryTables (RegistryOf m)
     )
  => WorkerConfig m payload
  -> m ()
runWorkerPool :: forall payload (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m),
 EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> m ()
runWorkerPool WorkerConfig m payload
config = do
  WorkerConfig m payload -> m ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
throwIfInvalid WorkerConfig m payload
config
  let workerCap :: Int
workerCap = WorkerConfig m payload -> Int
forall (m :: * -> *) payload. WorkerConfig m payload -> Int
workerCount WorkerConfig m payload
config
      queueName :: Text
queueName = forall payload (m :: * -> *).
KnownSymbol (TableForPayload payload (RegistryOf m)) =>
Text
Arb.queueTable @payload @m
      -- Built once for the pool.
      consumeSpan :: ConsumeSpan
consumeSpan = Text -> ConsumeShape -> ConsumeSpan
consumeSpanFor Text
queueName (WorkerConfig m payload -> ConsumeShape
forall (m :: * -> *) payload.
WorkerConfig m payload -> ConsumeShape
poolSpanShape WorkerConfig m payload
config)

  schemaName <- m Text
forall (m :: * -> *). MonadArbiter m => m Text
getSchema
  workQueue <- newWorkQueue
  guard <- newHeartbeatGuard config
  let statements =
        forall payload.
JobPayload payload =>
Text
-> Text -> Int -> Int -> NominalDiffTime -> UUID -> JobStatements
Ops.mkJobStatements @payload
          Text
schemaName
          Text
queueName
          (WorkerConfig m payload -> Int
forall (m :: * -> *) payload. WorkerConfig m payload -> Int
handlerBatchSize WorkerConfig m payload
config)
          Int
workerCap
          (WorkerConfig m payload -> NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> NominalDiffTime
visibilityTimeout WorkerConfig m payload
config)
          (WorkerConfig m payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig m payload
config)
  mode <- poolMode config statements
  effectsFor <- poolEffects config statements consumeSpan

  tryAny (registerSelf config schemaName queueName)
    >>= either
      ( \SomeException
exception ->
          LogConfig -> Text -> SomeException -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> Text -> SomeException -> m ()
warnEx (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) Text
"Worker registry insert failed, starting paused" SomeException
exception
            m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (WorkerConfig m payload -> Bool -> STM ()
forall (n :: * -> *) payload.
WorkerConfig n payload -> Bool -> STM ()
writePause WorkerConfig m payload
config Bool
True)
      )
      (traverse_ (atomically . writePause config))

  dispatcherNotifVar <- newTVarIO Nothing
  cronRunVar <- newTVarIO False
  let handlers = WorkerConfig m payload
-> PoolGuard payload
-> Text
-> Text
-> TVar (Maybe Notification)
-> TVar Bool
-> [(ByteString, Notification -> m ())]
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload
-> PoolGuard payload
-> Text
-> Text
-> TVar (Maybe Notification)
-> TVar Bool
-> [(ByteString, Notification -> m ())]
poolChannels WorkerConfig m payload
config PoolGuard payload
guard Text
schemaName Text
queueName TVar (Maybe Notification)
dispatcherNotifVar TVar Bool
cronRunVar

  evalContT $ do
    withLivenessFile config
    mListener <- lift getListener
    listenerReady <- case mListener of
      Maybe Listener
Nothing -> do
        m () -> ContT () m ()
forall (m :: * -> *) a. Monad m => m a -> ContT () m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ContT () m ()) -> m () -> ContT () m ()
forall a b. (a -> b) -> a -> b
$ LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) LogLevel
Info Text
"No listen connection, running poll-only"
        STM Bool -> ContT () m (STM Bool)
forall a. a -> ContT () m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True)
      Just Listener
listener ->
        ((STM Bool -> m ()) -> m ()) -> ContT () m (STM Bool)
forall {k} (r :: k) (m :: k -> *) a.
((a -> m r) -> m r) -> ContT r m a
ContT (((STM Bool -> m ()) -> m ()) -> ContT () m (STM Bool))
-> ((STM Bool -> m ()) -> m ()) -> ContT () m (STM Bool)
forall a b. (a -> b) -> a -> b
$
          Listener
-> HubLog
-> [(ByteString, Notification -> m ())]
-> (STM Bool -> m ())
-> m ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
Listener
-> HubLog
-> [(ByteString, Notification -> m ())]
-> (STM Bool -> m a)
-> m a
Listen.withChannels Listener
listener (LogConfig -> HubLog
toHubLog (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config)) [(ByteString, Notification -> m ())]
handlers
    void . ContT $ Async.withAsync (publishListenerReady config listenerReady)
    let spawn = TVar WorkerState
-> LogConfig -> Text -> Text -> m () -> ContT () m (Async ())
forall (m :: * -> *) r.
MonadUnliftIO m =>
TVar WorkerState
-> LogConfig -> Text -> Text -> m () -> ContT r m (Async ())
spawnRetried (WorkerConfig m payload -> TVar WorkerState
forall (n :: * -> *) payload.
WorkerConfig n payload -> TVar WorkerState
workerStateVar WorkerConfig m payload
config) (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) Text
queueName
    heartbeat <-
      spawn "Worker heartbeat" $
        heartbeatLoop config schemaName queueName
    jobGuard <-
      spawn "Job heartbeat guard" $
        void (runHeartbeatGuard guard)
    dispatcher <-
      spawn "Dispatcher" $
        runDispatcher config workerCap statements workQueue dispatcherNotifVar
    workers <-
      replicateM workerCap
        $ spawn "Worker thread"
        $ workerLoop config guard mode effectsFor workQueue
    crons <-
      sequence
        [ spawn "Cron scheduler" $
            runCronScheduler (logConfig config) (workerStateVar config) cronRunVar schemaName queueName (cronJobs config)
        | not (null (cronJobs config))
        ]
    reaper <-
      spawn "Reaper" $
        reaperLoop (logConfig config) (onMaintenance config) (reaperPace config) (reaperTimeout config)

    (_, res) <- waitAnyCatch (dispatcher : jobGuard : reaper : heartbeat : crons <> workers)
    case res of
      Left SomeException
exception ->
        m () -> ContT () m ()
forall (m :: * -> *) a. Monad m => m a -> ContT () m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ContT () m ()) -> m () -> ContT () m ()
forall a b. (a -> b) -> a -> b
$ LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog (WorkerConfig m payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig m payload
config) LogLevel
Error (Text -> m ()) -> Text -> m ()
forall a b. (a -> b) -> a -> b
$ Text
"Thread pool exception: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> SomeException -> Text
displayEx SomeException
exception
      Right ()
_ -> () -> ContT () m ()
forall a. a -> ContT () m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

    lift $ shutdownPool config schemaName workQueue (void (Async.waitCatchSTM dispatcher))

-- | The pool's NOTIFY channel names and their handlers. The cron run channel is
-- subscribed only when the pool has cron jobs.
poolChannels
  :: (MonadIO m)
  => WorkerConfig n payload
  -> PoolGuard payload
  -> SchemaName
  -> Text
  -> TVar (Maybe Listen.Notification)
  -> TVar Bool
  -> [(ByteString, Listen.Notification -> m ())]
poolChannels :: forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload
-> PoolGuard payload
-> Text
-> Text
-> TVar (Maybe Notification)
-> TVar Bool
-> [(ByteString, Notification -> m ())]
poolChannels WorkerConfig n payload
config PoolGuard payload
guard Text
schemaName Text
queueName TVar (Maybe Notification)
dispatcherNotifVar TVar Bool
cronRunVar =
  [ (ByteString
createChannel, STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ())
-> (Notification -> STM ()) -> Notification -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TVar (Maybe Notification) -> Maybe Notification -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Maybe Notification)
dispatcherNotifVar (Maybe Notification -> STM ())
-> (Notification -> Maybe Notification) -> Notification -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Notification -> Maybe Notification
forall a. a -> Maybe a
Just)
  , (ByteString
pauseChannel, WorkerConfig n payload -> Notification -> m ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> Notification -> m ()
handlePauseNotif WorkerConfig n payload
config)
  , (ByteString
cancelChannel, WorkerConfig n payload -> PoolGuard payload -> Notification -> m ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> PoolGuard payload -> Notification -> m ()
handleCancelNotif WorkerConfig n payload
config PoolGuard payload
guard)
  ]
    [(ByteString, Notification -> m ())]
-> [(ByteString, Notification -> m ())]
-> [(ByteString, Notification -> m ())]
forall a. Semigroup a => a -> a -> a
<> [(ByteString
cronRunChannel, Set Text -> TVar Bool -> Notification -> m ()
forall (m :: * -> *).
MonadIO m =>
Set Text -> TVar Bool -> Notification -> m ()
handleCronRunNotif Set Text
cronNames TVar Bool
cronRunVar) | Bool -> Bool
not ([CronJob payload] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (WorkerConfig n payload -> [CronJob payload]
forall (m :: * -> *) payload.
WorkerConfig m payload -> [CronJob payload]
cronJobs WorkerConfig n payload
config))]
  where
    createChannel :: ByteString
createChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text
Schema.notificationChannelForTable Text
queueName)
    pauseChannel :: ByteString
pauseChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text -> Text
Schema.pauseNotifyChannel Text
schemaName Text
queueName)
    cancelChannel :: ByteString
cancelChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text -> Text
Schema.cancelNotifyChannel Text
schemaName Text
queueName)
    cronRunChannel :: ByteString
cronRunChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text
Schema.cronRunNotifyChannel Text
schemaName)
    cronNames :: Set Text
cronNames = [Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList ((CronJob payload -> Text) -> [CronJob payload] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map CronJob payload -> Text
forall payload. CronJob payload -> Text
name (WorkerConfig n payload -> [CronJob payload]
forall (m :: * -> *) payload.
WorkerConfig m payload -> [CronJob payload]
cronJobs WorkerConfig n payload
config))

-- | Flip 'listenerReadyVar' once the pool's channels are subscribed. Runs
-- alongside the pool. Startup does not wait on it.
publishListenerReady :: (MonadUnliftIO m) => WorkerConfig n payload -> STM Bool -> m ()
publishListenerReady :: forall (m :: * -> *) (n :: * -> *) payload.
MonadUnliftIO m =>
WorkerConfig n payload -> STM Bool -> m ()
publishListenerReady WorkerConfig n payload
config STM Bool
ready =
  STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    STM Bool
ready 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
    TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (WorkerConfig n payload -> TVar Bool
forall (n :: * -> *) payload. WorkerConfig n payload -> TVar Bool
listenerReadyVar WorkerConfig n payload
config) Bool
True

-- | Remove the liveness file when the pool exits, after the drain.
withLivenessFile :: (MonadUnliftIO m) => WorkerConfig n payload -> ContT r m ()
withLivenessFile :: forall (m :: * -> *) (n :: * -> *) payload r.
MonadUnliftIO m =>
WorkerConfig n payload -> ContT r m ()
withLivenessFile WorkerConfig n payload
config =
  (String -> ContT r m ()) -> Maybe String -> ContT r m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_
    (\String
path -> ((() -> m r) -> m r) -> ContT r m ()
forall {k} (r :: k) (m :: k -> *) a.
((a -> m r) -> m r) -> ContT r m a
ContT (((() -> m r) -> m r) -> ContT r m ())
-> ((() -> m r) -> m r) -> ContT r m ()
forall a b. (a -> b) -> a -> b
$ \() -> m r
continue -> () -> m r
continue () m r -> m () -> m r
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` (m (Either SomeException ()) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Either SomeException ()) -> m ())
-> (String -> m (Either SomeException ())) -> String -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m () -> m (Either SomeException ())
forall (m :: * -> *) a.
MonadUnliftIO m =>
m a -> m (Either SomeException a)
tryAny (m () -> m (Either SomeException ()))
-> (String -> m ()) -> String -> m (Either SomeException ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (String -> IO ()) -> String -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO ()
removeFile) String
path)
    (WorkerConfig n payload -> Maybe String
forall (m :: * -> *) payload.
WorkerConfig m payload -> Maybe String
livenessFile WorkerConfig n payload
config)

-- | Re-insert the worker's registry row, returning the effective paused state for the
-- caller to seed 'Arbiter.Worker.Config.pauseVar' with.
registerSelf :: (MonadArbiter m) => WorkerConfig n payload -> SchemaName -> Text -> m (Maybe Bool)
registerSelf :: forall (m :: * -> *) (n :: * -> *) payload.
MonadArbiter m =>
WorkerConfig n payload -> Text -> Text -> m (Maybe Bool)
registerSelf WorkerConfig n payload
config Text
schemaName Text
queueName =
  Text
-> UUID
-> Text
-> Maybe Text
-> Maybe Int32
-> NominalDiffTime
-> Maybe Value
-> m (Maybe Bool)
forall (m :: * -> *).
MonadArbiter m =>
Text
-> UUID
-> Text
-> Maybe Text
-> Maybe Int32
-> NominalDiffTime
-> Maybe Value
-> m (Maybe Bool)
Ops.registerWorker
    Text
schemaName
    (WorkerConfig n payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig n payload
config)
    Text
queueName
    (WorkerConfig n payload -> Maybe Text
forall (m :: * -> *) payload. WorkerConfig m payload -> Maybe Text
workerHost WorkerConfig n payload
config)
    (Int32 -> Maybe Int32
forall a. a -> Maybe a
Just (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (WorkerConfig n payload -> Int
forall (m :: * -> *) payload. WorkerConfig m payload -> Int
workerCount WorkerConfig n payload
config)))
    (WorkerConfig n payload -> NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> NominalDiffTime
workerStaleThreshold WorkerConfig n payload
config)
    (WorkerConfig n payload -> Maybe Value
forall (m :: * -> *) payload. WorkerConfig m payload -> Maybe Value
workerMetadata WorkerConfig n payload
config)

-- | Mark shutting-down, drain, then deregister. Every write is best-effort and logged
-- when it fails.
shutdownPool
  :: (MonadArbiter m)
  => WorkerConfig n payload
  -> SchemaName
  -> WorkQueue a
  -> STM ()
  -- ^ Retries until the dispatcher has stopped claiming.
  -> m ()
shutdownPool :: forall (m :: * -> *) (n :: * -> *) payload a.
MonadArbiter m =>
WorkerConfig n payload -> Text -> WorkQueue a -> STM () -> m ()
shutdownPool WorkerConfig n payload
config Text
schemaName WorkQueue a
workQueue STM ()
claimsStopped = do
  WorkerConfig n payload -> m ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
shutdownWorkerPool WorkerConfig n payload
config
  let wid :: UUID
wid = WorkerConfig n payload -> UUID
forall (m :: * -> *) payload. WorkerConfig m payload -> UUID
workerId WorkerConfig n payload
config
      logCfg :: LogConfig
logCfg = WorkerConfig n payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig n payload
config
  LogConfig -> Text -> m Int64 -> m ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig -> Text -> m a -> m ()
tryWarn LogConfig
logCfg Text
"Failed to mark worker shutting down" (Text -> UUID -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> UUID -> m Int64
Ops.markWorkerShuttingDown Text
schemaName UUID
wid)
  LogConfig -> Maybe NominalDiffTime -> WorkQueue a -> STM () -> m ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig -> Maybe NominalDiffTime -> WorkQueue a -> STM () -> m ()
drainPool LogConfig
logCfg (WorkerConfig n payload -> Maybe NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> Maybe NominalDiffTime
gracefulShutdownTimeout WorkerConfig n payload
config) WorkQueue a
workQueue STM ()
claimsStopped
  LogConfig -> Text -> m Int64 -> m ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig -> Text -> m a -> m ()
tryWarn LogConfig
logCfg Text
"Failed to deregister worker" (Text -> UUID -> m Int64
forall (m :: * -> *). MonadArbiter m => Text -> UUID -> m Int64
Ops.deregisterWorker Text
schemaName UUID
wid)

-- | How often an unbounded drain logs its progress.
drainProgressInterval :: NominalDiffTime
drainProgressInterval :: NominalDiffTime
drainProgressInterval = NominalDiffTime
10

-- | Wait for the dispatcher to stop, the work queue to drain and all worker threads to go idle,
-- optionally bounded by a timeout. Logs the entry, progress every
-- 'drainProgressInterval' when no timeout is set, and the result.
drainPool
  :: (MonadUnliftIO m)
  => LogConfig
  -> Maybe NominalDiffTime
  -> WorkQueue a
  -> STM ()
  -> m ()
drainPool :: forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig -> Maybe NominalDiffTime -> WorkQueue a -> STM () -> m ()
drainPool LogConfig
logCfg Maybe NominalDiffTime
mTimeout WorkQueue a
workQueue STM ()
claimsStopped = do
  LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog LogConfig
logCfg LogLevel
Info Text
"Starting graceful shutdown. Draining in-flight jobs..."
  result <- case Maybe NominalDiffTime
mTimeout of
    Maybe NominalDiffTime
Nothing -> () -> Either () ()
forall a b. b -> Either a b
Right () Either () () -> m () -> m (Either () ())
forall a b. a -> m b -> m a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ m ()
drainLoop
    Just NominalDiffTime
timeoutSecs ->
      m () -> m () -> m (Either () ())
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> m b -> m (Either a b)
Async.race (Int -> m ()
forall (m :: * -> *). MonadIO m => Int -> m ()
threadDelay (NominalDiffTime -> Int
Ops.micros NominalDiffTime
timeoutSecs)) m ()
waitForDrain
  case result of
    Right () -> LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog LogConfig
logCfg LogLevel
Info Text
"All workers are now idle. Graceful shutdown complete."
    Left () -> LogConfig -> LogLevel -> Text -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
LogConfig -> LogLevel -> Text -> m ()
tryLog LogConfig
logCfg LogLevel
Warning Text
"Graceful shutdown timed out. Some jobs may still be in-flight."
  where
    waitForDrain :: m ()
waitForDrain = STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM ()
claimsStopped STM () -> STM () -> STM ()
forall a b. STM a -> STM b -> STM b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> (WorkQueue a -> STM Int
forall a. WorkQueue a -> STM Int
inFlight WorkQueue a
workQueue STM Int -> (Int -> 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 (Bool -> STM ()) -> (Int -> Bool) -> Int -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0)))
    drainLoop :: m ()
drainLoop = do
      drainOrTick <- m () -> m () -> m (Either () ())
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> m b -> m (Either a b)
Async.race (Int -> m ()
forall (m :: * -> *). MonadIO m => Int -> m ()
threadDelay (NominalDiffTime -> Int
Ops.micros NominalDiffTime
drainProgressInterval)) m ()
waitForDrain
      case drainOrTick of
        Right () -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Left () -> do
          (busy, qLen) <- STM (Int, Int) -> m (Int, Int)
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically ((,) (Int -> Int -> (Int, Int)) -> STM Int -> STM (Int -> (Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> WorkQueue a -> STM Int
forall a. WorkQueue a -> STM Int
busyCount WorkQueue a
workQueue STM (Int -> (Int, Int)) -> STM Int -> STM (Int, Int)
forall a b. STM (a -> b) -> STM a -> STM b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> WorkQueue a -> STM Int
forall a. WorkQueue a -> STM Int
queuedCount WorkQueue a
workQueue)
          tryLog logCfg Info $
            "Graceful shutdown: waiting for "
              <> showText busy
              <> " busy worker(s), "
              <> showText qLen
              <> " job(s) in queue..."
          drainLoop

-- | Ticks at most once per 'workerHeartbeatInterval', gated by a proof-of-work signal unless paused.
-- Each tick bumps @arbiter_workers.last_heartbeat@, reconciles the registry
-- @paused@ flag into local state, and re-registers if the sweeper deleted
-- the row.
heartbeatLoop
  :: (MonadArbiter m)
  => WorkerConfig n payload
  -> SchemaName
  -> Text
  -- ^ Queue name (used when the row needs re-registering).
  -> m ()
heartbeatLoop :: forall (m :: * -> *) (n :: * -> *) payload.
MonadArbiter m =>
WorkerConfig n payload -> Text -> Text -> m ()
heartbeatLoop WorkerConfig n payload
config Text
schemaName Text
queueName = do
  gate <- m FailureGate
forall (m :: * -> *). MonadIO m => m FailureGate
newFailureGate
  tick gate
  forever $ throttledWait *> tick gate
  where
    logCfg :: LogConfig
logCfg = WorkerConfig n payload -> LogConfig
forall (m :: * -> *) payload. WorkerConfig m payload -> LogConfig
logConfig WorkerConfig n payload
config
    sig :: TMVar ()
sig = WorkerConfig n payload -> TMVar ()
forall (n :: * -> *) payload. WorkerConfig n payload -> TMVar ()
heartbeatSignal WorkerConfig n payload
config
    readShuttingDown :: STM Bool
readShuttingDown = TVar WorkerState -> STM Bool
isShuttingDown (WorkerConfig n payload -> TVar WorkerState
forall (n :: * -> *) payload.
WorkerConfig n payload -> TVar WorkerState
workerStateVar WorkerConfig n payload
config)
    cadenceMicros :: Int
cadenceMicros = NominalDiffTime -> Int
Ops.micros (WorkerConfig n payload -> NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> NominalDiffTime
workerHeartbeatInterval WorkerConfig n payload
config)
    throttledWait :: m ()
throttledWait = do
      delayVar <- Int -> m (TVar Bool)
forall (m :: * -> *). MonadIO m => Int -> m (TVar Bool)
registerDelay Int
cadenceMicros
      atomically $ do
        readTVar delayVar >>= checkSTM
        paused <- readTVar (pauseVar config)
        unless paused $ takeTMVar sig
    tick :: FailureGate -> m ()
tick FailureGate
gate = do
      (String -> m ()) -> Maybe String -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_
        (\String
path -> LogConfig -> Text -> m () -> m ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig -> Text -> m a -> m ()
tryWarn LogConfig
logCfg Text
"Liveness probe write failed" (IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> String -> IO ()
writeFile String
path String
""))
        (WorkerConfig n payload -> Maybe String
forall (m :: * -> *) payload.
WorkerConfig m payload -> Maybe String
livenessFile WorkerConfig n payload
config)
      epoch <- TVar Word64 -> m Word64
forall (m :: * -> *) a. MonadIO m => TVar a -> m a
readTVarIO (WorkerConfig n payload -> TVar Word64
forall (n :: * -> *) payload. WorkerConfig n payload -> TVar Word64
pauseEpoch WorkerConfig n payload
config)
      result <-
        tryReported logCfg Warning gate "Worker registry heartbeat" $
          Ops.heartbeatWorker schemaName (workerId config)
      traverse_ (maybe (reregister epoch) (reconcile epoch)) result
    reregister :: Word64 -> m ()
reregister Word64
epoch = do
      shutting <- STM Bool -> m Bool
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically STM Bool
readShuttingDown
      unless shutting $ do
        tryLog logCfg Warning "Worker registry row missing, re-registering"
        tryWarnWith logCfg "Worker re-registration failed" Nothing (registerSelf config schemaName queueName)
          >>= traverse_ (reconcile epoch)
    reconcile :: Word64 -> Bool -> m ()
reconcile Word64
epoch Bool
registryPaused =
      STM () -> m ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        shutting <- STM Bool
readShuttingDown
        unless shutting $ writePauseIfCurrent config epoch registryPaused