deepline 0.3.89 → 0.3.91

package/dist/bundling-sources/shared_libs/play-runtime/scheduler-backends/postgres-rate-state.ts

@@ -0,0 +1,942 @@
+import type {
+  PlayRunnerRateStateAcquireInput,
+  PlayRunnerRateStateAcquireResult,
+  PlayRunnerRateStatePenalizeInput,
+  PlayRunnerRateStateReleaseInput,
+} from '../protocol';
+import type { PacingRule } from '../governor/rate-state-backend';
+import type { PostgresSchedulerQueryClient } from './postgres';
+import { resolveRuntimeSchedulerSchemaEnv } from './runtime-scheduler-schema';
+import { stringifyPostgresJson } from '../postgres-json';
+
+const DEFAULT_SCHEMA = 'runtime_scheduler';
+const RATE_BUCKET_TABLE = 'rate_bucket';
+const RATE_CONCURRENCY_LEASE_TABLE = 'rate_concurrency_leases';
+const MIN_CONCURRENCY_WAIT_MS = 10;
+const ACTIVE_LEASE_TTL_MS = 5 * 60 * 1000;
+const MIN_ADAPTIVE_RPS_FLOOR = 0.01;
+
+/**
+ * Internal cooldown cap. A 429 report (penalize) can carry an arbitrarily large
+ * server Retry-After; we store that verbatim in `claimed_retry_at_ms` for
+ * observability but cap the pacer's own cooldown at 45 minutes so a stuck
+ * provider is re-probed at least this often instead of freezing the bucket for
+ * the whole (possibly hours-long) Retry-After.
+ */
+const MAX_INTERNAL_COOLDOWN_MS = 45 * 60 * 1000;
+
+/**
+ * Additive-ramp multiplier applied to undeclared default-pacing buckets.
+ * Declared rules stay pinned at base unless they carry an explicit adaptive
+ * ceiling.
+ */
+const DEFAULT_PACING_MAX_RPS_MULTIPLIER = 25;
+
+/**
+ * Low-rate provider requests cross the sandbox -> gateway -> app boundary after
+ * their permit wakes. Small network-latency differences can compress two
+ * otherwise exact 100ms/200ms intervals at the provider and turn a mathematically
+ * valid edge into a rolling-window 429. Keep a 20% arrival margin below 100
+ * RPS. High-rate providers retain their microburst schedule and full rate.
+ */
+const LOW_RATE_DISPATCH_HEADROOM_FACTOR = 0.8;
+const LOW_RATE_DISPATCH_HEADROOM_THRESHOLD_RPS = 100;
+
+function quoteIdent(value: string): string {
+  return `"${value.replace(/"/g, '""')}"`;
+}
+
+function schemaName(options?: {
+  schema?: string | null;
+  schedulerSchema?: string | null;
+}): string {
+  return (
+    options?.schema?.trim() ||
+    options?.schedulerSchema?.trim() ||
+    resolveRuntimeSchedulerSchemaEnv() ||
+    DEFAULT_SCHEMA
+  );
+}
+
+function bucketTableName(options?: {
+  schema?: string | null;
+  schedulerSchema?: string | null;
+}): string {
+  return `${quoteIdent(schemaName(options))}.${quoteIdent(RATE_BUCKET_TABLE)}`;
+}
+
+function concurrencyLeaseTableName(options?: {
+  schema?: string | null;
+  schedulerSchema?: string | null;
+}): string {
+  return `${quoteIdent(schemaName(options))}.${quoteIdent(RATE_CONCURRENCY_LEASE_TABLE)}`;
+}
+
+/**
+ * Acquire input rules may carry an optional `adaptiveMaxRps` the client can send
+ * to widen a bucket's recovery ceiling above its declared base rate. When
+ * absent, `default:*` rules get base * 25 and declared providers pin at base.
+ */
+type AcquirePacingRule = PacingRule & {
+  adaptiveMaxRps?: number | null;
+  adaptiveMaxConfigured?: boolean;
+};
+
+function normalizeRules(
+  rules: readonly AcquirePacingRule[],
+): AcquirePacingRule[] {
+  return [...rules]
+    .map((rule) => {
+      const isDefaultPacing = rule.ruleId.startsWith('default:');
+      return {
+        ruleId: rule.ruleId,
+        // Rates are fractional; only dispatched requests must be whole.
+        requestsPerWindow: rule.requestsPerWindow,
+        windowMs: Math.max(0, Math.trunc(rule.windowMs)),
+        maxConcurrency:
+          rule.maxConcurrency == null
+            ? null
+            : Math.max(0, Math.trunc(rule.maxConcurrency)),
+        adaptiveMaxRps:
+          rule.adaptiveMaxRps == null || !Number.isFinite(rule.adaptiveMaxRps)
+            ? null
+            : Math.max(0, rule.adaptiveMaxRps),
+        adaptiveMaxConfigured:
+          isDefaultPacing ||
+          Object.prototype.hasOwnProperty.call(rule, 'adaptiveMaxRps'),
+      };
+    })
+    .filter(
+      (rule) =>
+        rule.ruleId &&
+        Number.isFinite(rule.requestsPerWindow) &&
+        rule.requestsPerWindow > 0 &&
+        rule.windowMs > 0,
+    )
+    .sort((a, b) => a.ruleId.localeCompare(b.ruleId));
+}
+
+/**
+ * Seed parameters for a rule row. `base_refill` is the declared rate in tokens
+ * per second; `capacity` is one window of burst. The adaptive `[min,max]` range
+ * is derived value-first so the same wire shape drives both declared and
+ * default-pacing buckets:
+ *   - min_rps = min(base, max(0.01, base/8))
+ *     (low-RPM rules must still be able to back down below their base)
+ *   - max_rps = base               (declared: recovery climbs back TO base only)
+ *   - max_rps = base * 25          (undeclared default:* buckets, or explicit
+ *                                   adaptiveMaxRps override)
+ */
+function seedForRule(rule: AcquirePacingRule): {
+  rule_id: string;
+  capacity: number;
+  base_refill_per_sec: number;
+  min_rps: number;
+  max_rps: number;
+  adaptive_max_configured: boolean;
+} {
+  // Even a sub-one-request window must accumulate one complete permit.
+  // Preserve its fractional refill rate, with dispatch times resolved in ms.
+  const capacity = Math.max(1, rule.requestsPerWindow);
+  const base = (rule.requestsPerWindow / rule.windowMs) * 1000;
+  const minRps = Math.min(base, Math.max(MIN_ADAPTIVE_RPS_FLOOR, base / 8));
+  const isDefaultPacing = rule.ruleId.startsWith('default:');
+  const maxRps =
+    rule.adaptiveMaxRps != null && rule.adaptiveMaxRps > base
+      ? rule.adaptiveMaxRps
+      : isDefaultPacing
+        ? base * DEFAULT_PACING_MAX_RPS_MULTIPLIER
+        : base;
+  return {
+    rule_id: rule.ruleId,
+    capacity,
+    base_refill_per_sec: base,
+    min_rps: minRps,
+    max_rps: Math.max(maxRps, base),
+    adaptive_max_configured: rule.adaptiveMaxConfigured === true,
+  };
+}
+
+function positiveInt(value: number, label: string): number {
+  if (!Number.isFinite(value) || value <= 0) {
+    throw new Error(`${label} must be a positive finite number.`);
+  }
+  return Math.trunc(value);
+}
+
+export async function ensurePostgresRateStateSchema(
+  client: PostgresSchedulerQueryClient,
+  options?: { schema?: string | null },
+): Promise<void> {
+  // Token-bucket + AIMD-in-the-row state. One row per (bucket_id, rule_id). The
+  // hot path is a single autocommit statement. It locks only the selected
+  // bucket rows, in rule-id order, while it refills, grants, schedules, deducts,
+  // and folds the ramp into the same UPDATE.
+  await client.query(`
+    CREATE TABLE IF NOT EXISTS ${bucketTableName(options)} (
+      bucket_id text NOT NULL,
+      rule_id text NOT NULL,
+      tokens float8 NOT NULL,
+      capacity float8 NOT NULL,
+      refill_per_sec float8 NOT NULL,
+      base_refill_per_sec float8 NOT NULL,
+      min_rps float8 NOT NULL,
+      max_rps float8 NOT NULL,
+      seen_backpressure boolean NOT NULL DEFAULT false,
+      cool_until_ms bigint NOT NULL DEFAULT 0,
+      claimed_retry_at_ms bigint NOT NULL DEFAULT 0,
+      last_ramp_at_ms bigint NOT NULL DEFAULT 0,
+      refilled_at_ms bigint NOT NULL,
+      PRIMARY KEY (bucket_id, rule_id)
+    )
+  `);
+  await client.query(`
+    ALTER TABLE ${bucketTableName(options)}
+      ADD COLUMN IF NOT EXISTS next_dispatch_at_ms bigint NOT NULL DEFAULT 0
+  `);
+  await client.query(`
+    ALTER TABLE ${bucketTableName(options)}
+      ADD COLUMN IF NOT EXISTS adaptive_max_configured boolean NOT NULL DEFAULT false
+  `);
+  await client.query(`
+    CREATE TABLE IF NOT EXISTS ${concurrencyLeaseTableName(options)} (
+      bucket_id text NOT NULL,
+      rule_id text NOT NULL,
+      lease_id text NOT NULL,
+      permits integer NOT NULL CHECK (permits > 0),
+      expires_at_ms bigint NOT NULL,
+      created_at timestamptz NOT NULL DEFAULT now(),
+      updated_at timestamptz NOT NULL DEFAULT now(),
+      PRIMARY KEY (bucket_id, rule_id, lease_id)
+    )
+  `);
+  await client.query(`
+    CREATE INDEX IF NOT EXISTS runtime_scheduler_rate_concurrency_expiry_idx
+      ON ${concurrencyLeaseTableName(options)} (expires_at_ms)
+  `);
+  await client.query(`
+    CREATE INDEX IF NOT EXISTS runtime_scheduler_rate_bucket_bucket_idx
+      ON ${bucketTableName(options)} (bucket_id)
+  `);
+}
+
+/**
+ * Seed missing rule rows and refresh changed configuration without erasing
+ * learned AIMD state. The conflict WHERE keeps the warm path a no-op, while a
+ * newly deployed adaptive ceiling reaches existing production buckets.
+ */
+async function seedBucketRows(
+  client: PostgresSchedulerQueryClient,
+  input: { bucketId: string } & {
+    schema?: string | null;
+    schedulerSchema?: string | null;
+  },
+  rules: AcquirePacingRule[],
+  nowMs: number,
+): Promise<void> {
+  const seeds = rules.map(seedForRule);
+  await client.query(
+    `
+      INSERT INTO ${bucketTableName(input)} AS existing (
+        bucket_id,
+        rule_id,
+        tokens,
+        capacity,
+        refill_per_sec,
+        base_refill_per_sec,
+        min_rps,
+        max_rps,
+        adaptive_max_configured,
+        seen_backpressure,
+        cool_until_ms,
+        claimed_retry_at_ms,
+        last_ramp_at_ms,
+        refilled_at_ms,
+        next_dispatch_at_ms
+      )
+      SELECT
+        $1,
+        seed.rule_id,
+        seed.capacity,
+        seed.capacity,
+        seed.base_refill_per_sec,
+        seed.base_refill_per_sec,
+        seed.min_rps,
+        seed.max_rps,
+        seed.adaptive_max_configured,
+        false,
+        0,
+        0,
+        0,
+        $3::bigint,
+        $3::bigint
+      FROM jsonb_to_recordset($2::jsonb) AS seed(
+        rule_id text,
+        capacity float8,
+        base_refill_per_sec float8,
+        min_rps float8,
+        max_rps float8,
+        adaptive_max_configured boolean
+      )
+      ON CONFLICT (bucket_id, rule_id) DO UPDATE SET
+        tokens = LEAST(existing.tokens, EXCLUDED.capacity),
+        capacity = EXCLUDED.capacity,
+        refill_per_sec = LEAST(
+          EXCLUDED.max_rps,
+          GREATEST(EXCLUDED.min_rps, existing.refill_per_sec)
+        ),
+        base_refill_per_sec = EXCLUDED.base_refill_per_sec,
+        min_rps = EXCLUDED.min_rps,
+        max_rps = CASE
+          WHEN EXCLUDED.adaptive_max_configured THEN EXCLUDED.max_rps
+          ELSE existing.max_rps
+        END,
+        adaptive_max_configured =
+          existing.adaptive_max_configured OR EXCLUDED.adaptive_max_configured
+      WHERE existing.capacity IS DISTINCT FROM EXCLUDED.capacity
+        OR existing.base_refill_per_sec IS DISTINCT FROM EXCLUDED.base_refill_per_sec
+        OR existing.min_rps IS DISTINCT FROM EXCLUDED.min_rps
+        OR (
+          EXCLUDED.adaptive_max_configured
+          AND existing.max_rps IS DISTINCT FROM EXCLUDED.max_rps
+        )
+    `,
+    [input.bucketId, stringifyPostgresJson(seeds), nowMs],
+  );
+}
+
+/**
+ * Pure-rate (no maxConcurrency) acquire. A deterministic row lock inside one
+ * autocommit statement makes every contender reserve from the latest committed
+ * token and dispatch cursor instead of reusing a stale CTE snapshot. It refills
+ * every rule row against clock_timestamp(), grants the integer floor of the min
+ * available tokens across the bucket's rules (multi-rule = most constrained),
+ * deducts, and folds the AIMD ramp in. Returns the granted count and, when
+ * granted is 0, the wait until one more token accrues on the binding rule (or
+ * the cooldown clears).
+ */
+async function acquirePureRate(
+  client: PostgresSchedulerQueryClient,
+  input: PlayRunnerRateStateAcquireInput & { nowMs?: number },
+  rules: AcquirePacingRule[],
+  requested: number,
+): Promise<PlayRunnerRateStateAcquireResult> {
+  const ruleIds = rules.map((rule) => rule.ruleId);
+  const result = await client.query<{
+    granted: string | number;
+    wait_ms: string | number;
+    server_now_ms: string | number;
+    cool_until_ms: string | number;
+    claimed_retry_at_ms: string | number;
+    scheduled_at_ms: Array<string | number> | null;
+    effective_rps: string | number;
+  }>(
+    `
+      WITH now_value AS (
+        SELECT COALESCE(
+          $3::bigint,
+          FLOOR(EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::bigint
+        ) AS now_ms
+      ),
+      refreshed AS (
+        SELECT
+          b.bucket_id,
+          b.rule_id,
+          b.capacity,
+          b.min_rps,
+          b.max_rps,
+          b.seen_backpressure,
+          b.base_refill_per_sec,
+          b.refill_per_sec,
+          b.cool_until_ms,
+          b.last_ramp_at_ms,
+          b.next_dispatch_at_ms,
+          nv.now_ms AS now_ms,
+          LEAST(
+            b.capacity,
+            b.tokens + GREATEST(0, nv.now_ms - b.refilled_at_ms) / 1000.0 * b.refill_per_sec
+          ) AS tokens_after_refill
+        FROM ${bucketTableName(input)} b
+        CROSS JOIN now_value nv
+        WHERE b.bucket_id = $1
+          AND b.rule_id = ANY($2::text[])
+        ORDER BY b.rule_id
+        FOR UPDATE OF b
+      ),
+      cooling AS (
+        SELECT COALESCE(MAX(cool_until_ms), 0) AS cool_until_ms
+        FROM refreshed
+      ),
+      grant_calc AS (
+        SELECT
+          -- Cooling buckets grant nothing; otherwise the integer floor of the
+          -- minimum post-refill tokens across every rule row (most constrained).
+          CASE
+            WHEN (SELECT cool_until_ms FROM cooling) > (SELECT now_ms FROM refreshed LIMIT 1)
+              THEN 0
+            ELSE GREATEST(0, LEAST($4::int, FLOOR(MIN(tokens_after_refill))::int))
+          END AS granted
+        FROM refreshed
+      ),
+      schedule_calc AS (
+        SELECT
+          GREATEST(
+            MAX(now_ms),
+            MAX(cool_until_ms),
+            MAX(next_dispatch_at_ms)
+          )::bigint AS first_at_ms,
+          CEIL(MAX(
+            1000.0 / GREATEST(
+              refill_per_sec * CASE
+                WHEN refill_per_sec < ${LOW_RATE_DISPATCH_HEADROOM_THRESHOLD_RPS}
+                  THEN ${LOW_RATE_DISPATCH_HEADROOM_FACTOR}
+                ELSE 1.0
+              END,
+              0.0001
+            )
+          ))::bigint AS spacing_ms,
+          GREATEST(1, FLOOR(MIN(refill_per_sec) / 100.0))::bigint AS burst_size
+        FROM refreshed
+      ),
+      updated AS (
+        UPDATE ${bucketTableName(input)} b
+        SET
+          tokens = r.tokens_after_refill - (SELECT granted FROM grant_calc),
+          refilled_at_ms = r.now_ms,
+          next_dispatch_at_ms = CASE
+            WHEN (SELECT granted FROM grant_calc) > 0
+              THEN (SELECT first_at_ms FROM schedule_calc)
+                + CEIL(
+                    (SELECT granted FROM grant_calc)::numeric
+                    / (SELECT burst_size FROM schedule_calc)
+                  )::bigint
+                  * (SELECT spacing_ms FROM schedule_calc)
+                  * (SELECT burst_size FROM schedule_calc)
+            ELSE b.next_dispatch_at_ms
+          END,
+          -- AIMD ramp, folded in. Only after confirmed provider success,
+          -- under real demand, not cooling, and at least 1s since the last
+          -- ramp. Slow-start doubles until the first 429, then adds
+          -- max(1, base/10) per step. Never exceeds max_rps.
+          refill_per_sec = CASE
+            WHEN $5::int > 0
+              AND r.now_ms - r.last_ramp_at_ms >= 1000
+              AND r.cool_until_ms <= r.now_ms
+              AND r.tokens_after_refill < r.capacity / 2
+            THEN LEAST(
+              r.max_rps,
+              CASE
+                WHEN r.seen_backpressure
+                  THEN r.refill_per_sec + GREATEST(1, r.base_refill_per_sec / 10)
+                ELSE r.refill_per_sec * 2
+              END
+            )
+            ELSE r.refill_per_sec
+          END,
+          last_ramp_at_ms = CASE
+            WHEN $5::int > 0
+              AND r.now_ms - r.last_ramp_at_ms >= 1000
+              AND r.cool_until_ms <= r.now_ms
+              AND r.tokens_after_refill < r.capacity / 2
+            THEN r.now_ms
+            ELSE r.last_ramp_at_ms
+          END
+        FROM refreshed r
+        WHERE b.bucket_id = r.bucket_id
+          AND b.rule_id = r.rule_id
+        RETURNING
+          b.rule_id,
+          r.tokens_after_refill,
+          b.refill_per_sec AS new_refill_per_sec,
+          r.now_ms AS now_ms
+      )
+      SELECT
+        (SELECT granted FROM grant_calc) AS granted,
+        (SELECT now_ms FROM refreshed LIMIT 1) AS server_now_ms,
+        (SELECT cool_until_ms FROM cooling) AS cool_until_ms,
+        (SELECT COALESCE(MAX(claimed_retry_at_ms), 0) FROM ${bucketTableName(input)}
+           WHERE bucket_id = $1 AND rule_id = ANY($2::text[])) AS claimed_retry_at_ms,
+        (SELECT COALESCE(MIN(new_refill_per_sec), 0) FROM updated) AS effective_rps,
+        CASE
+          WHEN (SELECT granted FROM grant_calc) > 0 THEN ARRAY(
+            SELECT
+              (SELECT first_at_ms FROM schedule_calc)
+                + FLOOR(
+                    (permit_index - 1)::numeric
+                    / (SELECT burst_size FROM schedule_calc)
+                  )::bigint
+                  * (SELECT spacing_ms FROM schedule_calc)
+                  * (SELECT burst_size FROM schedule_calc)
+            FROM generate_series(1, (SELECT granted FROM grant_calc)) AS permit_index
+          )
+          ELSE '{}'::bigint[]
+        END AS scheduled_at_ms,
+        -- Wait until either the cooldown clears or one token accrues on the
+        -- most-constrained rule (whichever a granted=0 caller must sleep for).
+        CASE
+          WHEN (SELECT granted FROM grant_calc) > 0 THEN 0
+          WHEN (SELECT cool_until_ms FROM cooling) > (SELECT now_ms FROM refreshed LIMIT 1)
+            THEN (SELECT cool_until_ms FROM cooling) - (SELECT now_ms FROM refreshed LIMIT 1)
+          ELSE (
+            SELECT CEIL(
+              MAX((1.0 - tokens_after_refill) / GREATEST(new_refill_per_sec, 0.0001) * 1000)
+            )::bigint
+            FROM updated
+          )
+        END AS wait_ms
+    `,
+    [
+      input.bucketId,
+      ruleIds,
+      input.nowMs ?? null,
+      requested,
+      Math.max(0, Math.trunc(input.observedSuccesses ?? 0)),
+    ],
+  );
+  const row = result.rows[0];
+  if (!row) {
+    // Seeding raced away or rows vanished; treat as a short retryable wait.
+    return { granted: 0, waitMs: MIN_CONCURRENCY_WAIT_MS, leaseIds: [] };
+  }
+  const granted = Math.max(0, Math.trunc(Number(row.granted)));
+  const serverNowMs = Math.max(0, Number(row.server_now_ms));
+  const coolUntilMs = Math.max(0, Number(row.cool_until_ms));
+  const claimedRetryAtMs = Math.max(0, Number(row.claimed_retry_at_ms));
+  const scheduledAtMs = Array.isArray(row.scheduled_at_ms)
+    ? row.scheduled_at_ms.map(Number)
+    : [];
+  const effectiveRequestsPerSecond = Math.max(0, Number(row.effective_rps));
+  const waitMs = granted > 0 ? 0 : Math.max(1, Math.trunc(Number(row.wait_ms)));
+  return {
+    granted,
+    waitMs,
+    serverNowMs,
+    leaseIds: [],
+    scheduledAtMs,
+    effectiveRequestsPerSecond,
+    coolUntilMs,
+    claimedRetryAtMs,
+  };
+}
+
+/**
+ * Mixed / concurrency-bearing acquire. The token bucket still governs the
+ * request rate; the durable lease table governs the simultaneous-in-flight cap.
+ * Expired-lease reclaim + active count + rate refill + grant + token deduction +
+ * lease minting all execute in ONE multi-CTE statement (one implicit
+ * transaction), so a crash can never deduct tokens without minting the matching
+ * lease rows. Lease ids are generated in SQL (gen_random_uuid) and returned to
+ * the client, one per granted permit, releasable independently.
+ */
+async function acquireWithConcurrency(
+  client: PostgresSchedulerQueryClient,
+  input: PlayRunnerRateStateAcquireInput & { nowMs?: number },
+  rules: AcquirePacingRule[],
+  requested: number,
+): Promise<PlayRunnerRateStateAcquireResult> {
+  const ruleIds = rules.map((rule) => rule.ruleId);
+  const concurrencyRules = rules.filter((rule) => rule.maxConcurrency != null);
+  const concurrencyRuleIds = concurrencyRules.map((rule) => rule.ruleId);
+  // Per-rule limits, NOT collapsed to a single min. When a bucket carries more
+  // than one concurrency rule, each rule mints its own lease row per permit, so
+  // SUM(permits) over all rules double-counts a single in-flight permit. The
+  // grant must instead be bounded by the MIN free slots computed PER RULE
+  // (limit_r - active_r), matching the per-rule enforcement on origin/main.
+  const concurrencyLimits = concurrencyRules.map((rule) => ({
+    rule_id: rule.ruleId,
+    max_concurrency: rule.maxConcurrency as number,
+  }));
+
+  // One multi-CTE statement: reclaim expired leases, refill the rate rows,
+  // count live leases, compute the grant bounded by rate AND free concurrency
+  // slots, deduct the tokens, AND mint the lease rows — atomically. Grant +
+  // deduction + lease minting share one implicit transaction, so a crash can
+  // never deduct tokens without the matching lease rows (which would otherwise
+  // over-admit the provider until the phantom leases "expired"). PgBouncer-safe:
+  // autocommit, deterministic row locking, no lock_timeout, no session state.
+  const result = await client.query<{
+    granted: string | number;
+    wait_ms: string | number;
+    server_now_ms: string | number;
+    cool_until_ms: string | number;
+    claimed_retry_at_ms: string | number;
+    lease_ids: string[] | null;
+    scheduled_at_ms: Array<string | number> | null;
+    effective_rps: string | number;
+  }>(
+    `
+      WITH now_value AS (
+        SELECT COALESCE(
+          $3::bigint,
+          FLOOR(EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::bigint
+        ) AS now_ms
+      ),
+      reclaim AS (
+        DELETE FROM ${concurrencyLeaseTableName(input)}
+        WHERE bucket_id = $1
+          AND rule_id = ANY($5::text[])
+          AND expires_at_ms <= (SELECT now_ms FROM now_value)
+        RETURNING 1
+      ),
+      limits AS (
+        -- One row per concurrency rule with its declared cap.
+        SELECT rule_id, max_concurrency
+        FROM jsonb_to_recordset($9::jsonb) AS lim(
+          rule_id text,
+          max_concurrency int
+        )
+      ),
+      active_by_rule AS (
+        -- Live in-flight permits PER RULE. A CTE SELECT sees the pre-statement
+        -- snapshot, so expired rows the reclaim DELETE removes are still visible
+        -- here; filter them out explicitly by expiry so the free-slot count
+        -- matches the post-reclaim table. Grouping by rule_id (not SUM across
+        -- all rules) is what keeps a single permit — which mints one lease row
+        -- per concurrency rule — from being counted once per rule.
+        SELECT l.rule_id,
+          COALESCE((
+            SELECT SUM(cl.permits)
+            FROM ${concurrencyLeaseTableName(input)} cl
+            WHERE cl.bucket_id = $1
+              AND cl.rule_id = l.rule_id
+              AND cl.expires_at_ms > (SELECT now_ms FROM now_value)
+          ), 0)::int AS active_count
+        FROM limits l
+      ),
+      free_slots AS (
+        -- Most-constrained free concurrency across all rules of the bucket.
+        SELECT COALESCE(
+          MIN(l.max_concurrency - a.active_count),
+          $6::int
+        )::int AS free
+        FROM limits l
+        JOIN active_by_rule a ON a.rule_id = l.rule_id
+      ),
+      refreshed AS (
+        SELECT
+          b.rule_id,
+          b.capacity,
+          b.max_rps,
+          b.seen_backpressure,
+          b.base_refill_per_sec,
+          b.cool_until_ms,
+          b.last_ramp_at_ms,
+          b.next_dispatch_at_ms,
+          nv.now_ms AS now_ms,
+          LEAST(
+            b.capacity,
+            b.tokens + GREATEST(0, nv.now_ms - b.refilled_at_ms) / 1000.0 * b.refill_per_sec
+          ) AS tokens_after_refill,
+          b.refill_per_sec
+        FROM ${bucketTableName(input)} b
+        CROSS JOIN now_value nv
+        WHERE b.bucket_id = $1
+          AND b.rule_id = ANY($2::text[])
+        ORDER BY b.rule_id
+        FOR UPDATE OF b
+      ),
+      cooling AS (
+        SELECT COALESCE(MAX(cool_until_ms), 0) AS cool_until_ms FROM refreshed
+      ),
+      grant_calc AS (
+        SELECT
+          CASE
+            WHEN (SELECT cool_until_ms FROM cooling) > (SELECT now_ms FROM refreshed LIMIT 1)
+              THEN 0
+            ELSE GREATEST(
+              0,
+              LEAST(
+                $4::int,
+                FLOOR(MIN(tokens_after_refill))::int,
+                (SELECT free FROM free_slots)
+              )
+            )
+          END AS granted
+        FROM refreshed
+      ),
+      schedule_calc AS (
+        SELECT
+          GREATEST(
+            MAX(now_ms),
+            MAX(cool_until_ms),
+            MAX(next_dispatch_at_ms)
+          )::bigint AS first_at_ms,
+          CEIL(MAX(
+            1000.0 / GREATEST(
+              refill_per_sec * CASE
+                WHEN refill_per_sec < ${LOW_RATE_DISPATCH_HEADROOM_THRESHOLD_RPS}
+                  THEN ${LOW_RATE_DISPATCH_HEADROOM_FACTOR}
+                ELSE 1.0
+              END,
+              0.0001
+            )
+          ))::bigint AS spacing_ms,
+          GREATEST(1, FLOOR(MIN(refill_per_sec) / 100.0))::bigint AS burst_size
+        FROM refreshed
+      ),
+      updated AS (
+        UPDATE ${bucketTableName(input)} b
+        SET
+          tokens = r.tokens_after_refill - (SELECT granted FROM grant_calc),
+          refilled_at_ms = r.now_ms,
+          next_dispatch_at_ms = CASE
+            WHEN (SELECT granted FROM grant_calc) > 0
+              THEN (SELECT first_at_ms FROM schedule_calc)
+                + CEIL(
+                    (SELECT granted FROM grant_calc)::numeric
+                    / (SELECT burst_size FROM schedule_calc)
+                  )::bigint
+                  * (SELECT spacing_ms FROM schedule_calc)
+                  * (SELECT burst_size FROM schedule_calc)
+            ELSE b.next_dispatch_at_ms
+          END,
+          -- Match the pure-rate path: concurrency leases change admission
+          -- accounting, not the AIMD learning contract.
+          refill_per_sec = CASE
+            WHEN $10::int > 0
+              AND r.now_ms - r.last_ramp_at_ms >= 1000
+              AND r.cool_until_ms <= r.now_ms
+              AND r.tokens_after_refill < r.capacity / 2
+            THEN LEAST(
+              r.max_rps,
+              CASE
+                WHEN r.seen_backpressure
+                  THEN r.refill_per_sec + GREATEST(1, r.base_refill_per_sec / 10)
+                ELSE r.refill_per_sec * 2
+              END
+            )
+            ELSE r.refill_per_sec
+          END,
+          last_ramp_at_ms = CASE
+            WHEN $10::int > 0
+              AND r.now_ms - r.last_ramp_at_ms >= 1000
+              AND r.cool_until_ms <= r.now_ms
+              AND r.tokens_after_refill < r.capacity / 2
+            THEN r.now_ms
+            ELSE r.last_ramp_at_ms
+          END
+        FROM refreshed r
+        WHERE b.bucket_id = $1
+          AND b.rule_id = r.rule_id
+        RETURNING
+          r.tokens_after_refill,
+          b.refill_per_sec AS new_refill_per_sec
+      ),
+      lease_tokens AS (
+        -- One uuid per granted permit. A WITH query is evaluated once (volatile
+        -- functions included), so each lease_id below stays stable across the
+        -- per-rule cross join in minted — every concurrency rule row shares the
+        -- same lease_id for a given permit, matching release semantics.
+        SELECT gen_random_uuid()::text AS lease_id
+        FROM generate_series(1, (SELECT granted FROM grant_calc))
+      ),
+      minted AS (
+        -- Mint the lease rows from grant_calc (NOT by re-reading the lease
+        -- table: all CTEs share the pre-statement snapshot). Same-statement =
+        -- same transaction as the token deduction above.
+        INSERT INTO ${concurrencyLeaseTableName(input)} (
+          bucket_id,
+          rule_id,
+          lease_id,
+          permits,
+          expires_at_ms
+        )
+        SELECT
+          $1,
+          rule_id,
+          lt.lease_id,
+          1,
+          (SELECT now_ms FROM now_value) + $8::bigint
+        FROM unnest($5::text[]) AS rule_id
+        CROSS JOIN lease_tokens lt
+        RETURNING lease_id
+      )
+      SELECT
+        (SELECT granted FROM grant_calc) AS granted,
+        (SELECT now_ms FROM refreshed LIMIT 1) AS server_now_ms,
+        (SELECT cool_until_ms FROM cooling) AS cool_until_ms,
+        (SELECT COALESCE(MAX(claimed_retry_at_ms), 0) FROM ${bucketTableName(input)}
+           WHERE bucket_id = $1 AND rule_id = ANY($2::text[])) AS claimed_retry_at_ms,
+        (SELECT COALESCE(MIN(new_refill_per_sec), 0) FROM updated) AS effective_rps,
+        CASE
+          WHEN (SELECT granted FROM grant_calc) > 0 THEN ARRAY(
+            SELECT
+              (SELECT first_at_ms FROM schedule_calc)
+                + FLOOR(
+                    (permit_index - 1)::numeric
+                    / (SELECT burst_size FROM schedule_calc)
+                  )::bigint
+                  * (SELECT spacing_ms FROM schedule_calc)
+                  * (SELECT burst_size FROM schedule_calc)
+            FROM generate_series(1, (SELECT granted FROM grant_calc)) AS permit_index
+          )
+          ELSE '{}'::bigint[]
+        END AS scheduled_at_ms,
+        (SELECT COALESCE(ARRAY_AGG(DISTINCT lease_id), '{}') FROM minted) AS lease_ids,
+        CASE
+          WHEN (SELECT granted FROM grant_calc) > 0 THEN 0
+          WHEN (SELECT cool_until_ms FROM cooling) > (SELECT now_ms FROM refreshed LIMIT 1)
+            THEN (SELECT cool_until_ms FROM cooling) - (SELECT now_ms FROM refreshed LIMIT 1)
+          WHEN (SELECT free FROM free_slots) <= 0 THEN $7::bigint
+          ELSE (
+            SELECT CEIL(
+              MAX((1.0 - tokens_after_refill) / GREATEST(new_refill_per_sec, 0.0001) * 1000)
+            )::bigint
+            FROM updated
+          )
+        END AS wait_ms
+    `,
+    [
+      input.bucketId,
+      ruleIds,
+      input.nowMs ?? null,
+      requested,
+      concurrencyRuleIds,
+      // $6 is only the free_slots COALESCE fallback (no concurrency rows). Use
+      // the tightest declared cap so it can never widen the grant.
+      Math.min(...concurrencyLimits.map((limit) => limit.max_concurrency)),
+      MIN_CONCURRENCY_WAIT_MS,
+      ACTIVE_LEASE_TTL_MS,
+      stringifyPostgresJson(concurrencyLimits),
+      Math.max(0, Math.trunc(input.observedSuccesses ?? 0)),
+    ],
+  );
+  const row = result.rows[0];
+  if (!row) {
+    return { granted: 0, waitMs: MIN_CONCURRENCY_WAIT_MS, leaseIds: [] };
+  }
+  const granted = Math.max(0, Math.trunc(Number(row.granted)));
+  const serverNowMs = Math.max(0, Number(row.server_now_ms));
+  const coolUntilMs = Math.max(0, Number(row.cool_until_ms));
+  const claimedRetryAtMs = Math.max(0, Number(row.claimed_retry_at_ms));
+  const scheduledAtMs = Array.isArray(row.scheduled_at_ms)
+    ? row.scheduled_at_ms.map(Number)
+    : [];
+  const effectiveRequestsPerSecond = Math.max(0, Number(row.effective_rps));
+  const waitMs = granted > 0 ? 0 : Math.max(1, Math.trunc(Number(row.wait_ms)));
+  const leaseIds = Array.isArray(row.lease_ids)
+    ? row.lease_ids.filter(
+        (leaseId): leaseId is string =>
+          typeof leaseId === 'string' && leaseId.length > 0,
+      )
+    : [];
+  if (granted > 0 && leaseIds.length !== granted) {
+    throw new Error(
+      `Rate-state lease minting drift for ${input.bucketId}: granted ${granted} with ${leaseIds.length} lease tokens.`,
+    );
+  }
+  return {
+    granted,
+    waitMs,
+    serverNowMs,
+    leaseIds: granted > 0 ? leaseIds : [],
+    scheduledAtMs,
+    effectiveRequestsPerSecond,
+    coolUntilMs,
+    claimedRetryAtMs,
+  };
+}
+
+export async function acquirePostgresRateState(
+  client: PostgresSchedulerQueryClient,
+  input: PlayRunnerRateStateAcquireInput & { nowMs?: number },
+): Promise<PlayRunnerRateStateAcquireResult> {
+  const rules = normalizeRules(input.rules as AcquirePacingRule[]);
+  if (rules.length === 0) return { granted: 0, waitMs: 0, leaseIds: [] };
+  const requested = positiveInt(input.requested, 'requested permits');
+  const nowMs =
+    input.nowMs ??
+    Number(
+      (
+        await client.query<{ now_ms: string | number }>(
+          `SELECT FLOOR(EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::bigint AS now_ms`,
+        )
+      ).rows[0]?.now_ms ?? Date.now(),
+    );
+  // Cold path: seed missing rows (ON CONFLICT DO NOTHING, no-op when warm).
+  await seedBucketRows(client, input, rules, nowMs);
+  const hasConcurrency = rules.some((rule) => rule.maxConcurrency != null);
+  if (!hasConcurrency) {
+    return acquirePureRate(client, { ...input, nowMs }, rules, requested);
+  }
+  return acquireWithConcurrency(client, { ...input, nowMs }, rules, requested);
+}
+
+export async function releasePostgresRateState(
+  client: PostgresSchedulerQueryClient,
+  input: PlayRunnerRateStateReleaseInput & { nowMs?: number },
+): Promise<void> {
+  const rules = normalizeRules(input.rules as AcquirePacingRule[]).filter(
+    (rule) => rule.maxConcurrency != null,
+  );
+  if (rules.length === 0) return;
+  const leaseIds = [
+    ...new Set(
+      input.leaseIds
+        .filter((leaseId): leaseId is string => typeof leaseId === 'string')
+        .map((leaseId) => leaseId.trim())
+        .filter(Boolean),
+    ),
+  ];
+  if (leaseIds.length === 0) return;
+  // DELETE by holder token is naturally idempotent, including HTTP retries.
+  await client.query(
+    `
+      DELETE FROM ${concurrencyLeaseTableName(input)}
+      WHERE bucket_id = $1
+        AND rule_id = ANY($2::text[])
+        AND lease_id = ANY($3::text[])
+    `,
+    [input.bucketId, rules.map((rule) => rule.ruleId), leaseIds],
+  );
+}
+
+export async function penalizePostgresRateState(
+  client: PostgresSchedulerQueryClient,
+  input: PlayRunnerRateStatePenalizeInput & { nowMs?: number },
+): Promise<void> {
+  if (!Number.isFinite(input.cooldownMs) || input.cooldownMs <= 0) return;
+  // Single autocommit statement applied to every rule row of the bucket:
+  //  - claimed_retry_at_ms: server Retry-After stored verbatim (observability).
+  //  - cool_until_ms: extended by the cooldown, internally capped at 45min.
+  //  - refill_per_sec: multiplicative decrease, but ONLY once per active
+  //    cooldown window (halve-once guard: skip if already cooling), floored at
+  //    min_rps.
+  //  - seen_backpressure flips the ramp from slow-start (x2) to additive.
+  //  - tokens drained: no more grants until refill re-accrues past the cooldown.
+  await client.query(
+    `
+      WITH now_value AS (
+        SELECT COALESCE(
+          $3::bigint,
+          FLOOR(EXTRACT(EPOCH FROM clock_timestamp()) * 1000)::bigint
+        ) AS now_ms
+      )
+      UPDATE ${bucketTableName(input)} b
+      SET
+        claimed_retry_at_ms = (SELECT now_ms FROM now_value) + $2::bigint,
+        refill_per_sec = CASE
+          WHEN (SELECT now_ms FROM now_value) >= b.cool_until_ms
+            THEN GREATEST(b.min_rps, b.refill_per_sec / 2)
+          ELSE b.refill_per_sec
+        END,
+        cool_until_ms = GREATEST(
+          b.cool_until_ms,
+          (SELECT now_ms FROM now_value) + LEAST($2::bigint, $4::bigint)
+        ),
+        seen_backpressure = true,
+        tokens = 0,
+        next_dispatch_at_ms = GREATEST(
+          b.next_dispatch_at_ms,
+          (SELECT now_ms FROM now_value) + LEAST($2::bigint, $4::bigint)
+        )
+      WHERE b.bucket_id = $1
+    `,
+    [
+      input.bucketId,
+      Math.ceil(input.cooldownMs),
+      input.nowMs ?? null,
+      MAX_INTERNAL_COOLDOWN_MS,
+    ],
+  );
+}