{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Naming, sizing, and coordinated lifecycle for multiple worker pools.
module Arbiter.Worker.MultiQueue
  ( NamedWorkerPool (..)
  , namedWorkerPool
  , WorkerPoolSelectionException (..)
  , shutdownWorkerPools
  , 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.Proxy (Proxy (..))
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 (..)
  , workerStateVar
  )
import Arbiter.Worker.EnabledQueues (getEnabledQueues, requestedQueues)
import Arbiter.Worker.Logger (withIdentity, (.=))
import Arbiter.Worker.Pool (runWorkerPool, throwIfInvalid)
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
  -- ^ The queue name from the registry.
  , ()
workerPoolConfig :: WorkerConfig m payload
  -- ^ The pool's config. Get it by pattern match.
  }

-- | 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 = NamedWorkerPool {workerPoolName :: Text
workerPoolName = forall payload (m :: * -> *).
KnownSymbol (TableForPayload payload (RegistryOf m)) =>
Text
Arb.queueTable @payload @m, workerPoolConfig :: WorkerConfig m payload
workerPoolConfig = WorkerConfig m payload
cfg}

-- | Run the pools selected by @ARBITER_ENABLED_QUEUES@, or every configured
-- pool when it is unset or blank. Throws 'Arbiter.Core.Exceptions.InternalException'
-- when @ARBITER_ENABLED_QUEUES@ has an unknown name or names no queues, and
-- 'WorkerPoolSelectionException' for a selected queue with no pool, or when no pool
-- is selected.
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
$ Proxy (RegistryOf m) -> IO (Maybe [Text])
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> IO (Maybe [Text])
requestedQueues (forall (t :: JobPayloadRegistry). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(RegistryOf m))
  runSelectedWorkerPools (fromMaybe (map workerPoolName pools) requested) pools

-- | Signal graceful shutdown to every pool atomically.
shutdownWorkerPools :: (MonadIO m) => [NamedWorkerPool m'] -> m ()
shutdownWorkerPools :: forall (m :: * -> *) (m' :: * -> *).
MonadIO m =>
[NamedWorkerPool m'] -> m ()
shutdownWorkerPools [NamedWorkerPool m']
pools =
  STM () -> m ()
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. After every peer has
-- been joined, the failure of the first failed pool in list order is rethrown. Throws
-- 'WorkerPoolSelectionException' for a name with no pool, or when no pool is selected.
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 ((NamedWorkerPool m -> Text) -> [NamedWorkerPool m] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map NamedWorkerPool m -> Text
forall (m :: * -> *). NamedWorkerPool m -> Text
workerPoolName [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 [Text] -> [NamedWorkerPool m] -> [NamedWorkerPool m]
forall (m :: * -> *).
[Text] -> [NamedWorkerPool m] -> [NamedWorkerPool m]
selectPools [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
        shutdownWorkerPools 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 = withIdentity ("pool" .= 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')

-- | Pool connections per handler thread of an enabled pool.
connectionsPerWorker :: Int
connectionsPerWorker :: Int
connectionsPerWorker = Int
2

-- | Floor on the worker share of the pool.
minWorkerConnections :: Int
minWorkerConnections :: Int
minWorkerConnections = Int
2

-- | Pool connections for the listener.
listenerConnections :: Int
listenerConnections :: Int
listenerConnections = Int
1

-- | A single-stripe connection pool sized for the enabled pools' threads and the listener.
-- Unset or blank @ARBITER_ENABLED_QUEUES@ enables all queues. Throws
-- 'Arbiter.Worker.Config.WorkerConfigException' for an invalid enabled pool config, and
-- 'Arbiter.Core.Exceptions.InternalException' when @ARBITER_ENABLED_QUEUES@ has an unknown
-- name or names no queues.
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 <- Proxy (RegistryOf m) -> IO [Text]
forall (registry :: JobPayloadRegistry).
RegistryTables registry =>
Proxy registry -> IO [Text]
getEnabledQueues (forall (t :: JobPayloadRegistry). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(RegistryOf m))
  let selected = [Text] -> [NamedWorkerPool m] -> [NamedWorkerPool m]
forall (m :: * -> *).
[Text] -> [NamedWorkerPool m] -> [NamedWorkerPool m]
selectPools [Text]
enabled [NamedWorkerPool m]
pools
      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
_ WorkerConfig m payload
cfg <- [NamedWorkerPool m]
selected]
  traverse_ (\(NamedWorkerPool Text
_ WorkerConfig m payload
cfg) -> WorkerConfig m payload -> IO ()
forall (m :: * -> *) (n :: * -> *) payload.
MonadIO m =>
WorkerConfig n payload -> m ()
throwIfInvalid WorkerConfig m payload
cfg) selected
  pure
    defaultPoolConfig {poolSize = max minWorkerConnections (connectionsPerWorker * enabledWorkers) + listenerConnections}

-- | The pools whose name is in the list, in pool order.
selectPools :: [Text] -> [NamedWorkerPool m] -> [NamedWorkerPool m]
selectPools :: forall (m :: * -> *).
[Text] -> [NamedWorkerPool m] -> [NamedWorkerPool m]
selectPools [Text]
enabled = (NamedWorkerPool m -> Bool)
-> [NamedWorkerPool m] -> [NamedWorkerPool m]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
enabled) (Text -> Bool)
-> (NamedWorkerPool m -> Text) -> NamedWorkerPool m -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NamedWorkerPool m -> Text
forall (m :: * -> *). NamedWorkerPool m -> Text
workerPoolName)