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

-- | Types and DDL for the @arbiter_workers@ table.
module Arbiter.Core.Worker
  ( WorkerRow (..)
  , WorkerHealth (..)
  , workerHealthFromText
  , arbiterWorkersTable
  , arbiterWorkersTableName
  , createWorkersTableSQL
  , addClaimedByColumnSQL
  , addCancelRequestedAtColumnSQL
  , addArchiveForColumnSQL
  ) where

import Data.Aeson (FromJSON (..), ToJSON (..), Value, withText)
import Data.Aeson qualified as Aeson
import Data.Int (Int32)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Time (UTCTime)
import Data.UUID.Types (UUID)
import GHC.Generics (Generic)

import Arbiter.Core.Job.Schema (SchemaName, TableName, jobQueueDLQTable, jobQueueTable)
import Arbiter.Core.SqlLiterals (quoteIdentifier)

-- | Heartbeat-derived health of a worker. Independent of its 'paused' flag.
data WorkerHealth
  = Live
  | Stale
  | Draining
  deriving stock (WorkerHealth -> WorkerHealth -> Bool
(WorkerHealth -> WorkerHealth -> Bool)
-> (WorkerHealth -> WorkerHealth -> Bool) -> Eq WorkerHealth
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WorkerHealth -> WorkerHealth -> Bool
== :: WorkerHealth -> WorkerHealth -> Bool
$c/= :: WorkerHealth -> WorkerHealth -> Bool
/= :: WorkerHealth -> WorkerHealth -> Bool
Eq, (forall x. WorkerHealth -> Rep WorkerHealth x)
-> (forall x. Rep WorkerHealth x -> WorkerHealth)
-> Generic WorkerHealth
forall x. Rep WorkerHealth x -> WorkerHealth
forall x. WorkerHealth -> Rep WorkerHealth x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. WorkerHealth -> Rep WorkerHealth x
from :: forall x. WorkerHealth -> Rep WorkerHealth x
$cto :: forall x. Rep WorkerHealth x -> WorkerHealth
to :: forall x. Rep WorkerHealth x -> WorkerHealth
Generic, Int -> WorkerHealth -> ShowS
[WorkerHealth] -> ShowS
WorkerHealth -> String
(Int -> WorkerHealth -> ShowS)
-> (WorkerHealth -> String)
-> ([WorkerHealth] -> ShowS)
-> Show WorkerHealth
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> WorkerHealth -> ShowS
showsPrec :: Int -> WorkerHealth -> ShowS
$cshow :: WorkerHealth -> String
show :: WorkerHealth -> String
$cshowList :: [WorkerHealth] -> ShowS
showList :: [WorkerHealth] -> ShowS
Show)

instance ToJSON WorkerHealth where
  toJSON :: WorkerHealth -> Value
toJSON = \case
    WorkerHealth
Live -> Text -> Value
Aeson.String Text
"live"
    WorkerHealth
Stale -> Text -> Value
Aeson.String Text
"stale"
    WorkerHealth
Draining -> Text -> Value
Aeson.String Text
"draining"

instance FromJSON WorkerHealth where
  parseJSON :: Value -> Parser WorkerHealth
parseJSON = String
-> (Text -> Parser WorkerHealth) -> Value -> Parser WorkerHealth
forall a. String -> (Text -> Parser a) -> Value -> Parser a
withText String
"WorkerHealth" ((Text -> Parser WorkerHealth) -> Value -> Parser WorkerHealth)
-> (Text -> Parser WorkerHealth) -> Value -> Parser WorkerHealth
forall a b. (a -> b) -> a -> b
$ (Text -> Parser WorkerHealth)
-> (WorkerHealth -> Parser WorkerHealth)
-> Either Text WorkerHealth
-> Parser WorkerHealth
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (String -> Parser WorkerHealth
forall a. HasCallStack => String -> Parser a
forall (m :: * -> *) a.
(MonadFail m, HasCallStack) =>
String -> m a
fail (String -> Parser WorkerHealth)
-> (Text -> String) -> Text -> Parser WorkerHealth
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack) WorkerHealth -> Parser WorkerHealth
forall a. a -> Parser a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either Text WorkerHealth -> Parser WorkerHealth)
-> (Text -> Either Text WorkerHealth)
-> Text
-> Parser WorkerHealth
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Either Text WorkerHealth
workerHealthFromText

-- | Decode the @health@ SQL token into a 'WorkerHealth'.
workerHealthFromText :: Text -> Either Text WorkerHealth
workerHealthFromText :: Text -> Either Text WorkerHealth
workerHealthFromText = \case
  Text
"live" -> WorkerHealth -> Either Text WorkerHealth
forall a b. b -> Either a b
Right WorkerHealth
Live
  Text
"stale" -> WorkerHealth -> Either Text WorkerHealth
forall a b. b -> Either a b
Right WorkerHealth
Stale
  Text
"draining" -> WorkerHealth -> Either Text WorkerHealth
forall a b. b -> Either a b
Right WorkerHealth
Draining
  Text
other -> Text -> Either Text WorkerHealth
forall a b. a -> Either a b
Left (Text
"unknown worker health: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
other)

-- | A row in the worker registry. One row per running worker pool.
data WorkerRow = WorkerRow
  { WorkerRow -> UUID
workerId :: UUID
  , WorkerRow -> Text
queueName :: Text
  , WorkerRow -> Maybe Text
hostName :: Maybe Text
  , WorkerRow -> Maybe Int32
workerCount :: Maybe Int32
  , WorkerRow -> UTCTime
startedAt :: UTCTime
  , WorkerRow -> UTCTime
lastHeartbeat :: UTCTime
  , WorkerRow -> Bool
shuttingDown :: Bool
  , WorkerRow -> Bool
paused :: Bool
  , WorkerRow -> Double
staleThresholdSecs :: Double
  , WorkerRow -> Maybe Value
metadata :: Maybe Value
  , WorkerRow -> WorkerHealth
health :: WorkerHealth
  }
  deriving stock (WorkerRow -> WorkerRow -> Bool
(WorkerRow -> WorkerRow -> Bool)
-> (WorkerRow -> WorkerRow -> Bool) -> Eq WorkerRow
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: WorkerRow -> WorkerRow -> Bool
== :: WorkerRow -> WorkerRow -> Bool
$c/= :: WorkerRow -> WorkerRow -> Bool
/= :: WorkerRow -> WorkerRow -> Bool
Eq, (forall x. WorkerRow -> Rep WorkerRow x)
-> (forall x. Rep WorkerRow x -> WorkerRow) -> Generic WorkerRow
forall x. Rep WorkerRow x -> WorkerRow
forall x. WorkerRow -> Rep WorkerRow x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. WorkerRow -> Rep WorkerRow x
from :: forall x. WorkerRow -> Rep WorkerRow x
$cto :: forall x. Rep WorkerRow x -> WorkerRow
to :: forall x. Rep WorkerRow x -> WorkerRow
Generic, Int -> WorkerRow -> ShowS
[WorkerRow] -> ShowS
WorkerRow -> String
(Int -> WorkerRow -> ShowS)
-> (WorkerRow -> String)
-> ([WorkerRow] -> ShowS)
-> Show WorkerRow
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> WorkerRow -> ShowS
showsPrec :: Int -> WorkerRow -> ShowS
$cshow :: WorkerRow -> String
show :: WorkerRow -> String
$cshowList :: [WorkerRow] -> ShowS
showList :: [WorkerRow] -> ShowS
Show)
  deriving anyclass (Maybe WorkerRow
Value -> Parser [WorkerRow]
Value -> Parser WorkerRow
(Value -> Parser WorkerRow)
-> (Value -> Parser [WorkerRow])
-> Maybe WorkerRow
-> FromJSON WorkerRow
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser WorkerRow
parseJSON :: Value -> Parser WorkerRow
$cparseJSONList :: Value -> Parser [WorkerRow]
parseJSONList :: Value -> Parser [WorkerRow]
$comittedField :: Maybe WorkerRow
omittedField :: Maybe WorkerRow
FromJSON, [WorkerRow] -> Value
[WorkerRow] -> Encoding
WorkerRow -> Bool
WorkerRow -> Value
WorkerRow -> Encoding
(WorkerRow -> Value)
-> (WorkerRow -> Encoding)
-> ([WorkerRow] -> Value)
-> ([WorkerRow] -> Encoding)
-> (WorkerRow -> Bool)
-> ToJSON WorkerRow
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: WorkerRow -> Value
toJSON :: WorkerRow -> Value
$ctoEncoding :: WorkerRow -> Encoding
toEncoding :: WorkerRow -> Encoding
$ctoJSONList :: [WorkerRow] -> Value
toJSONList :: [WorkerRow] -> Value
$ctoEncodingList :: [WorkerRow] -> Encoding
toEncodingList :: [WorkerRow] -> Encoding
$comitField :: WorkerRow -> Bool
omitField :: WorkerRow -> Bool
ToJSON)

-- | Qualified table name for the arbiter_workers table.
arbiterWorkersTable :: SchemaName -> Text
arbiterWorkersTable :: Text -> Text
arbiterWorkersTable Text
schemaName = Text -> Text
quoteIdentifier Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
arbiterWorkersTableName

-- | Bare name of the workers table, for catalog lookups by relname.
arbiterWorkersTableName :: Text
arbiterWorkersTableName :: Text
arbiterWorkersTableName = Text
"arbiter_workers"

-- | DDL for the @arbiter_workers@ table.
createWorkersTableSQL :: SchemaName -> Text
createWorkersTableSQL :: Text -> Text
createWorkersTableSQL Text
schemaName =
  [Text] -> Text
T.unlines
    [ Text
"CREATE TABLE IF NOT EXISTS " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
arbiterWorkersTable Text
schemaName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ("
    , Text
"  worker_id UUID PRIMARY KEY,"
    , Text
"  queue_name TEXT NOT NULL,"
    , Text
"  host_name TEXT,"
    , Text
"  worker_count INT,"
    , Text
"  started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),"
    , Text
"  last_heartbeat TIMESTAMPTZ NOT NULL DEFAULT NOW(),"
    , Text
"  shutting_down BOOLEAN NOT NULL DEFAULT FALSE,"
    , Text
"  paused BOOLEAN NOT NULL DEFAULT FALSE,"
    , Text
"  stale_threshold_secs DOUBLE PRECISION NOT NULL DEFAULT 300,"
    , Text
"  metadata JSONB"
    , Text
");"
    ]

-- | Idempotent migration adding the @claimed_by@ column to a queue's job and DLQ tables.
addClaimedByColumnSQL :: SchemaName -> TableName -> Text
addClaimedByColumnSQL :: Text -> Text -> Text
addClaimedByColumnSQL Text
schemaName Text
tableName =
  [Text] -> Text
T.unlines
    [ Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS claimed_by UUID;"
    , Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueDLQTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS claimed_by UUID;"
    ]

-- | Idempotent migration adding the @cancel_requested_at@ column to a queue's
-- job table, with a partial index backing the reaper's flagged-job sweep.
addCancelRequestedAtColumnSQL :: SchemaName -> TableName -> Text
addCancelRequestedAtColumnSQL :: Text -> Text -> Text
addCancelRequestedAtColumnSQL Text
schemaName Text
tableName =
  [Text] -> Text
T.unlines
    [ Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS cancel_requested_at TIMESTAMPTZ;"
    , Text
"CREATE INDEX IF NOT EXISTS "
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text
quoteIdentifier (Text
"idx_" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_cancel_requested")
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ON "
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
        Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" (id ASC) WHERE cancel_requested_at IS NOT NULL;"
    ]

-- | Idempotent migration adding the @archive_for@ column to a queue's job and DLQ tables.
addArchiveForColumnSQL :: SchemaName -> TableName -> Text
addArchiveForColumnSQL :: Text -> Text -> Text
addArchiveForColumnSQL Text
schemaName Text
tableName =
  [Text] -> Text
T.unlines
    [ Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS archive_for INT;"
    , Text
"ALTER TABLE " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text
jobQueueDLQTable Text
schemaName Text
tableName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" ADD COLUMN IF NOT EXISTS archive_for INT;"
    ]