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

-- | Naming, sizing, and coordinated lifecycle for multiple worker pools.
module Arbiter.Worker.MultiQueue
  ( NamedWorkerPool (..)
  , namedWorkerPool
  , WorkerPoolSelectionException (..)
  , shutdownPools
  , runWorkerPools
  , runSelectedWorkerPools
  , poolConfigForWorkers
  ) where

import Arbiter.Core.HighLevel (QueueOperation)
import Arbiter.Core.HighLevel qualified as Arb
import Arbiter.Core.Job.Types (RegistryAdmissionPolicies)
import Arbiter.Core.JobResult (EncodeJobResult)
import Arbiter.Core.MonadArbiter (MonadArbiter (..), ResultOf)
import Arbiter.Core.PoolConfig (PoolConfig (..), defaultPoolConfig)
import Arbiter.Core.QueueRegistry (RegistryTables)
import Arbiter.Core.Threads (labelArbiterThread)
import Control.Exception qualified as E
import Control.Monad (unless)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Cont (ContT (..), evalContT)
import Data.Foldable (traverse_)
import Data.Maybe (fromMaybe)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import UnliftIO (MonadUnliftIO)
import UnliftIO.Async qualified as Async
import UnliftIO.STM qualified as STM

import Arbiter.Worker.Config
  ( WorkerConfig (..)
  , WorkerConfigException (..)
  , validateWorkerConfig
  , workerStateVar
  )
import Arbiter.Worker.EnabledQueues (enabledQueuesForMonad, requestedQueuesForMonad)
import Arbiter.Worker.Logger (LogConfig (..), (.=))
import Arbiter.Worker.Pool (runWorkerPool)
import Arbiter.Worker.WorkerState (WorkerState (ShuttingDown))

-- | A requested queue has no configured pool, or selection is empty.
newtype WorkerPoolSelectionException = WorkerPoolSelectionException Text
  deriving stock (WorkerPoolSelectionException
-> WorkerPoolSelectionException -> Bool
(WorkerPoolSelectionException
 -> WorkerPoolSelectionException -> Bool)
-> (WorkerPoolSelectionException
    -> WorkerPoolSelectionException -> Bool)
-> Eq WorkerPoolSelectionException
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WorkerPoolSelectionException
-> WorkerPoolSelectionException -> Bool
== :: WorkerPoolSelectionException
-> WorkerPoolSelectionException -> Bool
$c/= :: WorkerPoolSelectionException
-> WorkerPoolSelectionException -> Bool
/= :: WorkerPoolSelectionException
-> WorkerPoolSelectionException -> Bool
Eq, Int -> WorkerPoolSelectionException -> ShowS
[WorkerPoolSelectionException] -> ShowS
WorkerPoolSelectionException -> String
(Int -> WorkerPoolSelectionException -> ShowS)
-> (WorkerPoolSelectionException -> String)
-> ([WorkerPoolSelectionException] -> ShowS)
-> Show WorkerPoolSelectionException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> WorkerPoolSelectionException -> ShowS
showsPrec :: Int -> WorkerPoolSelectionException -> ShowS
$cshow :: WorkerPoolSelectionException -> String
show :: WorkerPoolSelectionException -> String
$cshowList :: [WorkerPoolSelectionException] -> ShowS
showList :: [WorkerPoolSelectionException] -> ShowS
Show)
  deriving anyclass (Show WorkerPoolSelectionException
Typeable WorkerPoolSelectionException
(Typeable WorkerPoolSelectionException,
 Show WorkerPoolSelectionException) =>
(WorkerPoolSelectionException -> SomeException)
-> (SomeException -> Maybe WorkerPoolSelectionException)
-> (WorkerPoolSelectionException -> String)
-> (WorkerPoolSelectionException -> Bool)
-> Exception WorkerPoolSelectionException
SomeException -> Maybe WorkerPoolSelectionException
WorkerPoolSelectionException -> Bool
WorkerPoolSelectionException -> String
WorkerPoolSelectionException -> SomeException
forall e.
(Typeable e, Show e) =>
(e -> SomeException)
-> (SomeException -> Maybe e)
-> (e -> String)
-> (e -> Bool)
-> Exception e
$ctoException :: WorkerPoolSelectionException -> SomeException
toException :: WorkerPoolSelectionException -> SomeException
$cfromException :: SomeException -> Maybe WorkerPoolSelectionException
fromException :: SomeException -> Maybe WorkerPoolSelectionException
$cdisplayException :: WorkerPoolSelectionException -> String
displayException :: WorkerPoolSelectionException -> String
$cbacktraceDesired :: WorkerPoolSelectionException -> Bool
backtraceDesired :: WorkerPoolSelectionException -> Bool
E.Exception)

-- | A worker pool paired with its registry-derived queue name.
data NamedWorkerPool m
  = forall payload.
  ( EncodeJobResult (ResultOf m payload)
  , QueueOperation m payload
  , RegistryAdmissionPolicies (RegistryOf m)
  , RegistryTables (RegistryOf m)
  ) =>
  NamedWorkerPool
  { forall (m :: * -> *). NamedWorkerPool m -> Text
workerPoolName :: Text
  , ()
workerPoolConfig :: WorkerConfig m payload
  }

-- | Name a pool from its payload's registry entry.
namedWorkerPool
  :: forall payload m
   . ( EncodeJobResult (ResultOf m payload)
     , QueueOperation m payload
     , RegistryAdmissionPolicies (RegistryOf m)
     , RegistryTables (RegistryOf m)
     )
  => WorkerConfig m payload
  -> NamedWorkerPool m
namedWorkerPool :: forall payload (m :: * -> *).
(EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryAdmissionPolicies (RegistryOf m),
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> NamedWorkerPool m
namedWorkerPool WorkerConfig m payload
cfg = Text -> WorkerConfig m payload -> NamedWorkerPool m
forall (m :: * -> *) payload.
(EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryAdmissionPolicies (RegistryOf m),
 RegistryTables (RegistryOf m)) =>
Text -> WorkerConfig m payload -> NamedWorkerPool m
NamedWorkerPool (forall payload (m :: * -> *).
KnownSymbol (TableForPayload payload (RegistryOf m)) =>
Text
Arb.queueTable @payload @m) WorkerConfig m payload
cfg

-- | Run the pools selected by @ARBITER_ENABLED_QUEUES@, or every configured
-- pool when it is unset.
runWorkerPools
  :: forall m
   . (MonadUnliftIO m, RegistryTables (RegistryOf m))
  => [NamedWorkerPool m]
  -> m ()
runWorkerPools :: forall (m :: * -> *).
(MonadUnliftIO m, RegistryTables (RegistryOf m)) =>
[NamedWorkerPool m] -> m ()
runWorkerPools [NamedWorkerPool m]
pools = do
  requested <- IO (Maybe [Text]) -> m (Maybe [Text])
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe [Text]) -> m (Maybe [Text]))
-> IO (Maybe [Text]) -> m (Maybe [Text])
forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
RegistryTables (RegistryOf m) =>
IO (Maybe [Text])
requestedQueuesForMonad @m
  runSelectedWorkerPools (fromMaybe [name | NamedWorkerPool name _ <- pools] requested) pools

-- | Signal graceful shutdown to every pool atomically.
shutdownPools :: (MonadIO m) => [NamedWorkerPool m'] -> m ()
shutdownPools :: forall (m :: * -> *) (m' :: * -> *).
MonadIO m =>
[NamedWorkerPool m'] -> m ()
shutdownPools [NamedWorkerPool m']
pools =
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (STM () -> IO ()) -> STM () -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. STM () -> IO ()
forall (m :: * -> *) a. MonadIO m => STM a -> m a
STM.atomically (STM () -> m ()) -> STM () -> m ()
forall a b. (a -> b) -> a -> b
$
    (TVar WorkerState -> STM ()) -> [TVar WorkerState] -> STM ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (TVar WorkerState -> WorkerState -> STM ()
forall a. TVar a -> a -> STM ()
`STM.writeTVar` WorkerState
ShuttingDown) [WorkerConfig m' payload -> TVar WorkerState
forall (n :: * -> *) payload.
WorkerConfig n payload -> TVar WorkerState
workerStateVar WorkerConfig m' payload
cfg | NamedWorkerPool Text
_ WorkerConfig m' payload
cfg <- [NamedWorkerPool m']
pools]

-- | Run only named pools. A pool that exits winds down its peers. The first
-- failure among them is rethrown after every peer has been joined.
runSelectedWorkerPools
  :: forall m
   . (MonadUnliftIO m)
  => [Text]
  -> [NamedWorkerPool m]
  -> m ()
runSelectedWorkerPools :: forall (m :: * -> *).
MonadUnliftIO m =>
[Text] -> [NamedWorkerPool m] -> m ()
runSelectedWorkerPools [Text]
enabled [NamedWorkerPool m]
pools = do
  let available :: Set Text
available = [Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList [Text
name | NamedWorkerPool Text
name WorkerConfig m payload
_ <- [NamedWorkerPool m]
pools]
      missing :: [Text]
missing = Set Text -> [Text]
forall a. Set a -> [a]
Set.toList ([Text] -> Set Text
forall a. Ord a => [a] -> Set a
Set.fromList [Text]
enabled Set Text -> Set Text -> Set Text
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set Text
available)
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless ([Text] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Text]
missing)
    (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ 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
. WorkerPoolSelectionException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO (WorkerPoolSelectionException -> IO ())
-> (Text -> WorkerPoolSelectionException) -> Text -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> WorkerPoolSelectionException
WorkerPoolSelectionException
    (Text -> m ()) -> Text -> m ()
forall a b. (a -> b) -> a -> b
$ Text
"No worker pool configured for: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> [Text] -> Text
T.intercalate Text
", " [Text]
missing
  case (NamedWorkerPool m -> Bool)
-> [NamedWorkerPool m] -> [NamedWorkerPool m]
forall a. (a -> Bool) -> [a] -> [a]
filter (\(NamedWorkerPool Text
name WorkerConfig m payload
_) -> Text
name Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
enabled) [NamedWorkerPool m]
pools of
    [] -> 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
$ WorkerPoolSelectionException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
E.throwIO (Text -> WorkerPoolSelectionException
WorkerPoolSelectionException Text
"No worker pools selected")
    [NamedWorkerPool m]
selected -> ContT () m () -> m ()
forall (m :: * -> *) r. Monad m => ContT r m r -> m r
evalContT (ContT () m () -> m ()) -> ContT () m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      asyncs <- (NamedWorkerPool m -> ContT () m (Async ()))
-> [NamedWorkerPool m] -> ContT () m [Async ()]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse NamedWorkerPool m -> ContT () m (Async ())
withPoolAsync [NamedWorkerPool m]
selected
      lift $ do
        _ <- Async.waitAnyCatch asyncs
        shutdownPools selected
        results <- traverse Async.waitCatch asyncs
        either (liftIO . E.throwIO) pure (sequence_ results)
  where
    withPoolAsync :: NamedWorkerPool m -> ContT () m (Async.Async ())
    withPoolAsync :: NamedWorkerPool m -> ContT () m (Async ())
withPoolAsync (NamedWorkerPool Text
name WorkerConfig m payload
cfg) =
      let cfg' :: WorkerConfig m payload
cfg' = WorkerConfig m payload
cfg {logConfig = withPoolContext name (logConfig cfg)}
       in ((Async () -> m ()) -> m ()) -> ContT () m (Async ())
forall {k} (r :: k) (m :: k -> *) a.
((a -> m r) -> m r) -> ContT r m a
ContT (((Async () -> m ()) -> m ()) -> ContT () m (Async ()))
-> ((Async () -> m ()) -> m ()) -> ContT () m (Async ())
forall a b. (a -> b) -> a -> b
$ m () -> (Async () -> m ()) -> m ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
m a -> (Async a -> m b) -> m b
Async.withAsync (Text -> Maybe Text -> m ()
forall (m :: * -> *). MonadIO m => Text -> Maybe Text -> m ()
labelArbiterThread Text
"pool" (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
name) m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> WorkerConfig m payload -> m ()
forall payload (m :: * -> *).
(RegistryAdmissionPolicies (RegistryOf m),
 EncodeJobResult (ResultOf m payload), QueueOperation m payload,
 RegistryTables (RegistryOf m)) =>
WorkerConfig m payload -> m ()
runWorkerPool WorkerConfig m payload
cfg')

-- | Name the pool a message came from.
withPoolContext :: Text -> LogConfig -> LogConfig
withPoolContext :: Text -> LogConfig -> LogConfig
withPoolContext Text
poolName LogConfig
logCfg =
  LogConfig
logCfg {identityContext = identityContext logCfg <> ["pool" .= poolName]}

-- | A single-stripe pool sized at twice the enabled worker count plus one for
-- the listener, with a minimum size of three.
poolConfigForWorkers
  :: forall m
   . (RegistryTables (RegistryOf m))
  => [NamedWorkerPool m]
  -> IO PoolConfig
poolConfigForWorkers :: forall (m :: * -> *).
RegistryTables (RegistryOf m) =>
[NamedWorkerPool m] -> IO PoolConfig
poolConfigForWorkers [NamedWorkerPool m]
pools = do
  enabled <- forall (m :: * -> *). RegistryTables (RegistryOf m) => IO [Text]
enabledQueuesForMonad @m
  traverse_ (validateSelected enabled) pools
  let enabledWorkers = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [WorkerConfig m payload -> Int
forall (m :: * -> *) payload. WorkerConfig m payload -> Int
workerCount WorkerConfig m payload
cfg | NamedWorkerPool Text
name WorkerConfig m payload
cfg <- [NamedWorkerPool m]
pools, Text
name Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
enabled]
  pure defaultPoolConfig {poolSize = max 2 (2 * enabledWorkers) + 1}
  where
    validateSelected :: [Text] -> NamedWorkerPool m -> IO ()
    validateSelected :: [Text] -> NamedWorkerPool m -> IO ()
validateSelected [Text]
enabled (NamedWorkerPool Text
name WorkerConfig m payload
cfg)
      | Text
name Text -> [Text] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Text]
enabled = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      | Bool
otherwise = (Text -> IO ()) -> (() -> IO ()) -> Either Text () -> IO ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (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) () -> IO ()
forall a. a -> IO 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
cfg)