-- | 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.Text (Text)
import Data.Time (NominalDiffTime)
import Data.UUID.Types (UUID)

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

-- | Register or refresh a worker and return its effective pause state (worker or
-- queue). 'Nothing' when the upsert returns no row.
registerWorker
  :: (MonadArbiter m)
  => SchemaName
  -> UUID
  -- ^ Worker id
  -> Text
  -- ^ Queue name
  -> Maybe Text
  -- ^ Host name
  -> Maybe Int32
  -- ^ Worker thread count
  -> NominalDiffTime
  -- ^ Heartbeat age after which the worker counts as stale
  -> Maybe Value
  -- ^ Free-form metadata
  -> m (Maybe Bool)
registerWorker :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName
-> UUID
-> SchemaName
-> Maybe SchemaName
-> Maybe Int32
-> NominalDiffTime
-> Maybe Value
-> m (Maybe Bool)
registerWorker SchemaName
schemaName UUID
workerId SchemaName
queue Maybe SchemaName
host Maybe Int32
threads NominalDiffTime
staleThreshold Maybe Value
metadata =
  Query Bool -> m (Maybe Bool)
forall (m :: * -> *) a. MonadArbiter m => Query a -> m (Maybe a)
MA.queryMaybe
    (SchemaName
-> UUID
-> SchemaName
-> Maybe SchemaName
-> Maybe Int32
-> Double
-> Maybe Value
-> Query Bool
Sql.upsertWorkerSQL SchemaName
schemaName 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 pause state (worker or queue).
-- 'Nothing' when the worker has no registry row.
heartbeatWorker :: (MonadArbiter m) => SchemaName -> UUID -> m (Maybe Bool)
heartbeatWorker :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> UUID -> m (Maybe Bool)
heartbeatWorker SchemaName
schemaName UUID
workerId =
  Query Bool -> m (Maybe Bool)
forall (m :: * -> *) a. MonadArbiter m => Query a -> m (Maybe a)
MA.queryMaybe (SchemaName -> UUID -> Query Bool
Sql.heartbeatWorkerSQL SchemaName
schemaName UUID
workerId)

-- | Set a worker's pause flag and notify the worker.
setWorkerPaused :: (MonadArbiter m) => SchemaName -> UUID -> Bool -> m Int64
setWorkerPaused :: forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> UUID -> Bool -> m Int64
setWorkerPaused SchemaName
schemaName UUID
workerId Bool
paused =
  Query Int64 -> m Int64
forall (m :: * -> *). MonadArbiter m => Query Int64 -> m Int64
countOr0 (SchemaName -> Bool -> UUID -> Query Int64
Sql.setWorkerPausedSQL SchemaName
schemaName 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
schemaName 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
schemaName 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
schemaName 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
schemaName 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
schemaName UUID
workerId =
  Query Bool -> m Bool
forall (m :: * -> *). MonadArbiter m => Query Bool -> m Bool
MA.queryAny (SchemaName -> UUID -> Query Bool
Sql.workerRegisteredSQL SchemaName
schemaName 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
schemaName 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
schemaName 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 <- Either SchemaName WorkerHealth -> m WorkerHealth
forall (m :: * -> *) a. MonadIO m => Either SchemaName a -> m a
orThrowParsing (SchemaName -> Either SchemaName WorkerHealth
workerHealthFromText SchemaName
rawHealth)
      pure worker {health}

-- | Delete workers whose last heartbeat is older than their stale threshold.
sweepStaleWorkers :: (MonadArbiter m) => SchemaName -> m Int64
sweepStaleWorkers :: forall (m :: * -> *). MonadArbiter m => SchemaName -> m Int64
sweepStaleWorkers SchemaName
schemaName = 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
schemaName)