{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}

-- | Claim-time SQL: candidate selection, admission gates, and the batched claim.
module Arbiter.Core.Sql.Claim
  ( ClaimAdmission (..)
  , claimJobsBatchedSQL
  ) where

import Data.Text (Text)
import Data.Text qualified as T
import Data.Time (NominalDiffTime)
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.Job.Types (defaultMaxAttempts)
import Arbiter.Core.RateLimit.Schema (arbiterRateLimitPoliciesTable, arbiterRateLimitsTable, bucketSeedInsert)
import Arbiter.Core.Sql.Jobs (jobColumns)
import Arbiter.Core.Sql.Query (mwhen)
import Arbiter.Core.Sql.RateLimit (defaultThrottleWaitSeconds, refilledExpr)

-- | 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
  , ClaimAdmission -> Bool
admitConcurrent :: Bool
  }
  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. Shared by
-- admission, debit, and the deny-wait.
clampedCostExpr :: Text
clampedCostExpr :: Text
clampedCostExpr = Text
"LEAST(GREATEST(candidate.rate_limit_cost, 0), bucket.max_tokens)"

-- | The batch a gated group would take, reduced to the row the group cut ranks first,
-- and the headroom check on that row. Retried rows rank ahead of fresh ones. Each run
-- comes from an index.
gatedHeadLateral :: Text -> Text -> Text -> Text -> Text
gatedHeadLateral :: Text -> Text -> Text -> Text -> Text
gatedHeadLateral Text
tbl Text
dma Text
batchLimit Text
headGate =
  [text|
    CROSS JOIN LATERAL (
      SELECT head_batch.concurrency_key, head_batch.claimed_by
      FROM (
        SELECT merged.concurrency_key, merged.claimed_by, merged.priority, merged.id
        FROM (
          (
            SELECT job.concurrency_key, job.claimed_by, job.attempts, job.priority, job.id
            FROM ${tbl} job
            WHERE job.group_key = eligible.group_key
              AND NOT job.suspended
              AND job.cancel_requested_at IS NULL
              AND (job.not_visible_until IS NULL OR job.not_visible_until <= NOW())
              AND job.attempts < COALESCE(job.max_attempts, ${dma})
              AND job.attempts > 0
            ORDER BY job.attempts DESC, job.priority ASC, job.id ASC
            LIMIT ${batchLimit}
          )
          UNION ALL
          (
            SELECT job.concurrency_key, job.claimed_by, job.attempts, job.priority, job.id
            FROM ${tbl} job
            WHERE job.group_key = eligible.group_key
              AND NOT job.suspended
              AND job.cancel_requested_at IS NULL
              AND (job.not_visible_until IS NULL OR job.not_visible_until <= NOW())
              AND job.attempts < COALESCE(job.max_attempts, ${dma})
              AND job.attempts = 0
            ORDER BY job.priority ASC, job.id ASC
            LIMIT ${batchLimit}
          )
        ) merged
        ORDER BY merged.attempts DESC, merged.priority ASC, merged.id ASC
        LIMIT ${batchLimit}
      ) 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 -> Text
groupCandidateCtes :: Text -> Text -> Text -> Text -> Text -> Text
groupCandidateCtes Text
groupsTbl Text
tbl Text
overfetch Text
dma Text
gateLateral =
  [text|
    group_candidates AS (
      (
        SELECT group_key FROM ${groupsTbl}
        WHERE ready_count > 0 AND in_flight_until IS NULL
        ORDER BY min_priority ASC, min_id ASC
        LIMIT ${overfetch}
      )
      UNION
      (
        SELECT group_key FROM ${groupsTbl}
        WHERE next_due <= NOW()
        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 NOT job.suspended
          AND job.cancel_requested_at IS NULL
          AND (job.not_visible_until IS NULL OR job.not_visible_until <= NOW())
          AND job.attempts < COALESCE(job.max_attempts, ${dma})
        ORDER BY job.priority ASC, job.id ASC
        LIMIT 1
      ) head
      ${gateLateral}
    )
  |]

-- | Candidate stage two. The ungrouped ready and due pools, numbered into batches.
ungroupedPoolCtes :: Text -> Text -> Text -> Text -> Text -> Text
ungroupedPoolCtes :: Text -> Text -> Text -> Text -> Text -> Text
ungroupedPoolCtes Text
tbl Text
ungroupedLimit Text
batchLimit Text
dma Text
ccGate =
  [text|
    ungrouped_pool AS (
      (
        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 job.not_visible_until IS NULL
          AND job.attempts < COALESCE(job.max_attempts, ${dma})
          ${ccGate}
        ORDER BY job.priority ASC, job.id ASC
        LIMIT ${ungroupedLimit}
      )
      UNION ALL
      (
        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 job.not_visible_until <= NOW()
          AND job.attempts < COALESCE(job.max_attempts, ${dma})
          ${ccGate}
        ORDER BY job.not_visible_until ASC
        LIMIT ${ungroupedLimit}
      )
    ),
    ungrouped_numbered AS (
      SELECT id, priority,
        ((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
    )
  |]

-- | Candidate stage four. Resolves the allocated slots to rows and locks the claimable set.
lockedCandidateCtes :: Text -> Text -> Text -> Text -> Text
lockedCandidateCtes :: Text -> Text -> Text -> Text -> Text
lockedCandidateCtes Text
tbl Text
batchLimit Text
dma Text
ungroupedLimit =
  [text|
    grouped_candidates AS (
      SELECT batch.id, target_group.group_key AS expected_group
      FROM final_locked_groups target_group
      CROSS JOIN LATERAL (
        SELECT id
        FROM ${tbl}
        WHERE group_key = target_group.group_key
          AND NOT suspended
          AND cancel_requested_at IS NULL
          AND (not_visible_until IS NULL OR not_visible_until <= NOW())
          AND attempts < COALESCE(max_attempts, ${dma})
        ORDER BY attempts DESC, priority ASC, id ASC
        LIMIT ${batchLimit}
      ) batch
    ),
    ungrouped_candidates AS (
      SELECT id, NULL::text AS expected_group
      FROM ungrouped_numbered
      WHERE batch_num IN (
        SELECT ungrouped_batch
        FROM allocated_slots
        WHERE ungrouped_batch IS NOT NULL
      )
    ),
    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 (
        SELECT id, expected_group FROM grouped_candidates
        UNION ALL
        SELECT id, expected_group FROM ungrouped_candidates
      ) selected
      INNER JOIN ${tbl} job ON job.id = selected.id
      WHERE NOT job.suspended
        AND job.cancel_requested_at IS NULL
        AND (job.not_visible_until IS NULL OR job.not_visible_until <= NOW())
        AND job.group_key IS NOT DISTINCT FROM selected.expected_group
        AND job.attempts < COALESCE(job.max_attempts, ${dma})
      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. Without it the planner hashes the whole
-- count table once per claim.
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(pool 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
",\n 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
",\n MIN(CASE WHEN NOT rl_ok THEN grp_rank END) OVER (PARTITION BY group_key) AS rl_cut"
      admitCols :: Text
admitCols =
        Text -> [Text] -> Text
T.intercalate Text
",\n" ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$
          [Text
"(conc_ok AND (grp_rank IS NULL OR conc_cut IS NULL OR grp_rank < conc_cut)) AS conc_admit" | Bool
cc]
            [Text] -> [Text] -> [Text]
forall a. Semigroup a => a -> a -> a
<> [Text
"(rl_ok AND (grp_rank IS NULL OR rl_cut IS NULL OR grp_rank < rl_cut)) AS rl_admit" | Bool
rl]
      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.*,
              ${admitCols}
            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
        ),
      |]

-- | Rate-limit spend. Debits each held bucket by the cost admitted by every gate, banks
-- the accrued refill, then computes a jittered defer time for denied keyed jobs. A
-- stale-leased job on a count row this claim does not hold stays visible for a later
-- cycle, since parking it would make the update trigger lock that row.
rlSpendCtes :: ClaimAdmission -> Text -> Text
rlSpendCtes :: ClaimAdmission -> Text -> Text
rlSpendCtes ClaimAdmission
admission Text
buckets =
  let fallback :: Text
fallback = String -> Text
T.pack (Double -> String
forall a. Show a => a -> String
show Double
defaultThrottleWaitSeconds)
      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 * 0.5, 2.0))
                     * 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 keyed jobs
-- 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. Rate limiting splits it into an admit/defer decision.
-- Without it, a straight claim of the admitted ids. A defer clears the holder and
-- moves the token.
claimedCte :: ClaimAdmission -> Text -> Text -> Text
claimedCte :: ClaimAdmission -> Text -> Text -> Text
claimedCte ClaimAdmission
admission Text
tbl Text
timeout
  | ClaimAdmission -> Bool
admitRateLimited ClaimAdmission
admission =
      [text|
        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 ?::uuid ELSE NULL END)::uuid
          FROM decision verdict
          WHERE job.id = verdict.id
          RETURNING job.*, verdict._admit
        )
      |]
  | Bool
otherwise =
      [text|
        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 = ?::uuid
          FROM admitted admitted_row
          WHERE job.id = admitted_row.id
          RETURNING job.*
        )
      |]

-- | The single-CTE batched claim, which at batch size 1 is the single-job claim. Takes
-- any unsuspended visible job, rollup children and woken rollup parents included.
-- Each gate's CTEs render when the payload declares that kind of policy.
claimJobsBatchedSQL :: SchemaName -> TableName -> ClaimAdmission -> Int -> Int -> NominalDiffTime -> Text
claimJobsBatchedSQL :: Text
-> Text -> ClaimAdmission -> Int -> Int -> NominalDiffTime -> Text
claimJobsBatchedSQL Text
schema Text
tableName ClaimAdmission
admission Int
batchSize Int
maxBatches NominalDiffTime
timeoutSeconds =
  let tbl :: Text
tbl = Text -> Text -> Text
jobQueueTable Text
schema Text
tableName
      groupsTbl :: Text
groupsTbl = Text -> Text -> Text
jobQueueGroupsTable Text
schema Text
tableName
      buckets :: Text
buckets = Text -> Text
arbiterRateLimitsTable Text
schema
      concTbl :: Text
concTbl = Text -> Text
arbiterConcurrencyTable Text
schema
      concPolicies :: Text
concPolicies = Text -> Text
arbiterConcurrencyPoliciesTable Text
schema
      rlPolicies :: Text
rlPolicies = Text -> Text
arbiterRateLimitPoliciesTable Text
schema
      batchLimit :: Text
batchLimit = String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
batchSize)
      batchBudget :: Text
batchBudget = String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
maxBatches)
      timeout :: Text
timeout = String -> Text
T.pack (Double -> String
forall a. Show a => a -> String
show (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
timeoutSeconds :: Double))
      ungroupedLimit :: Text
ungroupedLimit = String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show (Int
maxBatches Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
batchSize))
      overfetch :: Text
overfetch = String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show (Int
maxBatches Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
10))
      dma :: Text
dma = String -> Text
T.pack (Int32 -> String
forall a. Show a => a -> String
show Int32
defaultMaxAttempts)
      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}|]
      gateLateral :: Text
gateLateral = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasConcurrency (Text -> Text -> Text -> Text -> Text
gatedHeadLateral Text
tbl Text
dma Text
batchLimit Text
headHeadroom)
      groupCandidates :: Text
groupCandidates = Text -> Text -> Text -> Text -> Text -> Text
groupCandidateCtes Text
groupsTbl Text
tbl Text
overfetch Text
dma Text
gateLateral
      ungroupedPool :: Text
ungroupedPool = Text -> Text -> Text -> Text -> Text -> Text
ungroupedPoolCtes Text
tbl Text
ungroupedLimit Text
batchLimit Text
dma Text
ccGate
      allocatedSlots :: Text
allocatedSlots = Text -> Text
allocatedSlotCtes Text
batchBudget
      lockedCandidates :: Text
lockedCandidates = Text -> Text -> Text -> Text -> Text
lockedCandidateCtes Text
tbl Text
batchLimit Text
dma Text
ungroupedLimit
      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
      claimed :: Text
claimed = ClaimAdmission -> Text -> Text -> Text
claimedCte ClaimAdmission
admission Text
tbl Text
timeout
      admitFilter :: Text
admitFilter = Bool -> Text -> Text
forall m. Monoid m => Bool -> m -> m
mwhen Bool
hasRateLimit Text
" WHERE _admit"
   in [text|
        WITH
        ${groupCandidates},
        ${ungroupedPool},
        ${allocatedSlots},
        ${lockedCandidates},
        ${concLocked}
        ${rlSeed}
        ${judged}
        ${admitted}
        ${rlSpend}
        ${decision}
        ${claimed}
        SELECT ${jobColumns} FROM claimed${admitFilter} ORDER BY priority ASC, id ASC
      |]