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

-- | Internal to the arbiter packages. Not covered by the PVP.
--
-- Claim-time SQL: candidate selection, admission gates, and the batched claim.
module Arbiter.Core.Sql.Claim
  ( ClaimAdmission (..)
  , claimJobsBatchedSQL
  , concHeadroomPred
  , groupHeadBatch
  ) where

import Data.Text (Text)
import Data.Text qualified as T
import Data.Time (NominalDiffTime)
import Data.UUID.Types (UUID)
import NeatInterpolation (text)

import Arbiter.Core.Admission (effectivePolicyCol)
import Arbiter.Core.Concurrency.Schema (arbiterConcurrencyPoliciesTable, arbiterConcurrencyTable)
import Arbiter.Core.Job.Schema (SchemaName, TableName, jobQueueGroupsTable, jobQueueTable)
import Arbiter.Core.RateLimit.Schema (arbiterRateLimitPoliciesTable, arbiterRateLimitsTable, bucketSeedInsert)
import Arbiter.Core.Sql.Jobs (claimablePred, jobReadColumns)
import Arbiter.Core.Sql.QQ (sql)
import Arbiter.Core.Sql.Query (Query, mwhen)
import Arbiter.Core.Sql.RateLimit (defaultThrottleWaitSeconds, refilledExpr)
import Arbiter.Core.SqlLiterals (attemptsLeftSQL, doubleLiteral, intLiteral)

-- | Which admission filters the claim SQL renders. A payload type that declares
-- no policy of a kind gets no filter for it.
data ClaimAdmission = ClaimAdmission
  { ClaimAdmission -> Bool
admitRateLimited :: Bool
  -- ^ Render the rate-limit gate.
  , ClaimAdmission -> Bool
admitConcurrent :: Bool
  -- ^ Render the concurrency gate.
  }
  deriving stock (ClaimAdmission -> ClaimAdmission -> Bool
(ClaimAdmission -> ClaimAdmission -> Bool)
-> (ClaimAdmission -> ClaimAdmission -> Bool) -> Eq ClaimAdmission
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ClaimAdmission -> ClaimAdmission -> Bool
== :: ClaimAdmission -> ClaimAdmission -> Bool
$c/= :: ClaimAdmission -> ClaimAdmission -> Bool
/= :: ClaimAdmission -> ClaimAdmission -> Bool
Eq, Int -> ClaimAdmission -> ShowS
[ClaimAdmission] -> ShowS
ClaimAdmission -> String
(Int -> ClaimAdmission -> ShowS)
-> (ClaimAdmission -> String)
-> ([ClaimAdmission] -> ShowS)
-> Show ClaimAdmission
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ClaimAdmission -> ShowS
showsPrec :: Int -> ClaimAdmission -> ShowS
$cshow :: ClaimAdmission -> String
show :: ClaimAdmission -> String
$cshowList :: [ClaimAdmission] -> ShowS
showList :: [ClaimAdmission] -> ShowS
Show)

-- | Per-group rank of candidate rows (alias @candidate@), which each gate's group cut ranks by.
grpRankExpr :: Text
grpRankExpr :: Text
grpRankExpr =
  [text|
    CASE WHEN candidate.group_key IS NOT NULL
         THEN ROW_NUMBER() OVER (
                PARTITION BY candidate.group_key ORDER BY candidate.priority ASC, candidate.id ASC
              )
         END AS grp_rank
  |]

-- | A job's cost, floored at 0 and capped at the bucket max.
clampedCostExpr :: Text
clampedCostExpr :: Text
clampedCostExpr = Text
"LEAST(GREATEST(candidate.rate_limit_cost, 0), bucket.max_tokens)"

-- | A group's claimable rows in claim order: the two index-served head runs,
-- merged with retried rows first. Projects @extras@, @priority@ and @id@ over alias @job@.
groupHeadBatch :: Text -> Text -> [Text] -> Text -> Text
groupHeadBatch :: Text -> Text -> [Text] -> Text -> Text
groupHeadBatch Text
tbl Text
groupKey [Text]
extras Text
limit =
  let claimable :: Text
claimable = Text -> Text
claimablePred Text
"job"
      runCols :: Text
runCols = Text -> [Text] -> Text
T.intercalate Text
", " ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text
"job." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) ([Text]
extras [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text
"attempts", Text
"priority", Text
"id"]))
      mergedCols :: Text
mergedCols = Text -> [Text] -> Text
T.intercalate Text
", " ((Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text
"merged." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) ([Text]
extras [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text
"priority", Text
"id"]))
   in [text|
    SELECT ${mergedCols}
    FROM (
      (
        SELECT ${runCols}
        FROM ${tbl} job
        WHERE job.group_key = ${groupKey}
          AND ${claimable}
          AND job.attempts > 0
        ORDER BY job.attempts DESC, job.priority ASC, job.id ASC
        LIMIT ${limit}
      )
      UNION ALL
      (
        SELECT ${runCols}
        FROM ${tbl} job
        WHERE job.group_key = ${groupKey}
          AND ${claimable}
          AND job.attempts = 0
        ORDER BY job.priority ASC, job.id ASC
        LIMIT ${limit}
      )
    ) merged
    ORDER BY merged.attempts DESC, merged.priority ASC, merged.id ASC
    LIMIT ${limit}
  |]

-- | The batch a gated group would take, reduced to the row the group cut ranks first,
-- and the headroom check on that row.
gatedHeadPred :: Text -> Text -> Text -> Text -> Text
gatedHeadPred :: Text -> Text -> Text -> Text -> Text
gatedHeadPred Text
tbl Text
groupKey Text
batchLimit Text
headGate =
  let headBatch :: Text
headBatch = Text -> Text -> [Text] -> Text -> Text
groupHeadBatch Text
tbl Text
groupKey [Text
"concurrency_key", Text
"claimed_by"] Text
batchLimit
   in [text|
    EXISTS (
      SELECT 1
      FROM (
        SELECT head_batch.concurrency_key, head_batch.claimed_by
        FROM (
          ${headBatch}
        ) head_batch
        ORDER BY head_batch.priority ASC, head_batch.id ASC
        LIMIT 1
      ) gated_head
      WHERE ${headGate}
    )
  |]

-- | Candidate stage one. Groups with ready or due work, locked and reduced
-- to each head row. A gated group is judged on the row its next batch would take.
groupCandidateCtes :: Text -> Text -> Text -> Text -> Text
groupCandidateCtes :: Text -> Text -> Text -> Text -> Text
groupCandidateCtes Text
groupsTbl Text
tbl Text
overfetch Text
headGate =
  let claimable :: Text
claimable = Text -> Text
claimablePred Text
"job"
   in [text|
    group_candidates AS (
      (
        SELECT candidate_group.group_key FROM ${groupsTbl} candidate_group
        WHERE ready_count > 0 AND in_flight_until IS NULL
          ${headGate}
        ORDER BY min_priority ASC, min_id ASC
        LIMIT ${overfetch}
      )
      UNION
      (
        SELECT candidate_group.group_key FROM ${groupsTbl} candidate_group
        WHERE next_due <= NOW()
          ${headGate}
        ORDER BY next_due ASC
        LIMIT ${overfetch}
      )
    ),
    eligible_groups AS (
      SELECT summary.group_key FROM ${groupsTbl} summary
      WHERE summary.group_key IN (SELECT group_key FROM group_candidates)
        AND summary.job_count > 0
        AND (summary.in_flight_until IS NULL OR summary.in_flight_until <= NOW())
      ORDER BY summary.min_priority ASC, summary.min_id ASC
      LIMIT ${overfetch}
      FOR UPDATE SKIP LOCKED
    ),
    eligible_heads AS (
      SELECT eligible.group_key, head.min_priority, head.min_id
      FROM eligible_groups eligible
      CROSS JOIN LATERAL (
        SELECT job.priority AS min_priority, job.id AS min_id
        FROM ${tbl} job
        WHERE job.group_key = eligible.group_key
          AND ${claimable}
        ORDER BY job.priority ASC, job.id ASC
        LIMIT 1
      ) head
    )
  |]

-- | One ungrouped pool in its index order, ready or due, over alias @job@.
ungroupedScan :: Text -> Text -> Text -> Text -> Text -> Text
ungroupedScan :: Text -> Text -> Text -> Text -> Text -> Text
ungroupedScan Text
tbl Text
ccGate Text
visible Text
order Text
limit =
  let attemptsLeft :: Text
attemptsLeft = Text -> Text
attemptsLeftSQL Text
"job."
   in [text|
    SELECT job.id, job.priority
    FROM ${tbl} job
    WHERE job.group_key IS NULL
      AND NOT job.suspended
      AND job.cancel_requested_at IS NULL
      AND ${visible}
      AND ${attemptsLeft}
      ${ccGate}
    ORDER BY ${order}
    LIMIT ${limit}
  |]

readyScan :: Text -> Text -> Text -> Text
readyScan :: Text -> Text -> Text -> Text
readyScan Text
tbl Text
ccGate = Text -> Text -> Text -> Text -> Text -> Text
ungroupedScan Text
tbl Text
ccGate Text
"job.not_visible_until IS NULL" Text
"job.priority ASC, job.id ASC"

dueScan :: Text -> Text -> Text -> Text
dueScan :: Text -> Text -> Text -> Text
dueScan Text
tbl Text
ccGate = Text -> Text -> Text -> Text -> Text -> Text
ungroupedScan Text
tbl Text
ccGate Text
"job.not_visible_until <= NOW()" Text
"job.not_visible_until ASC"

-- | Candidate stage two. The ungrouped ready and due pools, numbered into batches.
ungroupedPoolCtes :: Text -> Text -> Text -> Text -> Text
ungroupedPoolCtes :: Text -> Text -> Text -> Text -> Text
ungroupedPoolCtes Text
tbl Text
ungroupedLimit Text
batchLimit Text
ccGate =
  let ready :: Text
ready = Text -> Text -> Text -> Text
readyScan Text
tbl Text
ccGate Text
ungroupedLimit
      due :: Text
due = Text -> Text -> Text -> Text
dueScan Text
tbl Text
ccGate Text
ungroupedLimit
   in [text|
    ungrouped_pool AS (
      (SELECT ready_scan.*, TRUE AS ready FROM (${ready}) ready_scan)
      UNION ALL
      (SELECT due_scan.*, FALSE AS ready FROM (${due}) due_scan)
    ),
    ungrouped_numbered AS (
      SELECT id, priority, ready,
        ((ROW_NUMBER() OVER (ORDER BY priority ASC, id ASC) - 1)
          / ${batchLimit}) + 1 AS batch_num
      FROM ungrouped_pool
      ORDER BY priority ASC, id ASC
      LIMIT ${ungroupedLimit}
    ),
    ungrouped_batch_info AS (
      SELECT batch_num, MIN(priority) AS min_priority,
        (MIN(ARRAY[priority::bigint, id]))[2] AS min_id
      FROM ungrouped_numbered
      GROUP BY batch_num
    )
  |]

-- | Candidate stage three. Interleaves group heads and ungrouped batches by
-- priority, capped at the batch budget.
allocatedSlotCtes :: Text -> Text
allocatedSlotCtes :: Text -> Text
allocatedSlotCtes Text
batchBudget =
  [text|
    allocated_slots AS (
      SELECT slot.group_key, slot.ungrouped_batch
      FROM (
        SELECT group_key, NULL::bigint AS ungrouped_batch, min_priority, min_id
        FROM eligible_heads
        UNION ALL
        SELECT NULL::text, batch_num, min_priority, min_id
        FROM ungrouped_batch_info
        ORDER BY min_priority ASC, min_id ASC
      ) slot
      LIMIT ${batchBudget}
    ),
    final_locked_groups AS (
      SELECT group_key FROM allocated_slots WHERE group_key IS NOT NULL
    )
  |]

-- | Keeps @alias@'s scan on the primary key, bounded by the ids in @source@.
idsIn :: Text -> Text -> Text
idsIn :: Text -> Text -> Text
idsIn Text
alias Text
source = Text
alias Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
".id = ANY(ARRAY(SELECT id FROM " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
source Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"))"

-- | Candidate stage four. Resolves the allocated slots to rows and locks the claimable set.
-- Each ungrouped pool locks as many rows as its allocated batches hold, skipping locked rows.
lockedCandidateCtes :: Text -> Text -> Text -> Text -> Text
lockedCandidateCtes :: Text -> Text -> Text -> Text -> Text
lockedCandidateCtes Text
tbl Text
batchLimit Text
ungroupedLimit Text
ccGate =
  let claimable :: Text
claimable = Text -> Text
claimablePred Text
"job"
      heldIds :: Text
heldIds = Text -> Text -> Text
idsIn Text
"job" Text
"ungrouped_allocated"
      selectedIds :: Text
selectedIds = Text -> Text -> Text
idsIn Text
"job" Text
"selected"
      unheld :: Text
unheld = Text
ccGate Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
" AND job.id NOT IN (SELECT id FROM ungrouped_held)"
      ready :: Text
ready = Text -> Text -> Text -> Text
readyScan Text
tbl Text
unheld Text
"(SELECT ready_rows FROM ungrouped_shortfall)"
      due :: Text
due = Text -> Text -> Text -> Text
dueScan Text
tbl Text
unheld Text
"(SELECT due_rows FROM ungrouped_shortfall)"
   in [text|
    grouped_candidates AS (
      SELECT batch.id, target_group.group_key AS expected_group
      FROM final_locked_groups target_group
      CROSS JOIN LATERAL (
        SELECT job.id
        FROM ${tbl} job
        WHERE job.group_key = target_group.group_key
          AND ${claimable}
        ORDER BY job.attempts DESC, job.priority ASC, job.id ASC
        LIMIT ${batchLimit}
      ) batch
    ),
    ungrouped_allocated AS (
      SELECT id, ready
      FROM ungrouped_numbered
      WHERE batch_num IN (
        SELECT ungrouped_batch
        FROM allocated_slots
        WHERE ungrouped_batch IS NOT NULL
      )
    ),
    ungrouped_held AS MATERIALIZED (
      SELECT job.id, allocated.ready
      FROM ${tbl} job
      INNER JOIN ungrouped_allocated allocated ON allocated.id = job.id
      WHERE ${heldIds}
        AND job.group_key IS NULL
        AND ${claimable}
      FOR UPDATE OF job SKIP LOCKED
    ),
    ungrouped_shortfall AS (
      SELECT COUNT(*) FILTER (WHERE allocated.ready) AS ready_rows,
             COUNT(*) FILTER (WHERE NOT allocated.ready) AS due_rows
      FROM ungrouped_allocated allocated
      LEFT JOIN ungrouped_held held ON held.id = allocated.id
      WHERE held.id IS NULL
    ),
    ungrouped_ready_locked AS MATERIALIZED (
      ${ready}
      FOR UPDATE OF job SKIP LOCKED
    ),
    ungrouped_due_locked AS MATERIALIZED (
      ${due}
      FOR UPDATE OF job SKIP LOCKED
    ),
    ungrouped_candidates AS (
      SELECT id, NULL::text AS expected_group
      FROM (SELECT id FROM ungrouped_held LIMIT ${ungroupedLimit}) held
      UNION ALL
      SELECT id, NULL::text
      FROM (SELECT id FROM ungrouped_ready_locked LIMIT ${ungroupedLimit}) ready_locked
      UNION ALL
      SELECT id, NULL::text
      FROM (SELECT id FROM ungrouped_due_locked LIMIT ${ungroupedLimit}) due_locked
    ),
    selected AS (
      SELECT id, expected_group FROM grouped_candidates
      UNION ALL
      SELECT id, expected_group FROM ungrouped_candidates
    ),
    locked AS (
      SELECT job.id, job.priority, job.group_key,
             job.concurrency_key, job.claimed_by,
             job.rate_limit_key, job.rate_limit_prefix, job.rate_limit_cost
      FROM selected
      INNER JOIN ${tbl} job ON job.id = selected.id
      WHERE ${selectedIds}
        AND ${claimable}
        AND job.group_key IS NOT DISTINCT FROM selected.expected_group
      ORDER BY job.priority ASC, job.id ASC
      FOR UPDATE OF job SKIP LOCKED
      LIMIT ${ungroupedLimit}
    )
  |]

-- | Concurrency headroom over a candidate row alias. Keeps a full key off the
-- bounded candidate window, in the ungrouped pool and at each group's head.
-- OFFSET 0 keeps the probe correlated.
concHeadroomPred :: Text -> Text -> Text -> Text
concHeadroomPred :: Text -> Text -> Text -> Text
concHeadroomPred Text
concTbl Text
concPolicies Text
alias =
  let effLimit :: Text
effLimit = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"limit"
   in [text|
        (${alias}.claimed_by IS NOT NULL OR ${alias}.concurrency_key IS NULL OR EXISTS (
          SELECT 1 FROM ${concTbl} counts
          LEFT JOIN ${concPolicies} policy ON policy.prefix_id = counts.concurrency_prefix
          WHERE counts.concurrency_key = ${alias}.concurrency_key
            AND (${effLimit} IS NULL
                 OR counts.in_flight < ${effLimit})
          OFFSET 0
        ))
      |]

-- | Lock each referenced count row. Whoever holds the row gates that key this cycle.
-- eff_limit is COALESCE(policy override, default). An undeclared prefix leaves it NULL
-- and runs uncapped. Emits @conc_locked@.
concLockedCte :: Text -> Text -> Text
concLockedCte :: Text -> Text -> Text
concLockedCte Text
concTbl Text
concPolicies =
  let effLimit :: Text
effLimit = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"limit"
   in [text|
        conc_locked AS (
          SELECT counts.concurrency_key, counts.in_flight,
                 ${effLimit} AS eff_limit
          FROM ${concTbl} counts
          LEFT JOIN ${concPolicies} policy ON policy.prefix_id = counts.concurrency_prefix
          WHERE counts.concurrency_key IN (
            SELECT concurrency_key FROM locked WHERE concurrency_key IS NOT NULL
          )
          FOR UPDATE OF counts SKIP LOCKED
        ),
      |]

-- | Seed a bucket for each candidate key without one. This claim does not see it.
rlSeedCte :: Text -> Text -> Text
rlSeedCte :: Text -> Text -> Text
rlSeedCte Text
buckets Text
rlPolicies =
  let seed :: Text
seed = Text -> Text -> Text -> Text -> Text
bucketSeedInsert Text
buckets Text
rlPolicies Text
"locked n" Text
"n.rate_limit_key IS NOT NULL"
   in [text|
        rl_seed AS (
          ${seed}
        ),
      |]

-- | The candidate's bucket, locked and refilled under its policy. Absent when another
-- claimer holds it.
rlBucketLateral :: Text -> Text -> Text
rlBucketLateral :: Text -> Text -> Text
rlBucketLateral Text
buckets Text
rlPolicies =
  let maxCol :: Text
maxCol = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"max_tokens"
      refillCol :: Text
refillCol = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"refill_amount"
      intervalCol :: Text
intervalCol = Text -> Text -> Text
effectivePolicyCol Text
"policy" Text
"interval"
      availCol :: Text
availCol = Text -> Text -> Text -> Text -> Text -> Text
refilledExpr Text
maxCol Text
"stored_bucket.tokens" Text
"stored_bucket.last_refill" Text
refillCol Text
intervalCol
   in [text|
        LEFT JOIN LATERAL (
          SELECT stored_bucket.rate_limit_key,
                 ${maxCol} AS max_tokens, ${refillCol} AS refill_amount, ${intervalCol} AS refill_interval,
                 ${availCol} AS available
          FROM ${buckets} stored_bucket
          JOIN ${rlPolicies} policy ON policy.prefix_id = stored_bucket.policy_prefix
          WHERE stored_bucket.rate_limit_key = candidate.rate_limit_key
          FOR UPDATE OF stored_bucket SKIP LOCKED
        ) bucket ON TRUE
      |]

-- | The conjunction of every rendered gate's admission column.
admitExpr :: ClaimAdmission -> Text
admitExpr :: ClaimAdmission -> Text
admitExpr ClaimAdmission
admission =
  Text -> [Text] -> Text
T.intercalate Text
" AND " ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$
    [Text
"conc_admit" | ClaimAdmission -> Bool
admitConcurrent ClaimAdmission
admission] [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text
"rl_admit" | ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission]

-- | Every rendered gate's verdict on every locked row in one windowed pass. A row
-- carries its rank in its group, its own admission per gate, the group cut per gate,
-- and the cost its key admits. Emits @judged@.
judgedCte :: ClaimAdmission -> Text -> Text -> Text -> Text
judgedCte :: ClaimAdmission -> Text -> Text -> Text -> Text
judgedCte ClaimAdmission
admission Text
rlPolicies Text
concJoin Text
rlJoin =
  let cc :: Bool
cc = ClaimAdmission -> Bool
admitConcurrent ClaimAdmission
admission
      rl :: Bool
rl = ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission
      admit :: Text
admit = ClaimAdmission -> Text
admitExpr ClaimAdmission
admission
      concRow :: Text
concRow =
        Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen
          Bool
cc
          [text|,
            (pool.concurrency_key IS NOT NULL) AS pool_held,
            (candidate.concurrency_key IS NULL
              OR (pool.concurrency_key IS NOT NULL
                  AND (candidate.claimed_by IS NOT NULL OR pool.eff_limit IS NULL
                       OR pool.in_flight + SUM((candidate.claimed_by IS NULL)::int) OVER (
                            PARTITION BY candidate.concurrency_key ORDER BY candidate.priority ASC, candidate.id ASC
                            ROWS UNBOUNDED PRECEDING
                          ) <= pool.eff_limit))) AS conc_ok
          |]
      rlRow :: Text
rlRow =
        Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen
          Bool
rl
          [text|,
            (bucket.rate_limit_key IS NOT NULL) AS bucket_held,
            bucket.max_tokens, bucket.refill_amount, bucket.refill_interval, bucket.available,
            ${clampedCostExpr} AS clamped_cost,
            CASE WHEN bucket.rate_limit_key IS NOT NULL
                 THEN bucket.max_tokens > 0
                      AND SUM(${clampedCostExpr}) OVER (
                            PARTITION BY candidate.rate_limit_key ORDER BY candidate.priority ASC, candidate.id ASC
                            ROWS UNBOUNDED PRECEDING
                          ) <= bucket.available
                 ELSE candidate.rate_limit_key IS NULL
                      OR NOT EXISTS (
                        SELECT 1 FROM ${rlPolicies} policy WHERE policy.prefix_id = candidate.rate_limit_prefix
                      )
            END AS rl_ok
          |]
      concCut :: Text
concCut = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
cc Text
", MIN(CASE WHEN NOT conc_ok THEN grp_rank END) OVER (PARTITION BY group_key) AS conc_cut"
      rlCut :: Text
rlCut = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
rl Text
", MIN(CASE WHEN NOT rl_ok THEN grp_rank END) OVER (PARTITION BY group_key) AS rl_cut"
      concAdmit :: Text
concAdmit = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
cc Text
", (conc_ok AND (grp_rank IS NULL OR conc_cut IS NULL OR grp_rank < conc_cut)) AS conc_admit"
      rlAdmit :: Text
rlAdmit = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
rl Text
", (rl_ok AND (grp_rank IS NULL OR rl_cut IS NULL OR grp_rank < rl_cut)) AS rl_admit"
      keyCost :: Text
keyCost =
        Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen
          Bool
rl
          [text|,
            SUM(CASE WHEN ${admit} THEN clamped_cost ELSE 0 END) OVER (PARTITION BY rate_limit_key) AS key_admitted_cost
          |]
   in [text|
        judged AS MATERIALIZED (
          SELECT judged_admit.*${keyCost}
          FROM (
            SELECT judged_cut.*${concAdmit}${rlAdmit}
            FROM (
              SELECT judged_row.*${concCut}${rlCut}
              FROM (
                SELECT candidate.id, candidate.priority, candidate.group_key, candidate.claimed_by,
                       candidate.concurrency_key, candidate.rate_limit_key, candidate.rate_limit_prefix,
                       candidate.rate_limit_cost,
                       ${grpRankExpr}${concRow}${rlRow}
                FROM locked candidate
                ${concJoin}
                ${rlJoin}
              ) judged_row
            ) judged_cut
          ) judged_admit
        ),
      |]

-- | The share of a denied job's wait that its defer adds at most as jitter.
deferJitterShare :: Double
deferJitterShare :: Double
deferJitterShare = Double
0.5

-- | The most jitter a denied job's defer adds, in seconds.
maxDeferJitterSeconds :: Double
maxDeferJitterSeconds :: Double
maxDeferJitterSeconds = Double
2.0

-- | Rate-limit spend. Debits each held bucket by the cost admitted by every gate, only
-- where that cost is above 0, banks the accrued refill, then computes a jittered defer
-- time for denied jobs whose bucket this claim holds. A stale-leased job on a count row
-- this claim does not hold stays visible for a later cycle.
rlSpendCtes :: ClaimAdmission -> Text -> Text
rlSpendCtes :: ClaimAdmission -> Text -> Text
rlSpendCtes ClaimAdmission
admission Text
buckets =
  let fallback :: Text
fallback = Double -> Text
doubleLiteral Double
defaultThrottleWaitSeconds
      jitterShare :: Text
jitterShare = Double -> Text
doubleLiteral Double
deferJitterShare
      maxJitter :: Text
maxJitter = Double -> Text
doubleLiteral Double
maxDeferJitterSeconds
      admit :: Text
admit = ClaimAdmission -> Text
admitExpr ClaimAdmission
admission
      deferGuard :: Text
deferGuard = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen (ClaimAdmission -> Bool
admitConcurrent ClaimAdmission
admission) Text
" AND (claimed_by IS NULL OR concurrency_key IS NULL OR pool_held)"
      -- Time to accrue the token deficit when the policy refills, else one
      -- interval, else the fallback.
      rlWaitExpr :: Text
rlWaitExpr =
        [text|
          CASE WHEN judged.refill_amount > 0 AND judged.refill_interval > 0 AND judged.max_tokens > 0
               THEN (judged.refill_interval / judged.refill_amount)
                    * GREATEST(1.0, judged.clamped_cost - (judged.available - judged.key_admitted_cost))
               ELSE COALESCE(NULLIF(judged.refill_interval, 0), ${fallback})
          END
        |]
   in [text|
        rl_spend_agg AS (
          SELECT rate_limit_key, MAX(available) AS available, SUM(clamped_cost) AS admitted_cost
          FROM judged
          WHERE ${admit} AND rate_limit_key IS NOT NULL AND bucket_held
          GROUP BY rate_limit_key
        ),
        rl_spent AS (
          UPDATE ${buckets} stored_bucket
          SET tokens = spend.available - spend.admitted_cost, last_refill = NOW()
          FROM rl_spend_agg spend
          WHERE stored_bucket.rate_limit_key = spend.rate_limit_key AND spend.admitted_cost > 0
          RETURNING stored_bucket.rate_limit_key
        ),
        rl_deferred AS (
          SELECT id,
            NOW() + ((wait + random() * LEAST(wait * ${jitterShare}, ${maxJitter}))
                     * interval '1 second') AS defer_until
          FROM (
            SELECT judged.id, GREATEST(${rlWaitExpr}, 0) AS wait
            FROM judged
            WHERE NOT rl_admit AND bucket_held${deferGuard}
          ) denied
        ),
      |]

-- | The claimable set. The rows every rendered gate admits, or the full locked set
-- when no gate is rendered.
admittedCte :: ClaimAdmission -> Text
admittedCte :: ClaimAdmission -> Text
admittedCte ClaimAdmission
admission =
  [text|
    admitted AS (
      ${body}
    ),
  |]
  where
    body :: Text
body
      | ClaimAdmission -> Bool
admitConcurrent ClaimAdmission
admission Bool -> Bool -> Bool
|| ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission =
          let admit :: Text
admit = ClaimAdmission -> Text
admitExpr ClaimAdmission
admission
           in [text|SELECT id FROM judged WHERE ${admit}|]
      | Bool
otherwise = Text
"SELECT id FROM locked"

-- | With rate limiting the claim splits into an admit/defer decision. Denied jobs whose
-- bucket this claim holds are parked with a defer time alongside the claimed ids.
decisionCte :: ClaimAdmission -> Text
decisionCte :: ClaimAdmission -> Text
decisionCte ClaimAdmission
admission =
  Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen
    (ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission)
    [text|
      decision AS (
        SELECT id, TRUE AS _admit, NULL::timestamptz AS _defer FROM admitted
        UNION ALL
        SELECT id, FALSE, defer_until FROM rl_deferred
      ),
    |]

-- | The @claimed@ UPDATE and the final SELECT, over the CTEs rendered before them.
claimStatement :: ClaimAdmission -> Text -> Text -> Text -> UUID -> Query ()
claimStatement :: ClaimAdmission -> Text -> Text -> Text -> UUID -> Query ()
claimStatement ClaimAdmission
admission Text
tbl Text
timeout Text
ctes UUID
claimant =
  let claimed :: Query ()
claimed = ClaimAdmission -> Text -> Text -> UUID -> Query ()
claimedCte ClaimAdmission
admission Text
tbl Text
timeout UUID
claimant
      admitFilter :: Text
admitFilter = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen (ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission) (Text
" WHERE _admit" :: Text)
   in [sql|
        WITH
        ${ctes}
        ${claimed}
        SELECT ${jobReadColumns} FROM claimed${admitFilter} ORDER BY priority ASC, id ASC
      |]

-- | The @claimed@ UPDATE. Under rate limiting a deferred row moves its claim token
-- and parks, taking no attempt.
claimedCte :: ClaimAdmission -> Text -> Text -> UUID -> Query ()
claimedCte :: ClaimAdmission -> Text -> Text -> UUID -> Query ()
claimedCte ClaimAdmission
admission Text
tbl Text
timeout UUID
claimant
  | ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission =
      [sql|
        claimed AS (
          UPDATE ${tbl} job
          SET not_visible_until = CASE
                WHEN verdict._admit THEN NOW() + (${timeout} * interval '1 second')
                ELSE verdict._defer
              END,
              attempts = CASE
                WHEN verdict._admit THEN job.attempts + 1
                ELSE job.attempts
              END,
              claim_seq = job.claim_seq + 1,
              last_attempted_at = CASE
                WHEN verdict._admit THEN NOW()
                ELSE job.last_attempted_at
              END,
              updated_at = NOW(),
              throttled_until = CASE
                WHEN verdict._admit THEN NULL
                ELSE verdict._defer
              END,
              claimed_by =
                (CASE WHEN verdict._admit THEN #{claimant :: CUuid} ELSE NULL END)::uuid
          FROM decision verdict
          WHERE job.id = verdict.id
            AND ${decisionIds}
          RETURNING job.*, verdict._admit
        )
      |]
  | Bool
otherwise =
      [sql|
        claimed AS (
          UPDATE ${tbl} job
          SET not_visible_until = NOW() + (${timeout} * interval '1 second'),
              attempts = job.attempts + 1,
              claim_seq = job.claim_seq + 1,
              last_attempted_at = NOW(),
              updated_at = NOW(),
              claimed_by = #{claimant :: CUuid}::uuid
          WHERE ${admittedIds}
          RETURNING job.*
        )
      |]
  where
    decisionIds :: Text
decisionIds = Text -> Text -> Text
idsIn Text
"job" Text
"decision"
    admittedIds :: Text
admittedIds = Text -> Text -> Text
idsIn Text
"job" Text
"admitted"

-- | Group candidates fetched per batch of budget, ahead of the admission filters.
groupOverfetchFactor :: Int
groupOverfetchFactor :: Int
groupOverfetchFactor = Int
10

-- | The single-CTE batched claim, which at batch size 1 is the single-job claim. Takes
-- any unsuspended, uncancelled, visible job within its attempt budget, rollup children
-- and woken rollup parents included. Each gate's CTEs render from the 'ClaimAdmission' flags.
-- The caller attaches the row decoder.
claimJobsBatchedSQL :: SchemaName -> TableName -> ClaimAdmission -> Int -> Int -> NominalDiffTime -> UUID -> Query ()
claimJobsBatchedSQL :: Text
-> Text
-> ClaimAdmission
-> Int
-> Int
-> NominalDiffTime
-> UUID
-> Query ()
claimJobsBatchedSQL Text
schemaName Text
tableName ClaimAdmission
admission Int
batchSize Int
maxBatches NominalDiffTime
timeoutSeconds =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schemaName Text
tableName
      groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schemaName Text
tableName
      buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schemaName
      concTbl :: Text
concTbl = Text -> Text
arbiterConcurrencyTable Text
schemaName
      concPolicies :: Text
concPolicies = Text -> Text
arbiterConcurrencyPoliciesTable Text
schemaName
      rlPolicies :: Text
rlPolicies = Text -> Text
arbiterRateLimitPoliciesTable Text
schemaName
      batchLimit :: Text
batchLimit = Int -> Text
forall a. Integral a => a -> Text
intLiteral Int
batchSize
      batchBudget :: Text
batchBudget = Int -> Text
forall a. Integral a => a -> Text
intLiteral Int
maxBatches
      timeout :: Text
timeout = Double -> Text
doubleLiteral (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
timeoutSeconds)
      ungroupedLimit :: Text
ungroupedLimit = Int -> Text
forall a. Integral a => a -> Text
intLiteral (Int
maxBatches Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
batchSize)
      overfetch :: Text
overfetch = Int -> Text
forall a. Integral a => a -> Text
intLiteral (Int
maxBatches Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
groupOverfetchFactor)
      hasRateLimit :: Bool
hasRateLimit = ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission
      hasConcurrency :: Bool
hasConcurrency = ClaimAdmission -> Bool
admitConcurrent ClaimAdmission
admission
      jobHeadroom :: Text
jobHeadroom = Text -> Text -> Text -> Text
concHeadroomPred Text
concTbl Text
concPolicies Text
"job"
      headHeadroom :: Text
headHeadroom = Text -> Text -> Text -> Text
concHeadroomPred Text
concTbl Text
concPolicies Text
"gated_head"
      ccGate :: Text
ccGate = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasConcurrency [text|AND ${jobHeadroom}|]
      -- Admission precedes the bounded candidate window, so a full key cannot
      -- pin every poll to the same prefix of blocked groups.
      headGate :: Text
headGate = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasConcurrency (Text
"AND " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text -> Text -> Text -> Text -> Text
gatedHeadPred Text
tbl Text
"candidate_group.group_key" Text
batchLimit Text
headHeadroom)
      groupCandidates :: Text
groupCandidates = Text -> Text -> Text -> Text -> Text
groupCandidateCtes Text
groupsTbl Text
tbl Text
overfetch Text
headGate
      ungroupedPool :: Text
ungroupedPool = Text -> Text -> Text -> Text -> Text
ungroupedPoolCtes Text
tbl Text
ungroupedLimit Text
batchLimit Text
ccGate
      allocatedSlots :: Text
allocatedSlots = Text -> Text
allocatedSlotCtes Text
batchBudget
      lockedCandidates :: Text
lockedCandidates = Text -> Text -> Text -> Text -> Text
lockedCandidateCtes Text
tbl Text
batchLimit Text
ungroupedLimit Text
ccGate
      concLocked :: Text
concLocked = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasConcurrency (Text -> Text -> Text
concLockedCte Text
concTbl Text
concPolicies)
      rlSeed :: Text
rlSeed = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasRateLimit (Text -> Text -> Text
rlSeedCte Text
buckets Text
rlPolicies)
      concJoin :: Text
concJoin = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasConcurrency Text
"LEFT JOIN conc_locked pool ON pool.concurrency_key = candidate.concurrency_key"
      rlJoin :: Text
rlJoin = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasRateLimit (Text -> Text -> Text
rlBucketLateral Text
buckets Text
rlPolicies)
      judged :: Text
judged = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen (Bool
hasConcurrency Bool -> Bool -> Bool
|| Bool
hasRateLimit) (ClaimAdmission -> Text -> Text -> Text -> Text
judgedCte ClaimAdmission
admission Text
rlPolicies Text
concJoin Text
rlJoin)
      admitted :: Text
admitted = ClaimAdmission -> Text
admittedCte ClaimAdmission
admission
      rlSpend :: Text
rlSpend = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasRateLimit (ClaimAdmission -> Text -> Text
rlSpendCtes ClaimAdmission
admission Text
buckets)
      decision :: Text
decision = ClaimAdmission -> Text
decisionCte ClaimAdmission
admission
      ctes :: Text
ctes =
        [text|
          ${groupCandidates},
          ${ungroupedPool},
          ${allocatedSlots},
          ${lockedCandidates},
          ${concLocked}
          ${rlSeed}
          ${judged}
          ${admitted}
          ${rlSpend}
          ${decision}
        |]
   in ClaimAdmission -> Text -> Text -> Text -> UUID -> Query ()
claimStatement ClaimAdmission
admission Text
tbl Text
timeout Text
ctes