{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# OPTIONS_HADDOCK not-home #-}

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Stats SQL templates.
module Arbiter.Core.Sql.Stats
  ( getQueueStatsSQL
  , allQueueStatsSQL
  , countChildrenBatchSQL
  , listGroupsSQL
  , countGroupsSQL
  ) where

import Data.Int (Int64)
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Data.Text qualified as T
import NeatInterpolation (text)

import Arbiter.Core.Admission (effectivePolicyCol)
import Arbiter.Core.Codec (RowCodec)
import Arbiter.Core.Concurrency.Schema (arbiterConcurrencyPoliciesTable, arbiterConcurrencyTable)
import Arbiter.Core.Job.Schema (SchemaName, TableName, jobQueueDLQTable, jobQueueGroupsTable, jobQueueTable)
import Arbiter.Core.Job.Schema.Groups (inFlightPredicate)
import Arbiter.Core.Job.Status (JobStatus (Ready), jobStatusToText)
import Arbiter.Core.Queues (arbiterQueuesTable)
import Arbiter.Core.RateLimit.Schema (arbiterRateLimitPoliciesTable, arbiterRateLimitsTable)
import Arbiter.Core.Sql.Claim (concHeadroomPred, groupHeadBatch)
import Arbiter.Core.Sql.Jobs (claimablePred, jobStatusCaseSQL, unionAllOverQueueTables)
import Arbiter.Core.Sql.QQ (sql)
import Arbiter.Core.Sql.Query (Query, rawRows, rows)
import Arbiter.Core.Sql.RateLimit (refilledBucketTokens)
import Arbiter.Core.Sql.Workers (workerStaleSQL)
import Arbiter.Core.SqlLiterals (textLiteral)
import Arbiter.Core.Worker (arbiterWorkersTable)

-- | Per-status queue counts plus the age of the oldest @ready@ and @in_flight@ job.
-- The ready age covers ready and blocked rows and runs from when the row became visible.
-- Counts follow the 'jobStatusCaseSQL' taxonomy, with blocked rows split out of ready,
-- and sum to @total_jobs@.
getQueueStatsSQL :: RowCodec a -> SchemaName -> TableName -> [Text] -> Query a
getQueueStatsSQL :: forall a. RowCodec a -> Text -> Text -> [Text] -> Query a
getQueueStatsSQL RowCodec a
codec Text
schemaName Text
tableName [Text]
kinds = RowCodec a -> Text -> Query a
forall a. RowCodec a -> Text -> Query a
rawRows RowCodec a
codec (DlqCount -> Text -> Text -> [Text] -> Text
queueStatsSelect DlqCount
DlqByKind Text
schemaName Text
tableName [Text]
kinds)

-- | How a stats row counts the DLQ.
data DlqCount = DlqByKind | DlqTotal

-- | One queue's stats row. The classified rows are aggregated once per kind and once
-- over the whole table in one pass. @total_row@ marks the whole-table row. A ready
-- row a claim would skip counts as blocked instead.
queueStatsSelect :: DlqCount -> SchemaName -> TableName -> [Text] -> Text
queueStatsSelect :: DlqCount -> Text -> Text -> [Text] -> Text
queueStatsSelect DlqCount
dlqCount Text
schemaName Text
tableName [Text]
kinds =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      dlqTbl :: Text
dlqTbl = Text -> Text -> Text
jobQueueDLQTable Text
schemaName Text
tableName
      kindExpr :: Text
kindExpr = [Text] -> Text
declaredKindSQL [Text]
kinds
      heads :: Text
heads = Text -> Text -> Text
groupHeadsSQL Text
schemaName Text
tableName
      blocked :: Text
blocked = Text -> Text
blockedExpr Text
schemaName
      (Text
dlqCte, Text
dlqTotal, Text
dlqKinds) = case DlqCount
dlqCount of
        DlqCount
DlqByKind ->
          ( [text|WITH dlq_by_kind AS (SELECT ${kindExpr} AS kind, COUNT(*)::int8 AS dlq_count FROM ${dlqTbl} GROUP BY 1)|]
          , Text
"(SELECT COALESCE(SUM(dlq_count), 0)::int8 FROM dlq_by_kind)"
          , Text
"(SELECT jsonb_object_agg(kind, dlq_count) FILTER (WHERE kind IS NOT NULL) FROM dlq_by_kind)"
          )
        DlqCount
DlqTotal -> (Text
"", [text|(SELECT COUNT(*)::int8 FROM ${dlqTbl})|], Text
"NULL::jsonb")
   in [text|
        ${dlqCte}
        SELECT MAX(total_jobs) FILTER (WHERE total_row = 1) AS total_jobs,
               MAX(ready_jobs) FILTER (WHERE total_row = 1) AS ready_jobs,
               MAX(in_flight_jobs) FILTER (WHERE total_row = 1) AS in_flight_jobs,
               MAX(scheduled_jobs) FILTER (WHERE total_row = 1) AS scheduled_jobs,
               MAX(backoff_jobs) FILTER (WHERE total_row = 1) AS backoff_jobs,
               MAX(throttled_jobs) FILTER (WHERE total_row = 1) AS throttled_jobs,
               MAX(suspended_jobs) FILTER (WHERE total_row = 1) AS suspended_jobs,
               MAX(cancelled_jobs) FILTER (WHERE total_row = 1) AS cancelled_jobs,
               MAX(exhausted_jobs) FILTER (WHERE total_row = 1) AS exhausted_jobs,
               MAX(blocked_jobs) FILTER (WHERE total_row = 1) AS blocked_jobs,
               MAX(oldest_ready_age_seconds) FILTER (WHERE total_row = 1) AS oldest_ready_age_seconds,
               MAX(oldest_in_flight_age_seconds) FILTER (WHERE total_row = 1) AS oldest_in_flight_age_seconds,
               ${dlqTotal} AS dlq_jobs,
               jsonb_object_agg(kind, total_jobs) FILTER (WHERE total_row = 0 AND kind IS NOT NULL) AS kind_counts,
               ${dlqKinds} AS dlq_kind_counts
        FROM (
          SELECT kind, GROUPING(kind) AS total_row,
                 COUNT(*)::int8 AS total_jobs,
                 COUNT(*) FILTER (WHERE status = 'ready' AND NOT blocked) AS ready_jobs,
                 COUNT(*) FILTER (WHERE status = 'in_flight') AS in_flight_jobs,
                 COUNT(*) FILTER (WHERE status = 'scheduled') AS scheduled_jobs,
                 COUNT(*) FILTER (WHERE status = 'backoff') AS backoff_jobs,
                 COUNT(*) FILTER (WHERE status = 'throttled') AS throttled_jobs,
                 COUNT(*) FILTER (WHERE status = 'suspended') AS suspended_jobs,
                 COUNT(*) FILTER (WHERE status = 'cancelled') AS cancelled_jobs,
                 COUNT(*) FILTER (WHERE status = 'exhausted') AS exhausted_jobs,
                 COUNT(*) FILTER (WHERE blocked) AS blocked_jobs,
                 EXTRACT(EPOCH FROM (
                   clock_timestamp() - MIN(GREATEST(inserted_at, not_visible_until, updated_at)) FILTER (WHERE status = 'ready')
                 ))::float8 AS oldest_ready_age_seconds,
                 EXTRACT(EPOCH FROM (
                   clock_timestamp() - MIN(last_attempted_at) FILTER (WHERE status = 'in_flight')
                 ))::float8 AS oldest_in_flight_age_seconds
          FROM (
            SELECT inserted_at, not_visible_until, updated_at, last_attempted_at, ${kindExpr} AS kind, ${jobStatusCaseSQL} AS status,
                   ${blocked} AS blocked
            FROM ${tbl} job
            LEFT JOIN (${heads}) head ON head.id = job.id
          ) classified
          GROUP BY GROUPING SETS ((), (kind))
        ) rollup
      |]

-- | Whether a row is one a claim would skip, over alias @job@ joined to its group
-- head as @head@: behind the head, or behind a full concurrency or rate-limit key.
-- The CASE keeps the key probes off every row a claim would not consider.
blockedExpr :: SchemaName -> Text
blockedExpr :: Text -> Text
blockedExpr Text
schemaName =
  let concTbl :: Text
concTbl = Text -> Text
arbiterConcurrencyTable Text
schemaName
      concPolicies :: Text
concPolicies = Text -> Text
arbiterConcurrencyPoliciesTable Text
schemaName
      buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schemaName
      rlPolicies :: Text
rlPolicies = Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName
      claimable :: Text
claimable = Text -> Text
claimablePred Text
"job"
      concOk :: Text
concOk = Text -> Text -> Text -> Text
concHeadroomPred Text
concTbl Text
concPolicies Text
"job"
      rlOk :: Text
rlOk = Text -> Text -> Text -> Text
rateLimitHeadroomPred Text
buckets Text
rlPolicies Text
"job"
   in [text|
        CASE WHEN ${claimable}
             THEN (job.group_key IS NOT NULL AND head.id IS NULL) OR NOT ${concOk} OR NOT ${rlOk}
             ELSE FALSE END
      |]

-- | A batch limit of one row. The head is judged as a single claim takes it.
headOnly :: Text
headOnly :: Text
headOnly = Text
"1"

-- | The job columns the gate predicates read.
gateColumns :: [Text]
gateColumns :: [Text]
gateColumns = [Text
"claimed_by", Text
"concurrency_key", Text
"rate_limit_key", Text
"rate_limit_prefix", Text
"rate_limit_cost"]

-- | The id of each open group's head, the row a claim of that group takes first.
groupHeadsSQL :: SchemaName -> TableName -> Text
groupHeadsSQL :: Text -> Text -> Text
groupHeadsSQL Text
schemaName Text
tableName =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      headBatch :: Text
headBatch = Text -> Text -> [Text] -> Text -> Text
groupHeadBatch Text
tbl Text
"summary.group_key" [] Text
headOnly
   in [text|
        SELECT head.id
        FROM ${groupsTbl} summary
        CROSS JOIN LATERAL (
          ${headBatch}
        ) head
        WHERE summary.job_count > 0
          AND ((summary.ready_count > 0 AND summary.in_flight_until IS NULL) OR summary.next_due <= NOW())
          AND (summary.in_flight_until IS NULL OR summary.in_flight_until <= NOW())
      |]

-- | Whether a job's rate-limit policy admits its cost, over a row alias. A key without
-- a policy is admitted. A key without a bucket is held only to the policy cap. The first
-- claim seeds a full bucket for it. OFFSET 0 keeps the probe correlated.
rateLimitHeadroomPred :: Text -> Text -> Text -> Text
rateLimitHeadroomPred :: Text -> Text -> Text -> Text
rateLimitHeadroomPred Text
buckets Text
rlPolicies Text
alias =
  let effMax :: Text
effMax = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"max_tokens"
   in [text|
        (${alias}.rate_limit_key IS NULL OR NOT EXISTS (
          SELECT 1 FROM ${rlPolicies} policy
          LEFT JOIN ${buckets} bucket ON bucket.rate_limit_key = ${alias}.rate_limit_key
          WHERE policy.prefix_id = ${alias}.rate_limit_prefix
            AND NOT (${effMax} > 0
                     AND (bucket.rate_limit_key IS NULL
                          OR LEAST(GREATEST(${alias}.rate_limit_cost, 0), ${effMax}) <= ${refilledBucketTokens}))
          OFFSET 0
        ))
      |]

-- | A stored label the payload declares, and NULL for anything else.
declaredKindSQL :: [Text] -> Text
declaredKindSQL :: [Text] -> Text
declaredKindSQL [] = Text
"NULL::text"
declaredKindSQL [Text]
kinds =
  let literals :: Text
literals = Text -> [Text] -> Text
T.intercalate Text
", " ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
textLiteral [Text]
kinds)
   in [text|CASE WHEN kind IN (${literals}) THEN kind END|]

-- | Every queue's stats in one query, tagged by name. Caller guards the empty list.
allQueueStatsSQL :: RowCodec a -> SchemaName -> [(TableName, [Text])] -> Query a
allQueueStatsSQL :: forall a. RowCodec a -> Text -> [(Text, [Text])] -> Query a
allQueueStatsSQL RowCodec a
codec Text
schemaName [(Text, [Text])]
queueKinds =
  let queuesTbl :: Text
queuesTbl = Text -> Text
arbiterQueuesTable Text
schemaName
      workersTbl :: Text
workersTbl = Text -> Text
arbiterWorkersTable Text
schemaName
   in RowCodec a -> Text -> Query a
forall a. RowCodec a -> Text -> Query a
rawRows RowCodec a
codec (Text -> Query a) -> Text -> Query a
forall a b. (a -> b) -> a -> b
$ Text -> [Text] -> (Text -> Text -> Text) -> Text
unionAllOverQueueTables Text
schemaName (((Text, [Text]) -> Text) -> [(Text, [Text])] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text, [Text]) -> Text
forall a b. (a, b) -> a
fst [(Text, [Text])]
queueKinds) ((Text -> Text -> Text) -> Text) -> (Text -> Text -> Text) -> Text
forall a b. (a -> b) -> a -> b
$ \Text
tableName Text
_ ->
        let stats :: Text
stats = DlqCount -> Text -> Text -> [Text] -> Text
queueStatsSelect DlqCount
DlqTotal Text
schemaName Text
tableName ([Text] -> Maybe [Text] -> [Text]
forall a. a -> Maybe a -> a
fromMaybe [] (Text -> [(Text, [Text])] -> Maybe [Text]
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Text
tableName [(Text, [Text])]
queueKinds))
            queue :: Text
queue = Text -> Text
textLiteral Text
tableName
         in [text|
          SELECT ${queue} AS queue, stats.*,
                 COALESCE((SELECT paused FROM ${queuesTbl} WHERE queue_name = ${queue}), FALSE) AS queue_paused,
                 worker_counts.workers_live, worker_counts.workers_paused
          FROM (${stats}) stats
          CROSS JOIN (
            -- Live matches the worker health CASE: a fresh heartbeat and not draining.
            SELECT COUNT(*)::int8 AS workers_live, COUNT(*) FILTER (WHERE paused)::int8 AS workers_paused
            FROM ${workersTbl}
            WHERE queue_name = ${queue}
              AND NOT (${workerStaleSQL})
              AND NOT shutting_down
          ) worker_counts
        |]

-- ---------------------------------------------------------------------------
-- Parent-Child Operations
-- ---------------------------------------------------------------------------

-- | Child counts as @(parent_id, total, suspended)@ per parent, over a set of job ids.
-- The caller attaches the row decoder.
countChildrenBatchSQL :: SchemaName -> TableName -> [Int64] -> Query ()
countChildrenBatchSQL :: Text -> Text -> [Int64] -> Query ()
countChildrenBatchSQL Text
schemaName Text
tableName [Int64]
jobIds =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
   in [sql|
        SELECT parent_id, COUNT(*) AS count,
               COUNT(*) FILTER (WHERE suspended) AS count_suspended
        FROM ${tbl}
        WHERE parent_id = ANY(#{jobIds :: [CInt8]})
        GROUP BY parent_id
      |]

-- ---------------------------------------------------------------------------
-- Group Listing
-- ---------------------------------------------------------------------------

-- | A page of a queue's open groups from the summary table, largest first, each with
-- its head: the job that holds the group or that its next claim takes.
listGroupsSQL :: RowCodec a -> SchemaName -> TableName -> Maybe Text -> Int64 -> Int64 -> Query a
listGroupsSQL :: forall a.
RowCodec a
-> Text -> Text -> Maybe Text -> Int64 -> Int64 -> Query a
listGroupsSQL RowCodec a
codec Text
schemaName Text
tableName Maybe Text
groupKey Int64
limit Int64
offset =
  let groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      keyFilter :: Query ()
keyFilter = Maybe Text -> Query ()
openGroupKeyFilter Maybe Text
groupKey
      heads :: Text
heads = Text -> Text -> Text
groupListHeadSQL Text
schemaName Text
tableName
   in RowCodec a -> Query () -> Query a
forall a. RowCodec a -> Query () -> Query a
rows
        RowCodec a
codec
        [sql|
          SELECT summary.group_key, summary.job_count::int8 AS job_count,
                 summary.ready_count::int8 AS ready_count, summary.next_due, summary.in_flight_until,
                 COALESCE(summary.in_flight_until > NOW(), FALSE) AS in_flight,
                 head.id AS head_id, COALESCE(head.blocked, FALSE) AS head_blocked, head.status AS head_status
          FROM (
            SELECT * FROM ${groupsTbl} summary
            WHERE summary.job_count > 0 ${keyFilter}
            ORDER BY summary.job_count DESC, summary.group_key ASC
            LIMIT #{limit :: CInt8} OFFSET #{offset :: CInt8}
          ) summary
          LEFT JOIN LATERAL (${heads}) head ON TRUE
          ORDER BY summary.job_count DESC, summary.group_key ASC
        |]

-- | The number of open groups 'listGroupsSQL' pages through.
countGroupsSQL :: SchemaName -> TableName -> Maybe Text -> Query Int64
countGroupsSQL :: Text -> Text -> Maybe Text -> Query Int64
countGroupsSQL Text
schemaName Text
tableName Maybe Text
groupKey =
  let groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      keyFilter :: Query ()
keyFilter = Maybe Text -> Query ()
openGroupKeyFilter Maybe Text
groupKey
   in [sql|SELECT COUNT(*)::int8 AS @{count :: CInt8} FROM ${groupsTbl} summary WHERE summary.job_count > 0 ${keyFilter}|]

openGroupKeyFilter :: Maybe Text -> Query ()
openGroupKeyFilter :: Maybe Text -> Query ()
openGroupKeyFilter = (Text -> Query ()) -> Maybe Text -> Query ()
forall m a. Monoid m => (a -> m) -> Maybe a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Text
key -> [sql|AND summary.group_key = #{key :: CText}|])

-- | A group's head with its status: the job that holds the group, else the first row
-- a claim takes, and whether a gate holds that row back.
groupListHeadSQL :: SchemaName -> TableName -> Text
groupListHeadSQL :: Text -> Text -> Text
groupListHeadSQL Text
schemaName Text
tableName =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      inFlight :: Text
inFlight = Text -> Text
inFlightPredicate Text
"job."
      headBatch :: Text
headBatch = Text -> Text -> [Text] -> Text -> Text
groupHeadBatch Text
tbl Text
"summary.group_key" [Text]
gateColumns Text
headOnly
      ready :: Text
ready = Text -> Text
textLiteral (JobStatus -> Text
jobStatusToText JobStatus
Ready)
      concOk :: Text
concOk = Text -> Text -> Text -> Text
concHeadroomPred (Text -> Text
arbiterConcurrencyTable Text
schemaName) (Text -> Text
arbiterConcurrencyPoliciesTable Text
schemaName) Text
"head_batch"
      rlOk :: Text
rlOk = Text -> Text -> Text -> Text
rateLimitHeadroomPred (Text -> Text
arbiterRateLimitsTable Text
schemaName) (Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName) Text
"head_batch"
   in [text|
        SELECT head.id, head.status, head.blocked
        FROM (
          (
            SELECT 0 AS rank, job.id, ${jobStatusCaseSQL} AS status, FALSE AS blocked
            FROM ${tbl} job
            WHERE summary.in_flight_until > NOW() AND job.group_key = summary.group_key AND ${inFlight}
            ORDER BY job.not_visible_until DESC NULLS LAST LIMIT 1
          )
          UNION ALL
          (
            SELECT 1, head_batch.id, ${ready}, NOT (${concOk} AND ${rlOk})
            FROM (
              ${headBatch}
            ) head_batch
          )
        ) head
        ORDER BY head.rank
        LIMIT 1
      |]