-- | Worker-registry database operations.
module Arbiter.Core.Operations.Workers
  ( registerWorker
  , heartbeatWorker
  , setWorkerPaused
  , markWorkerShuttingDown
  , deregisterWorker
  , workerRegistered
  , listWorkers
  , sweepStaleWorkers
  ) where

import Data.Aeson (Value)
import Data.Int (Int32, Int64)
import Data.Maybe (listToMaybe)
import Data.Text (Text)
import Data.Time (NominalDiffTime)
import Data.UUID.Types (UUID)

import Arbiter.Core.Exceptions (throwParsing)
import Arbiter.Core.Job.Schema (SchemaName)
import Arbiter.Core.MonadArbiter (MonadArbiter)
import Arbiter.Core.MonadArbiter qualified as MA
import Arbiter.Core.Sql.Query (Query)
import Arbiter.Core.Sql.Workers qualified as Sql
import Arbiter.Core.Worker (WorkerRow (..), workerHealthFromText)

-- | Register or refresh a worker and return its effective paused state.
registerWorker
  :: (MonadArbiter m)
  => SchemaName
  -> UUID
  -> Text
  -> Maybe Text
  -> Maybe Int32
  -> NominalDiffTime
  -> Maybe Value
  -> m (Maybe Bool)
registerWorker :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName
-> UUID
-> SchemaName
-> Maybe SchemaName
-> Maybe Int32
-> NominalDiffTime
-> Maybe Value
-> m (Maybe Bool)
registerWorker SchemaName
schema UUID
workerId SchemaName
queue Maybe SchemaName
host Maybe Int32
threads NominalDiffTime
staleThreshold Maybe Value
metadata =
  [Bool] -> Maybe Bool
forall a. [a] -> Maybe a
listToMaybe
    ([Bool] -> Maybe Bool) -> m [Bool] -> m (Maybe Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Query Bool -> m [Bool]
forall a. Query a -> m [a]
forall (m :: * -> *) a. MonadArbiter m => Query a -> m [a]
MA.executeQuery
      (SchemaName
-> UUID
-> SchemaName
-> Maybe SchemaName
-> Maybe Int32
-> Double
-> Maybe Value
-> Query Bool
Sql.upsertWorkerSQL SchemaName
schema UUID
workerId SchemaName
queue Maybe SchemaName
host Maybe Int32
threads (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
staleThreshold) Maybe Value
metadata)

-- | Record a heartbeat and return the worker's effective paused state.
heartbeatWorker :: (MonadArbiter m) => SchemaName -> UUID -> m (Maybe Bool)
heartbeatWorker :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> UUID -> m (Maybe Bool)
heartbeatWorker SchemaName
schema UUID
workerId =
  [Bool] -> Maybe Bool
forall a. [a] -> Maybe a
listToMaybe ([Bool] -> Maybe Bool) -> m [Bool] -> m (Maybe Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Query Bool -> m [Bool]
forall a. Query a -> m [a]
forall (m :: * -> *) a. MonadArbiter m => Query a -> m [a]
MA.executeQuery (SchemaName -> UUID -> Query Bool
Sql.heartbeatWorkerSQL SchemaName
schema UUID
workerId)

-- | Set a worker's pause flag.
setWorkerPaused :: (MonadArbiter m) => SchemaName -> UUID -> Bool -> m Int64
setWorkerPaused :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> UUID -> Bool -> m Int64
setWorkerPaused SchemaName
schema UUID
workerId Bool
paused =
  Query Int64 -> m Int64
forall (m :: * -> *). MonadArbiter m => Query Int64 -> m Int64
countOrZero (SchemaName -> Bool -> UUID -> Query Int64
Sql.setWorkerPausedSQL SchemaName
schema Bool
paused UUID
workerId)

-- | Mark a worker as gracefully draining.
markWorkerShuttingDown :: (MonadArbiter m) => SchemaName -> UUID -> m Int64
markWorkerShuttingDown :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> UUID -> m Int64
markWorkerShuttingDown SchemaName
schema UUID
workerId =
  Query () -> m Int64
forall a. Query a -> m Int64
forall (m :: * -> *) a. MonadArbiter m => Query a -> m Int64
MA.executeStatement (SchemaName -> UUID -> Query ()
Sql.markWorkerShuttingDownSQL SchemaName
schema UUID
workerId)

-- | Remove a worker registry row.
deregisterWorker :: (MonadArbiter m) => SchemaName -> UUID -> m Int64
deregisterWorker :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> UUID -> m Int64
deregisterWorker SchemaName
schema UUID
workerId =
  Query () -> m Int64
forall a. Query a -> m Int64
forall (m :: * -> *) a. MonadArbiter m => Query a -> m Int64
MA.executeStatement (SchemaName -> UUID -> Query ()
Sql.deleteWorkerSQL SchemaName
schema UUID
workerId)

-- | Whether the worker registry holds this identity.
workerRegistered :: (MonadArbiter m) => SchemaName -> UUID -> m Bool
workerRegistered :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> UUID -> m Bool
workerRegistered SchemaName
schema UUID
workerId =
  [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ([Bool] -> Bool) -> m [Bool] -> m Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Query Bool -> m [Bool]
forall a. Query a -> m [a]
forall (m :: * -> *) a. MonadArbiter m => Query a -> m [a]
MA.executeQuery (SchemaName -> UUID -> Query Bool
Sql.workerRegisteredSQL SchemaName
schema UUID
workerId)

-- | List workers, optionally filtered by queue and heartbeat age.
listWorkers
  :: (MonadArbiter m)
  => SchemaName
  -> Maybe Text
  -> Maybe NominalDiffTime
  -> m [WorkerRow]
listWorkers :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName
-> Maybe SchemaName -> Maybe NominalDiffTime -> m [WorkerRow]
listWorkers SchemaName
schema Maybe SchemaName
queue Maybe NominalDiffTime
liveSecs = do
  rows <- Query (WorkerRow, SchemaName) -> m [(WorkerRow, SchemaName)]
forall a. Query a -> m [a]
forall (m :: * -> *) a. MonadArbiter m => Query a -> m [a]
MA.executeQuery (SchemaName
-> Maybe SchemaName
-> Maybe Double
-> Query (WorkerRow, SchemaName)
Sql.listWorkersSQL SchemaName
schema Maybe SchemaName
queue (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (NominalDiffTime -> Double)
-> Maybe NominalDiffTime -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe NominalDiffTime
liveSecs))
  traverse decodeHealth rows
  where
    decodeHealth :: (WorkerRow, SchemaName) -> m WorkerRow
decodeHealth (WorkerRow
worker, SchemaName
rawHealth) = do
      health <- (SchemaName -> m WorkerHealth)
-> (WorkerHealth -> m WorkerHealth)
-> Either SchemaName WorkerHealth
-> m WorkerHealth
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either SchemaName -> m WorkerHealth
forall (m :: * -> *) a. MonadIO m => SchemaName -> m a
throwParsing WorkerHealth -> m WorkerHealth
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SchemaName -> Either SchemaName WorkerHealth
workerHealthFromText SchemaName
rawHealth)
      pure worker {health}

-- | Delete workers older than their recorded stale threshold.
sweepStaleWorkers :: (MonadArbiter m) => SchemaName -> m Int64
sweepStaleWorkers :: forall (m :: * -> *). MonadArbiter m => SchemaName -> m Int64
sweepStaleWorkers SchemaName
schema = Query () -> m Int64
forall a. Query a -> m Int64
forall (m :: * -> *) a. MonadArbiter m => Query a -> m Int64
MA.executeStatement (SchemaName -> Query ()
Sql.deleteStaleWorkersSQL SchemaName
schema)

countOrZero :: (MonadArbiter m) => Query Int64 -> m Int64
countOrZero :: forall (m :: * -> *). MonadArbiter m => Query Int64 -> m Int64
countOrZero Query Int64
query = do
  rows <- Query Int64 -> m [Int64]
forall a. Query a -> m [a]
forall (m :: * -> *) a. MonadArbiter m => Query a -> m [a]
MA.executeQuery Query Int64
query
  pure $ case rows of
    [Int64
count] -> Int64
count
    [Int64]
_ -> Int64
0