{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

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

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)
import Control.Exception qualified as E
import Control.Monad (forever, replicateM, unless, void)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Cont (ContT (..), evalContT)
import Data.Bool (bool)
import Data.Foldable (traverse_)
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Time (NominalDiffTime)
import System.Directory (removeFile)
import UnliftIO
  ( MonadUnliftIO
  , atomically
  , checkSTM
  , finally
  , isEmptyTBQueue
  , lengthTBQueue
  , newTBQueueIO
  , newTVarIO
  , readTVar
  , tryAny
  , waitAnyCatch
  , writeTVar
  )
import UnliftIO.Async qualified as Async
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.STM (STM, TBQueue, TVar)
import UnliftIO.STM qualified as STM

import Arbiter.Worker.ChannelHandlers
  ( handleCancelNotif
  , handleCronRunNotif
  , handlePauseNotif
  )
import Arbiter.Worker.Config
import Arbiter.Worker.Cron (CronJob (..), runCronScheduler)
import Arbiter.Worker.Dispatcher
import Arbiter.Worker.Logger
import Arbiter.Worker.Logger.Internal (tryWarn, tryWarnWith)
import Arbiter.Worker.Processing (workerLoop)
import Arbiter.Worker.Reaper (MaintenancePace (..), reaperLoop, runReaperOp)
import Arbiter.Worker.Retry (spawnRetried)

-- ---------------------------------------------------------------------------
-- Worker Pool
-- ---------------------------------------------------------------------------

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

-- | The pace the pool's reaper keeps.
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: a dispatcher and its worker threads.
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
  (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 (WorkerConfig m payload -> Either Text ()
forall (m :: * -> *) payload.
WorkerConfig m payload -> Either Text ()
validateWorkerConfig 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 <- newTBQueueIO (fromIntegral workerCap)
  busyWorkerCount <- newTVarIO 0
  workerFinishedVar <- newTVarIO False
  runningJobs <- STM.newTVarIO Map.empty

  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 <- STM.newTVarIO Nothing
  cronRunVar <- STM.newTVarIO False
  let createChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text
Schema.notificationChannelForTable Text
queueName)
      pauseChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text -> Text
Schema.pauseNotifyChannel Text
schemaName Text
queueName)
      cancelChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text -> Text
Schema.cancelNotifyChannel Text
schemaName Text
queueName)
      cronRunChannel = Text -> ByteString
TE.encodeUtf8 (Text -> Text
Schema.cronRunNotifyChannel Text
schemaName)
      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 m payload -> [CronJob payload]
forall (m :: * -> *) payload.
WorkerConfig m payload -> [CronJob payload]
cronJobs WorkerConfig m payload
config))
      cronHandlers =
        if [CronJob payload] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (WorkerConfig m payload -> [CronJob payload]
forall (m :: * -> *) payload.
WorkerConfig m payload -> [CronJob payload]
cronJobs WorkerConfig m payload
config)
          then []
          else [(ByteString
cronRunChannel, Set Text -> TVar Bool -> Notification -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
Set Text -> TVar Bool -> Notification -> m ()
handleCronRunNotif Set Text
cronNames TVar Bool
cronRunVar)]
      handlers =
        [ (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 ()
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 m payload -> Notification -> m ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadUnliftIO m =>
WorkerConfig n payload -> Notification -> m ()
handlePauseNotif WorkerConfig m payload
config)
        , (ByteString
cancelChannel, WorkerConfig m payload
-> TVar (Map Int64 (Async ())) -> Notification -> m ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadUnliftIO m =>
WorkerConfig n payload
-> TVar (Map Int64 (Async ())) -> Notification -> m ()
handleCancelNotif WorkerConfig m payload
config TVar (Map Int64 (Async ()))
runningJobs)
        ]
          [(ByteString, Notification -> m ())]
-> [(ByteString, Notification -> m ())]
-> [(ByteString, Notification -> m ())]
forall a. Semigroup a => a -> a -> a
<> [(ByteString, Notification -> m ())]
cronHandlers

  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
hubLogFor (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
    dispatcher <-
      spawn "Dispatcher" $
        runDispatcher config workerCap workQueue busyWorkerCount workerFinishedVar dispatcherNotifVar
    workers <-
      replicateM workerCap
        $ spawn "Worker thread"
        $ workerLoop config consumeSpan runningJobs workQueue busyWorkerCount workerFinishedVar
    crons <-
      unlessNull (cronJobs config)
        $ spawn "Cron scheduler"
        $ runCronScheduler (workerStateVar config) cronRunVar (logConfig config) schemaName queueName (cronJobs config)
    reaper <-
      spawn "Reaper" $
        reaperLoop (logConfig config) (onMaintenance config) (reaperPace config) (reaperTimeout config)

    (_, res) <- waitAnyCatch (dispatcher : 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 busyWorkerCount

-- | 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 '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
  -> TBQueue a
  -> TVar Int
  -> m ()
shutdownPool :: forall (m :: * -> *) (n :: * -> *) payload a.
MonadArbiter m =>
WorkerConfig n payload -> Text -> TBQueue a -> TVar Int -> m ()
shutdownPool WorkerConfig n payload
config Text
schemaName TBQueue a
workQueue TVar Int
busyCount = do
  WorkerConfig n payload -> m ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
shutdownWorker 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 -> TBQueue a -> TVar Int -> m ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig -> Maybe NominalDiffTime -> TBQueue a -> TVar Int -> m ()
drainPool LogConfig
logCfg (WorkerConfig n payload -> Maybe NominalDiffTime
forall (m :: * -> *) payload.
WorkerConfig m payload -> Maybe NominalDiffTime
gracefulShutdownTimeout WorkerConfig n payload
config) TBQueue a
workQueue TVar Int
busyCount
  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)

-- | Wait for the work queue to drain and all worker threads to go idle,
-- optionally bounded by a timeout. Logs the entry, periodic progress (every
-- 10s) when no timeout is set, and the result.
drainPool
  :: (MonadUnliftIO m)
  => LogConfig
  -> Maybe NominalDiffTime
  -> TBQueue a
  -> TVar Int
  -> m ()
drainPool :: forall (m :: * -> *) a.
MonadUnliftIO m =>
LogConfig -> Maybe NominalDiffTime -> TBQueue a -> TVar Int -> m ()
drainPool LogConfig
logCfg Maybe NominalDiffTime
mTimeout TBQueue a
workQueue TVar Int
busyCount = 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 () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      qEmpty <- TBQueue a -> STM Bool
forall a. TBQueue a -> STM Bool
isEmptyTBQueue TBQueue a
workQueue
      checkSTM qEmpty
      busy <- readTVar busyCount
      checkSTM (busy == 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 Int
10_000_000) 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 (STM (Int, Int) -> m (Int, Int)) -> STM (Int, Int) -> m (Int, Int)
forall a b. (a -> b) -> a -> b
$
              (,)
                (Int -> Int -> (Int, Int)) -> STM Int -> STM (Int -> (Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
busyCount
                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
<*> (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Natural -> Int) -> STM Natural -> STM Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TBQueue a -> STM Natural
forall a. TBQueue a -> STM Natural
lengthTBQueue TBQueue a
workQueue)
          tryLog logCfg Info $
            "Graceful shutdown: waiting for "
              <> T.pack (show (busy :: Int))
              <> " busy worker(s), "
              <> T.pack (show (qLen :: Int))
              <> " job(s) in queue..."
          drainLoop

-- | Ticks at @pollInterval@, 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 = (WorkerState -> WorkerState -> Bool
forall a. Eq a => a -> a -> Bool
== WorkerState
ShuttingDown) (WorkerState -> Bool) -> STM WorkerState -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar WorkerState -> STM WorkerState
forall a. TVar a -> STM a
STM.readTVar (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)
STM.registerDelay Int
cadenceMicros
      atomically $ do
        STM.readTVar delayVar >>= checkSTM
        paused <- STM.readTVar (pauseVar config)
        unless paused $ STM.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 <- STM Word64 -> m Word64
forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically (STM Word64 -> m Word64) -> STM Word64 -> m Word64
forall a b. (a -> b) -> a -> b
$ TVar Word64 -> STM Word64
forall a. TVar a -> STM a
STM.readTVar (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
STM.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
STM.atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        shutting <- STM Bool
readShuttingDown
        unless shutting $ writePauseIfCurrent config epoch registryPaused

-- | @[]@ when the list is empty, else a singleton holding @act@'s result.
unlessNull :: (Applicative f) => [a] -> f b -> f [b]
unlessNull :: forall (f :: * -> *) a b. Applicative f => [a] -> f b -> f [b]
unlessNull [a]
items f b
act = if [a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
items then [b] -> f [b]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [] else (b -> [b] -> [b]
forall a. a -> [a] -> [a]
: []) (b -> [b]) -> f b -> f [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f b
act