use crate::admin::{CallbackConfig, CallbackPollResult}; use crate::audited_sql; use crate::dlq::{ListDlqFilter, RetryFromDlqOpts}; use crate::error::AwaError; use crate::insert::prepare_row_raw; use crate::output::AttemptError; use crate::storage_footprint_sql::brin_fsm_empty_pages_sql; use crate::{InsertParams, JobRow, JobState}; use chrono::TimeDelta; use chrono::{DateTime, Utc}; use sqlx::Acquire; use sqlx::{PgPool, Postgres}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; use std::sync::Mutex; use std::time::{Duration, Instant}; use uuid::Uuid; /// Transaction-local setting that marks a transaction holding a receipt /// attempt lock, so a later cancel in it does not wait for another. const ATTEMPT_LOCKS_HELD_SETTING: &str = "awa.attempt_locks_held"; const DEFAULT_SCHEMA: &str = "awa"; const DEFAULT_QUEUE_SLOT_COUNT: usize = 16; const DEFAULT_LEASE_SLOT_COUNT: usize = 8; const DEFAULT_CLAIM_SLOT_COUNT: usize = 8; const DEFAULT_QUEUE_STRIPE_COUNT: usize = 1; const QUEUE_STRIPE_DELIMITER: &str = "#"; const COPY_NULL_SENTINEL: &str = "__AWA_NULL__"; const COPY_CHUNK_TARGET_BYTES: usize = 256 * 1024; /// Batches smaller than this are written with a single `unnest` INSERT even on /// the COPY entry points; below it the INSERT is the faster path. pub(crate) const COPY_MIN_ROWS: usize = 2048; /// Claim and claim-cursor transactions commit without waiting for the WAL /// flush. A claim lost to a server crash leaves its job claimable again, and /// any later synchronous commit flushes every earlier claim first. const CLAIM_TRANSACTION_BEGIN: &str = "BEGIN; SET LOCAL synchronous_commit = off"; const TERMINAL_COUNTER_BUCKETS: i16 = 256; const RECEIPT_RESCUE_BATCH_LIMIT: i64 = 500; const DRAIN_BATCH_SIZE: i64 = 1_000; // The planner overestimates the drain statements enough to JIT-compile them, // which takes seconds against milliseconds of execution. const DRAIN_TRANSACTION_BEGIN: &str = "BEGIN; SET LOCAL jit = off"; const DRAIN_CLAIM_CHECK_ROWS: i64 = 64; const RECEIPT_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000; const RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000; const DEFAULT_PRUNE_LOCK_TIMEOUT: Duration = Duration::from_millis(10); // Physical bytes above empty-index overhead on PostgreSQL 17/18. Installed // btree/GiST main forks retain one page; BRIN retains three main // and three FSM pages (four below 4 KiB block size) when empty. Unknown // access methods conservatively count any allocation as reclaimable. Include // heap/TOAST/index forks: VACUUM can leave an empty heap with expanded indexes. // Accept one child or a partitioned parent (the lease-ring idle probe). const RECLAIMABLE_CHILD_STORAGE_SQL: &str = concat!( r#" WITH heaps AS ( SELECT oid, reltoastrelid FROM pg_class WHERE oid = $1::regclass AND relkind = 'r' UNION ALL SELECT c.oid, c.reltoastrelid FROM pg_inherits i JOIN pg_class c ON c.oid = i.inhrelid WHERE i.inhparent = $1::regclass AND c.relkind = 'r' ), heap_relations AS ( SELECT oid FROM heaps UNION ALL SELECT reltoastrelid FROM heaps WHERE reltoastrelid <> 0 ), storage_relations AS ( SELECT oid, NULL::name AS amname FROM heap_relations UNION ALL SELECT i.indexrelid, am.amname FROM pg_index i JOIN pg_class c ON c.oid = i.indexrelid JOIN pg_am am ON am.oid = c.relam WHERE i.indrelid IN (SELECT oid FROM heap_relations) ) SELECT EXISTS ( SELECT 1 FROM storage_relations CROSS JOIN (SELECT current_setting('block_size')::bigint AS block_size) settings WHERE pg_relation_size(oid, 'fsm') > current_setting('block_size')::bigint * CASE amname WHEN 'brin' THEN "#, brin_fsm_empty_pages_sql!(), r#" ELSE 0 END OR pg_relation_size(oid, 'vm') > 0 OR pg_relation_size(oid) > current_setting('block_size')::bigint * CASE amname WHEN 'btree' THEN 1 WHEN 'gist' THEN 1 WHEN 'brin' THEN 3 ELSE 0 END ) "# ); /// Minimum age of another backend's active snapshot before the horizon-gated /// folds (ledger trim, terminal-delta fold, delta rollup) treat it as pinning /// the MVCC horizon and stand down (#371, ADR-040). /// /// The folds DELETE rows whose versions are only immediately reclaimable if no /// snapshot older than the DELETE still needs them. The gate exists to skip /// while such a snapshot is live. Its first implementation stood down whenever /// *any* backend had `backend_xmin` set — but `backend_xmin` is set for the /// duration of *every* ordinary query, including sub-millisecond hot-path /// claims and enqueues. Under any continuous traffic there is always such a /// backend, so the folds skipped on essentially every tick and the ledgers and /// `queue_terminal_rollup_deltas` grew without bound for the whole run. That is /// the opposite of the intent: a transient autocommit snapshot is gone long /// before the next 30s fold tick, so it can never strand the tuples a fold /// deletes now. /// /// We therefore only count an *active* backend as pinning once its transaction /// has been open longer than this threshold — comfortably above normal /// query/claim latency (single-digit ms) yet well under the 30s fold cadence, /// so a genuinely long-lived reader still stands the folds down before their /// deletions accumulate. Idle-in-transaction backends holding a live xid are /// always treated as pinning regardless of age (they hold their snapshot /// indefinitely with no query in flight to bound it). const MVCC_HORIZON_PIN_MIN_AGE: Duration = Duration::from_secs(5); /// Portable 64-bit hash over raw ordering-key bytes. /// /// This is intentionally simple enough to implement byte-for-byte in /// `awa.insert_job_compat`, so SQL, Rust, and Python producers that /// pass the same ordering-key bytes can derive the same routing facts. pub fn ordering_key_hash64(ordering_key: &[u8]) -> u64 { let mut hash: u128 = 14_695_981_039_346_656_037; const PRIME: u128 = 1_099_511_628_211; const MASK: u128 = u64::MAX as u128; for byte in ordering_key { hash = hash.wrapping_mul(PRIME).wrapping_add(*byte as u128) & MASK; } hash as u64 } /// Deterministically map an ordering key to a shard in `[0, shards)`. /// /// Inputs sharing a key always produce the same shard, which is what /// lets producers preserve FIFO within a key when the destination /// queue is sharded. `shards <= 1` returns shard 0 unconditionally. /// /// This deliberately remains the raw `ordering_key_hash64(key) % shards` /// mapping used by SQL compatibility producers. pub fn shard_for_ordering_key(ordering_key: &[u8], shards: i16) -> i16 { if shards <= 1 { return 0; } (ordering_key_hash64(ordering_key) % shards as u64) as i16 } fn terminal_counter_bucket(job_id: i64) -> i16 { job_id.rem_euclid(TERMINAL_COUNTER_BUCKETS as i64) as i16 } #[derive(Debug, Clone)] pub struct QueueStorageConfig { pub schema: String, pub queue_slot_count: usize, pub lease_slot_count: usize, /// Number of child partitions the receipt ring splits /// `lease_claims`, `lease_claim_batches`, and receipt-closure evidence /// across (ADR-023). /// Mirrors `lease_slot_count`: a small fixed set of slots /// reclaimed by rotation + TRUNCATE rather than by row-level /// DELETE. pub claim_slot_count: usize, pub queue_stripe_count: usize, /// Use the receipt-plane short path: claim writes compact batches into /// `lease_claim_batches` for both zero-deadline and deadline-backed /// attempts (#246) — a batch shares one `claimed_at`, hence one /// `deadline_at`, and expired batches are swept by the compact batch /// deadline-rescue cursor. Row-local `lease_claims` evidence remains /// readable for legacy in-flight claims mid-upgrade. Compact claims /// can still materialize into `leases` when the worker later needs mutable /// attempt state. Successful compact /// completion writes durable terminal history through /// `receipt_completion_batches` and compact claim-local closure evidence /// in `lease_claim_closure_batches`, or falls back to `done_entries`, while /// non-success exits and cold terminal-delete paths materialize explicit /// closures in `lease_claim_closures`. /// Receipt claim evidence is reclaimed by claim-ring rotation + /// TRUNCATE. /// Default `true`. /// Set to `false` to force every claim through the legacy /// `leases` materialization path. pub lease_claim_receipts: bool, } impl Default for QueueStorageConfig { fn default() -> Self { Self { schema: DEFAULT_SCHEMA.to_string(), queue_slot_count: DEFAULT_QUEUE_SLOT_COUNT, lease_slot_count: DEFAULT_LEASE_SLOT_COUNT, claim_slot_count: DEFAULT_CLAIM_SLOT_COUNT, queue_stripe_count: DEFAULT_QUEUE_STRIPE_COUNT, lease_claim_receipts: true, } } } #[derive(Debug, Clone, PartialEq, Eq, sqlx::FromRow)] pub struct ClaimedEntry { pub queue: String, pub priority: i16, pub lane_seq: i64, pub ready_slot: i32, pub ready_generation: i64, pub lease_slot: i32, pub lease_generation: i64, /// ADR-023: the `claim_slot` partition this attempt's receipt /// evidence landed in. Receipt-closure paths use this to co-locate /// explicit `lease_claim_closures`; compact successful receipt /// completions keep claim-local evidence in /// `lease_claim_closure_batches` instead. pub claim_slot: i32, /// Stable identity for immutable receipt claim evidence. Present for /// receipt-backed claims and used by the compact completion path to /// validate the exact claim attempt without row-locking it. pub receipt_id: Option, /// Compact claim batch row identity for receipt claims (zero-deadline /// and deadline-backed alike since #246). Row-local receipt claims /// leave this unset. pub claim_batch_id: Option, /// One-based item position inside `lease_claim_batches.*_ids` arrays. /// Present with `claim_batch_id` and used as an O(1) compact proof. pub claim_batch_index: Option, pub lease_claim_receipt: bool, /// The enqueue shard the row was claimed from. Routes the /// terminal `done_entries` write onto the correct shard's /// `(ready_slot, queue, priority, enqueue_shard, lane_seq)` key /// and is the join predicate for receipt and admin lookups that /// touch `queue_claim_heads` / `ready_entries` / `leases`. pub enqueue_shard: i16, } #[derive(Debug, Clone)] pub struct ClaimedRuntimeJob { pub claim: ClaimedEntry, pub job: JobRow, pub unique_states: Option, } /// Identity of a compact receipt claim for heartbeat purposes: addresses /// `lease_claim_batches` by primary key and the attempt by item index, so a /// heartbeat does not have to scan every open claim batch. Built with /// [`ClaimHeartbeatRef::from_claim`]. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[non_exhaustive] pub struct ClaimHeartbeatRef { pub job_id: i64, pub run_lease: i64, pub claim_slot: i32, pub claim_batch_id: i64, pub claim_batch_index: i32, pub receipt_id: i64, } impl ClaimHeartbeatRef { /// The heartbeat identity of a compact receipt claim, or `None` for /// legacy or row-local claims. pub fn from_claim(job_id: i64, run_lease: i64, claim: &ClaimedEntry) -> Option { if !claim.lease_claim_receipt { return None; } Some(Self { job_id, run_lease, claim_slot: claim.claim_slot, claim_batch_id: claim.claim_batch_id?, claim_batch_index: claim.claim_batch_index?, receipt_id: claim.receipt_id?, }) } } #[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)] pub struct QueueClaimerLease { pub claimer_slot: i16, pub lease_epoch: i64, } /// How long a claim may reuse this instance's cached claimer lease before /// re-reading the lease row and the queue's claimer target. const CLAIMER_GATE_CACHE_TTL: Duration = Duration::from_millis(250); #[derive(Debug, Clone, Copy)] struct CachedClaimerGate { checked_at: Instant, lease: QueueClaimerLeaseRow, } #[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)] struct QueueClaimerLeaseRow { claimer_slot: i16, lease_epoch: i64, last_claimed_at: DateTime, expires_at: DateTime, } impl QueueClaimerLeaseRow { fn lease(self) -> QueueClaimerLease { QueueClaimerLease { claimer_slot: self.claimer_slot, lease_epoch: self.lease_epoch, } } fn needs_refresh( self, now: DateTime, lease_ttl: Duration, idle_threshold: Duration, ) -> bool { let Ok(idle_refresh_delta) = TimeDelta::from_std(idle_threshold / 2) else { return true; }; let Ok(expiry_refresh_delta) = TimeDelta::from_std(lease_ttl / 2) else { return true; }; self.last_claimed_at <= now - idle_refresh_delta || self.expires_at <= now + expiry_refresh_delta } } #[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)] pub struct QueueClaimerState { pub target_claimers: i16, } impl ClaimedRuntimeJob { fn into_done_row(self, finalized_at: DateTime) -> Result { let payload = QueueStorage::payload_from_parts( self.job.metadata, self.job.tags, self.job.errors, None, )?; Ok(DoneJobRow { ready_slot: self.claim.ready_slot, ready_generation: self.claim.ready_generation, job_id: self.job.id, kind: self.job.kind, queue: self.job.queue, args: self.job.args, state: JobState::Completed, priority: self.claim.priority, attempt: self.job.attempt, run_lease: self.job.run_lease, max_attempts: self.job.max_attempts, lane_seq: self.claim.lane_seq, enqueue_shard: self.claim.enqueue_shard, run_at: self.job.run_at, attempted_at: self.job.attempted_at, finalized_at, created_at: self.job.created_at, unique_key: self.job.unique_key, unique_states: self.unique_states, payload, output: self.job.output, }) } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct QueueCounts { pub available: i64, pub running: i64, /// Count of rows in *any* terminal state (`completed`, `failed`, or /// `cancelled`) for the queue. The name reflects what the field /// actually counts: it is `{schema}.terminal_jobs` semantics, not /// `count(*) WHERE state = 'completed'`. Physical terminal facts may /// live in `done_entries` or compact receipt completion batches. /// The historical name `completed` was a misnomer — /// `queue_counts_exact` has always included failed and cancelled /// terminals; renamed in #290 along with the counter-backed read /// path. pub terminal: i64, /// Cumulative count of `failed` terminal rows pruned past the /// failed-retention floor, summed from /// `queue_terminal_rollups.pruned_failed_count`. These rows no /// longer exist in `done_entries` and cannot be retried. /// Monotonically non-decreasing — rollups never shrink. pub pruned_failed: i64, } /// Cheap available-only signal used by the dispatcher's claimer-sizing /// control loop. Derives the count from enqueue and claim sequence cursors /// summed over the queue's physical stripes — two PK reads per lane, O(few /// rows) regardless of backlog size. /// /// This is intentionally a separate type from [`QueueCounts`]: the /// dispatcher claim hot path only consumes the available count, and /// returning a `QueueCounts` with two perpetually-zero fields would /// invite future code to read `.running` or `.terminal` and silently /// get wrong answers. Code that legitimately needs the full counts /// should call [`QueueStorage::queue_counts`]. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct AvailableSignal { pub available: i64, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RotateOutcome { Rotated { slot: i32, generation: i64, }, /// Target slot has live state; rotation deferred. `busy` carries a /// bounded per-table presence indicator observed at the gate (only fields /// relevant to the ring being rotated are populated). SkippedBusy { slot: i32, busy: BusyCounts, }, /// Ring is idle: every child table of the *current* slot is empty, so /// there is nothing to seal and rotation was skipped without updating /// any ring-state row. Issue #371: timer-driven rotation of an idle /// ring UPDATEs the singleton ring-state row every tick, and under a /// pinned MVCC horizon each of those dead tuples stays unreclaimable. /// `slot` is the current slot, which stays open. SkippedIdle { slot: i32, /// Cursor generation left unchanged by the idle skip. generation: i64, }, /// The current slot holds rows, is smaller than /// [`RingRotationCadence::target_slot_bytes`] and younger than /// [`RingRotationCadence::max_slot_age`], and is either younger than /// [`RingRotationCadence::min_slot_age`] or holds fewer than /// [`RingRotationCadence::min_slot_rows`], so it stays open and keeps /// filling. Sealing it now would spend a prune (`TRUNCATE` of every /// child that received rows, one `pg_class` row version per relation) /// on a slot that has barely been used. `slot` is the current slot. SkippedFilling { slot: i32, /// Cursor generation left unchanged by the skip. generation: i64, }, } /// When a maintenance tick may seal the current slot of a ring. /// /// A non-empty slot becomes eligible to seal once it is `min_slot_age` past /// this store's first observation of rows and holds `min_slot_rows`, when its /// total physical storage reaches `target_slot_bytes`, or once it is /// `max_slot_age` old. The minimum age uses a local monotonic clock, and a new /// store starts a fresh observation window. The maximum age also counts from /// when the slot opened on the database clock, so a change of maintenance /// leader does not restart it. Observing an empty slot clears its timer and /// lets the existing busy, idle and reclamation gates decide rotation. /// /// The byte target includes heap, indexes, TOAST and auxiliary forks. It is /// sampled on rotate calls and is not a storage limit: batches, delayed ticks, /// blocked reuse and writes to sealed slots can exceed it. /// /// [`RingRotationCadence::EAGER`] restores sealing on every tick. Build other /// cadences from [`RingRotationCadence::DEFAULT`] or `EAGER` with the `with_*` /// methods, so that later fields do not break callers. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[non_exhaustive] pub struct RingRotationCadence { /// Time since the first non-empty observation before a tick may seal. pub min_slot_age: Duration, /// Physical size that makes a non-empty slot eligible regardless of age. /// Zero disables the size trigger. pub target_slot_bytes: u64, /// Rows the slot must hold before `min_slot_age` alone seals it, so a /// trickle of writes does not seal a nearly empty slot every window. /// Zero disables the row floor. pub min_slot_rows: u64, /// Age after which a non-empty slot seals regardless of `min_slot_rows`, /// bounding how long a trickle keeps a slot open. Zero disables it. pub max_slot_age: Duration, } impl RingRotationCadence { /// Seal on every tick: the pre-cadence behaviour. pub const EAGER: Self = Self { min_slot_age: Duration::ZERO, target_slot_bytes: 0, min_slot_rows: 0, max_slot_age: Duration::ZERO, }; /// Eligible after five seconds of observed filling once the slot holds /// 256 rows, at 2 MiB regardless of age, or after 60 seconds regardless /// of rows. pub const DEFAULT: Self = Self { min_slot_age: Duration::from_secs(5), target_slot_bytes: 2 * 1024 * 1024, min_slot_rows: 256, max_slot_age: Duration::from_secs(60), }; /// Set [`Self::min_slot_age`]. pub const fn with_min_slot_age(mut self, min_slot_age: Duration) -> Self { self.min_slot_age = min_slot_age; self } /// Set [`Self::target_slot_bytes`]. pub const fn with_target_slot_bytes(mut self, target_slot_bytes: u64) -> Self { self.target_slot_bytes = target_slot_bytes; self } /// Set [`Self::min_slot_rows`]. pub const fn with_min_slot_rows(mut self, min_slot_rows: u64) -> Self { self.min_slot_rows = min_slot_rows; self } /// Set [`Self::max_slot_age`]. pub const fn with_max_slot_age(mut self, max_slot_age: Duration) -> Self { self.max_slot_age = max_slot_age; self } } impl Default for RingRotationCadence { fn default() -> Self { Self::DEFAULT } } /// Per-table presence observed at a rotation gate. Each non-zero value means /// "this table was non-empty"; it is not an exact row count. Each ring populates /// only the fields meaningful for it; unused fields stay zero. The /// maintenance loop emits one OTel metric label per non-zero field so /// dashboards can attribute "rotation pinned" to the responsible side. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct BusyCounts { /// Queue ring: rows in the next `ready_entries` child. pub queue_ready: i64, /// Queue ring: rows in the next `ready_claim_attempt_batches` child. pub queue_claim_attempt_batches: i64, /// Queue ring: rows in the next `done_entries` child. pub queue_done: i64, /// Queue ring: rows in the next `ready_tombstones` child. pub queue_tombstones: i64, /// Queue ring: rows in the next `ready_segments` child. pub queue_ready_segments: i64, /// Queue ring: rows in the next `receipt_completion_batches` child. pub queue_receipt_completion_batches: i64, /// Queue ring: rows in the next `receipt_completion_tombstones` child. pub queue_receipt_completion_tombstones: i64, /// Queue ring: rows in the next `queue_terminal_count_deltas` child. pub queue_terminal_deltas: i64, /// Lease ring: rows in the next `leases` child. pub leases: i64, /// Claim ring: rows in the next `lease_claims` child. pub claims: i64, /// Claim ring: rows in the next `lease_claim_closures` child. pub closures: i64, /// Claim ring: rows in the next compact closure-batch child. pub closure_batches: i64, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PruneOutcome { Noop, Pruned { slot: i32, /// Generation of the sealed slot that was reclaimed. generation: i64, /// Failed terminal rows inside the failed-retention floor that /// the queue prune re-homed into the live `done_entries` /// segment instead of dropping. Always zero for the lease and /// claim rings. carried_failed_rows: u64, /// Time spent in the database phases that can force a WAL flush. durations: PruneDurations, }, /// Lock acquisition timed out (held-tx, lock contention). Blocked { slot: i32, }, /// Target slot still has live state. `reason` discriminates which gate /// fired and `count` gives its magnitude. SkippedActive { slot: i32, reason: SkipReason, count: i64, }, } /// Database phase timings for a successful destructive ring prune. /// /// These phases are separated because managed Postgres can spend most of a /// prune inside a durability wait in `TRUNCATE` or `COMMIT`, even after the /// bounded ACCESS EXCLUSIVE acquisition has succeeded. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct PruneDurations { /// Time acquiring `ACCESS EXCLUSIVE` on the child relations. pub lock: Duration, /// Time executing the destructive `TRUNCATE` statement. pub truncate: Duration, /// Time committing the prune transaction. pub commit: Duration, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct TerminalDeltaRollupOutcome { pub rolled_slots: usize, pub delta_rows: i64, pub grouped_keys: i64, pub skipped_active_slots: usize, pub blocked_slots: usize, pub skipped_mvcc_pinned: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum TerminalDeltaSlotRollup { Empty, Rolled { delta_rows: i64, grouped_keys: i64 }, SkippedActive, SkippedMvccPinned, Blocked, } /// Outcome of [`QueueStorage::fold_terminal_rollup_deltas`] (#371). #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct TerminalRollupFoldOutcome { /// Delta rows consumed (deleted) by the fold. pub folded_delta_rows: usize, /// Distinct (queue, priority) keys upserted into the rollups. pub folded_keys: usize, /// True when the fold stood down because another backend pinned the /// MVCC horizon; the delta rows stay visible to exact readers. pub skipped_mvcc_pinned: bool, } /// Outcome of [`QueueStorage::fold_ring_rotation_ledgers`] (#371). #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct RingLedgerFoldOutcome { /// Ledger rows trimmed across the queue, lease, and claim rings. pub trimmed_rows: u64, /// True when the fold stood down because another backend pinned the /// MVCC horizon; the ledgers keep their full history until the next /// clear-horizon tick. pub skipped_mvcc_pinned: bool, } /// Discriminator for [`PruneOutcome::SkippedActive`]. /// /// Multiple gates can fire `SkippedActive` for the same ring (e.g. queue /// prune checks both `active_leases` and `pending_ready`). Carrying the /// reason separately from `count` lets dashboards split out "ring saturated /// because backlog never drained" from "leases lingering on prior /// generation" without re-parsing log lines. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SkipReason { /// Queue prune: leases on the prior generation persist. QueueActiveLeases, /// Queue prune: receipt claims still need same-slot terminal evidence. QueueUnclosedClaimRefs, /// Queue prune: ready rows without matching done or tombstone evidence. QueuePendingReady, /// Lease prune: target slot equals the current slot (rotator race). LeaseCurrent, /// Lease prune: pending leases on target slot. LeaseActive, /// Claim prune: target slot equals the current slot (rotator race). ClaimCurrent, /// Claim prune: open claims on target slot (no closure evidence). ClaimOpen, } impl SkipReason { /// Stable, low-cardinality label suitable for OTel metric attributes. pub fn as_str(self) -> &'static str { match self { Self::QueueActiveLeases => "queue.active_leases", Self::QueueUnclosedClaimRefs => "queue.unclosed_claim_refs", Self::QueuePendingReady => "queue.pending_ready", Self::LeaseCurrent => "lease.current", Self::LeaseActive => "lease.active", Self::ClaimCurrent => "claim.current", Self::ClaimOpen => "claim.open", } } } fn map_sqlx_error(err: sqlx::Error) -> AwaError { if let sqlx::Error::Database(ref db_err) = err { if db_err.code().as_deref() == Some("23505") { return AwaError::UniqueConflict { constraint: db_err.constraint().map(|c| c.to_string()), }; } } AwaError::Database(err) } fn is_lock_contention_error(err: &sqlx::Error) -> bool { matches!( err, sqlx::Error::Database(db_err) if db_err.code().as_deref() == Some("55P03") ) } async fn set_prune_lock_timeout_tx( tx: &mut sqlx::Transaction<'_, Postgres>, timeout: Duration, ) -> Result<(), AwaError> { let millis = timeout.as_millis().max(1).min(i64::MAX as u128); let timeout = format!("{millis}ms"); sqlx::query("SELECT set_config('lock_timeout', $1, true)") .bind(timeout) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } fn validate_ident(ident: &str) -> Result<(), AwaError> { let mut chars = ident.chars(); match chars.next() { Some(first) if first.is_ascii_lowercase() || first == '_' => {} _ => { return Err(AwaError::Validation(format!( "invalid SQL identifier: {ident}" ))); } } if chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_') { Ok(()) } else { Err(AwaError::Validation(format!( "invalid SQL identifier: {ident}" ))) } } /// Gate a schema name that came back from the database before it is /// interpolated into query text. Mirrors the [`validate_ident`] check /// [`QueueStorage::new`] applies to configured names. fn validate_active_schema(schema: String) -> Result { validate_ident(&schema).map_err(|_| { AwaError::Validation(format!( "active queue-storage schema name is not a valid SQL identifier: {schema:?}" )) })?; Ok(schema) } fn ready_child_name(schema: &str, slot: usize) -> String { format!("{schema}.ready_entries_{slot}") } fn ready_claim_attempt_batch_child_name(schema: &str, slot: usize) -> String { format!("{schema}.ready_claim_attempt_batches_{slot}") } fn done_child_name(schema: &str, slot: usize) -> String { format!("{schema}.done_entries_{slot}") } fn ready_tombstone_child_name(schema: &str, slot: usize) -> String { format!("{schema}.ready_tombstones_{slot}") } fn ready_segment_child_name(schema: &str, slot: usize) -> String { format!("{schema}.ready_segments_{slot}") } fn receipt_completion_batch_child_name(schema: &str, slot: usize) -> String { format!("{schema}.receipt_completion_batches_{slot}") } fn receipt_completion_tombstone_child_name(schema: &str, slot: usize) -> String { format!("{schema}.receipt_completion_tombstones_{slot}") } fn terminal_delta_child_name(schema: &str, slot: usize) -> String { format!("{schema}.queue_terminal_count_deltas_{slot}") } fn done_ready_join(schema: &str, done_alias: &str, ready_alias: &str) -> String { format!( r#" LEFT JOIN {schema}.ready_entries AS {ready_alias} ON {ready_alias}.ready_slot = {done_alias}.ready_slot AND {ready_alias}.ready_generation = {done_alias}.ready_generation AND {ready_alias}.queue = {done_alias}.queue AND {ready_alias}.priority = {done_alias}.priority AND {ready_alias}.enqueue_shard = {done_alias}.enqueue_shard AND {ready_alias}.lane_seq = {done_alias}.lane_seq "# ) } fn done_row_projection(done_alias: &str, ready_alias: &str) -> String { format!( r#" {done_alias}.ready_slot, {done_alias}.ready_generation, {done_alias}.job_id, {done_alias}.kind, {done_alias}.queue, COALESCE({done_alias}.args, {ready_alias}.args, '{{}}'::jsonb) AS args, {done_alias}.state, {done_alias}.priority, {done_alias}.attempt, {done_alias}.run_lease, COALESCE({done_alias}.max_attempts, {ready_alias}.max_attempts, 25::smallint) AS max_attempts, {done_alias}.lane_seq, {done_alias}.enqueue_shard, COALESCE({done_alias}.run_at, {ready_alias}.run_at, {done_alias}.finalized_at) AS run_at, COALESCE({done_alias}.attempted_at, {ready_alias}.attempted_at) AS attempted_at, {done_alias}.finalized_at, COALESCE({done_alias}.created_at, {ready_alias}.created_at, {done_alias}.finalized_at) AS created_at, COALESCE({done_alias}.unique_key, {ready_alias}.unique_key) AS unique_key, COALESCE({done_alias}.unique_states, {ready_alias}.unique_states) AS unique_states, COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) AS payload "# ) } fn lease_child_name(schema: &str, slot: usize) -> String { format!("{schema}.leases_{slot}") } fn claim_child_name(schema: &str, slot: usize) -> String { format!("{schema}.lease_claims_{slot}") } /// SQL true when `column` names queue `param` or one of its stripes /// (`queue#N`). fn queue_and_stripes_sql(column: &str, param: &str) -> String { format!( "({column} = {param} \ OR (left({column}, char_length({param}) + 1) = {param} || '{QUEUE_STRIPE_DELIMITER}' \ AND substr({column}, char_length({param}) + 2) ~ '^[0-9]+$'))" ) } fn claim_batch_child_name(schema: &str, slot: usize) -> String { format!("{schema}.lease_claim_batches_{slot}") } fn closure_child_name(schema: &str, slot: usize) -> String { format!("{schema}.lease_claim_closures_{slot}") } fn claim_closure_batch_child_name(schema: &str, slot: usize) -> String { format!("{schema}.lease_claim_closure_batches_{slot}") } fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result { let index = usize::try_from(slot).map_err(|_| { AwaError::Validation(format!( "invalid {ring} ring slot {slot}: slot must be non-negative" )) })?; if index >= slot_count { return Err(AwaError::Validation(format!( "invalid {ring} ring slot {slot}: configured slot count is {slot_count}" ))); } Ok(index) } fn receipt_closed_evidence_sql( schema: &str, closure_rel: &str, closure_batch_rel: &str, claims_alias: &str, ) -> String { // `receipt_id` is allocated from a global sequence, so it identifies a // claim without a `claim_slot` predicate. Compact closure batches retain // the raw receipt array for audit/debugging, but hot membership checks use // the per-child GiST index on the derived multirange. format!( r#" ( {claims_alias}.closed_at IS NOT NULL OR EXISTS ( SELECT 1 FROM {closure_rel} AS closures WHERE closures.claim_slot = {claims_alias}.claim_slot AND closures.job_id = {claims_alias}.job_id AND closures.run_lease = {claims_alias}.run_lease ) OR EXISTS ( SELECT 1 FROM {closure_batch_rel} AS closure_batches WHERE closure_batches.receipt_ranges @> {claims_alias}.receipt_id ) OR EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = {claims_alias}.job_id AND done.run_lease = {claims_alias}.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = {claims_alias}.job_id AND deferred.run_lease = {claims_alias}.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = {claims_alias}.job_id AND dlq.run_lease = {claims_alias}.run_lease ) ) "# ) } /// Each queue's available rows, exactly, and the oldest of them, as /// `(queue, available, oldest_run_at)`. /// /// Each lane's unclaimed rows are found the way a claim finds them: through /// `ready_segments` from the lane's claim cursor, then along the ready /// primary key within each segment's ring slot. The ready key leads with the /// ring slot, so a lane range without it cannot use the index and would scan /// every retained ready row. The cost follows the backlog. /// /// `oldest_run_at` is the earliest `run_at` among those rows, not the first /// row's: a queue move or priority aging appends an older job behind newer /// ones. pub fn available_by_queue_sql(schema: &str) -> String { format!( r#" SELECT heads.queue, COALESCE(sum(lane.available), 0)::bigint AS available, min(lane.oldest_run_at) AS oldest_run_at FROM ( SELECT claims.queue, claims.priority, claims.enqueue_shard, {schema}.sequence_next_value(claims.seq_name) AS claim_seq, {schema}.sequence_next_value(enqueues.seq_name) AS next_seq FROM {schema}.queue_claim_heads AS claims JOIN {schema}.queue_enqueue_heads AS enqueues ON enqueues.queue = claims.queue AND enqueues.priority = claims.priority AND enqueues.enqueue_shard = claims.enqueue_shard ) AS heads CROSS JOIN LATERAL ( {lane_rows} ) AS lane GROUP BY heads.queue "#, lane_rows = available_lane_rows_sql( schema, "heads", "count(*)::bigint AS available, min(ready.run_at) AS oldest_run_at", "TRUE", ), ) } /// A `SELECT {columns} FROM ready_segments JOIN ready_entries` over the /// unclaimed, untombstoned rows of the lane in `heads` (which must carry /// `queue`, `priority`, `enqueue_shard`, `claim_seq` and `next_seq`) that /// satisfy `condition` on `ready`. See [`available_by_queue_sql`] for why it /// walks segments. pub(crate) fn available_lane_rows_sql( schema: &str, heads: &str, columns: &str, condition: &str, ) -> String { format!( r#" SELECT {columns} FROM ( -- The lane's segments from its claim cursor, in lane order -- along the lane index. Without the order the planner can -- prefer walking every segment by ring slot. SELECT segment.ready_slot, segment.ready_generation, GREATEST(segment.first_lane_seq, {heads}.claim_seq) AS from_seq, LEAST(segment.next_lane_seq, {heads}.next_seq) AS to_seq FROM {schema}.ready_segments AS segment WHERE segment.queue = {heads}.queue AND segment.priority = {heads}.priority AND segment.enqueue_shard = {heads}.enqueue_shard AND segment.next_lane_seq > {heads}.claim_seq AND segment.first_lane_seq < {heads}.next_seq ORDER BY segment.next_lane_seq ) AS segment -- Each segment's rows along the ready key within its ring slot. -- OFFSET 0 keeps this a per-segment probe: flattened into a -- join, a filter such as a creation bound can lead the planner -- to scan every retained ready row and match segments after. CROSS JOIN LATERAL ( SELECT ready.* FROM {schema}.ready_entries AS ready WHERE ready.ready_slot = segment.ready_slot AND ready.ready_generation = segment.ready_generation AND ready.queue = {heads}.queue AND ready.priority = {heads}.priority AND ready.enqueue_shard = {heads}.enqueue_shard AND ready.lane_seq >= segment.from_seq AND ready.lane_seq < segment.to_seq OFFSET 0 ) AS ready WHERE ({condition}) AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation AND tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq ) "# ) } /// What [`queue_signals_sql`] reads. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum QueueSignalReads { /// No receipt-claim table, so claim-ring rotation's `ACCESS EXCLUSIVE` /// locks never block it: `running` counts materialised leases only, and /// `completed_last_hour` is 0. For gauges polled by the maintenance loop. Gauge, /// Counts available rows exactly, and also counts receipt-backed running /// claims and completions in the last hour from completion batches and /// `done_entries`. For the admin queue views. Admin, } /// Per-queue counts that never expand terminal history, ending in a /// `signals` CTE with `queue`, `available`, `running`, `waiting_external`, /// `scheduled`, `retryable`, `failed_done`, `failed_dlq`, `lag_seconds` and /// `completed_last_hour`. /// /// With [`QueueSignalReads::Gauge`], `available` is the sequence-head signal, /// an upper bound as in [`QueueStorage::queue_counts_fast`]; with /// [`QueueSignalReads::Admin`] it is exact. Nothing here reads `terminal_jobs`, /// whose completion batches expand and join back to `ready_entries`, so the /// cost follows live work rather than retained history. pub fn queue_signals_sql(schema: &str, reads: QueueSignalReads) -> String { let (ready, lag, receipt_running, completed_recent) = match reads { QueueSignalReads::Gauge => ( // Enqueued minus claimed: an upper bound, since a tombstoned // row or a sequence gap still counts until the claimer passes. r#" SELECT head_signal.queue, COALESCE( sum(GREATEST(head_signal.next_seq - head_signal.claim_seq, 0)), 0 )::bigint AS available FROM head_signal GROUP BY head_signal.queue "# .to_string(), // The first unclaimed row per lane, read through its segment: two // index probes a lane, cheap enough for a gauge polled by every // maintenance leader. An older job appended behind newer ones by a // queue move or priority aging doesn't show until it is reached. format!( r#" SELECT head_signal.queue, EXTRACT(EPOCH FROM clock_timestamp() - min(next_ready.run_at))::double precision AS lag_seconds FROM head_signal JOIN LATERAL ( SELECT next_in_segment.run_at FROM {schema}.ready_segments AS segment CROSS JOIN LATERAL ( SELECT ready.run_at FROM {schema}.ready_entries AS ready WHERE ready.ready_slot = segment.ready_slot AND ready.ready_generation = segment.ready_generation AND ready.queue = segment.queue AND ready.priority = segment.priority AND ready.enqueue_shard = segment.enqueue_shard AND ready.lane_seq >= GREATEST(segment.first_lane_seq, head_signal.claim_seq) AND ready.lane_seq < segment.next_lane_seq AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation AND tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq ) ORDER BY ready.lane_seq LIMIT 1 ) AS next_in_segment WHERE segment.queue = head_signal.queue AND segment.priority = head_signal.priority AND segment.enqueue_shard = head_signal.enqueue_shard AND segment.next_lane_seq > head_signal.claim_seq ORDER BY segment.next_lane_seq LIMIT 1 ) AS next_ready ON TRUE GROUP BY head_signal.queue "# ), "SELECT NULL::text AS queue, 0::bigint AS running WHERE false".to_string(), "SELECT NULL::text AS queue, 0::bigint AS completed_last_hour WHERE false".to_string(), ), QueueSignalReads::Admin => ( available_by_queue_sql(schema), // The oldest available job, as before queue storage: it reads the // same backlog rows the exact count already reads. r#" SELECT queue, EXTRACT(EPOCH FROM clock_timestamp() - oldest_run_at)::double precision AS lag_seconds FROM ready WHERE oldest_run_at IS NOT NULL "# .to_string(), format!( "SELECT queue, count(*)::bigint AS running FROM ({}) AS open_claims GROUP BY queue", open_receipt_running_claims_sql(schema) ), format!( r#" WITH recent_tombstones AS MATERIALIZED ( SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, job_id, run_lease FROM {schema}.receipt_completion_tombstones WHERE deleted_at > clock_timestamp() - interval '1 hour' ), -- Recent batches in a lane with a recent tombstone: the only -- ones whose members need reading. Selected before expanding -- any member array, so retained batches are never expanded. tombstoned_batches AS MATERIALIZED ( SELECT batch.queue, batch.ready_slot, batch.job_ids, batch.run_leases FROM {schema}.receipt_completion_batches AS batch WHERE batch.finalized_at > clock_timestamp() - interval '1 hour' AND EXISTS ( SELECT 1 FROM recent_tombstones AS tomb WHERE tomb.ready_slot = batch.ready_slot AND tomb.ready_generation = batch.ready_generation AND tomb.queue = batch.queue AND tomb.priority = batch.priority AND tomb.enqueue_shard = batch.enqueue_shard ) ) SELECT queue, sum(completed)::bigint AS completed_last_hour FROM ( -- A batch carries its queue and size, so it is counted -- without unnesting its members. SELECT queue, completed_count::bigint AS completed FROM {schema}.receipt_completion_batches WHERE finalized_at > clock_timestamp() - interval '1 hour' UNION ALL -- Less its members deleted or retried since, as -- `terminal_jobs` drops them. A tombstone on a member -- finalised within the hour is itself newer than that. -- Each tombstoned batch is expanded once and its members -- matched by hash, so the work follows the last hour's -- completions and deletions, never one batch scan per -- tombstone. SELECT batch.queue, -1::bigint FROM tombstoned_batches AS batch CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS member(job_id, run_lease) JOIN recent_tombstones AS tomb ON tomb.ready_slot = batch.ready_slot AND tomb.job_id = member.job_id AND tomb.run_lease = member.run_lease UNION ALL SELECT queue, 1::bigint FROM {schema}.done_entries WHERE state = 'completed' AND finalized_at > clock_timestamp() - interval '1 hour' ) AS recent GROUP BY queue "# ), ), }; format!( r#" WITH head_signal AS ( SELECT enqueues.queue, enqueues.priority, enqueues.enqueue_shard, {schema}.sequence_next_value(enqueues.seq_name) AS next_seq, {schema}.sequence_next_value(claims.seq_name) AS claim_seq FROM {schema}.queue_enqueue_heads AS enqueues JOIN {schema}.queue_claim_heads AS claims ON claims.queue = enqueues.queue AND claims.priority = enqueues.priority AND claims.enqueue_shard = enqueues.enqueue_shard ), queues AS ( SELECT DISTINCT queue FROM ( SELECT queue FROM awa.queue_meta UNION ALL SELECT queue FROM head_signal UNION ALL SELECT queue FROM {schema}.leases UNION ALL SELECT queue FROM {schema}.deferred_jobs UNION ALL SELECT queue FROM {schema}.queue_terminal_live_counts UNION ALL SELECT queue FROM {schema}.queue_terminal_rollups UNION ALL SELECT queue FROM {schema}.dlq_entries ) queues ), ready AS ( {ready} ), lag AS ( {lag} ), leases AS ( -- Both legs are `leases`-only because this gauge must not touch -- the receipt tables (see the function-level lock note). The -- `running` leg is therefore an approximation (misses compact -- lease_claim_batches claims, #246). `waiting_external` is exact -- even so: a receipt claim cannot be `waiting_external` off the -- `leases` plane — QueueStorage::register_callback runs -- ensure_mutable_running_attempt_tx (materialises the claim into -- `leases`) before enter_callback_wait_in_tx flips -- `leases.state`. (Same waiting_external invariant, exact form, -- at admin::state_counts.) SELECT queue, count(*) FILTER (WHERE state = 'running')::bigint AS running, count(*) FILTER (WHERE state = 'waiting_external')::bigint AS waiting_external FROM {schema}.leases GROUP BY queue ), deferred AS ( SELECT queue, count(*) FILTER (WHERE state = 'scheduled')::bigint AS scheduled, count(*) FILTER (WHERE state = 'retryable')::bigint AS retryable FROM {schema}.deferred_jobs GROUP BY queue ), terminal AS ( SELECT queue, count(*)::bigint AS failed_done FROM {schema}.done_entries WHERE state = 'failed' GROUP BY queue ), dlq AS ( SELECT queue, count(*)::bigint AS failed_dlq FROM {schema}.dlq_entries GROUP BY queue ), receipt_running AS ( {receipt_running} ), completed_recent AS ( {completed_recent} ), signals AS ( SELECT queues.queue, COALESCE(ready.available, 0)::bigint AS available, (COALESCE(leases.running, 0) + COALESCE(receipt_running.running, 0))::bigint AS running, COALESCE(leases.waiting_external, 0)::bigint AS waiting_external, COALESCE(deferred.scheduled, 0)::bigint AS scheduled, COALESCE(deferred.retryable, 0)::bigint AS retryable, COALESCE(terminal.failed_done, 0)::bigint AS failed_done, COALESCE(dlq.failed_dlq, 0)::bigint AS failed_dlq, lag.lag_seconds, COALESCE(completed_recent.completed_last_hour, 0)::bigint AS completed_last_hour FROM queues LEFT JOIN ready ON ready.queue = queues.queue LEFT JOIN lag ON lag.queue = queues.queue LEFT JOIN leases ON leases.queue = queues.queue LEFT JOIN receipt_running ON receipt_running.queue = queues.queue LEFT JOIN deferred ON deferred.queue = queues.queue LEFT JOIN terminal ON terminal.queue = queues.queue LEFT JOIN dlq ON dlq.queue = queues.queue LEFT JOIN completed_recent ON completed_recent.queue = queues.queue ) "# ) } /// A subquery that yields the identity `(job_id, run_lease, queue)` of every /// receipt claim currently open (i.e. `running`) but not yet materialised /// into `leases`. Since #246 receipt claims — zero-deadline and /// deadline-backed alike — are written as compact `lease_claim_batches` rows /// rather than one `lease_claims` row per job, so the "open receipt claims" /// set is the union of legacy row-local `lease_claims` and the unnested /// members of `lease_claim_batches`, each anti-joined against every durable /// closure/terminal/materialised-lease evidence shape. /// /// This is the canonical open-receipt-claim shape (mirrors `queue_counts`'s /// `live_running` derivation). Compose it wherever a running enumeration must /// see receipt claims: wrap in `(SELECT count(*) FROM () …)`, join it, or /// `UNION ALL` its rows into a `current_jobs`-style view. `waiting_external` /// deliberately does NOT read this set — a claim can only reach /// `waiting_external` after `register_callback` materialises it into `leases` /// (see `enter_callback_wait_in_tx`), so those legs stay `leases`-only. /// /// `leases` here is materialised-claim evidence; a claim present in `leases` /// is reported by the caller's own `leases` leg, so it is excluded here to /// avoid double counting. pub fn open_receipt_running_claims_sql(schema: &str) -> String { open_receipt_running_claims_sql_for(schema, None) } /// [`open_receipt_running_claims_sql`] restricted to the jobs in `job_ids`, /// an SQL expression yielding `bigint[]`. The restriction applies before the /// closure and materialisation checks. Row-local claims are found through the /// `job_id` index. Compact claim batches have no index on `job_ids`, so every /// retained batch header is read, but only batches holding one of the jobs /// are unnested and checked. pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option<&str>) -> String { let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let row_closed = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let (row_filter, batch_filter) = match job_ids { Some(ids) => ( format!("claims.job_id = ANY({ids}) AND "), format!("batches.job_ids && {ids} AND items.job_id = ANY({ids}) AND "), ), None => (String::new(), String::new()), }; format!( r#" SELECT claims.job_id, claims.run_lease, claims.queue FROM {schema}.lease_claims AS claims WHERE {row_filter}NOT {row_closed} AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) UNION ALL SELECT items.job_id, items.run_lease, batches.queue FROM {schema}.lease_claim_batches AS batches CROSS JOIN LATERAL unnest( batches.job_ids, batches.run_leases, batches.receipt_ids ) AS items(job_id, run_lease, receipt_id) -- Skips, before unnesting, a batch whose receipts closure batches -- have all closed: most retained batches, once their jobs finish. WHERE {batch_filter}NOT ( batches.receipt_ranges <@ COALESCE(( SELECT range_agg(closed) FROM {schema}.lease_claim_closure_batches AS closure_batches CROSS JOIN LATERAL unnest(closure_batches.receipt_ranges) AS closed WHERE closure_batches.claim_slot = batches.claim_slot AND closure_batches.receipt_ranges && batches.receipt_ranges ), '{{}}'::int8multirange) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = items.job_id AND lease.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) "# ) } fn busy_indicator(has_rows: bool) -> i64 { if has_rows { 1 } else { 0 } } async fn queue_prune_has_active_leases_tx( tx: &mut sqlx::Transaction<'_, Postgres>, schema: &str, slot: i32, generation: i64, ) -> Result { sqlx::query_scalar(audited_sql(format!( r#" SELECT EXISTS ( SELECT 1 FROM {schema}.leases WHERE ready_slot = $1 AND ready_generation = $2 LIMIT 1 ) "# ))) .bind(slot) .bind(generation) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error) } /// Whether the ready slot still holds a row at or after its lane's claim /// cursor. The `ready_slot` predicate lets each per-lane probe use the /// partition primary key `(ready_slot, queue, priority, enqueue_shard, /// lane_seq)`. async fn queue_prune_has_pending_ready_tx( tx: &mut sqlx::Transaction<'_, Postgres>, schema: &str, slot: i32, generation: i64, ) -> Result { let ready_child = ready_child_name(schema, slot as usize); sqlx::query_scalar(audited_sql(format!( r#" WITH claim_cursors AS MATERIALIZED ( SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS claim_seq FROM {schema}.queue_claim_heads ) SELECT EXISTS ( SELECT 1 FROM claim_cursors AS claims CROSS JOIN LATERAL ( SELECT 1 FROM {ready_child} AS ready WHERE ready.ready_slot = $2 AND ready.ready_generation = $1 AND ready.queue = claims.queue AND ready.priority = claims.priority AND ready.enqueue_shard = claims.enqueue_shard AND ready.lane_seq >= claims.claim_seq LIMIT 1 ) AS pending_ready LIMIT 1 ) "# ))) .bind(generation) .bind(slot) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error) } async fn queue_prune_has_unclosed_claim_refs_tx( tx: &mut sqlx::Transaction<'_, Postgres>, schema: &str, slot: i32, generation: i64, ) -> Result { let count_proves_claim_refs_closed: bool = sqlx::query_scalar(audited_sql(format!( r#" WITH claim_count AS ( SELECT count(*)::bigint AS total FROM {schema}.lease_claims AS claims WHERE claims.ready_slot = $1 AND claims.ready_generation = $2 ), compact_claim_count AS ( SELECT COALESCE(sum(claimed_count), 0)::bigint AS total FROM {schema}.lease_claim_batches AS batches WHERE batches.ready_slot = $1 AND batches.ready_generation = $2 ), explicit_count AS ( SELECT count(*)::bigint AS total FROM {schema}.lease_claims AS claims JOIN {schema}.lease_claim_closures AS closures ON closures.claim_slot = claims.claim_slot AND closures.job_id = claims.job_id AND closures.run_lease = claims.run_lease WHERE claims.ready_slot = $1 AND claims.ready_generation = $2 ), compact_count AS ( SELECT COALESCE(sum(closed_count), 0)::bigint AS total FROM {schema}.lease_claim_closure_batches AS batches WHERE batches.ready_slot = $1 AND batches.ready_generation = $2 ) SELECT claim_count.total + compact_claim_count.total = explicit_count.total + compact_count.total FROM claim_count, compact_claim_count, explicit_count, compact_count "# ))) .bind(slot) .bind(generation) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; if count_proves_claim_refs_closed { return Ok(false); } // The exact anti-join over compact claim batches requires unnesting // retained claim history. Under a hot MVCC horizon that proof can become // more expensive than the prune it guards. A count mismatch is enough to // make truncate unsafe, so skip and try again after closures catch up. Ok(true) } async fn claim_prune_has_open_claims_tx( tx: &mut sqlx::Transaction<'_, Postgres>, _schema: &str, claim_child: &str, claim_batch_child: &str, closure_child: &str, closure_batch_child: &str, ) -> Result { let count_proves_claims_closed: bool = sqlx::query_scalar(audited_sql(format!( r#" WITH claim_count AS ( SELECT count(*)::bigint AS total FROM {claim_child} ), compact_claim_count AS ( SELECT COALESCE(sum(claimed_count), 0)::bigint AS total FROM {claim_batch_child} ), explicit_count AS ( SELECT count(*)::bigint AS total FROM {closure_child} ), compact_count AS ( SELECT COALESCE(sum(closed_count), 0)::bigint AS total FROM {closure_batch_child} ) SELECT claim_count.total + compact_claim_count.total = explicit_count.total + compact_count.total FROM claim_count, compact_claim_count, explicit_count, compact_count "# ))) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; if count_proves_claims_closed { return Ok(false); } // A count mismatch means at least one claim is not proven closed by the // append-only closure ledgers. Returning SkippedActive is conservative and // avoids an unbounded anti-join over retained compact claim batches. Ok(true) } fn oldest_initialized_ring_slot( current_slot: i32, generation: i64, slot_count: i32, ) -> Option<(i32, i64)> { if slot_count <= 1 { return None; } let initialized_slots = (generation + 1).min(slot_count as i64) as i32; if initialized_slots <= 1 { return None; } let offset = initialized_slots - 1; let oldest_slot = (current_slot - offset).rem_euclid(slot_count); let oldest_generation = generation - offset as i64; if oldest_generation < 0 { return None; } Some((oldest_slot, oldest_generation)) } /// Ring families that share the append-only rotation-ledger pattern /// (#371). The current cursor of each ring is the max-generation row of /// `{schema}.{ring}_ring_rotations`; the `{ring}_ring_state` singleton /// holds only cold config (`slot_count`, plus the #290 trust marker on /// the queue ring). #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] enum RingFamily { Queue, Lease, Claim, } #[derive(Debug, Clone, Copy)] struct SlotFillObservation { cursor: (i32, i64), first_nonempty_at: Instant, } impl RingFamily { fn name(self) -> &'static str { match self { RingFamily::Queue => "queue", RingFamily::Lease => "lease", RingFamily::Claim => "claim", } } fn ledger_relname(self) -> &'static str { match self { RingFamily::Queue => "queue_ring_rotations", RingFamily::Lease => "lease_ring_rotations", RingFamily::Claim => "claim_ring_rotations", } } fn state_relname(self) -> &'static str { match self { RingFamily::Queue => "queue_ring_state", RingFamily::Lease => "lease_ring_state", RingFamily::Claim => "claim_ring_state", } } /// Per-slot table holding the row-lock targets and (in compat mode) the /// per-slot `generation` a 0.6 rotator stamps on each advance. fn slots_relname(self) -> &'static str { match self { RingFamily::Queue => "queue_ring_slots", RingFamily::Lease => "lease_ring_slots", RingFamily::Claim => "claim_ring_slots", } } /// Advisory-lock name serializing rotate ↔ prune ↔ delta-rollup for /// one ring. Replaces the pre-#371 `FOR UPDATE` on the ring-state /// singleton: the row is gone (cursor lives in the append-only /// ledger), but the mutual exclusion it provided is still required — /// on a fully-initialized ring the rotation target and the prune /// target are the same slot, and rotating into a slot mid-truncate /// would route fresh writes into a partition being wiped. fn rotation_lock_name(self, schema: &str) -> String { format!("awa.queue_storage.ring_rotation:{schema}:{}", self.name()) } } /// Which representation of a ring cursor is authoritative for a schema /// during the #371 staged rolling upgrade. Read once per rotate/prune tick /// from the per-schema `ring_cursor_authority` control row and threaded /// through the cursor/lock/append seams so a whole rotate transaction picks /// exactly one mode. The flip is one-way (`Columns -> Ledger`) and takes /// the ring-state singletons `FOR UPDATE`, so it cannot interleave with a /// compat rotator mid-critical-section (modelled in AwaStorageLockOrder). #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RingAuthority { /// Compat: the pre-0.7 `{ring}_ring_state` `current_slot` / `generation` /// singleton columns are authoritative (a live 0.6 rotator keeps them /// current). 0.7 rotators CAS the columns AND shadow the ledger. Columns, /// The append-only `{ring}_ring_rotations` ledger is authoritative /// (dead-tuple-free rotation; today's merged behaviour). Ledger, } /// Latest generation at which `slot` was the ring's open slot, derived /// from the current cursor. Rotation advances the cursor slot and the /// generation in lock-step by exactly one (`slot = generation mod /// slot_count`, genesis at `(0, 0)`), so the ledger's max row determines /// every older slot's last-open generation without per-slot state. /// Returns `None` for slots the ring has not rotated through yet. fn ring_slot_generation( current_slot: i32, generation: i64, slot_count: i32, slot: i32, ) -> Option { if slot_count <= 0 || slot < 0 || slot >= slot_count { return None; } let offset = (current_slot - slot).rem_euclid(slot_count) as i64; let slot_generation = generation - offset; (slot_generation >= 0).then_some(slot_generation) } /// All sealed (initialized, non-current) ring slots with their last-open /// generations, ordered by generation ascending — the same order the /// pre-#371 `queue_ring_slots` scan (`WHERE generation >= 0 AND slot <> /// current ORDER BY generation`) produced. fn initialized_sealed_ring_slots( current_slot: i32, generation: i64, slot_count: i32, ) -> Vec<(i32, i64)> { if slot_count <= 1 { return Vec::new(); } let initialized_slots = (generation + 1).min(slot_count as i64) as i32; (1..initialized_slots) .rev() .map(|offset| { ( (current_slot - offset).rem_euclid(slot_count), generation - offset as i64, ) }) .collect() } #[cfg(test)] mod identifier_tests { use super::{validate_ident, QueueStorage, QueueStorageConfig}; #[test] fn queue_storage_schema_identifiers_are_lowercase_unquoted_names() { for ident in ["awa", "awa_queue_storage", "_awa123"] { validate_ident(ident).expect("identifier should be accepted"); } for ident in ["Awa", "awa-queue", "123awa", "awa.queue"] { assert!( validate_ident(ident).is_err(), "identifier should be rejected: {ident}" ); } } #[test] fn active_schema_read_path_rejects_names_the_constructor_would_reject() { // `active_schema` interpolates its result into query text, so the // same gate the constructor applies has to hold on the read path. for schema in ["Awa", "awa-queue", "123awa", "awa.queue", "awa\"; DROP"] { assert!( super::validate_active_schema(schema.to_string()).is_err(), "active schema name should be rejected: {schema}" ); } for schema in ["awa", "awa_shadow", "_awa123"] { assert_eq!( super::validate_active_schema(schema.to_string()).expect("should be accepted"), schema ); } } #[test] fn default_queue_storage_schema_requires_default_physical_shape() { for config in [ QueueStorageConfig { queue_slot_count: 32, ..Default::default() }, QueueStorageConfig { lease_slot_count: 4, ..Default::default() }, QueueStorageConfig { claim_slot_count: 4, ..Default::default() }, QueueStorageConfig { lease_claim_receipts: false, ..Default::default() }, ] { let err = QueueStorage::new(config).expect_err("default awa schema shape must reject"); assert!( err.to_string() .contains("default `awa` queue-storage schema"), "unexpected error: {err}" ); } QueueStorage::new(QueueStorageConfig { schema: "awa_custom".to_string(), queue_slot_count: 4, lease_slot_count: 2, claim_slot_count: 2, lease_claim_receipts: false, ..Default::default() }) .expect("custom schema should allow custom physical shape"); } } #[cfg(test)] mod shard_routing_tests { use super::shard_for_ordering_key; use std::collections::HashSet; #[test] fn shards_le_one_collapse_to_zero() { assert_eq!(shard_for_ordering_key(b"customer-42", 1), 0); assert_eq!(shard_for_ordering_key(b"", 1), 0); assert_eq!(shard_for_ordering_key(b"customer-42", 0), 0); } #[test] fn same_key_lands_on_same_shard() { let key = b"customer-42"; let first = shard_for_ordering_key(key, 8); for _ in 0..100 { assert_eq!(shard_for_ordering_key(key, 8), first); } } #[test] fn shard_is_within_range() { for n in 0..256u32 { let key = format!("order-{n}"); let shard = shard_for_ordering_key(key.as_bytes(), 8); assert!((0..8).contains(&shard)); } } #[test] fn distinct_keys_spread_across_shards() { let mut hit: HashSet = HashSet::new(); for n in 0..1024u32 { let key = format!("order-{n}"); hit.insert(shard_for_ordering_key(key.as_bytes(), 8)); } assert_eq!(hit.len(), 8, "1024 distinct keys should cover all 8 shards"); } } #[cfg(test)] mod ring_slot_tests { use super::oldest_initialized_ring_slot; #[test] fn oldest_initialized_ring_slot_is_none_until_second_slot_exists() { assert_eq!(oldest_initialized_ring_slot(0, 0, 8), None); } #[test] fn oldest_initialized_ring_slot_tracks_partial_ring_startup() { assert_eq!(oldest_initialized_ring_slot(1, 1, 8), Some((0, 0))); assert_eq!(oldest_initialized_ring_slot(2, 2, 8), Some((0, 0))); assert_eq!(oldest_initialized_ring_slot(3, 3, 8), Some((0, 0))); } #[test] fn oldest_initialized_ring_slot_wraps_after_full_rotation() { assert_eq!(oldest_initialized_ring_slot(7, 7, 8), Some((0, 0))); assert_eq!(oldest_initialized_ring_slot(0, 8, 8), Some((1, 1))); assert_eq!(oldest_initialized_ring_slot(1, 9, 8), Some((2, 2))); } } #[cfg(test)] mod ring_ledger_derivation_tests { use super::{ initialized_sealed_ring_slots, oldest_initialized_ring_slot, ring_slot_generation, }; /// Rotation advances slot and generation in lock-step from genesis /// (0, 0), so the cursor always satisfies `slot = generation mod /// slot_count` and each slot's last-open generation is derivable. #[test] fn cursor_slot_is_generation_mod_slot_count() { for slot_count in [2_i32, 4, 8, 16] { for generation in 0..(slot_count as i64 * 3) { let slot = (generation % slot_count as i64) as i32; assert_eq!( ring_slot_generation(slot, generation, slot_count, slot), Some(generation), "current slot derives to the cursor generation" ); } } } #[test] fn slot_generation_derives_older_slots() { // Cursor at slot 3, generation 19, slot_count 16: slot 3 opened // at 19, slot 2 at 18, ..., slot 0 at 16, slot 15 at 15, etc. assert_eq!(ring_slot_generation(3, 19, 16, 3), Some(19)); assert_eq!(ring_slot_generation(3, 19, 16, 2), Some(18)); assert_eq!(ring_slot_generation(3, 19, 16, 0), Some(16)); assert_eq!(ring_slot_generation(3, 19, 16, 15), Some(15)); assert_eq!(ring_slot_generation(3, 19, 16, 4), Some(4)); } #[test] fn slot_generation_is_none_for_never_opened_slots() { // Cursor at (2, 2) in an 8-slot ring: slots 3..7 never opened. for slot in 3..8 { assert_eq!(ring_slot_generation(2, 2, 8, slot), None); } // Out-of-range slots are never derivable. assert_eq!(ring_slot_generation(2, 2, 8, -1), None); assert_eq!(ring_slot_generation(2, 2, 8, 8), None); } #[test] fn sealed_slots_match_slot_generation_and_ordering() { for slot_count in [2_i32, 4, 8, 16] { for generation in 0..(slot_count as i64 * 2 + 3) { let current_slot = (generation % slot_count as i64) as i32; let sealed = initialized_sealed_ring_slots(current_slot, generation, slot_count); // Ascending generations, all below the cursor. let mut previous = -1_i64; for &(slot, slot_generation) in &sealed { assert!(slot_generation > previous); assert!(slot_generation < generation); assert_ne!(slot, current_slot); assert_eq!( ring_slot_generation(current_slot, generation, slot_count, slot), Some(slot_generation) ); previous = slot_generation; } // The oldest sealed slot agrees with the prune target. assert_eq!( sealed.first().copied(), oldest_initialized_ring_slot(current_slot, generation, slot_count) ); // Ring keeps at most slot_count - 1 sealed slots. assert!(sealed.len() < slot_count as usize); } } } } #[cfg(test)] mod claim_cursor_advance_tests { use super::{ClaimCursorAdvance, QueueStorage}; fn advance(next_seq: i64, only_if_current: Option) -> ClaimCursorAdvance { ClaimCursorAdvance { queue: "queue".to_string(), priority: 2, enqueue_shard: 0, next_seq, only_if_current, evidence: None, } } #[test] fn normalize_claim_cursor_advances_sorts_conditional_lane_updates() { let normalized = QueueStorage::normalize_claim_cursor_advances(&[ advance(7, Some(6)), advance(6, Some(5)), advance(8, Some(7)), ]); let ordered: Vec<(i64, i64)> = normalized .iter() .map(|advance| (advance.only_if_current.unwrap(), advance.next_seq)) .collect(); assert_eq!(ordered, vec![(5, 6), (6, 7), (7, 8)]); } #[test] fn normalize_claim_cursor_advances_coalesces_unconditional_lane_updates() { let normalized = QueueStorage::normalize_claim_cursor_advances(&[ advance(3, None), advance(5, None), advance(4, Some(3)), ]); assert_eq!(normalized.len(), 1); assert_eq!(normalized[0].next_seq, 5); assert_eq!(normalized[0].only_if_current, None); } } fn default_payload_metadata() -> serde_json::Value { serde_json::json!({}) } fn is_empty_json_object(value: &serde_json::Value) -> bool { value.as_object().is_some_and(serde_json::Map::is_empty) } fn is_compact_receipt_completion_metadata(value: &serde_json::Value) -> bool { let Some(metadata) = value.as_object() else { return false; }; metadata .keys() .all(|key| key == "_awa_original_priority" || key == "_awa_original_queue") } #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] struct RuntimePayload { #[serde( default = "default_payload_metadata", skip_serializing_if = "is_empty_json_object" )] metadata: serde_json::Value, #[serde(default, skip_serializing_if = "Vec::is_empty")] tags: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] errors: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] progress: Option, } impl Default for RuntimePayload { fn default() -> Self { Self { metadata: default_payload_metadata(), tags: Vec::new(), errors: Vec::new(), progress: None, } } } impl RuntimePayload { fn from_json(value: serde_json::Value) -> Result { if value.is_null() { return Ok(Self::default()); } let payload: Self = serde_json::from_value(value)?; if !payload.metadata.is_object() { return Err(AwaError::Validation( "queue storage payload metadata must be a JSON object".to_string(), )); } Ok(payload) } fn into_json(self) -> serde_json::Value { serde_json::to_value(self).expect("runtime payload serializes") } fn errors_option(&self) -> Option> { (!self.errors.is_empty()).then(|| self.errors.clone()) } fn push_error(&mut self, error: serde_json::Value) { // ADR-048: at most one entry keeps an output, the latest to record // one. Retries copy the payload into every later ready row, so keeping // them all would grow retained bytes with the square of the attempt // count. An entry without an output leaves the earlier one in place. if error.get(crate::output::ERROR_ENTRY_OUTPUT_KEY).is_some() { crate::output::strip_error_outputs(&mut self.errors); } self.errors.push(error); } fn set_progress(&mut self, progress: Option) { self.progress = progress; } fn insert_callback_result(&mut self, payload: Option) { let metadata = self .metadata .as_object_mut() .expect("runtime payload metadata object"); metadata.insert( "_awa_callback_result".to_string(), payload.unwrap_or(serde_json::Value::Null), ); } } #[cfg(test)] mod runtime_payload_tests { use super::{ is_compact_receipt_completion_metadata, storage_payload, terminal_storage_payload, RuntimePayload, }; #[test] fn default_runtime_payload_serializes_compactly() { assert_eq!( RuntimePayload::default().into_json(), serde_json::json!({}), "default payloads should not write empty metadata/tags/errors/progress" ); assert_eq!( storage_payload(&RuntimePayload::default().into_json()), None ); } #[test] fn missing_runtime_payload_fields_round_trip_with_defaults() { let payload = RuntimePayload::from_json(serde_json::json!({})).unwrap(); assert_eq!(payload.metadata, serde_json::json!({})); assert!(payload.tags.is_empty()); assert!(payload.errors.is_empty()); assert_eq!(payload.progress, None); assert_eq!(payload.into_json(), serde_json::json!({})); } #[test] fn null_runtime_payload_round_trips_with_defaults() { let payload = RuntimePayload::from_json(serde_json::Value::Null).unwrap(); assert_eq!(payload.metadata, serde_json::json!({})); assert!(payload.tags.is_empty()); assert!(payload.errors.is_empty()); assert_eq!(payload.progress, None); assert_eq!(storage_payload(&payload.into_json()), None); } #[test] fn compact_receipt_completion_metadata_only_allows_awa_provenance() { assert!(is_compact_receipt_completion_metadata(&serde_json::json!( {} ))); assert!(is_compact_receipt_completion_metadata( &serde_json::json!({ "_awa_original_priority": 4 }) )); assert!(is_compact_receipt_completion_metadata(&serde_json::json!({ "_awa_original_queue": "default", "_awa_original_priority": 4 }))); assert!(!is_compact_receipt_completion_metadata( &serde_json::json!({ "tenant": "acme" }) )); assert!(!is_compact_receipt_completion_metadata(&serde_json::json!( null ))); } #[test] fn legacy_expanded_runtime_payload_round_trips_to_compact_form() { let payload = RuntimePayload::from_json(serde_json::json!({ "metadata": {}, "tags": [], "errors": [], "progress": null })) .unwrap(); assert_eq!(payload.metadata, serde_json::json!({})); assert!(payload.tags.is_empty()); assert!(payload.errors.is_empty()); assert_eq!(payload.progress, None); assert_eq!(payload.into_json(), serde_json::json!({})); } #[test] fn non_default_runtime_payload_fields_are_preserved() { let payload = RuntimePayload::from_json(serde_json::json!({ "metadata": { "source": "test" }, "tags": ["fast"], "errors": [{ "message": "boom" }], "progress": { "step": 1 } })) .unwrap(); assert_eq!( payload.into_json(), serde_json::json!({ "metadata": { "source": "test" }, "tags": ["fast"], "errors": [{ "message": "boom" }], "progress": { "step": 1 } }) ); } #[test] fn unchanged_terminal_payload_elides_storage_copy() { let payload = serde_json::json!({ "metadata": { "source": "test" }, "tags": ["fast"] }); assert_eq!(terminal_storage_payload(&payload, Some(&payload)), None); let changed = serde_json::json!({ "metadata": { "source": "test" }, "tags": ["fast"], "errors": [{ "message": "boom" }] }); assert_eq!( terminal_storage_payload(&changed, Some(&payload)), Some(&changed) ); } } fn unique_state_claims(unique_states: Option<&str>, state: JobState) -> bool { let Some(bitmask) = unique_states else { return false; }; let idx = state.bit_position() as usize; bitmask.as_bytes().get(idx).is_some_and(|bit| *bit == b'1') } fn write_copy_field(buf: &mut Vec, value: &str) { if value.contains(',') || value.contains('"') || value.contains('\n') || value.contains('\r') || value.contains('\\') || value == COPY_NULL_SENTINEL { buf.push(b'"'); for byte in value.bytes() { if byte == b'"' { buf.push(b'"'); } buf.push(byte); } buf.push(b'"'); } else { buf.extend_from_slice(value.as_bytes()); } } fn write_copy_json(buf: &mut Vec, value: &serde_json::Value) { let json = serde_json::to_string(value).expect("JSON serialization should not fail"); write_copy_field(buf, &json); } fn storage_payload(value: &serde_json::Value) -> Option<&serde_json::Value> { (!is_storage_payload_empty(value)).then_some(value) } fn terminal_storage_payload<'a>( value: &'a serde_json::Value, ready_payload: Option<&serde_json::Value>, ) -> Option<&'a serde_json::Value> { if is_storage_payload_empty(value) || ready_payload.is_some_and(|ready| ready == value) { None } else { Some(value) } } fn terminal_payloads<'r>( rows: &[&'r DoneJobRow], ready_payloads: &HashMap<(i32, i64, &str, i16, i16, i64), serde_json::Value>, ) -> Vec> { rows.iter() .map(|row| { let ready_key = ( row.ready_slot, row.ready_generation, row.queue.as_str(), row.priority, row.enqueue_shard, row.lane_seq, ); terminal_storage_payload(&row.payload, ready_payloads.get(&ready_key)) }) .collect() } fn is_storage_payload_empty(value: &serde_json::Value) -> bool { value.is_null() || is_empty_json_object(value) } fn write_copy_storage_payload(buf: &mut Vec, value: &serde_json::Value) { match storage_payload(value) { Some(value) => write_copy_json(buf, value), None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()), } } fn write_copy_datetime(buf: &mut Vec, value: DateTime) { write_copy_field(buf, &value.to_rfc3339()); } fn write_copy_optional_datetime(buf: &mut Vec, value: Option>) { match value { Some(value) => write_copy_datetime(buf, value), None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()), } } fn write_copy_optional_bytes(buf: &mut Vec, value: &Option>) { match value { Some(bytes) => { let bytea_hex = format!("\\x{}", hex::encode(bytes)); write_copy_field(buf, &bytea_hex); } None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()), } } fn write_copy_optional_string(buf: &mut Vec, value: Option<&str>) { match value { Some(value) => write_copy_field(buf, value), None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()), } } fn write_ready_copy_row( buf: &mut Vec, ready_slot: i32, ready_generation: i64, row: &RuntimeReadyInsert, ) { buf.extend_from_slice(ready_slot.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(ready_generation.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.job_id.to_string().as_bytes()); buf.push(b','); write_copy_field(buf, &row.kind); buf.push(b','); write_copy_field(buf, &row.queue); buf.push(b','); write_copy_json(buf, &row.args); buf.push(b','); buf.extend_from_slice(row.priority.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.attempt.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.run_lease.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.max_attempts.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.lane_seq.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.enqueue_shard.to_string().as_bytes()); buf.push(b','); write_copy_datetime(buf, row.run_at); buf.push(b','); write_copy_optional_datetime(buf, row.attempted_at); buf.push(b','); write_copy_datetime(buf, row.created_at); buf.push(b','); write_copy_optional_bytes(buf, &row.unique_key); buf.push(b','); write_copy_optional_string(buf, row.unique_states.as_deref()); buf.push(b','); write_copy_storage_payload(buf, &row.payload); buf.push(b'\n'); } fn write_deferred_copy_row(buf: &mut Vec, row: &DeferredJobRow) { buf.extend_from_slice(row.job_id.to_string().as_bytes()); buf.push(b','); write_copy_field(buf, &row.kind); buf.push(b','); write_copy_field(buf, &row.queue); buf.push(b','); write_copy_json(buf, &row.args); buf.push(b','); write_copy_field(buf, &row.state.to_string()); buf.push(b','); buf.extend_from_slice(row.priority.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.attempt.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.run_lease.to_string().as_bytes()); buf.push(b','); buf.extend_from_slice(row.max_attempts.to_string().as_bytes()); buf.push(b','); write_copy_datetime(buf, row.run_at); buf.push(b','); write_copy_optional_datetime(buf, row.attempted_at); buf.push(b','); write_copy_optional_datetime(buf, row.finalized_at); buf.push(b','); write_copy_datetime(buf, row.created_at); buf.push(b','); write_copy_optional_bytes(buf, &row.unique_key); buf.push(b','); write_copy_optional_string(buf, row.unique_states.as_deref()); buf.push(b','); write_copy_storage_payload(buf, &row.payload); buf.push(b'\n'); } /// Text bytes of recorded outputs a compact completion batch may hold before /// the flush starts another batch row. Reading one job's output detoasts its /// batch's whole array (ADR-048). const COMPACT_BATCH_OUTPUT_BUDGET: usize = 64 * 1024; /// Split a completion batch into contiguous runs whose recorded outputs stay /// within [`COMPACT_BATCH_OUTPUT_BUDGET`]; a run always holds at least one job. fn split_by_output_budget<'a>( claimed: &'a [&'a ClaimedRuntimeJob], ) -> Vec<&'a [&'a ClaimedRuntimeJob]> { let mut chunks = Vec::new(); let (mut start, mut bytes) = (0, 0); for (index, entry) in claimed.iter().enumerate() { let size = entry .job .output .as_ref() .and_then(|output| crate::output::serialized_len(output).ok()) .unwrap_or(0); if index > start && bytes + size > COMPACT_BATCH_OUTPUT_BUDGET { chunks.push(&claimed[start..index]); (start, bytes) = (index, 0); } bytes += size; } if start < claimed.len() { chunks.push(&claimed[start..]); } chunks } /// An error entry for a failed attempt, carrying the output the attempt /// recorded, if any (ADR-048). fn lifecycle_entry(error: AttemptError, attempt: i16, terminal: bool) -> serde_json::Value { let mut value = lifecycle_error(error.message, attempt, terminal); if let Some(output) = error.output { value[crate::output::ERROR_ENTRY_OUTPUT_KEY] = output; } value } fn lifecycle_error(error: impl Into, attempt: i16, terminal: bool) -> serde_json::Value { let mut value = serde_json::json!({ "error": error.into(), "attempt": attempt, "at": Utc::now().to_rfc3339(), }); if terminal { value["terminal"] = serde_json::Value::Bool(true); } value } fn transition_timestamp(job: &JobRow) -> DateTime { job.finalized_at .or(job.heartbeat_at) .or(job.deadline_at) .or(job.attempted_at) .unwrap_or(job.run_at) } fn state_rank(state: JobState) -> u8 { match state { JobState::Running | JobState::WaitingExternal => 4, JobState::Retryable | JobState::Scheduled => 3, JobState::Available => 2, JobState::Completed | JobState::Failed | JobState::Cancelled => 1, } } #[derive(Debug, Clone, sqlx::FromRow)] struct ReadyJobRow { job_id: i64, kind: String, queue: String, args: serde_json::Value, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, attempted_at: Option>, created_at: DateTime, unique_key: Option>, payload: serde_json::Value, } impl ReadyJobRow { fn into_job_row(self) -> Result { let payload = RuntimePayload::from_json(self.payload)?; Ok(JobRow { id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: JobState::Available, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, heartbeat_at: None, deadline_at: None, attempted_at: self.attempted_at, finalized_at: None, created_at: self.created_at, errors: payload.errors_option(), metadata: payload.metadata, tags: payload.tags, unique_key: self.unique_key, unique_states: None, callback_id: None, callback_timeout_at: None, callback_filter: None, callback_on_complete: None, callback_on_fail: None, callback_transform: None, progress: payload.progress, output: None, }) } } #[derive(Debug, Clone, sqlx::FromRow)] struct ReadyTransitionRow { ready_slot: i32, ready_generation: i64, job_id: i64, kind: String, queue: String, args: serde_json::Value, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, lane_seq: i64, enqueue_shard: i16, run_at: DateTime, attempted_at: Option>, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, } #[derive(Debug, Clone, sqlx::FromRow)] struct DrainLane { queue: String, priority: i16, enqueue_shard: i16, from_seq: i64, to_seq: i64, } #[derive(Debug, Clone)] struct ClaimCursorAdvance { queue: String, priority: i16, enqueue_shard: i16, next_seq: i64, only_if_current: Option, /// Claim batch row that must exist for the advance to apply: the cursor /// only moves past a lane position whose claim is visible on the server /// the advance runs against. evidence: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct ClaimEvidence { ready_slot: i32, ready_generation: i64, lane_seq: i64, legacy_attempt: Option<(i64, i64)>, } type ClaimCursorLaneKey = (String, i16, i16); type ConditionalClaimCursorAdvances = BTreeMap; type GroupedClaimCursorAdvances = BTreeMap< ClaimCursorLaneKey, ( Option<(i64, Option)>, ConditionalClaimCursorAdvances, ), >; type TerminalCounterKey = (i32, i64, String, i16, i16, i16); /// Test hook awaited at named points of a running-attempt cancel: /// `"lease_lookup"` after the first lease read, `"attempt_locked"` once the /// attempt lock is held, and `"attempt_busy"` before a retry pause. #[doc(hidden)] pub type CancelPhaseHook = std::sync::Arc< dyn Fn(&'static str) -> std::pin::Pin + Send>> + Send + Sync, >; struct CancelPhaseHookSlot(Option); impl std::fmt::Debug for CancelPhaseHookSlot { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str(if self.0.is_some() { "Some(..)" } else { "None" }) } } /// How long a cancel keeps retrying a running attempt that other /// transactions keep busy before reporting it as not cancelled. const CANCEL_BUSY_BUDGET: Duration = Duration::from_secs(5); /// What cancelling a job's running attempt did. enum RunningAttemptCancel { Cancelled(Box), /// No running attempt: not claimed, or already closed. NotRunning, /// A running attempt stayed busy, or could not be hydrated; left as is. Left, } /// One savepoint-scoped try at cancelling a running attempt. enum CancelStep { Done(RunningAttemptCancel), /// A lock was held elsewhere; roll back and retry after a pause. Busy, /// The attempt changed representation; roll back and look again. Again, } struct CancelJobTxResult { row: JobRow, claim_cursor_advance: Option, } struct ReadyBatchMoveResult { moved: bool, } impl ReadyTransitionRow { fn into_existing_ready_row( self, queue: String, priority: i16, payload: serde_json::Value, ) -> ExistingReadyRow { ExistingReadyRow { job_id: self.job_id, kind: self.kind, queue, args: self.args, priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, attempted_at: self.attempted_at, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, } } fn into_done_row( self, state: JobState, finalized_at: DateTime, payload: serde_json::Value, ) -> DoneJobRow { DoneJobRow { ready_slot: self.ready_slot, ready_generation: self.ready_generation, job_id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, lane_seq: self.lane_seq, enqueue_shard: self.enqueue_shard, run_at: self.run_at, attempted_at: self.attempted_at, finalized_at, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, output: None, } } } #[derive(Debug, Clone, sqlx::FromRow)] struct ReadyJobLeaseRow { ready_slot: i32, ready_generation: i64, lane_seq: i64, enqueue_shard: i16, lease_slot: i32, lease_generation: i64, claim_slot: i32, receipt_id: Option, claim_batch_id: Option, claim_batch_index: Option, job_id: i64, kind: String, queue: String, args: serde_json::Value, lane_priority: i16, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, heartbeat_at: Option>, deadline_at: Option>, attempted_at: Option>, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, } impl ReadyJobLeaseRow { fn claim_ref(&self, lease_claim_receipt: bool) -> ClaimedEntry { ClaimedEntry { queue: self.queue.clone(), priority: self.lane_priority, lane_seq: self.lane_seq, ready_slot: self.ready_slot, ready_generation: self.ready_generation, lease_slot: self.lease_slot, lease_generation: self.lease_generation, claim_slot: self.claim_slot, receipt_id: self.receipt_id, claim_batch_id: self.claim_batch_id, claim_batch_index: self.claim_batch_index, lease_claim_receipt, enqueue_shard: self.enqueue_shard, } } fn into_job_row(self) -> Result { let mut payload = RuntimePayload::from_json(self.payload)?; if self.priority < self.lane_priority { let metadata = payload.metadata.as_object_mut().ok_or_else(|| { AwaError::Validation( "queue storage payload metadata must be a JSON object".to_string(), ) })?; metadata .entry("_awa_original_priority".to_string()) .or_insert_with(|| serde_json::Value::from(i64::from(self.lane_priority))); } Ok(JobRow { id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: JobState::Running, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, heartbeat_at: self.heartbeat_at, deadline_at: self.deadline_at, attempted_at: self.attempted_at, finalized_at: None, created_at: self.created_at, errors: payload.errors_option(), metadata: payload.metadata, tags: payload.tags, unique_key: self.unique_key, unique_states: None, callback_id: None, callback_timeout_at: None, callback_filter: None, callback_on_complete: None, callback_on_fail: None, callback_transform: None, progress: payload.progress, output: None, }) } fn into_claimed_runtime_job( self, lease_claim_receipt: bool, ) -> Result { let claim = self.claim_ref(lease_claim_receipt); let unique_states = self.unique_states.clone(); let job = self.into_job_row()?; Ok(ClaimedRuntimeJob { claim, job, unique_states, }) } } #[derive(Debug, Clone)] struct RuntimeReadyRow { job_id: Option, kind: String, queue: String, args: serde_json::Value, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, attempted_at: Option>, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, /// Optional caller-supplied key. When the destination queue has /// `enqueue_shards > 1`, all rows sharing the same key route to /// the same shard so FIFO within the key is preserved. `None` /// falls back to the per-store rotor. ordering_key: Option>, } #[derive(Debug, Clone)] struct RuntimeReadyInsert { job_id: i64, kind: String, queue: String, args: serde_json::Value, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, attempted_at: Option>, lane_seq: i64, /// Enqueue shard the row belongs to. Selected per row by /// `shard_for_enqueue`; defaults to 0 when the queue's /// `enqueue_shards` is 1 (the default). enqueue_shard: i16, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, } #[derive(Debug)] struct ReadySegmentInsert { queue: String, priority: i16, enqueue_shard: i16, first_lane_seq: i64, next_lane_seq: i64, /// `run_at` for every row in this segment. Segment construction splits on /// `run_at` so claim-time priority aging stays exact after the claim cursor /// advances inside a multi-row segment. first_run_at: DateTime, } #[cfg(test)] mod ready_segment_tests { use super::{QueueStorage, RuntimeReadyInsert}; use chrono::{Duration, TimeZone, Utc}; fn ready_row(lane_seq: i64, run_at: chrono::DateTime) -> RuntimeReadyInsert { RuntimeReadyInsert { job_id: lane_seq, kind: "segment_test".to_string(), queue: "segment-q".to_string(), args: serde_json::json!({}), priority: 2, attempt: 0, run_lease: 0, max_attempts: 25, run_at, attempted_at: None, lane_seq, enqueue_shard: 0, created_at: run_at, unique_key: None, unique_states: None, payload: serde_json::json!({}), } } #[test] fn ready_segments_split_on_run_at_boundaries() { let first = Utc .with_ymd_and_hms(2026, 6, 14, 12, 0, 0) .single() .expect("valid test timestamp"); let second = first + Duration::seconds(1); let rows = vec![ ready_row(1, first), ready_row(2, first), ready_row(3, second), ready_row(4, first), ]; let segments = QueueStorage::ready_segments_from_rows(&rows); let ranges: Vec<_> = segments .iter() .map(|segment| { ( segment.first_lane_seq, segment.next_lane_seq, segment.first_run_at, ) }) .collect(); assert_eq!(ranges, vec![(1, 3, first), (3, 4, second), (4, 5, first)]); } } #[derive(Debug, Clone, sqlx::FromRow)] struct DoneJobRow { ready_slot: i32, ready_generation: i64, job_id: i64, kind: String, queue: String, args: serde_json::Value, state: JobState, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, lane_seq: i64, /// Enqueue shard the row was claimed from. Part of the /// `done_entries` primary key so two shards' terminal rows at /// the same `(ready_slot, queue, priority, lane_seq)` do not /// collide. enqueue_shard: i16, run_at: DateTime, attempted_at: Option>, finalized_at: DateTime, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, /// The completed attempt's output (ADR-048). Absent from projections /// that do not load it. #[sqlx(default)] output: Option, } impl DoneJobRow { fn into_job_row(self) -> Result { let payload = RuntimePayload::from_json(self.payload)?; Ok(JobRow { id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: self.state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, heartbeat_at: None, deadline_at: None, attempted_at: self.attempted_at, finalized_at: Some(self.finalized_at), created_at: self.created_at, errors: payload.errors_option(), metadata: payload.metadata, tags: payload.tags, unique_key: self.unique_key, unique_states: None, callback_id: None, callback_timeout_at: None, callback_filter: None, callback_on_complete: None, callback_on_fail: None, callback_transform: None, progress: payload.progress, output: self.output, }) } fn into_dlq_row(self, dlq_reason: String, dlq_at: DateTime) -> DlqJobRow { DlqJobRow { job_id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: self.state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, attempted_at: self.attempted_at, finalized_at: self.finalized_at, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload: self.payload, dlq_reason, dlq_at, original_run_lease: self.run_lease, } } } #[derive(Debug, Clone, sqlx::FromRow)] struct DlqJobRow { job_id: i64, kind: String, queue: String, args: serde_json::Value, state: JobState, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, attempted_at: Option>, finalized_at: DateTime, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, dlq_reason: String, dlq_at: DateTime, original_run_lease: i64, } impl DlqJobRow { fn into_job_row(self) -> Result { let payload = RuntimePayload::from_json(self.payload)?; Ok(JobRow { id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: self.state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, heartbeat_at: None, deadline_at: None, attempted_at: self.attempted_at, finalized_at: Some(self.finalized_at), created_at: self.created_at, errors: payload.errors_option(), metadata: payload.metadata, tags: payload.tags, unique_key: self.unique_key, unique_states: None, callback_id: None, callback_timeout_at: None, callback_filter: None, callback_on_complete: None, callback_on_fail: None, callback_transform: None, progress: payload.progress, output: None, }) } fn into_retry_ready_row( self, queue: String, priority: i16, run_at: DateTime, payload: serde_json::Value, ) -> ExistingReadyRow { ExistingReadyRow { job_id: self.job_id, kind: self.kind, queue, args: self.args, priority, attempt: 0, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at, attempted_at: None, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, } } fn into_retry_deferred_row( self, queue: String, priority: i16, run_at: DateTime, payload: serde_json::Value, ) -> DeferredJobRow { DeferredJobRow { job_id: self.job_id, kind: self.kind, queue, args: self.args, state: JobState::Scheduled, priority, attempt: 0, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at, attempted_at: None, finalized_at: None, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, } } } #[derive(Debug, Clone)] struct ExistingReadyRow { job_id: i64, kind: String, queue: String, args: serde_json::Value, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, attempted_at: Option>, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, } #[derive(Debug, Clone, sqlx::FromRow)] struct DeletedLeaseRow { ready_slot: i32, ready_generation: i64, job_id: i64, queue: String, state: JobState, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, lane_seq: i64, enqueue_shard: i16, attempted_at: Option>, } #[derive(Debug, Clone, sqlx::FromRow)] struct ReadySnapshotRow { ready_slot: i32, ready_generation: i64, job_id: i64, kind: String, queue: String, args: serde_json::Value, lane_seq: i64, enqueue_shard: i16, run_at: DateTime, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, } #[derive(Debug, Clone, sqlx::FromRow)] struct AttemptStateRow { job_id: i64, run_lease: i64, progress: Option, callback_filter: Option, callback_on_complete: Option, callback_on_fail: Option, callback_transform: Option, callback_result: Option, } #[derive(Debug, Clone)] struct LeaseTransitionRow { ready_slot: i32, ready_generation: i64, job_id: i64, kind: String, queue: String, args: serde_json::Value, state: JobState, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, lane_seq: i64, enqueue_shard: i16, run_at: DateTime, attempted_at: Option>, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, progress: Option, } impl LeaseTransitionRow { fn into_done_row( self, state: JobState, finalized_at: DateTime, payload: serde_json::Value, ) -> DoneJobRow { DoneJobRow { ready_slot: self.ready_slot, ready_generation: self.ready_generation, job_id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, lane_seq: self.lane_seq, enqueue_shard: self.enqueue_shard, run_at: self.run_at, attempted_at: self.attempted_at, finalized_at, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, output: None, } } fn into_deferred_row( self, state: JobState, run_at: DateTime, finalized_at: Option>, payload: serde_json::Value, ) -> DeferredJobRow { DeferredJobRow { job_id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at, attempted_at: self.attempted_at, finalized_at, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, } } fn into_ready_row(self, run_at: DateTime, payload: serde_json::Value) -> ExistingReadyRow { ExistingReadyRow { job_id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at, attempted_at: self.attempted_at, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, } } fn into_dlq_row( self, finalized_at: DateTime, payload: serde_json::Value, dlq_reason: String, dlq_at: DateTime, ) -> DlqJobRow { DlqJobRow { job_id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: JobState::Failed, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, attempted_at: self.attempted_at, finalized_at, created_at: self.created_at, unique_key: self.unique_key, unique_states: self.unique_states, payload, dlq_reason, dlq_at, original_run_lease: self.run_lease, } } } #[derive(Debug, Clone, sqlx::FromRow)] struct LeaseJobRow { job_id: i64, kind: String, queue: String, args: serde_json::Value, state: JobState, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, heartbeat_at: Option>, deadline_at: Option>, attempted_at: Option>, finalized_at: Option>, created_at: DateTime, unique_key: Option>, callback_id: Option, callback_timeout_at: Option>, callback_filter: Option, callback_on_complete: Option, callback_on_fail: Option, callback_transform: Option, payload: serde_json::Value, progress: Option, callback_result: Option, } impl LeaseJobRow { fn into_job_row(self) -> Result { let payload = QueueStorage::materialize_runtime_payload( self.payload, self.progress, self.callback_result, )?; Ok(JobRow { id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: self.state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, heartbeat_at: self.heartbeat_at, deadline_at: self.deadline_at, attempted_at: self.attempted_at, finalized_at: self.finalized_at, created_at: self.created_at, errors: payload.errors_option(), metadata: payload.metadata, tags: payload.tags, unique_key: self.unique_key, unique_states: None, callback_id: self.callback_id, callback_timeout_at: self.callback_timeout_at, callback_filter: self.callback_filter, callback_on_complete: self.callback_on_complete, callback_on_fail: self.callback_on_fail, callback_transform: self.callback_transform, progress: payload.progress, output: None, }) } } #[derive(Debug, Clone, sqlx::FromRow)] struct DeferredJobRow { job_id: i64, kind: String, queue: String, args: serde_json::Value, state: JobState, priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, run_at: DateTime, attempted_at: Option>, finalized_at: Option>, created_at: DateTime, unique_key: Option>, unique_states: Option, payload: serde_json::Value, } impl DeferredJobRow { fn into_job_row(self) -> Result { let payload = RuntimePayload::from_json(self.payload)?; Ok(JobRow { id: self.job_id, kind: self.kind, queue: self.queue, args: self.args, state: self.state, priority: self.priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, run_at: self.run_at, heartbeat_at: None, deadline_at: None, attempted_at: self.attempted_at, finalized_at: self.finalized_at, created_at: self.created_at, errors: payload.errors_option(), metadata: payload.metadata, tags: payload.tags, unique_key: self.unique_key, unique_states: None, callback_id: None, callback_timeout_at: None, callback_filter: None, callback_on_complete: None, callback_on_fail: None, callback_transform: None, progress: payload.progress, output: None, }) } } /// Segmented queue storage backend. /// /// Design goals: /// - append-only queue segments in a rotated ring /// - append-only completion segments keyed back to the queue segment /// - a separate, faster rotating lease ring so delete churn is bounded by the /// lease cycle rather than by queue retention /// - hot mutable state restricted to queue cursors, narrow leases, and /// optional per-attempt runtime state only when needed /// #[derive(Debug)] pub struct QueueStorage { config: QueueStorageConfig, next_stripe_probe: AtomicUsize, /// Per-store rotor that selects the enqueue shard for the next batch. /// Rotated once per `shard_for_enqueue` call so producers spread their /// writes across the per-`(queue, priority)` shard rows. shard_rotor: AtomicU16, /// Cache of `awa.queue_meta.enqueue_shards` per queue. Populated lazily /// on the first `shard_for_enqueue` call for a queue and invalidated by /// `reset()`. With the default `enqueue_shards = 1`, the cache holds 1 /// and `pick_shard` returns 0 unconditionally. enqueue_shards_cache: Mutex>, /// This instance's claimer lease per queue, as last read or written, so a /// claim does not re-read the lease and claimer target on every call. claimer_gates: Mutex>, /// Lane-presence cache: `(physical_queue, priority, enqueue_shard)` /// triples whose three lane rows (queue_lanes, queue_enqueue_heads, /// queue_claim_heads) we have previously inserted. Skips the three /// `INSERT … ON CONFLICT DO NOTHING` round-trips on subsequent enqueue /// batches to a known lane/shard. Cleared on `reset()` because reset /// TRUNCATEs queue_lanes; `advance_enqueue_head` repairs a stale entry /// (head row gone after a rolled-back ensure_lane). ensured_lanes: Mutex>, slot_fill_observations: Mutex>, prune_lock_timeout: Duration, cancel_phase_hook: CancelPhaseHookSlot, /// Whether this schema has migration v052's output columns (ADR-048), /// and when that was last checked. job_output_columns: Mutex>, } /// How long a negative answer about the output columns is trusted before /// the catalog is checked again, so a binary that started before the /// migration picks the columns up without a restart. const JOB_OUTPUT_PROBE_INTERVAL: Duration = Duration::from_secs(30); impl QueueStorage { pub fn new(config: QueueStorageConfig) -> Result { if config.queue_slot_count < 4 { return Err(AwaError::Validation( "queue storage requires at least 4 queue slots".into(), )); } if config.lease_slot_count < 2 { return Err(AwaError::Validation( "queue storage requires at least 2 lease slots".into(), )); } if config.claim_slot_count < 2 { return Err(AwaError::Validation( "queue storage requires at least 2 claim slots".into(), )); } if config.queue_stripe_count == 0 { return Err(AwaError::Validation( "queue storage requires at least 1 queue stripe".into(), )); } if config.schema == DEFAULT_SCHEMA && (config.queue_slot_count != DEFAULT_QUEUE_SLOT_COUNT || config.lease_slot_count != DEFAULT_LEASE_SLOT_COUNT || config.claim_slot_count != DEFAULT_CLAIM_SLOT_COUNT || !config.lease_claim_receipts) { return Err(AwaError::Validation( "the default `awa` queue-storage schema must use the default slot counts and \ lease_claim_receipts=true" .into(), )); } validate_ident(&config.schema)?; Ok(Self { config, next_stripe_probe: AtomicUsize::new(0), shard_rotor: AtomicU16::new(0), enqueue_shards_cache: Mutex::new(HashMap::new()), claimer_gates: Mutex::new(HashMap::new()), ensured_lanes: Mutex::new(HashSet::new()), slot_fill_observations: Mutex::new(HashMap::new()), prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT, cancel_phase_hook: CancelPhaseHookSlot(None), job_output_columns: Mutex::new(None), }) } /// Whether this schema can store job outputs (ADR-048, migration v052). /// /// A positive answer is cached for the life of the store. A negative one /// is re-checked at most every [`JOB_OUTPUT_PROBE_INTERVAL`]; until then /// writers drop outputs and readers project NULL instead of naming the /// missing columns. /// Discard the cached [`Self::supports_job_output`] answer, after a /// statement found the output columns missing, so the next call /// probes the catalog again. #[doc(hidden)] pub fn forget_job_output_columns(&self) { *self .job_output_columns .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = None; } pub async fn supports_job_output<'e, E>(&self, executor: E) -> Result where E: sqlx::PgExecutor<'e>, { if let Some((supported, checked_at)) = *self .job_output_columns .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) { if supported || checked_at.elapsed() < JOB_OUTPUT_PROBE_INTERVAL { return Ok(supported); } } let supported: bool = sqlx::query_scalar( r#" SELECT count(*) = 2 FROM pg_attribute WHERE NOT attisdropped AND ( (attrelid = to_regclass(format('%I.done_entries', $1)) AND attname = 'output') OR (attrelid = to_regclass(format('%I.receipt_completion_batches', $1)) AND attname = 'outputs') ) "#, ) .bind(self.schema()) .fetch_one(executor) .await .map_err(map_sqlx_error)?; *self .job_output_columns .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = Some((supported, std::time::Instant::now())); Ok(supported) } /// Await `hook` at named points of a running-attempt cancel, so tests can /// commit a concurrent change at an exact step. Not for production use. #[doc(hidden)] pub fn with_cancel_phase_hook(mut self, hook: CancelPhaseHook) -> Self { self.cancel_phase_hook = CancelPhaseHookSlot(Some(hook)); self } async fn cancel_phase(&self, phase: &'static str) { if let Some(hook) = &self.cancel_phase_hook.0 { hook(phase).await; } } /// Set the PostgreSQL `lock_timeout` applied to the exclusive-lock steps of /// ring prune (queue, lease and claim rings) and terminal-delta rollup. /// /// Queue ring prune waits only on its `ready_segments` lock; its other /// partition locks use `NOWAIT`, so a reader holding one of them makes it /// `Blocked` immediately whatever the timeout. /// /// A prune that cannot take its locks within the timeout rolls back and /// reports [`PruneOutcome::Blocked`]; the maintenance loop retries on its /// next tick. The default is 10 ms, which /// keeps prune from queueing behind readers and so from stalling claims /// and enqueues behind its pending `ACCESS EXCLUSIVE` request. Raise it /// when prune is persistently `Blocked` by short-lived readers of the /// ring partitions. Keep it well below `deadlock_timeout` (default 1 s) and /// any `statement_timeout`: a prune that waits as long as `deadlock_timeout` /// triggers deadlock detection, which may cancel a concurrent claim instead /// of the prune. The timeout is applied at millisecond granularity. /// /// Returns a validation error for a zero timeout or one above PostgreSQL's /// `lock_timeout` maximum (`i32::MAX` milliseconds, about 24.8 days). pub fn with_prune_lock_timeout(mut self, timeout: Duration) -> Result { self.set_prune_lock_timeout(timeout)?; Ok(self) } /// In-place form of [`Self::with_prune_lock_timeout`]. pub fn set_prune_lock_timeout(&mut self, timeout: Duration) -> Result<(), AwaError> { if timeout.is_zero() { return Err(AwaError::Validation( "queue storage prune lock timeout must be greater than zero".into(), )); } if timeout.as_millis() > i32::MAX as u128 { return Err(AwaError::Validation(format!( "queue storage prune lock timeout must not exceed {} ms", i32::MAX ))); } self.prune_lock_timeout = timeout; Ok(()) } /// The `lock_timeout` applied to prune's exclusive-lock steps. pub fn prune_lock_timeout(&self) -> Duration { self.prune_lock_timeout } pub fn from_existing_schema(schema: impl Into) -> Result { Self::new(QueueStorageConfig { schema: schema.into(), ..Default::default() }) } pub fn schema(&self) -> &str { &self.config.schema } pub fn slot_count(&self) -> usize { self.queue_slot_count() } pub fn queue_slot_count(&self) -> usize { self.config.queue_slot_count } pub fn lease_slot_count(&self) -> usize { self.config.lease_slot_count } pub fn claim_slot_count(&self) -> usize { self.config.claim_slot_count } pub fn queue_stripe_count(&self) -> usize { self.config.queue_stripe_count } pub fn lease_claim_receipts(&self) -> bool { self.config.lease_claim_receipts } /// Start fresh fill-age windows after acquiring maintenance leadership. pub fn reset_rotation_observations(&self) { self.slot_fill_observations .lock() .expect("slot fill observations poisoned") .clear(); } fn uses_queue_striping(&self) -> bool { self.queue_stripe_count() > 1 } fn is_physical_stripe_queue(&self, queue: &str) -> bool { self.uses_queue_striping() && queue .rsplit_once(QUEUE_STRIPE_DELIMITER) .is_some_and(|(_, suffix)| suffix.parse::().is_ok()) } fn physical_queue_for_stripe(&self, queue: &str, stripe: usize) -> String { format!("{queue}{QUEUE_STRIPE_DELIMITER}{stripe}") } pub(crate) fn physical_queues_for_logical(&self, queue: &str) -> Vec { if !self.uses_queue_striping() || self.is_physical_stripe_queue(queue) { return vec![queue.to_string()]; } (0..self.queue_stripe_count()) .map(|stripe| self.physical_queue_for_stripe(queue, stripe)) .collect() } fn stripe_probe_start(&self, stripe_count: usize) -> usize { if stripe_count <= 1 { return 0; } self.next_stripe_probe.fetch_add(1, Ordering::Relaxed) % stripe_count } fn logical_queue_name<'a>(&self, queue: &'a str) -> &'a str { if !self.uses_queue_striping() { return queue; } queue .rsplit_once(QUEUE_STRIPE_DELIMITER) .and_then(|(prefix, suffix)| suffix.parse::().ok().map(|_| prefix)) .unwrap_or(queue) } fn queue_stripe_for_enqueue( &self, queue: &str, unique_key: &Option>, salt: i64, ) -> String { if !self.uses_queue_striping() || self.is_physical_stripe_queue(queue) { return queue.to_string(); } let stripe = if let Some(key) = unique_key { let mut hasher = DefaultHasher::new(); key.hash(&mut hasher); (hasher.finish() as usize) % self.queue_stripe_count() } else { salt.rem_euclid(self.queue_stripe_count() as i64) as usize }; self.physical_queue_for_stripe(queue, stripe) } fn use_lease_claim_receipts_for_runtime(&self, _deadline_duration: Duration) -> bool { // Receipts mode supports per-claim deadlines via // `lease_claim_batches.deadline_at` (#246; rescued by // `rescue_expired_receipt_deadlines_tx`), so receipts is the // live path whenever the engine is configured for receipts — // the queue's `deadline_duration` no longer disqualifies it. self.lease_claim_receipts() } pub fn ready_child_relname(&self, slot: usize) -> String { format!("ready_entries_{slot}") } pub fn done_child_relname(&self, slot: usize) -> String { format!("done_entries_{slot}") } pub fn leases_relname(&self) -> &'static str { "leases" } pub fn lease_claims_relname(&self) -> &'static str { "lease_claims" } pub fn lease_claim_closures_relname(&self) -> &'static str { "lease_claim_closures" } pub fn leases_child_relname(&self, slot: usize) -> String { format!("leases_{slot}") } pub fn attempt_state_relname(&self) -> &'static str { "attempt_state" } /// Read the active queue-storage schema name from the transition state. /// /// The name is re-validated with [`validate_ident`] before it is handed /// back. Callers interpolate it into query text, so the check has to /// happen on the read path too and not only in [`Self::new`] — a name /// that never passed the constructor (hand-edited row, restore from an /// older or foreign database) must not reach a `format!`. No supported /// deployment can store a rejected name, since every writer of this row /// goes through the constructor first. pub async fn active_schema(pool: &PgPool) -> Result, AwaError> { let schema: Option = sqlx::query_scalar( "SELECT schema_name FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'", ) .fetch_optional(pool) .await .map_err(map_sqlx_error)?; schema.map(validate_active_schema).transpose() } /// Transaction-aware variant of [`Self::active_schema`] — read the /// active queue-storage schema name inside the caller's transaction /// rather than acquiring a separate pool connection. pub async fn active_schema_in_tx( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, ) -> Result, AwaError> { let schema: Option = sqlx::query_scalar( "SELECT schema_name FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'", ) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; schema.map(validate_active_schema).transpose() } fn materialize_runtime_payload( payload: serde_json::Value, progress: Option, callback_result: Option, ) -> Result { let mut payload = RuntimePayload::from_json(payload)?; if let Some(progress) = progress { payload.set_progress(Some(progress)); } if let Some(callback_result) = callback_result { payload.insert_callback_result(Some(callback_result)); } Ok(payload) } fn payload_with_attempt_state( payload: serde_json::Value, progress: Option, ) -> Result { let mut payload = RuntimePayload::from_json(payload)?; if let Some(progress) = progress { payload.set_progress(Some(progress)); } Ok(payload.into_json()) } fn payload_from_parts( metadata: serde_json::Value, tags: Vec, errors: Option>, progress: Option, ) -> Result { Ok(RuntimePayload { metadata, tags, errors: errors.unwrap_or_default(), progress, } .into_json()) } async fn sync_unique_claim<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, unique_key: &Option>, unique_states: Option<&str>, old_state: Option, new_state: Option, ) -> Result<(), AwaError> { let old_claim = old_state.is_some_and(|state| unique_state_claims(unique_states, state)); let new_claim = new_state.is_some_and(|state| unique_state_claims(unique_states, state)); if old_claim && !new_claim { if let Some(key) = unique_key { sqlx::query( "DELETE FROM awa.job_unique_claims WHERE unique_key = $1 AND job_id = $2", ) .bind(key) .bind(job_id) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } } if new_claim && !old_claim { if let Some(key) = unique_key { let result = sqlx::query( r#" INSERT INTO awa.job_unique_claims (unique_key, job_id) VALUES ($1, $2) ON CONFLICT (unique_key) DO UPDATE SET job_id = EXCLUDED.job_id WHERE awa.job_unique_claims.job_id = EXCLUDED.job_id "#, ) .bind(key) .bind(job_id) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; if result.rows_affected() == 0 { return Err(AwaError::UniqueConflict { constraint: Some("idx_awa_jobs_unique".to_string()), }); } } } Ok(()) } async fn sync_enqueue_unique_claims<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, claims: Vec<(Vec, i64)>, ) -> Result<(), AwaError> { if claims.is_empty() { return Ok(()); } let mut seen: HashSet<&[u8]> = HashSet::with_capacity(claims.len()); for (key, _) in &claims { if !seen.insert(key.as_slice()) { return Err(AwaError::UniqueConflict { constraint: Some("idx_awa_jobs_unique".to_string()), }); } } let (keys, job_ids): (Vec>, Vec) = claims.into_iter().unzip(); let (requested, applied): (i64, i64) = sqlx::query_as( r#" WITH input(unique_key, job_id) AS ( SELECT * FROM unnest($1::bytea[], $2::bigint[]) ), inserted AS ( INSERT INTO awa.job_unique_claims (unique_key, job_id) SELECT unique_key, job_id FROM input ORDER BY unique_key ON CONFLICT (unique_key) DO UPDATE SET job_id = EXCLUDED.job_id WHERE awa.job_unique_claims.job_id = EXCLUDED.job_id RETURNING unique_key ) SELECT (SELECT count(*)::bigint FROM input) AS requested, (SELECT count(*)::bigint FROM inserted) AS applied "#, ) .bind(keys) .bind(job_ids) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; if applied != requested { return Err(AwaError::UniqueConflict { constraint: Some("idx_awa_jobs_unique".to_string()), }); } Ok(()) } // Enqueue inserts have no prior storage state, so uniqueness only needs to // add claims for states included in the row's unique-state bitmask. State // transitions still use `sync_unique_claim`, which can release old claims. async fn sync_ready_enqueue_unique_claims<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[RuntimeReadyInsert], ) -> Result<(), AwaError> { let claims = rows .iter() .filter(|row| unique_state_claims(row.unique_states.as_deref(), JobState::Available)) .filter_map(|row| row.unique_key.as_ref().map(|key| (key.clone(), row.job_id))) .collect(); self.sync_enqueue_unique_claims(tx, claims).await } async fn sync_deferred_enqueue_unique_claims<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[DeferredJobRow], ) -> Result<(), AwaError> { let claims = rows .iter() .filter(|row| unique_state_claims(row.unique_states.as_deref(), row.state)) .filter_map(|row| row.unique_key.as_ref().map(|key| (key.clone(), row.job_id))) .collect(); self.sync_enqueue_unique_claims(tx, claims).await } fn batch_enqueue_unique_claims( ready_rows: &[RuntimeReadyRow], deferred_rows: &[DeferredJobRow], ) -> Vec<(Vec, i64)> { let ready_claims = ready_rows .iter() .filter(|row| unique_state_claims(row.unique_states.as_deref(), JobState::Available)) .filter_map(|row| { row.unique_key.as_ref().map(|key| { ( key.clone(), row.job_id.expect("batch ready job id is assigned"), ) }) }); let deferred_claims = deferred_rows .iter() .filter(|row| unique_state_claims(row.unique_states.as_deref(), row.state)) .filter_map(|row| row.unique_key.as_ref().map(|key| (key.clone(), row.job_id))); ready_claims.chain(deferred_claims).collect() } /// Idempotently install the queue-storage substrate (schema, tables, /// indexes, functions) for this configuration. /// /// Runs entirely inside one transaction on a single pooled connection, /// guarded by `pg_advisory_xact_lock` so concurrent worker startups /// serialize cleanly rather than fighting over pool slots. The lock /// auto-releases on COMMIT/ROLLBACK. /// /// **Note on index builds:** the `CREATE INDEX IF NOT EXISTS` calls /// below are not `CONCURRENTLY` (Postgres bans `CONCURRENTLY` inside a /// transaction). On a fresh install that's a no-op cost. On a redeploy /// against a database that already has rows in the queue tables, each /// fresh-index build takes `ShareUpdateExclusiveLock` on its table for /// the duration of the build, and concurrent writers to that table /// queue up. For typical operator scenarios this is bounded by the /// `IF NOT EXISTS` guard — only the *new* indexes added by a runtime /// upgrade get built; the rest are no-ops. #[tracing::instrument(skip(self, pool), name = "queue_storage.prepare_schema")] pub async fn prepare_schema(&self, pool: &PgPool) -> Result<(), AwaError> { let schema = self.schema(); let install_lock_name = format!("awa.queue_storage.install:{schema}"); let mut install_tx = pool.begin().await.map_err(map_sqlx_error)?; sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") .bind(&install_lock_name) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; let install_result = async { // Helper-owned DDL. The helper takes the same per-schema advisory // xact lock we just acquired (re-entrant on the same session) and // performs the entire forward-only substrate install: sequences, // ring-state singletons, partitioned ready/done/lease tables, // lane indexes, claim_ready_runtime() function, and seed rows. // Validation inside the helper rejects non-default configuration // against the default `awa` schema (see migration v023 and #308). // // What stays here (constraint 7 of #308 PR 1): // * `open_receipt_claims` non-empty rejection + drop (ADR-023). // * Legacy rename of `lease_claims` / `lease_claim_closures` // before the helper creates the partitioned parents, plus the // post-helper copy + drop of the renamed-aside rows. // * `queue_count_snapshots` legacy drop. // // `awa.runtime_storage_backends` is owned by migrations (v012); // activation paths only seed/update its row. The helper does NOT // touch it and does NOT change storage-transition state, so a call // to `prepare_schema` remains activation-neutral. sqlx::query(audited_sql(format!("CREATE SCHEMA IF NOT EXISTS {schema}"))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; // The hot path reads "currently open" by anti-joining the // partitioned `lease_claims` / `lease_claim_closures` pair, so // `open_receipt_claims` is unused (see ADR-023). Drop it on every // install. Refuse to drop a non-empty table — non-empty here // means an operator rolled forward from an older build that // still wrote rows we don't want to silently delete. let open_receipt_claims_exists: bool = sqlx::query_scalar( r#" SELECT EXISTS ( SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = 'open_receipt_claims' ) "#, ) .bind(schema) .fetch_one(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; if open_receipt_claims_exists { let row_count: i64 = sqlx::query_scalar(audited_sql(format!( "SELECT count(*)::bigint FROM {schema}.open_receipt_claims" ))) .fetch_one(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; if row_count > 0 { return Err(AwaError::Validation(format!( "{schema}.open_receipt_claims has {row_count} rows but the runtime no \ longer reads or writes this table. Run the ADR-023 reverse migration \ (recreate from lease_claims minus durable closure evidence) to drain it, \ then re-run prepare_schema." ))); } sqlx::query(audited_sql(format!( "DROP TABLE IF EXISTS {schema}.open_receipt_claims CASCADE" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; } // Detect the current shape of the legacy claim/closure parents. // The partitioned parents have relkind 'p'; regular tables have // 'r' and need to be renamed aside so the helper can create the // partitioned parents under the canonical name. Copying the data // back happens after the helper returns, all inside this single // transaction so a crash leaves the schema in one of exactly two // states (pre-migration or post-migration). let lease_claims_relkind: Option = sqlx::query_scalar( r#" SELECT c.relkind::text FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = 'lease_claims' "#, ) .bind(schema) .fetch_optional(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; let closures_relkind: Option = sqlx::query_scalar( r#" SELECT c.relkind::text FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = 'lease_claim_closures' "#, ) .bind(schema) .fetch_optional(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; if lease_claims_relkind.as_deref() == Some("r") { sqlx::query(audited_sql(format!( "ALTER TABLE {schema}.lease_claims RENAME TO lease_claims_legacy" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; } if closures_relkind.as_deref() == Some("r") { sqlx::query(audited_sql(format!( "ALTER TABLE {schema}.lease_claim_closures RENAME TO lease_claim_closures_legacy" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; } // queue_count_snapshots was a staleness-cached counterpart of // queue_counts_exact. The dispatcher now derives the available // count directly from the head tables and nothing else needs the // snapshot. Drop it on every prepare_schema so an upgrade from an // older install reclaims the storage. Done before the helper // runs because the helper does not touch this legacy table. sqlx::query(audited_sql(format!( "DROP TABLE IF EXISTS {schema}.queue_count_snapshots" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; // The single forward-only DDL call. See the migration file // `awa-model/migrations/v023_install_queue_storage_substrate.sql` // for the function body. Default values match the constants in // this file (DEFAULT_QUEUE_SLOT_COUNT etc); the helper validates // that the default `awa` schema only ever gets default-shaped // installs. sqlx::query("SELECT awa.install_queue_storage_substrate($1, $2, $3, $4, $5)") .bind(schema) .bind(self.queue_slot_count() as i32) .bind(self.lease_slot_count() as i32) .bind(self.claim_slot_count() as i32) .bind(self.lease_claim_receipts()) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; // #169: apply receipt-plane fillfactor tunings. The install // helper above creates structure only; tunings live in // their own SQL function (see v024) and the orchestrator // here is the canonical place that composes them. This // means future perf knobs land additively in // `apply_receipt_plane_fillfactor` (or a sibling helper) // without touching the bigger install helper, and every // path that creates substrate — Rust prepare_schema, fresh // `awa migrate` (via v024's sweep), explicit operator call // against a custom schema — composes the same two pieces. sqlx::query("SELECT awa.apply_receipt_plane_fillfactor($1)") .bind(schema) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; // Post-helper legacy fixups: copy any renamed-aside rows into the // newly partitioned parents created by the helper, then drop the // legacy table. ON CONFLICT DO NOTHING so a re-run after a // partial copy is idempotent. The outer transaction keeps the // copy + drop atomic so a crash between them leaves the schema // in one of exactly two states (pre or post migration); without // that the next `prepare_schema`'s ON CONFLICT would silently // mask the inconsistency. let lease_claims_legacy_exists: bool = sqlx::query_scalar( r#" SELECT EXISTS ( SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = 'lease_claims_legacy' ) "#, ) .bind(schema) .fetch_one(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; let closures_legacy_exists: bool = sqlx::query_scalar( r#" SELECT EXISTS ( SELECT 1 FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace WHERE n.nspname = $1 AND c.relname = 'lease_claim_closures_legacy' ) "#, ) .bind(schema) .fetch_one(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; let legacy_claim_slot: Option = if lease_claims_legacy_exists || closures_legacy_exists { Some( sqlx::query_scalar(audited_sql(format!( "SELECT slot FROM {schema}.claim_ring_rotations ORDER BY generation DESC LIMIT 1" ))) .fetch_one(install_tx.as_mut()) .await .map_err(map_sqlx_error)?, ) } else { None }; if lease_claims_legacy_exists { sqlx::query(audited_sql(format!( "ALTER TABLE {schema}.lease_claims_legacy ADD COLUMN IF NOT EXISTS enqueue_shard SMALLINT NOT NULL DEFAULT 0" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( "ALTER TABLE {schema}.lease_claims_legacy ADD COLUMN IF NOT EXISTS deadline_at TIMESTAMPTZ" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.lease_claims ( claim_slot, job_id, run_lease, ready_slot, ready_generation, queue, priority, attempt, max_attempts, lane_seq, enqueue_shard, claimed_at, materialized_at, deadline_at ) SELECT $1, job_id, run_lease, ready_slot, ready_generation, queue, priority, attempt, max_attempts, lane_seq, enqueue_shard, claimed_at, materialized_at, deadline_at FROM {schema}.lease_claims_legacy ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING "# ))) .bind(legacy_claim_slot.expect("legacy claim slot should be present")) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( "DROP TABLE {schema}.lease_claims_legacy" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; } if closures_legacy_exists { sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.lease_claim_closures ( claim_slot, job_id, run_lease, outcome, closed_at ) SELECT $1, job_id, run_lease, outcome, closed_at FROM {schema}.lease_claim_closures_legacy ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING "# ))) .bind(legacy_claim_slot.expect("legacy claim slot should be present")) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( "DROP TABLE {schema}.lease_claim_closures_legacy" ))) .execute(install_tx.as_mut()) .await .map_err(map_sqlx_error)?; } Ok(()) } .await; match install_result { Ok(()) => install_tx.commit().await.map_err(map_sqlx_error), Err(err) => { let _ = install_tx.rollback().await; Err(err) } } } #[tracing::instrument(skip(self, pool), name = "queue_storage.activate_backend")] pub async fn activate_backend(&self, pool: &PgPool) -> Result<(), AwaError> { let schema = self.schema(); let details = serde_json::json!({ "schema": schema }); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // Mark queue storage active only after the full schema, partitions, // indexes, and helper functions are in place. The explicit install // helper is used by tests and queue-storage-only setups, so it must // flip both the migration-owned routing registry row and the storage // transition state. sqlx::query( r#" INSERT INTO awa.runtime_storage_backends (backend, schema_name, updated_at) VALUES ('queue_storage', $1, now()) ON CONFLICT (backend) DO UPDATE SET schema_name = EXCLUDED.schema_name, updated_at = EXCLUDED.updated_at "#, ) .bind(schema) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; let activation_result = sqlx::query( r#" UPDATE awa.storage_transition_state AS sts SET current_engine = 'queue_storage', prepared_engine = NULL, state = 'active', transition_epoch = CASE WHEN sts.current_engine = 'queue_storage' AND sts.prepared_engine IS NULL AND sts.state = 'active' AND sts.details = $1 THEN sts.transition_epoch ELSE sts.transition_epoch + 1 END, details = $1, entered_at = CASE WHEN sts.current_engine = 'queue_storage' AND sts.prepared_engine IS NULL AND sts.state = 'active' AND sts.details = $1 THEN sts.entered_at ELSE now() END, updated_at = now(), finalized_at = CASE WHEN sts.current_engine = 'queue_storage' AND sts.prepared_engine IS NULL AND sts.state = 'active' AND sts.details = $1 THEN COALESCE(sts.finalized_at, now()) ELSE now() END WHERE sts.singleton "#, ) .bind(details) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; if activation_result.rows_affected() != 1 { return Err(AwaError::Validation( "queue storage activation requires the storage transition state row".into(), )); } tx.commit().await.map_err(map_sqlx_error)?; Ok(()) } #[tracing::instrument(skip(self, pool), name = "queue_storage.install")] pub async fn install(&self, pool: &PgPool) -> Result<(), AwaError> { self.prepare_schema(pool).await?; self.activate_backend(pool).await } pub async fn reset(&self, pool: &PgPool) -> Result<(), AwaError> { let schema = self.schema(); let install_lock_name = format!("awa.queue_storage.install:{schema}"); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") .bind(&install_lock_name) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; self.lock_ring_serializers_tx( &mut tx, &[RingFamily::Queue, RingFamily::Lease, RingFamily::Claim], ) .await?; for family in [RingFamily::Queue, RingFamily::Lease, RingFamily::Claim] { let state = family.state_relname(); sqlx::query(audited_sql(format!( "SELECT 1 FROM {schema}.{state} WHERE singleton = TRUE FOR UPDATE" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } // Drop any partial-migration leftover tables before the main // TRUNCATE. If `prepare_schema` crashed mid-migration, the // schema may contain `lease_claims_legacy` / // `lease_claim_closures_legacy` alongside the partitioned // parents. `reset()` must clean these out, otherwise the next // `prepare_schema()` runs the legacy migration again on top of // the freshly-emptied parent and silently re-inserts old rows. sqlx::query(audited_sql(format!( "DROP TABLE IF EXISTS {schema}.lease_claims_legacy" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( "DROP TABLE IF EXISTS {schema}.lease_claim_closures_legacy" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( r#" TRUNCATE {schema}.ready_entries, {schema}.ready_claim_attempt_batches, {schema}.ready_tombstones, {schema}.ready_segments, {schema}.done_entries, {schema}.receipt_completion_batches, {schema}.receipt_completion_tombstones, {schema}.dlq_entries, {schema}.leases, {schema}.lease_claims, {schema}.lease_claim_batches, {schema}.lease_claim_closures, {schema}.lease_claim_closure_batches, {schema}.attempt_state, {schema}.deferred_jobs, {schema}.queue_lanes, {schema}.queue_terminal_rollups, {schema}.queue_terminal_rollup_deltas, {schema}.queue_terminal_live_counts, {schema}.queue_terminal_count_deltas, {schema}.queue_claimer_leases, {schema}.queue_claimer_state, {schema}.queue_ring_slots, {schema}.lease_ring_slots, {schema}.claim_ring_slots, {schema}.queue_ring_rotations, {schema}.lease_ring_rotations, {schema}.claim_ring_rotations "# ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( "ALTER SEQUENCE {schema}.job_id_seq RESTART WITH 1" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; for (family, slot_count) in [ (RingFamily::Queue, self.queue_slot_count()), (RingFamily::Lease, self.lease_slot_count()), (RingFamily::Claim, self.claim_slot_count()), ] { let state = family.state_relname(); sqlx::query(audited_sql(format!( r#" UPDATE {schema}.{state} SET slot_count = $1 WHERE singleton = TRUE "# ))) .bind(slot_count as i32) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; // The rotation ledgers were truncated above; re-seed the // genesis cursor (slot 0, generation 0). let ledger = family.ledger_relname(); sqlx::query(audited_sql(format!( "INSERT INTO {schema}.{ledger} (generation, slot) VALUES (0, 0)" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } for slot in 0..self.queue_slot_count() { sqlx::query(audited_sql(format!( "INSERT INTO {schema}.queue_ring_slots (slot) VALUES ($1)" ))) .bind(slot as i32) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } for slot in 0..self.lease_slot_count() { sqlx::query(audited_sql(format!( "INSERT INTO {schema}.lease_ring_slots (slot) VALUES ($1)" ))) .bind(slot as i32) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } for slot in 0..self.claim_slot_count() { sqlx::query(audited_sql(format!( "INSERT INTO {schema}.claim_ring_slots (slot) VALUES ($1)" ))) .bind(slot as i32) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } tx.commit().await.map_err(map_sqlx_error)?; // queue_lanes and queue_claimer_leases were TRUNCATEd above and // queue_meta may have a new shard configuration for the next round. // Clear the caches so the next ensure_lane / shard_for_enqueue / // claim calls re-observe DB state. self.clear_lane_cache(); self.enqueue_shards_cache .lock() .expect("enqueue_shards_cache poisoned") .clear(); self.claimer_gates .lock() .expect("claimer gate cache poisoned") .clear(); self.reset_rotation_observations(); Ok(()) } async fn ensure_lane<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, queue: &str, priority: i16, enqueue_shard: i16, ) -> Result<(), AwaError> { // Fast path: this store has previously written the three lane // rows for this `(queue, priority, shard)` triple. The cache is // optimistic: another transaction may have marked the lane before // commit, or the marking transaction may later roll back. Verify the // head row is visible in this transaction before trusting the cache. if self.lane_is_cached(queue, priority, enqueue_shard) { let schema = self.schema(); let visible: bool = sqlx::query_scalar(audited_sql(format!( r#" SELECT EXISTS ( SELECT 1 FROM {schema}.queue_enqueue_heads WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 ) "# ))) .bind(queue) .bind(priority) .bind(enqueue_shard) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; if visible { return Ok(()); } self.invalidate_cached_lane(queue, priority, enqueue_shard); } self.ensure_lane_inserts(tx, queue, priority, enqueue_shard) .await } /// Provision every priority and shard of a queue before deploying a runtime /// role without schema CREATE. Run as the schema owner/migrator. This does /// not change enqueue routing or the configured shard count. pub async fn prepare_queue( &self, pool: &PgPool, queue: &str, enqueue_shards: i16, ) -> Result<(), AwaError> { if queue.is_empty() || !(1..=64).contains(&enqueue_shards) { return Err(AwaError::Validation( "prepare_queue requires a nonempty queue and 1..=64 enqueue shards".into(), )); } let mut transaction = pool.begin().await.map_err(map_sqlx_error)?; let mut physical_queues = self.physical_queues_for_logical(queue); physical_queues.sort_unstable(); for physical_queue in physical_queues { for priority in 1..=4 { for shard in 0..enqueue_shards { self.ensure_lane_inserts(&mut transaction, &physical_queue, priority, shard) .await?; } } } transaction.commit().await.map_err(map_sqlx_error)?; Ok(()) } /// Run the three lane-row inserts unconditionally and mark the /// `(queue, priority)` pair as cached on success. Skips the cache /// fast path so callers in the rollback-recovery path can force a /// re-insert without racing another transaction that has marked /// the lane but not yet committed its inserts. async fn ensure_lane_inserts<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, queue: &str, priority: i16, enqueue_shard: i16, ) -> Result<(), AwaError> { let schema = self.schema(); let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") .bind(lane_lock_key) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.queue_lanes (queue, priority) VALUES ($1, $2) ON CONFLICT (queue, priority) DO NOTHING "# ))) .bind(queue) .bind(priority) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; // The head tables' BEFORE INSERT triggers run even for a row that // ON CONFLICT then skips, and they read and set the lane's sequence. // Skipping rows that exist keeps provisioning a live lane off a // sequence another role owns, such as one a 0.6 runtime created. sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard) SELECT $1, $2, $3 WHERE NOT EXISTS ( SELECT 1 FROM {schema}.queue_enqueue_heads WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 ) ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING "# ))) .bind(queue) .bind(priority) .bind(enqueue_shard) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard) SELECT $1, $2, $3 WHERE NOT EXISTS ( SELECT 1 FROM {schema}.queue_claim_heads WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 ) ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING "# ))) .bind(queue) .bind(priority) .bind(enqueue_shard) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( r#" SELECT {schema}.ensure_lane_sequences($1, $2, $3) "# ))) .bind(queue) .bind(priority) .bind(enqueue_shard) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; self.mark_lane_ensured(queue, priority, enqueue_shard); Ok(()) } /// Pick the enqueue shard for a row on `queue`. /// /// Reads `awa.queue_meta.enqueue_shards` (default 1, cached /// in-process). With `enqueue_shards = 1` the result is always 0. /// With `enqueue_shards > 1`: /// /// - If `ordering_key` is `Some(k)`, the shard is a stable hash /// of `k` modulo the shard count. Two rows sharing an /// ordering key always land on the same shard, which preserves /// FIFO within the key. /// - If `ordering_key` is `None`, the shard comes from the /// per-store rotor, which spreads consecutive picks across /// shards. async fn shard_for_enqueue( &self, pool_executor: impl sqlx::PgExecutor<'_>, queue: &str, ordering_key: Option<&[u8]>, ) -> Result { if let Some(cached) = self .enqueue_shards_cache .lock() .expect("enqueue_shards_cache poisoned") .get(queue) .copied() { return Ok(self.pick_shard(cached, ordering_key)); } let shards: i16 = sqlx::query_scalar( r#" SELECT COALESCE(MAX(enqueue_shards), 1)::smallint FROM awa.queue_meta WHERE queue = $1 "#, ) .bind(queue) .fetch_one(pool_executor) .await .map_err(map_sqlx_error)?; let shards = shards.max(1); self.enqueue_shards_cache .lock() .expect("enqueue_shards_cache poisoned") .insert(queue.to_string(), shards); Ok(self.pick_shard(shards, ordering_key)) } /// Map a row (or a no-key sub-batch) to its shard. /// /// At `shards <= 1` every row goes to shard 0. Otherwise the /// caller-supplied `ordering_key` hashes deterministically into /// `[0, shards)` via [`shard_for_ordering_key`]; an absent key /// advances the per-store rotor once and returns the next shard. /// Callers route a whole no-key sub-batch through a single /// `pick_shard(_, None)` so the rotor amortises across the batch /// rather than firing per row. fn pick_shard(&self, shards: i16, ordering_key: Option<&[u8]>) -> i16 { if shards <= 1 { return 0; } match ordering_key { Some(key) => shard_for_ordering_key(key, shards), None => { let raw = self.shard_rotor.fetch_add(1, Ordering::Relaxed) as i32; raw.rem_euclid(shards as i32) as i16 } } } fn lane_is_cached(&self, queue: &str, priority: i16, enqueue_shard: i16) -> bool { let cache = self.ensured_lanes.lock().expect("ensured_lanes mutex"); cache.contains(&(queue.to_string(), priority, enqueue_shard)) } fn mark_lane_ensured(&self, queue: &str, priority: i16, enqueue_shard: i16) { self.ensured_lanes .lock() .expect("ensured_lanes mutex") .insert((queue.to_string(), priority, enqueue_shard)); } fn invalidate_cached_lane(&self, queue: &str, priority: i16, enqueue_shard: i16) { self.ensured_lanes .lock() .expect("ensured_lanes mutex") .remove(&(queue.to_string(), priority, enqueue_shard)); } fn clear_lane_cache(&self) { self.ensured_lanes .lock() .expect("ensured_lanes mutex") .clear(); } // Reserve lane sequence numbers for a specific // `(queue, priority, shard)` triple and return the lane sequence at // which the caller's range starts. If the head row is missing — // typically because a previous ensure_lane ran inside a transaction // that ultimately rolled back, leaving a stale cache entry behind — // the cache entry is invalidated, ensure_lane is re-run, and the reserve // is retried exactly once. A second miss surfaces as // RowNotFound. async fn advance_enqueue_head<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, queue: &str, priority: i16, enqueue_shard: i16, count: i64, ) -> Result { let schema = self.schema(); let sql = format!( r#" SELECT {schema}.reserve_enqueue_seq($1, $2, $3, $4) FROM {schema}.queue_enqueue_heads WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 "# ); let maybe_start: Option = sqlx::query_scalar(audited_sql(sql.clone())) .bind(queue) .bind(priority) .bind(enqueue_shard) .bind(count) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if let Some(start) = maybe_start { return Ok(start); } // Recovery path: a prior ensure_lane marked the cache and // then rolled back, leaving the head row missing. Bypass the // cache fast path here and run the inserts unconditionally — // calling `ensure_lane` would re-take the fast path if // another concurrent transaction has re-marked the cache but // not yet committed its inserts, leaving us in the same // failure state. self.invalidate_cached_lane(queue, priority, enqueue_shard); self.ensure_lane_inserts(tx, queue, priority, enqueue_shard) .await?; let start: i64 = sqlx::query_scalar(audited_sql(sql)) .bind(queue) .bind(priority) .bind(enqueue_shard) .bind(count) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(start) } /// Which cursor representation is authoritative for this schema (#371 /// staged upgrade). Read once per rotate/prune tick and threaded through /// the seams below so one transaction picks a single mode. async fn ring_authority_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ) -> Result { let schema = self.schema(); let authority: String = sqlx::query_scalar(audited_sql(format!( "SELECT authority FROM {schema}.ring_cursor_authority WHERE singleton = TRUE" ))) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(match authority.as_str() { "ledger" => RingAuthority::Ledger, _ => RingAuthority::Columns, }) } /// Current ring cursor. In ledger authority it is the max-generation row /// of the append-only ledger (backward PK scan, O(1)); in compat /// authority it is the pre-0.7 mutable singleton columns, exactly as a /// live 0.6 rotator maintains them (#371 staged upgrade). async fn ring_cursor_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, family: RingFamily, authority: RingAuthority, ) -> Result<(i32, i64), AwaError> { let schema = self.schema(); let sql = match authority { RingAuthority::Ledger => { let ledger = family.ledger_relname(); format!( "SELECT slot, generation FROM {schema}.{ledger} \ ORDER BY generation DESC LIMIT 1" ) } RingAuthority::Columns => { let state = family.state_relname(); format!( "SELECT current_slot, generation FROM {schema}.{state} \ WHERE singleton = TRUE" ) } }; sqlx::query_as(audited_sql(sql)) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error) } /// Ring width from the demoted cold-config singleton. async fn ring_slot_count_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, family: RingFamily, ) -> Result { let schema = self.schema(); let state = family.state_relname(); sqlx::query_scalar(audited_sql(format!( "SELECT slot_count FROM {schema}.{state} WHERE singleton = TRUE" ))) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error) } /// Wait for the per-ring rotation advisory locks, in a fixed ring order. /// Whole-schema operations (reset, counter rebuild) take these before any /// row or table lock, so in ledger authority, where rotate, prune and the /// delta rollup serialize only on these locks, they wait for an /// in-flight tick instead of interleaving table locks with it. async fn lock_ring_serializers_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, families: &[RingFamily], ) -> Result<(), AwaError> { let schema = self.schema(); for family in families { sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") .bind(family.rotation_lock_name(schema)) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } Ok(()) } /// Acquire the per-ring rotate ↔ prune ↔ delta-rollup serializer and /// resolve the authority under it. /// /// In compat (`columns`) authority the ring-state singleton `FOR UPDATE` /// is the serializer. It is taken before the authority is re-read: the /// one-way flip (`awa.flip_ring_authority`) takes the same three /// singletons before promoting the authority, so a rotator that loses the /// race to the flip reads `ledger` under the lock and runs the ledger /// path, and a flip that loses waits for the rotator to commit. The same /// row lock is the one a 0.6 rotator takes, so mixed-version rotators /// serialize. /// /// In ledger authority the non-blocking per-ring /// `pg_try_advisory_xact_lock` is the whole serializer, so a periodic /// rotate skips under contention with prune or rollup instead of /// queueing. `Ok((_, false))` means "skip this tick". The singleton row /// is not locked: a row lock writes a heap lock record and a commit /// record per tick even when nothing rotates, and because the flip is /// one-way a `ledger` reading taken without the row lock is final. async fn acquire_ring_serializer_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, family: RingFamily, ) -> Result<(RingAuthority, bool), AwaError> { let schema = self.schema(); let state = family.state_relname(); let authority = match self.ring_authority_tx(tx).await? { RingAuthority::Ledger => RingAuthority::Ledger, RingAuthority::Columns => { // Serialize against the flip and against 0.6 rotators, then // re-read the authority under the lock: a rotator that // loses the race to the flip sees `ledger` here. sqlx::query(audited_sql(format!( "SELECT 1 FROM {schema}.{state} WHERE singleton = TRUE FOR UPDATE" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; self.ring_authority_tx(tx).await? } }; match authority { RingAuthority::Columns => Ok((authority, true)), RingAuthority::Ledger => { let acquired: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock(hashtextextended($1, 0))") .bind(family.rotation_lock_name(schema)) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok((authority, acquired)) } } } /// Advance the ring cursor by one. /// /// In ledger authority this is a bare ledger append; the `generation` /// primary key makes the insert a compare-and-swap: a conflict means /// another rotator observed the same cursor and won — the caller treats /// `false` as a lost race, never retrying with the same generation. /// /// In compat authority (#371 staged upgrade) it CASes the pre-0.7 /// singleton columns exactly as a 0.6 rotator does (`SET current_slot, /// generation WHERE singleton AND generation = next-1`), updates the /// incoming slot's per-slot generation as 0.6 did, AND reconciles the /// shadow ledger: it inserts any generations the ledger is missing up to /// `next_generation` (an interleaved 0.6 rotator may have advanced the /// columns since our last append) with `ON CONFLICT DO NOTHING`. The CAS /// result (columns advanced by exactly this tick) is the return value. async fn append_ring_rotation_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, family: RingFamily, next_generation: i64, next_slot: i32, authority: RingAuthority, ) -> Result { let schema = self.schema(); let ledger = family.ledger_relname(); match authority { RingAuthority::Ledger => { let inserted = sqlx::query(audited_sql(format!( "INSERT INTO {schema}.{ledger} (generation, slot) \ VALUES ($1, $2) ON CONFLICT (generation) DO NOTHING" ))) .bind(next_generation) .bind(next_slot) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(inserted.rows_affected() == 1) } RingAuthority::Columns => { let state = family.state_relname(); let slots = family.slots_relname(); // CAS the authoritative compat columns (0.6 semantics). let advanced = sqlx::query(audited_sql(format!( "UPDATE {schema}.{state} \ SET current_slot = $1, generation = $2 \ WHERE singleton = TRUE AND generation = $2 - 1" ))) .bind(next_slot) .bind(next_generation) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; if advanced.rows_affected() != 1 { // Lost the CAS to an interleaved rotator (0.6 or 0.7). return Ok(false); } // Stamp the incoming slot's per-slot generation, as 0.6 did, // so a mixed-fleet 0.6 binary's sealed-slot logic stays // consistent. sqlx::query(audited_sql(format!( "UPDATE {schema}.{slots} SET generation = $1 WHERE slot = $2" ))) .bind(next_generation) .bind(next_slot) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; // Reconcile the shadow ledger up to the new column cursor. // An interleaved 0.6 rotator advances only the columns, so // the ledger can lag by more than one generation; backfill // every missing generation (deriving each slot) so the // ledger is a faithful shadow ready for the flip. let slot_count = self.ring_slot_count_tx(tx, family).await?; sqlx::query(audited_sql(format!( "INSERT INTO {schema}.{ledger} (generation, slot) \ SELECT g, ((g % $2) + $2) % $2 \ FROM generate_series( \ COALESCE((SELECT max(generation) FROM {schema}.{ledger}), -1) + 1, \ $1 \ ) AS g \ ON CONFLICT (generation) DO NOTHING" ))) .bind(next_generation) .bind(slot_count) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(true) } } } /// Current queue-ring cursor for the Rust-side enqueue paths, honouring /// the schema's authority (#371 staged upgrade). Reads the authority row /// (single cached page) then the O(1) cursor; callers that already hold /// an authority for the transaction should use `ring_cursor_tx` directly. async fn current_queue_ring<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ) -> Result<(i32, i64), AwaError> { let authority = self.ring_authority_tx(tx).await?; self.ring_cursor_tx(tx, RingFamily::Queue, authority).await } async fn next_job_ids<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, count: usize, ) -> Result, AwaError> { if count == 0 { return Ok(Vec::new()); } let query = format!( "SELECT nextval('{}')::bigint FROM generate_series(1, $1::int)", self.job_id_sequence() ); sqlx::query_scalar(audited_sql(query)) .bind(count as i32) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error) } async fn current_timestamp_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ) -> Result, AwaError> { sqlx::query_scalar("SELECT clock_timestamp()") .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error) } async fn claim_ready_rows_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, queue: &str, max_batch: i64, deadline_duration: Duration, aging_interval: Duration, ) -> Result, AwaError> { let schema = self.schema(); sqlx::query_as(audited_sql(format!( r#" SELECT ready_slot, ready_generation, lane_seq, enqueue_shard, lease_slot, lease_generation, claim_slot, receipt_id, claim_batch_id, claim_batch_index, job_id, kind, queue, args, lane_priority, priority, attempt, run_lease, max_attempts, run_at, heartbeat_at, deadline_at, attempted_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM {schema}.claim_ready_runtime($1, $2, $3, $4) "# ))) .bind(queue) .bind(max_batch) .bind(deadline_duration.as_secs_f64()) .bind(aging_interval.as_secs_f64()) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error) } fn claim_cursor_advances(rows: &[ReadyJobLeaseRow]) -> Vec { let mut last_by_lane: BTreeMap = BTreeMap::new(); for row in rows { let key = (row.queue.clone(), row.lane_priority, row.enqueue_shard); last_by_lane .entry(key) .and_modify(|current| { if row.lane_seq > current.lane_seq { *current = row; } }) .or_insert(row); } last_by_lane .into_iter() .map( |((queue, priority, enqueue_shard), last)| ClaimCursorAdvance { queue, priority, enqueue_shard, next_seq: last.lane_seq + 1, only_if_current: None, evidence: Some(ClaimEvidence { ready_slot: last.ready_slot, ready_generation: last.ready_generation, lane_seq: last.lane_seq, legacy_attempt: last .claim_batch_id .is_none() .then_some((last.job_id, last.run_lease)), }), }, ) .collect() } fn normalize_claim_cursor_advances(advances: &[ClaimCursorAdvance]) -> Vec { let mut grouped: GroupedClaimCursorAdvances = BTreeMap::new(); for advance in advances { let key = ( advance.queue.clone(), advance.priority, advance.enqueue_shard, ); let (unconditional, conditional) = grouped.entry(key).or_default(); if let Some(only_if_current) = advance.only_if_current { conditional .entry(only_if_current) .and_modify(|next| *next = (*next).max(advance.next_seq)) .or_insert(advance.next_seq); } else if unconditional.is_none_or(|(next, _)| advance.next_seq > next) { *unconditional = Some((advance.next_seq, advance.evidence)); } } let mut normalized = Vec::with_capacity(advances.len()); for ((queue, priority, enqueue_shard), (unconditional, conditional)) in grouped { if let Some((next_seq, evidence)) = unconditional { normalized.push(ClaimCursorAdvance { queue, priority, enqueue_shard, next_seq, only_if_current: None, evidence, }); continue; } for (only_if_current, next_seq) in conditional { normalized.push(ClaimCursorAdvance { queue: queue.clone(), priority, enqueue_shard, next_seq, only_if_current: Some(only_if_current), evidence: None, }); } } normalized } /// Replay the post-commit advance for an actual claim on the current server. #[doc(hidden)] pub async fn replay_claim_cursor_advance( &self, pool: &PgPool, claimed: &ClaimedRuntimeJob, ) -> Result<(), AwaError> { let claim = &claimed.claim; self.advance_claim_cursors_strict( pool, &[ClaimCursorAdvance { queue: claim.queue.clone(), priority: claim.priority, enqueue_shard: claim.enqueue_shard, next_seq: claim.lane_seq + 1, only_if_current: None, evidence: Some(ClaimEvidence { ready_slot: claim.ready_slot, ready_generation: claim.ready_generation, lane_seq: claim.lane_seq, legacy_attempt: claim .claim_batch_id .is_none() .then_some((claimed.job.id, claimed.job.run_lease)), }), }], ) .await } async fn advance_claim_cursors(&self, pool: &PgPool, advances: &[ClaimCursorAdvance]) { let advances = Self::normalize_claim_cursor_advances(advances); // PostgreSQL sequence state is not rolled back with the surrounding // transaction. Keep claim cursors lagging rather than risking a cursor // that gets ahead of ready rows whose claim/cancel transaction aborts. for attempt in 1..=3 { match self.advance_claim_cursors_strict(pool, &advances).await { Ok(()) => return, Err(err) if attempt < 3 => { tracing::warn!( error = ?err, lanes = advances.len(), attempt, "failed to advance queue-storage claim cursors after committed state change; retrying" ); tokio::time::sleep(Duration::from_millis(25 * attempt as u64)).await; } Err(err) => { tracing::warn!( error = ?err, lanes = advances.len(), attempts = attempt, "failed to advance queue-storage claim cursors after committed state change" ); return; } } } } async fn advance_claim_cursors_strict( &self, pool: &PgPool, advances: &[ClaimCursorAdvance], ) -> Result<(), AwaError> { if advances.is_empty() { return Ok(()); } let schema = self.schema(); let mut tx = pool .begin_with(CLAIM_TRANSACTION_BEGIN) .await .map_err(map_sqlx_error)?; for advance in advances { sqlx::query(audited_sql(format!( r#" WITH head AS MATERIALIZED ( SELECT seq_name FROM {schema}.queue_claim_heads WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3 FOR UPDATE ) SELECT {schema}.set_sequence_next(seq_name, $4) FROM head WHERE ($5::bigint IS NULL OR {schema}.sequence_next_value(seq_name) = $5) AND ( $6::int IS NULL OR ($9::bigint IS NULL AND EXISTS ( SELECT 1 FROM {schema}.ready_claim_attempt_batches AS attempts WHERE attempts.ready_slot = $6 AND attempts.ready_generation = $7 AND attempts.queue = $1 AND attempts.priority = $2 AND attempts.enqueue_shard = $3 AND attempts.first_lane_seq <= $8 AND attempts.next_lane_seq > $8 AND attempts.lane_ranges @> int8range($8, $8 + 1, '[)') )) OR EXISTS ( SELECT 1 FROM {schema}.leases AS leases WHERE $9::bigint IS NOT NULL AND leases.job_id = $9 AND leases.run_lease = $10 AND leases.ready_slot = $6 AND leases.ready_generation = $7 AND leases.queue = $1 AND leases.enqueue_shard = $3 AND leases.lane_seq = $8 ) ) "# ))) .bind(&advance.queue) .bind(advance.priority) .bind(advance.enqueue_shard) .bind(advance.next_seq) .bind(advance.only_if_current) .bind(advance.evidence.map(|evidence| evidence.ready_slot)) .bind(advance.evidence.map(|evidence| evidence.ready_generation)) .bind(advance.evidence.map(|evidence| evidence.lane_seq)) .bind( advance .evidence .and_then(|evidence| evidence.legacy_attempt.map(|attempt| attempt.0)), ) .bind( advance .evidence .and_then(|evidence| evidence.legacy_attempt.map(|attempt| attempt.1)), ) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } tx.commit().await.map_err(map_sqlx_error)?; Ok(()) } async fn execute_ready_inserts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ring: (i32, i64), rows: &[RuntimeReadyInsert], ) -> Result { if rows.is_empty() { return Ok(0); } let schema = self.schema(); sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.ready_entries (ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, payload) SELECT $1, $2, rows.* FROM unnest( $3::bigint[], $4::text[], $5::text[], $6::jsonb[], $7::smallint[], $8::smallint[], $9::bigint[], $10::smallint[], $11::bigint[], $12::smallint[], $13::timestamptz[], $14::timestamptz[], $15::timestamptz[], $16::bytea[], $17::text[], $18::jsonb[] ) AS rows "# ))) .bind(ring.0) .bind(ring.1) .bind(rows.iter().map(|row| row.job_id).collect::>()) .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind(rows.iter().map(|row| row.queue.as_str()).collect::>()) .bind(rows.iter().map(|row| &row.args).collect::>()) .bind(rows.iter().map(|row| row.priority).collect::>()) .bind(rows.iter().map(|row| row.attempt).collect::>()) .bind(rows.iter().map(|row| row.run_lease).collect::>()) .bind(rows.iter().map(|row| row.max_attempts).collect::>()) .bind(rows.iter().map(|row| row.lane_seq).collect::>()) .bind(rows.iter().map(|row| row.enqueue_shard).collect::>()) .bind(rows.iter().map(|row| row.run_at).collect::>()) .bind(rows.iter().map(|row| row.attempted_at).collect::>()) .bind(rows.iter().map(|row| row.created_at).collect::>()) .bind(rows.iter().map(|row| row.unique_key.as_deref()).collect::>()) .bind(rows.iter().map(|row| row.unique_states.as_deref()).collect::>()) .bind(rows.iter().map(|row| storage_payload(&row.payload)).collect::>()) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(rows.len()) } fn ready_segments_from_rows(rows: &[RuntimeReadyInsert]) -> Vec { if rows.is_empty() { return Vec::new(); } let mut segments = Vec::new(); let mut start_index = 0usize; for index in 1..=rows.len() { let split = if index == rows.len() { true } else { let previous = &rows[index - 1]; let current = &rows[index]; previous.queue != current.queue || previous.priority != current.priority || previous.enqueue_shard != current.enqueue_shard || previous.lane_seq + 1 != current.lane_seq || previous.run_at != current.run_at }; if split { let first = &rows[start_index]; let last = &rows[index - 1]; segments.push(ReadySegmentInsert { queue: first.queue.clone(), priority: first.priority, enqueue_shard: first.enqueue_shard, first_lane_seq: first.lane_seq, next_lane_seq: last.lane_seq + 1, first_run_at: first.run_at, }); start_index = index; } } segments } async fn insert_ready_segments_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ring: (i32, i64), rows: &[RuntimeReadyInsert], ) -> Result { let segments = Self::ready_segments_from_rows(rows); if segments.is_empty() { return Ok(0); } let schema = self.schema(); sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.ready_segments (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq, next_lane_seq, first_run_at) SELECT $1, $2, segments.* FROM unnest($3::text[], $4::smallint[], $5::smallint[], $6::bigint[], $7::bigint[], $8::timestamptz[]) AS segments ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) DO NOTHING "# ))) .bind(ring.0) .bind(ring.1) .bind(segments.iter().map(|segment| segment.queue.as_str()).collect::>()) .bind(segments.iter().map(|segment| segment.priority).collect::>()) .bind(segments.iter().map(|segment| segment.enqueue_shard).collect::>()) .bind(segments.iter().map(|segment| segment.first_lane_seq).collect::>()) .bind(segments.iter().map(|segment| segment.next_lane_seq).collect::>()) .bind(segments.iter().map(|segment| segment.first_run_at).collect::>()) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(segments.len()) } async fn execute_ready_copy_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ring: (i32, i64), rows: &[RuntimeReadyInsert], ) -> Result { if rows.is_empty() { return Ok(0); } if rows.len() < COPY_MIN_ROWS { return self.execute_ready_inserts_tx(tx, ring, rows).await; } let schema = self.schema(); // `copy_in_raw` is outside sqlx 0.9's `SqlSafeStr` guard, so this // site cannot opt in through `audited_sql`. It carries the same // invariant by hand: `schema` passed `validate_ident`, the sentinel // is a const, and every row value is CSV-encoded into the COPY // stream rather than into the statement. let copy_sql = format!( "COPY {schema}.ready_entries (ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, payload) FROM STDIN WITH (FORMAT csv, NULL '{COPY_NULL_SENTINEL}')" ); let mut copy_in = tx .as_mut() .copy_in_raw(©_sql) .await .map_err(map_sqlx_error)?; // 320 bytes/row is only a rough starting point; large JSON payloads // are bounded by chunked COPY sends below rather than by this reserve. let mut csv_buf = Vec::with_capacity(rows.len().min(1024) * 320); for row in rows { write_ready_copy_row(&mut csv_buf, ring.0, ring.1, row); if csv_buf.len() >= COPY_CHUNK_TARGET_BYTES { let chunk = std::mem::replace(&mut csv_buf, Vec::with_capacity(COPY_CHUNK_TARGET_BYTES)); copy_in.send(chunk).await.map_err(map_sqlx_error)?; } } if !csv_buf.is_empty() { copy_in.send(csv_buf).await.map_err(map_sqlx_error)?; } copy_in.finish().await.map_err(map_sqlx_error)?; Ok(rows.len()) } async fn insert_ready_rows_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, unique_claims_synced: bool, ) -> Result { if rows.is_empty() { return Ok(0); } let grouped = self.group_ready_rows_by_shard(tx, rows).await?; let total_rows: usize = grouped.values().map(Vec::len).sum(); let missing_job_ids = grouped .values() .flatten() .filter(|row| row.job_id.is_none()) .count(); let job_ids = self.next_job_ids(tx, missing_job_ids).await?; let mut job_id_iter = job_ids.into_iter(); let mut ready_rows = Vec::with_capacity(total_rows); let mut lane_ranges = Vec::with_capacity(grouped.len()); for ((queue, priority, enqueue_shard), lane_rows) in grouped { let range_start = ready_rows.len(); for row in lane_rows { let job_id = match row.job_id { Some(job_id) => job_id, None => job_id_iter.next().ok_or_else(|| { AwaError::Validation( "queue storage job id allocation underflow".to_string(), ) })?, }; ready_rows.push(RuntimeReadyInsert { job_id, kind: row.kind, queue: row.queue, args: row.args, priority: row.priority, attempt: row.attempt, run_lease: row.run_lease, max_attempts: row.max_attempts, run_at: row.run_at, attempted_at: row.attempted_at, lane_seq: 0, enqueue_shard, created_at: row.created_at, unique_key: row.unique_key, unique_states: row.unique_states, payload: row.payload, }); } lane_ranges.push(( queue, priority, enqueue_shard, range_start, ready_rows.len(), )); } if !unique_claims_synced { self.sync_ready_enqueue_unique_claims(tx, &ready_rows) .await?; } for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges { self.ensure_lane(tx, &queue, priority, enqueue_shard) .await?; let count = (range_end - range_start) as i64; let start_seq = self .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count) .await?; for (offset, row) in ready_rows[range_start..range_end].iter_mut().enumerate() { row.lane_seq = start_seq + offset as i64; } } let ring = self.current_queue_ring(tx).await?; self.execute_ready_inserts_tx(tx, ring, &ready_rows).await?; self.insert_ready_segments_tx(tx, ring, &ready_rows).await?; Ok(total_rows) } /// Re-group rows by `(queue, priority, enqueue_shard)` so each /// resulting bucket targets exactly one head row. /// /// The two routing modes are amortised differently: /// /// - **Rows with `ordering_key`** are hashed per row via /// `shard_for_ordering_key`, so jobs that share a key always /// land on the same shard. Each distinct key may produce its /// own sub-bucket inside one batch. /// - **Rows without `ordering_key`** share **one rotor pick per /// `(queue, priority)` per call**. A batch of 500 rotor-routed /// rows produces one `advance_enqueue_head` UPDATE and one /// INSERT, not 500. The rotor still advances per batch, so /// successive batches spread across shards; the per-batch /// amortisation is what makes `enqueue_shards > 1` net-faster /// than `S = 1` at moderate concurrency. /// /// Mixing keyed and rotor rows in one call is supported — the /// keyed rows fan out by key while the rotor rows collapse into /// a single shard-sub-batch. /// /// The map's key order, `(queue, priority, shard)` ascending, is the /// order callers reserve lanes in. It is the canonical lane-lock order /// that `insert_many` and `insert_many_copy` also follow, so batches /// from either path cannot deadlock on each other's lane locks. async fn group_ready_rows_by_shard<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, ) -> Result>, AwaError> { // Partition by (queue, priority) first so the rotor pick for // the no-key sub-batch can amortise over every row that shares // its destination lane. let mut by_queue_priority: BTreeMap<(String, i16), Vec> = BTreeMap::new(); for row in rows { by_queue_priority .entry((row.queue.clone(), row.priority)) .or_default() .push(row); } let mut grouped: BTreeMap<(String, i16, i16), Vec> = BTreeMap::new(); for ((queue, priority), bucket) in by_queue_priority { let mut rotor_rows: Vec = Vec::with_capacity(bucket.len()); for row in bucket { if row.ordering_key.is_some() { let shard = self .shard_for_enqueue(tx.as_mut(), &queue, row.ordering_key.as_deref()) .await?; grouped .entry((queue.clone(), priority, shard)) .or_default() .push(row); } else { rotor_rows.push(row); } } if !rotor_rows.is_empty() { let shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?; grouped .entry((queue.clone(), priority, shard)) .or_default() .extend(rotor_rows); } } Ok(grouped) } async fn insert_ready_rows_copy_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, unique_claims_synced: bool, ) -> Result { if rows.is_empty() { return Ok(0); } let grouped = self.group_ready_rows_by_shard(tx, rows).await?; let total_rows: usize = grouped.values().map(Vec::len).sum(); let missing_job_ids = grouped .values() .flatten() .filter(|row| row.job_id.is_none()) .count(); let job_ids = self.next_job_ids(tx, missing_job_ids).await?; let mut job_id_iter = job_ids.into_iter(); let mut ready_rows = Vec::with_capacity(total_rows); let mut lane_ranges = Vec::with_capacity(grouped.len()); for ((queue, priority, enqueue_shard), lane_rows) in grouped { let range_start = ready_rows.len(); for row in lane_rows { let job_id = match row.job_id { Some(job_id) => job_id, None => job_id_iter.next().ok_or_else(|| { AwaError::Validation( "queue storage job id allocation underflow".to_string(), ) })?, }; ready_rows.push(RuntimeReadyInsert { job_id, kind: row.kind, queue: row.queue, args: row.args, priority: row.priority, attempt: row.attempt, run_lease: row.run_lease, max_attempts: row.max_attempts, run_at: row.run_at, attempted_at: row.attempted_at, lane_seq: 0, enqueue_shard, created_at: row.created_at, unique_key: row.unique_key, unique_states: row.unique_states, payload: row.payload, }); } lane_ranges.push(( queue, priority, enqueue_shard, range_start, ready_rows.len(), )); } if !unique_claims_synced { self.sync_ready_enqueue_unique_claims(tx, &ready_rows) .await?; } for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges { self.ensure_lane(tx, &queue, priority, enqueue_shard) .await?; let count = (range_end - range_start) as i64; let start_seq = self .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count) .await?; for (offset, row) in ready_rows[range_start..range_end].iter_mut().enumerate() { row.lane_seq = start_seq + offset as i64; } } let ring = self.current_queue_ring(tx).await?; self.execute_ready_copy_tx(tx, ring, &ready_rows).await?; self.insert_ready_segments_tx(tx, ring, &ready_rows).await?; Ok(total_rows) } async fn insert_existing_ready_rows_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, old_state: Option, ) -> Result { if rows.is_empty() { return Ok(0); } let mut unique_rows: Vec<&ExistingReadyRow> = rows.iter().collect(); unique_rows.sort_unstable_by(|left, right| { left.unique_key .cmp(&right.unique_key) .then(left.job_id.cmp(&right.job_id)) }); for row in unique_rows { self.sync_unique_claim( tx, row.job_id, &row.unique_key, row.unique_states.as_deref(), old_state, Some(JobState::Available), ) .await?; } let mut grouped: BTreeMap<(String, i16), Vec> = BTreeMap::new(); for row in rows { grouped .entry((row.queue.clone(), row.priority)) .or_default() .push(row); } let total_rows: usize = grouped.values().map(Vec::len).sum(); let mut ready_rows = Vec::with_capacity(total_rows); for ((queue, priority), lane_rows) in grouped { // Re-enqueue paths (retry-after, age-waiting, DLQ retry, // callback resume) do not carry the producer's original // `ordering_key`: the row came back from terminal storage // or from a deferred row, where the key was not retained. // Fall back to the rotor so the retry batch picks a shard // uniformly. Workloads that need ordering preserved across // retries must re-enqueue with the original key from the // application layer. let enqueue_shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?; self.ensure_lane(tx, &queue, priority, enqueue_shard) .await?; let count = lane_rows.len() as i64; let start_seq = self .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count) .await?; for (offset, row) in lane_rows.into_iter().enumerate() { ready_rows.push(RuntimeReadyInsert { job_id: row.job_id, kind: row.kind, queue: row.queue, args: row.args, priority: row.priority, attempt: row.attempt, run_lease: row.run_lease, max_attempts: row.max_attempts, run_at: row.run_at, attempted_at: row.attempted_at, lane_seq: start_seq + offset as i64, enqueue_shard, created_at: row.created_at, unique_key: row.unique_key, unique_states: row.unique_states, payload: row.payload, }); } } let ring = self.current_queue_ring(tx).await?; self.execute_ready_inserts_tx(tx, ring, &ready_rows).await?; self.insert_ready_segments_tx(tx, ring, &ready_rows).await?; Ok(total_rows) } async fn insert_deferred_rows_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, old_state: Option, ) -> Result { if rows.is_empty() { return Ok(0); } if old_state.is_none() { self.sync_deferred_enqueue_unique_claims(tx, &rows).await?; } else { let mut unique_rows: Vec<&DeferredJobRow> = rows.iter().collect(); unique_rows.sort_unstable_by(|left, right| { left.unique_key .cmp(&right.unique_key) .then(left.job_id.cmp(&right.job_id)) }); for row in unique_rows { self.sync_unique_claim( tx, row.job_id, &row.unique_key, row.unique_states.as_deref(), old_state, Some(row.state), ) .await?; } } self.execute_deferred_inserts_tx(tx, &rows).await } async fn execute_deferred_inserts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[DeferredJobRow], ) -> Result { if rows.is_empty() { return Ok(0); } let schema = self.schema(); sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.deferred_jobs (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload) SELECT job_id, kind, queue, args, state::awa.job_state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload FROM unnest( $1::bigint[], $2::text[], $3::text[], $4::jsonb[], $5::text[], $6::smallint[], $7::smallint[], $8::bigint[], $9::smallint[], $10::timestamptz[], $11::timestamptz[], $12::timestamptz[], $13::timestamptz[], $14::bytea[], $15::text[], $16::jsonb[] ) AS rows (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload) "# ))) .bind(rows.iter().map(|row| row.job_id).collect::>()) .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind(rows.iter().map(|row| row.queue.as_str()).collect::>()) .bind(rows.iter().map(|row| &row.args).collect::>()) .bind(rows.iter().map(|row| row.state.as_str()).collect::>()) .bind(rows.iter().map(|row| row.priority).collect::>()) .bind(rows.iter().map(|row| row.attempt).collect::>()) .bind(rows.iter().map(|row| row.run_lease).collect::>()) .bind(rows.iter().map(|row| row.max_attempts).collect::>()) .bind(rows.iter().map(|row| row.run_at).collect::>()) .bind(rows.iter().map(|row| row.attempted_at).collect::>()) .bind(rows.iter().map(|row| row.finalized_at).collect::>()) .bind(rows.iter().map(|row| row.created_at).collect::>()) .bind(rows.iter().map(|row| row.unique_key.as_deref()).collect::>()) .bind(rows.iter().map(|row| row.unique_states.as_deref()).collect::>()) .bind(rows.iter().map(|row| storage_payload(&row.payload)).collect::>()) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(rows.len()) } async fn insert_deferred_rows_copy_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, unique_claims_synced: bool, ) -> Result { if rows.is_empty() { return Ok(0); } if !unique_claims_synced { self.sync_deferred_enqueue_unique_claims(tx, &rows).await?; } if rows.len() < COPY_MIN_ROWS { return self.execute_deferred_inserts_tx(tx, &rows).await; } let schema = self.schema(); // `copy_in_raw` is outside sqlx 0.9's `SqlSafeStr` guard, so this // site cannot opt in through `audited_sql`. It carries the same // invariant by hand: `schema` passed `validate_ident`, the sentinel // is a const, and every row value is CSV-encoded into the COPY // stream rather than into the statement. let copy_sql = format!( "COPY {schema}.deferred_jobs (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload) FROM STDIN WITH (FORMAT csv, NULL '{COPY_NULL_SENTINEL}')" ); let mut copy_in = tx .as_mut() .copy_in_raw(©_sql) .await .map_err(map_sqlx_error)?; // 320 bytes/row is only a rough starting point; large JSON payloads // are bounded by chunked COPY sends below rather than by this reserve. let mut csv_buf = Vec::with_capacity(rows.len().min(1024) * 320); for row in &rows { write_deferred_copy_row(&mut csv_buf, row); if csv_buf.len() >= COPY_CHUNK_TARGET_BYTES { let chunk = std::mem::replace(&mut csv_buf, Vec::with_capacity(COPY_CHUNK_TARGET_BYTES)); copy_in.send(chunk).await.map_err(map_sqlx_error)?; } } if !csv_buf.is_empty() { copy_in.send(csv_buf).await.map_err(map_sqlx_error)?; } copy_in.finish().await.map_err(map_sqlx_error)?; Ok(rows.len()) } async fn insert_done_rows_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[DoneJobRow], old_state: Option, ) -> Result { if rows.is_empty() { return Ok(0); } let mut unique_rows: Vec<&DoneJobRow> = rows.iter().collect(); unique_rows.sort_unstable_by(|left, right| { left.unique_key .cmp(&right.unique_key) .then(left.job_id.cmp(&right.job_id)) }); for row in unique_rows { self.sync_unique_claim( tx, row.job_id, &row.unique_key, row.unique_states.as_deref(), old_state, Some(row.state), ) .await?; } let schema = self.schema(); let keep_ready_backing = matches!( old_state, Some(JobState::Running | JobState::WaitingExternal) ); let ready_payloads = if rows .iter() .any(|row| !is_storage_payload_empty(&row.payload)) { self.ready_payloads_for_done_rows_tx(tx, rows).await? } else { HashMap::new() }; let mut ordered_rows: Vec<&DoneJobRow> = rows.iter().collect(); ordered_rows.sort_unstable_by_key(|row| { ( row.ready_slot, row.ready_generation, row.queue.as_str(), row.priority, row.enqueue_shard, row.lane_seq, row.job_id, ) }); let (ready_backed, synthetic): (Vec<_>, Vec<_>) = ordered_rows .into_iter() .partition(|row| keep_ready_backing && row.lane_seq >= 0); // Name the output column only when a row has an output, so a schema // without migration v052 never sees it; the executor drops outputs it // cannot store before they reach here (ADR-048). let has_output = |rows: &[&DoneJobRow]| rows.iter().any(|row| row.output.is_some()); if !ready_backed.is_empty() { let (output_column, output_select, output_array, output_alias) = if has_output(&ready_backed) { (", output", ", output::json", ", $15::text[]", ", output") } else { ("", "", "", "") }; let mut query = sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.done_entries (ready_slot, ready_generation, job_id, kind, queue, state, priority, attempt, run_lease, lane_seq, enqueue_shard, attempted_at, finalized_at, payload{output_column}) SELECT ready_slot, ready_generation, job_id, kind, queue, state::awa.job_state, priority, attempt, run_lease, lane_seq, enqueue_shard, attempted_at, finalized_at, payload{output_select} FROM unnest( $1::int[], $2::bigint[], $3::bigint[], $4::text[], $5::text[], $6::text[], $7::smallint[], $8::smallint[], $9::bigint[], $10::bigint[], $11::smallint[], $12::timestamptz[], $13::timestamptz[], $14::jsonb[]{output_array} ) AS rows (ready_slot, ready_generation, job_id, kind, queue, state, priority, attempt, run_lease, lane_seq, enqueue_shard, attempted_at, finalized_at, payload{output_alias}) "# ))) .bind(ready_backed.iter().map(|row| row.ready_slot).collect::>()) .bind(ready_backed.iter().map(|row| row.ready_generation).collect::>()) .bind(ready_backed.iter().map(|row| row.job_id).collect::>()) .bind(ready_backed.iter().map(|row| row.kind.as_str()).collect::>()) .bind(ready_backed.iter().map(|row| row.queue.as_str()).collect::>()) .bind(ready_backed.iter().map(|row| row.state.as_str()).collect::>()) .bind(ready_backed.iter().map(|row| row.priority).collect::>()) .bind(ready_backed.iter().map(|row| row.attempt).collect::>()) .bind(ready_backed.iter().map(|row| row.run_lease).collect::>()) .bind(ready_backed.iter().map(|row| row.lane_seq).collect::>()) .bind(ready_backed.iter().map(|row| row.enqueue_shard).collect::>()) .bind(ready_backed.iter().map(|row| row.attempted_at).collect::>()) .bind(ready_backed.iter().map(|row| row.finalized_at).collect::>()) .bind(terminal_payloads(&ready_backed, &ready_payloads)); if !output_column.is_empty() { // Sent as text and stored as `json`, so the value keeps its // exact text and skips jsonb's parse (ADR-048). query = query.bind( ready_backed .iter() .map(|row| row.output.as_ref().map(|output| output.to_string())) .collect::>(), ); } query.execute(tx.as_mut()).await.map_err(map_sqlx_error)?; } if !synthetic.is_empty() { let (output_column, output_select, output_array, output_alias) = if has_output(&synthetic) { (", output", ", output::json", ", $21::text[]", ", output") } else { ("", "", "", "") }; let mut query = sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.done_entries (ready_slot, ready_generation, job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload{output_column}) SELECT ready_slot, ready_generation, job_id, kind, queue, args, state::awa.job_state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload{output_select} FROM unnest( $1::int[], $2::bigint[], $3::bigint[], $4::text[], $5::text[], $6::jsonb[], $7::text[], $8::smallint[], $9::smallint[], $10::bigint[], $11::smallint[], $12::bigint[], $13::smallint[], $14::timestamptz[], $15::timestamptz[], $16::timestamptz[], $17::timestamptz[], $18::bytea[], $19::text[], $20::jsonb[]{output_array} ) AS rows (ready_slot, ready_generation, job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload{output_alias}) "# ))) .bind(synthetic.iter().map(|row| row.ready_slot).collect::>()) .bind(synthetic.iter().map(|row| row.ready_generation).collect::>()) .bind(synthetic.iter().map(|row| row.job_id).collect::>()) .bind(synthetic.iter().map(|row| row.kind.as_str()).collect::>()) .bind(synthetic.iter().map(|row| row.queue.as_str()).collect::>()) .bind(synthetic.iter().map(|row| &row.args).collect::>()) .bind(synthetic.iter().map(|row| row.state.as_str()).collect::>()) .bind(synthetic.iter().map(|row| row.priority).collect::>()) .bind(synthetic.iter().map(|row| row.attempt).collect::>()) .bind(synthetic.iter().map(|row| row.run_lease).collect::>()) .bind(synthetic.iter().map(|row| row.max_attempts).collect::>()) .bind(synthetic.iter().map(|row| row.lane_seq).collect::>()) .bind(synthetic.iter().map(|row| row.enqueue_shard).collect::>()) .bind(synthetic.iter().map(|row| row.run_at).collect::>()) .bind(synthetic.iter().map(|row| row.attempted_at).collect::>()) .bind(synthetic.iter().map(|row| row.finalized_at).collect::>()) .bind(synthetic.iter().map(|row| row.created_at).collect::>()) .bind(synthetic.iter().map(|row| row.unique_key.as_deref()).collect::>()) .bind(synthetic.iter().map(|row| row.unique_states.as_deref()).collect::>()) .bind(terminal_payloads(&synthetic, &ready_payloads)); if !output_column.is_empty() { query = query.bind( synthetic .iter() .map(|row| row.output.as_ref().map(|output| output.to_string())) .collect::>(), ); } query.execute(tx.as_mut()).await.map_err(map_sqlx_error)?; } // Terminal counts stay exact without mutating the hot live-counter // rows on every completion. The hot path appends signed deltas; the // maintenance rollup folds them into queue_terminal_live_counts later. self.increment_live_terminal_counters_tx(tx, rows).await?; Ok(rows.len()) } /// Append positive terminal-count deltas for newly inserted `done_entries` /// terminal rows. /// /// For done-entry terminal rows, the exact invariant is: /// /// `SUM(queue_terminal_live_counts) + SUM(queue_terminal_count_deltas) /// == count(*) FROM done_entries`. /// /// Compact receipt completions are already append-only batches, so /// `queue_counts_exact` counts retained `receipt_completion_batches` /// directly instead of adding another hot-path count-delta write. /// /// Maintenance asynchronously folds delta rows into the live-counter table /// and truncates the delta segment. The hot path therefore performs only /// append-only writes. async fn increment_live_terminal_counters_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[DoneJobRow], ) -> Result<(), AwaError> { if rows.is_empty() { return Ok(()); } let mut by_group: BTreeMap = BTreeMap::new(); for row in rows { let key = ( row.ready_slot, row.ready_generation, row.queue.clone(), row.priority, row.enqueue_shard, terminal_counter_bucket(row.job_id), ); *by_group.entry(key).or_insert(0) += 1; } self.append_terminal_count_deltas_tx(tx, by_group).await } /// Append negative terminal-count deltas for deleted `done_entries` /// terminal rows. /// /// This is used by retry-from-terminal, DLQ moves, discard paths, and the /// SQL compatibility delete function. The negative row may cancel a /// not-yet-rolled positive delta or reduce an already-folded live counter; /// exact reads sum both sources, so either order is correct. async fn decrement_live_terminal_counters_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[TerminalCounterKey], ) -> Result<(), AwaError> { if rows.is_empty() { return Ok(()); } let mut by_group: BTreeMap = BTreeMap::new(); for key in rows { *by_group.entry(key.clone()).or_insert(0) -= 1; } self.append_terminal_count_deltas_tx(tx, by_group).await } /// Append terminal-count deltas. Before the ring authority flip they keep /// their job-id bucket, because N-1 maintenance folds each delta into the /// bucket of the original row; after it, folds collapse each lane, so the /// statement writes one row per lane. The authority is read inside the /// statement, so the choice follows the caller's transaction, including a /// flip made and rolled back in it. async fn append_terminal_count_deltas_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, by_group: BTreeMap, ) -> Result<(), AwaError> { if by_group.is_empty() { return Ok(()); } let mut ready_slots: Vec = Vec::with_capacity(by_group.len()); let mut ready_generations: Vec = Vec::with_capacity(by_group.len()); let mut queues: Vec = Vec::with_capacity(by_group.len()); let mut priorities: Vec = Vec::with_capacity(by_group.len()); let mut enqueue_shards: Vec = Vec::with_capacity(by_group.len()); let mut counter_buckets: Vec = Vec::with_capacity(by_group.len()); let mut deltas: Vec = Vec::with_capacity(by_group.len()); for ((slot, generation, queue, prio, shard, bucket), delta) in by_group { if delta == 0 { continue; } ready_slots.push(slot); ready_generations.push(generation); queues.push(queue); priorities.push(prio); enqueue_shards.push(shard); counter_buckets.push(bucket); deltas.push(delta); } if deltas.is_empty() { return Ok(()); } let schema = self.schema(); sqlx::query(audited_sql(format!( r#" WITH authority AS ( SELECT COALESCE( ( SELECT authority = 'ledger' FROM {schema}.ring_cursor_authority WHERE singleton = TRUE ), FALSE ) AS per_lane ) INSERT INTO {schema}.queue_terminal_count_deltas ( ready_slot, ready_generation, queue, priority, enqueue_shard, counter_bucket, terminal_delta ) SELECT d.ready_slot, d.ready_generation, d.queue, d.priority, d.enqueue_shard, CASE WHEN authority.per_lane THEN 0::smallint ELSE d.counter_bucket END, sum(d.terminal_delta)::bigint FROM unnest( $1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::smallint[], $6::smallint[], $7::bigint[] ) AS d( ready_slot, ready_generation, queue, priority, enqueue_shard, counter_bucket, terminal_delta ) CROSS JOIN authority GROUP BY 1, 2, 3, 4, 5, 6 HAVING sum(d.terminal_delta) <> 0 ORDER BY 1, 2, 3, 4, 5, 6 "# ))) .bind(&ready_slots) .bind(&ready_generations) .bind(&queues) .bind(&priorities) .bind(&enqueue_shards) .bind(&counter_buckets) .bind(&deltas) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Build the row-key vector for `decrement_live_terminal_counters_tx` /// from a slice of `DoneJobRow` (used by retry-from-terminal, /// DLQ move, and discard, all of which DELETE FROM done_entries /// with `RETURNING *` materialised into `DoneJobRow`). fn done_rows_to_counter_keys(rows: &[DoneJobRow]) -> Vec { rows.iter() .map(|row| { ( row.ready_slot, row.ready_generation, row.queue.clone(), row.priority, row.enqueue_shard, terminal_counter_bucket(row.job_id), ) }) .collect() } async fn ensure_terminal_removed_receipt_closures_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[DoneJobRow], ) -> Result<(), AwaError> { if rows.is_empty() || !self.lease_claim_receipts() { return Ok(()); } let receipt_pairs: Vec<(i64, i64)> = rows.iter().map(|row| (row.job_id, row.run_lease)).collect(); self.lock_receipt_attempts_tx(tx, &receipt_pairs).await?; let schema = self.schema(); let job_ids: Vec = rows.iter().map(|row| row.job_id).collect(); let run_leases: Vec = rows.iter().map(|row| row.run_lease).collect(); sqlx::query(audited_sql(format!( r#" WITH refs(job_id, run_lease) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) ), inserted AS ( INSERT INTO {schema}.lease_claim_closures ( claim_slot, job_id, run_lease, outcome, closed_at ) SELECT claims.claim_slot, claims.job_id, claims.run_lease, 'terminal_removed', clock_timestamp() FROM {schema}.lease_claims AS claims JOIN refs ON refs.job_id = claims.job_id AND refs.run_lease = claims.run_lease ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), marked AS ( UPDATE {schema}.lease_claims AS claims SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) FROM inserted WHERE claims.claim_slot = inserted.claim_slot AND claims.job_id = inserted.job_id AND claims.run_lease = inserted.run_lease RETURNING claims.job_id ) SELECT count(*) FROM marked "# ))) .bind(&job_ids) .bind(&run_leases) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } async fn ready_payloads_for_done_rows_tx<'a, 'r>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &'r [DoneJobRow], ) -> Result, AwaError> { if rows.is_empty() { return Ok(HashMap::new()); } let schema = self.schema(); let ready_slots: Vec = rows.iter().map(|row| row.ready_slot).collect(); let ready_generations: Vec = rows.iter().map(|row| row.ready_generation).collect(); let queues: Vec<&str> = rows.iter().map(|row| row.queue.as_str()).collect(); let priorities: Vec = rows.iter().map(|row| row.priority).collect(); let enqueue_shards: Vec = rows.iter().map(|row| row.enqueue_shard).collect(); let lane_seqs: Vec = rows.iter().map(|row| row.lane_seq).collect(); let payload_rows: Vec<(i32, i64, String, i16, i16, i64, serde_json::Value)> = sqlx::query_as(audited_sql(format!( r#" WITH refs(ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq) AS ( SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::smallint[], $6::bigint[]) ) SELECT ready.ready_slot, ready.ready_generation, ready.queue, ready.priority, ready.enqueue_shard, ready.lane_seq, COALESCE(ready.payload, '{{}}'::jsonb) AS payload FROM refs JOIN {schema}.ready_entries AS ready ON ready.ready_slot = refs.ready_slot AND ready.ready_generation = refs.ready_generation AND ready.queue = refs.queue AND ready.priority = refs.priority AND ready.enqueue_shard = refs.enqueue_shard AND ready.lane_seq = refs.lane_seq "# ))) .bind(&ready_slots) .bind(&ready_generations) .bind(&queues) .bind(&priorities) .bind(&enqueue_shards) .bind(&lane_seqs) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; let mut payload_by_owned_key = HashMap::with_capacity(payload_rows.len()); for (ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, payload) in payload_rows { payload_by_owned_key.insert( ( ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, ), payload, ); } let mut payloads = HashMap::with_capacity(payload_by_owned_key.len()); for row in rows { if let Some(payload) = payload_by_owned_key.remove(&( row.ready_slot, row.ready_generation, row.queue.clone(), row.priority, row.enqueue_shard, row.lane_seq, )) { payloads.insert( ( row.ready_slot, row.ready_generation, row.queue.as_str(), row.priority, row.enqueue_shard, row.lane_seq, ), payload, ); } } Ok(payloads) } async fn insert_dlq_rows_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: &[DlqJobRow], old_state: Option, ) -> Result { if rows.is_empty() { return Ok(0); } let mut unique_rows: Vec<&DlqJobRow> = rows.iter().collect(); unique_rows.sort_unstable_by(|left, right| { left.unique_key .cmp(&right.unique_key) .then(left.job_id.cmp(&right.job_id)) }); for row in unique_rows { self.sync_unique_claim( tx, row.job_id, &row.unique_key, row.unique_states.as_deref(), old_state, Some(JobState::Failed), ) .await?; } let schema = self.schema(); sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.dlq_entries (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload, dlq_reason, dlq_at, original_run_lease) SELECT job_id, kind, queue, args, state::awa.job_state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload, dlq_reason, dlq_at, original_run_lease FROM unnest( $1::bigint[], $2::text[], $3::text[], $4::jsonb[], $5::text[], $6::smallint[], $7::smallint[], $8::bigint[], $9::smallint[], $10::timestamptz[], $11::timestamptz[], $12::timestamptz[], $13::timestamptz[], $14::bytea[], $15::text[], $16::jsonb[], $17::text[], $18::timestamptz[], $19::bigint[] ) AS rows (job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload, dlq_reason, dlq_at, original_run_lease) "# ))) .bind(rows.iter().map(|row| row.job_id).collect::>()) .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind(rows.iter().map(|row| row.queue.as_str()).collect::>()) .bind(rows.iter().map(|row| &row.args).collect::>()) .bind(rows.iter().map(|row| row.state.as_str()).collect::>()) .bind(rows.iter().map(|row| row.priority).collect::>()) .bind(rows.iter().map(|row| row.attempt).collect::>()) .bind(rows.iter().map(|row| row.run_lease).collect::>()) .bind(rows.iter().map(|row| row.max_attempts).collect::>()) .bind(rows.iter().map(|row| row.run_at).collect::>()) .bind(rows.iter().map(|row| row.attempted_at).collect::>()) .bind(rows.iter().map(|row| row.finalized_at).collect::>()) .bind(rows.iter().map(|row| row.created_at).collect::>()) .bind(rows.iter().map(|row| row.unique_key.as_deref()).collect::>()) .bind(rows.iter().map(|row| row.unique_states.as_deref()).collect::>()) .bind(rows.iter().map(|row| storage_payload(&row.payload)).collect::>()) .bind(rows.iter().map(|row| row.dlq_reason.as_str()).collect::>()) .bind(rows.iter().map(|row| row.dlq_at).collect::>()) .bind(rows.iter().map(|row| row.original_run_lease).collect::>()) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(rows.len()) } /// Upsert pruned terminal counts into the permanent /// `queue_terminal_rollups` denormaliser, clamping both columns at /// zero. Since #371 this runs only from the horizon-gated fold /// ([`Self::fold_terminal_rollup_deltas`]) — the prune transaction /// appends delta rows instead, so a pinned MVCC horizon cannot pile /// up dead rollup versions. Keys are folded through a `BTreeMap`, so /// the upsert always touches (queue, priority) keys in deterministic /// order. async fn adjust_terminal_rollups_batch<'a, I>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, deltas: I, ) -> Result<(), AwaError> where I: IntoIterator, { let mut grouped: BTreeMap<(String, i16), (i64, i64)> = BTreeMap::new(); for (queue, priority, pruned_completed_delta, pruned_failed_delta) in deltas { if pruned_completed_delta == 0 && pruned_failed_delta == 0 { continue; } let entry = grouped.entry((queue, priority)).or_insert((0_i64, 0_i64)); entry.0 += pruned_completed_delta; entry.1 += pruned_failed_delta; } if grouped.is_empty() { return Ok(()); } let schema = self.schema(); let mut queues = Vec::with_capacity(grouped.len()); let mut priorities = Vec::with_capacity(grouped.len()); let mut pruned_completed_deltas = Vec::with_capacity(grouped.len()); let mut pruned_failed_deltas = Vec::with_capacity(grouped.len()); for ((queue, priority), (pruned_completed_delta, pruned_failed_delta)) in grouped { queues.push(queue); priorities.push(priority); pruned_completed_deltas.push(pruned_completed_delta); pruned_failed_deltas.push(pruned_failed_delta); } sqlx::query(audited_sql(format!( r#" WITH deltas(queue, priority, pruned_completed_delta, pruned_failed_delta) AS ( SELECT * FROM unnest( $1::text[], $2::smallint[], $3::bigint[], $4::bigint[] ) ) INSERT INTO {schema}.queue_terminal_rollups AS rollups ( queue, priority, pruned_completed_count, pruned_failed_count ) SELECT deltas.queue, deltas.priority, deltas.pruned_completed_delta, deltas.pruned_failed_delta FROM deltas ON CONFLICT (queue, priority) DO UPDATE SET pruned_completed_count = GREATEST( 0, rollups.pruned_completed_count + EXCLUDED.pruned_completed_count ), pruned_failed_count = GREATEST( 0, rollups.pruned_failed_count + EXCLUDED.pruned_failed_count ) "# ))) .bind(&queues) .bind(&priorities) .bind(&pruned_completed_deltas) .bind(&pruned_failed_deltas) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Append pruned terminal counts as delta rows (#371). Runs inside /// the prune transaction in place of the pre-#371 rollup upsert; /// deltas are folded into `queue_terminal_rollups` by /// [`Self::fold_terminal_rollup_deltas`] when the MVCC horizon is /// clear. async fn append_terminal_rollup_deltas_tx<'a, I>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, deltas: I, ) -> Result<(), AwaError> where I: IntoIterator, { let mut grouped: BTreeMap<(String, i16), (i64, i64)> = BTreeMap::new(); for (queue, priority, pruned_completed_delta, pruned_failed_delta) in deltas { if pruned_completed_delta == 0 && pruned_failed_delta == 0 { continue; } let entry = grouped.entry((queue, priority)).or_insert((0_i64, 0_i64)); entry.0 += pruned_completed_delta; entry.1 += pruned_failed_delta; } if grouped.is_empty() { return Ok(()); } let schema = self.schema(); let mut queues = Vec::with_capacity(grouped.len()); let mut priorities = Vec::with_capacity(grouped.len()); let mut pruned_completed_deltas = Vec::with_capacity(grouped.len()); let mut pruned_failed_deltas = Vec::with_capacity(grouped.len()); for ((queue, priority), (pruned_completed_delta, pruned_failed_delta)) in grouped { queues.push(queue); priorities.push(priority); pruned_completed_deltas.push(pruned_completed_delta); pruned_failed_deltas.push(pruned_failed_delta); } sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.queue_terminal_rollup_deltas ( queue, priority, pruned_completed_delta, pruned_failed_delta ) SELECT * FROM unnest( $1::text[], $2::smallint[], $3::bigint[], $4::bigint[] ) "# ))) .bind(&queues) .bind(&priorities) .bind(&pruned_completed_deltas) .bind(&pruned_failed_deltas) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } async fn enqueue_runtime_rows( &self, pool: &PgPool, rows: Vec, ) -> Result { if rows.is_empty() { return Ok(0); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let total_rows = self .insert_ready_rows_tx(&mut tx, rows.clone(), false) .await?; let queues_to_notify: Vec = rows.iter().map(|row| row.queue.clone()).collect(); self.notify_queues_tx(&mut tx, queues_to_notify).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(total_rows) } pub async fn enqueue_batch( &self, pool: &PgPool, queue: &str, priority: i16, count: i64, ) -> Result { if count <= 0 { return Ok(0); } let rows: Vec<_> = (0..count) .map(|seq| RuntimeReadyRow { job_id: None, kind: "bench_job".to_string(), queue: if self.uses_queue_striping() && !self.is_physical_stripe_queue(queue) { self.physical_queue_for_stripe( queue, seq.rem_euclid(self.queue_stripe_count() as i64) as usize, ) } else { queue.to_string() }, args: serde_json::json!({ "seq": seq }), priority, attempt: 0, run_lease: 0, max_attempts: 25, run_at: Utc::now(), attempted_at: None, created_at: Utc::now(), unique_key: None, unique_states: None, payload: RuntimePayload::default().into_json(), ordering_key: None, }) .collect(); self.enqueue_runtime_rows(pool, rows) .await .map(|count| count as i64) } pub async fn enqueue_params_batch( &self, pool: &PgPool, jobs: &[InsertParams], ) -> Result { if jobs.is_empty() { return Ok(0); } let now = Utc::now(); let mut ready_rows = Vec::new(); let mut deferred_rows = Vec::new(); for (idx, job) in jobs.iter().enumerate() { let prepared = prepare_row_raw(job.kind.clone(), job.args.clone(), job.opts.clone())?; let payload = Self::payload_from_parts(prepared.metadata, prepared.tags, None, None)?; let queue = self.queue_stripe_for_enqueue(&prepared.queue, &prepared.unique_key, idx as i64); let ready_row = RuntimeReadyRow { job_id: None, kind: prepared.kind, queue: queue.clone(), args: prepared.args, priority: prepared.priority, attempt: 0, run_lease: 0, max_attempts: prepared.max_attempts, run_at: prepared.run_at.unwrap_or(now), attempted_at: None, created_at: now, unique_key: prepared.unique_key, unique_states: prepared.unique_states, payload: payload.clone(), ordering_key: prepared.ordering_key, }; match prepared.state { JobState::Available => ready_rows.push(ready_row), JobState::Scheduled => deferred_rows.push(DeferredJobRow { job_id: 0, kind: ready_row.kind, queue, args: ready_row.args, state: JobState::Scheduled, priority: ready_row.priority, attempt: ready_row.attempt, run_lease: ready_row.run_lease, max_attempts: ready_row.max_attempts, run_at: ready_row.run_at, attempted_at: ready_row.attempted_at, finalized_at: None, created_at: ready_row.created_at, unique_key: ready_row.unique_key, unique_states: ready_row.unique_states, payload: payload.clone(), }), other => { return Err(AwaError::Validation(format!( "queue storage does not support initial state {other}" ))); } } } let queues_to_notify: Vec = ready_rows.iter().map(|row| row.queue.clone()).collect(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let job_ids = self .next_job_ids(&mut tx, ready_rows.len() + deferred_rows.len()) .await?; let (ready_job_ids, deferred_job_ids) = job_ids.split_at(ready_rows.len()); for (row, job_id) in ready_rows.iter_mut().zip(ready_job_ids) { row.job_id = Some(*job_id); } let deferred_rows: Vec<_> = deferred_rows .into_iter() .zip(deferred_job_ids.iter().copied()) .map(|(row, job_id)| DeferredJobRow { job_id, ..row }) .collect(); let unique_claims = Self::batch_enqueue_unique_claims(&ready_rows, &deferred_rows); self.sync_enqueue_unique_claims(&mut tx, unique_claims) .await?; let mut total = 0usize; if !ready_rows.is_empty() { total += self .insert_ready_rows_tx(&mut tx, ready_rows.clone(), true) .await?; } if !deferred_rows.is_empty() { total += self .execute_deferred_inserts_tx(&mut tx, &deferred_rows) .await?; } self.notify_queues_tx(&mut tx, queues_to_notify).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(total) } /// Enqueue prepared jobs into queue storage using PostgreSQL COPY. /// /// This follows the same preparation, queue striping, lane sequencing, /// uniqueness, and notification semantics as [`Self::enqueue_params_batch`], /// but streams materialized rows directly into `ready_entries` and /// `deferred_jobs` instead of building multi-row `INSERT` statements. #[tracing::instrument(skip(self, pool, jobs), fields(job.count = jobs.len()), name = "queue_storage.enqueue_params_copy")] pub async fn enqueue_params_copy( &self, pool: &PgPool, jobs: &[InsertParams], ) -> Result { if jobs.is_empty() { return Ok(0); } let now = Utc::now(); let mut ready_rows = Vec::new(); let mut deferred_rows = Vec::new(); for (idx, job) in jobs.iter().enumerate() { let prepared = prepare_row_raw(job.kind.clone(), job.args.clone(), job.opts.clone())?; let payload = Self::payload_from_parts(prepared.metadata, prepared.tags, None, None)?; let queue = self.queue_stripe_for_enqueue(&prepared.queue, &prepared.unique_key, idx as i64); let ready_row = RuntimeReadyRow { job_id: None, kind: prepared.kind, queue: queue.clone(), args: prepared.args, priority: prepared.priority, attempt: 0, run_lease: 0, max_attempts: prepared.max_attempts, run_at: prepared.run_at.unwrap_or(now), attempted_at: None, created_at: now, unique_key: prepared.unique_key, unique_states: prepared.unique_states, payload: payload.clone(), ordering_key: prepared.ordering_key, }; match prepared.state { JobState::Available => ready_rows.push(ready_row), JobState::Scheduled => deferred_rows.push(DeferredJobRow { job_id: 0, kind: ready_row.kind, queue, args: ready_row.args, state: JobState::Scheduled, priority: ready_row.priority, attempt: ready_row.attempt, run_lease: ready_row.run_lease, max_attempts: ready_row.max_attempts, run_at: ready_row.run_at, attempted_at: ready_row.attempted_at, finalized_at: None, created_at: ready_row.created_at, unique_key: ready_row.unique_key, unique_states: ready_row.unique_states, payload: payload.clone(), }), other => { return Err(AwaError::Validation(format!( "queue storage does not support initial state {other}" ))); } } } let queues_to_notify: Vec = ready_rows.iter().map(|row| row.queue.clone()).collect(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let mut total = 0usize; let job_ids = self .next_job_ids(&mut tx, ready_rows.len() + deferred_rows.len()) .await?; let (ready_job_ids, deferred_job_ids) = job_ids.split_at(ready_rows.len()); for (row, job_id) in ready_rows.iter_mut().zip(ready_job_ids) { row.job_id = Some(*job_id); } let deferred_rows: Vec<_> = deferred_rows .into_iter() .zip(deferred_job_ids.iter().copied()) .map(|(row, job_id)| DeferredJobRow { job_id, ..row }) .collect(); let unique_claims = Self::batch_enqueue_unique_claims(&ready_rows, &deferred_rows); self.sync_enqueue_unique_claims(&mut tx, unique_claims) .await?; if !ready_rows.is_empty() { total += self .insert_ready_rows_copy_tx(&mut tx, ready_rows, true) .await?; } if !deferred_rows.is_empty() { total += self .insert_deferred_rows_copy_tx(&mut tx, deferred_rows, true) .await?; } self.notify_queues_tx(&mut tx, queues_to_notify).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(total) } #[tracing::instrument(skip(self, pool), name = "queue_storage.claim_batch")] pub async fn claim_batch( &self, pool: &PgPool, queue: &str, max_batch: i64, ) -> Result, AwaError> { if max_batch <= 0 { return Ok(Vec::new()); } let mut tx = pool .begin_with(CLAIM_TRANSACTION_BEGIN) .await .map_err(map_sqlx_error)?; let mut claimed_rows = Vec::new(); let stripe_queues = self.physical_queues_for_logical(queue); let start = self.stripe_probe_start(stripe_queues.len()); for offset in 0..stripe_queues.len() { if claimed_rows.len() >= max_batch as usize { break; } let stripe_queue = &stripe_queues[(start + offset) % stripe_queues.len()]; let remaining = max_batch - claimed_rows.len() as i64; claimed_rows.extend( self.claim_ready_rows_tx( &mut tx, stripe_queue, remaining, Duration::ZERO, Duration::ZERO, ) .await?, ); } let claim_cursor_advances = Self::claim_cursor_advances(&claimed_rows); let claimed = claimed_rows .into_iter() .map(|row| row.claim_ref(self.lease_claim_receipts())) .collect(); tx.commit().await.map_err(map_sqlx_error)?; self.advance_claim_cursors(pool, &claim_cursor_advances) .await; Ok(claimed) } #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue), name = "queue_storage.claim_runtime_batch")] pub async fn claim_runtime_batch( &self, pool: &PgPool, queue: &str, max_batch: i64, deadline_duration: Duration, ) -> Result, AwaError> { self.claim_runtime_batch_with_aging( pool, queue, max_batch, deadline_duration, Duration::ZERO, ) .await } #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue), name = "queue_storage.claim_runtime_batch_with_aging")] pub async fn claim_runtime_batch_with_aging( &self, pool: &PgPool, queue: &str, max_batch: i64, deadline_duration: Duration, aging_interval: Duration, ) -> Result, AwaError> { if max_batch <= 0 { return Ok(Vec::new()); } let stripe_queues = self.physical_queues_for_logical(queue); if stripe_queues.len() > 1 { let mut claimed = Vec::new(); let start = self.stripe_probe_start(stripe_queues.len()); for offset in 0..stripe_queues.len() { if claimed.len() >= max_batch as usize { break; } let stripe_queue = &stripe_queues[(start + offset) % stripe_queues.len()]; let remaining = max_batch - claimed.len() as i64; match self .claim_runtime_batch_with_aging_physical( pool, stripe_queue, remaining, deadline_duration, aging_interval, ) .await { Ok(stripe_claims) => claimed.extend(stripe_claims), Err(err) if claimed.is_empty() => return Err(err), Err(err) => { tracing::warn!( queue = %queue, messaging.destination.name = %queue, stripe_queue = %stripe_queue, claimed = claimed.len(), error = ?err, "returning already-claimed runtime jobs after striped claim error" ); break; } } } return Ok(claimed); } self.claim_runtime_batch_with_aging_physical( pool, &stripe_queues[0], max_batch, deadline_duration, aging_interval, ) .await } async fn claim_runtime_batch_with_aging_physical( &self, pool: &PgPool, queue: &str, max_batch: i64, deadline_duration: Duration, aging_interval: Duration, ) -> Result, AwaError> { if max_batch <= 0 { return Ok(Vec::new()); } let mut tx = pool .begin_with(CLAIM_TRANSACTION_BEGIN) .await .map_err(map_sqlx_error)?; let mut claimed = Vec::new(); claimed.extend( self.claim_ready_rows_tx(&mut tx, queue, max_batch, deadline_duration, aging_interval) .await?, ); let mut unique_claimed: Vec<&ReadyJobLeaseRow> = claimed.iter().collect(); unique_claimed.sort_unstable_by(|left, right| { left.unique_key .cmp(&right.unique_key) .then(left.job_id.cmp(&right.job_id)) }); for row in unique_claimed { self.sync_unique_claim( &mut tx, row.job_id, &row.unique_key, row.unique_states.as_deref(), Some(JobState::Available), Some(JobState::Running), ) .await?; } let use_lease_claim_receipts = self.use_lease_claim_receipts_for_runtime(deadline_duration); if !use_lease_claim_receipts && deadline_duration.is_zero() { // Legacy zero-deadline claims have no heartbeat/deadline rescue // timestamp, so a post-commit conversion error would strand the // materialized lease indefinitely. Keep the old rollback semantics // for that path. let converted = claimed .iter() .cloned() .map(|row| row.into_claimed_runtime_job(use_lease_claim_receipts)) .collect::, _>>()?; let claim_cursor_advances = Self::claim_cursor_advances(&claimed); tx.commit().await.map_err(map_sqlx_error)?; self.advance_claim_cursors(pool, &claim_cursor_advances) .await; return Ok(converted); } let claim_cursor_advances = Self::claim_cursor_advances(&claimed); // Release claim locks before doing Rust-side payload conversion; this // keeps the hot claim transaction focused on database state changes. tx.commit().await.map_err(map_sqlx_error)?; self.advance_claim_cursors(pool, &claim_cursor_advances) .await; claimed .into_iter() .map(|row| row.into_claimed_runtime_job(use_lease_claim_receipts)) .collect() } #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue, instance_id = %instance_id), name = "queue_storage.acquire_queue_claimer")] pub async fn acquire_queue_claimer( &self, pool: &PgPool, queue: &str, instance_id: Uuid, max_claimers: i16, lease_ttl: Duration, idle_threshold: Duration, ) -> Result, AwaError> { Ok(self .acquire_queue_claimer_row( pool, queue, instance_id, max_claimers, lease_ttl, idle_threshold, ) .await? .map(QueueClaimerLeaseRow::lease)) } async fn acquire_queue_claimer_row( &self, pool: &PgPool, queue: &str, instance_id: Uuid, max_claimers: i16, lease_ttl: Duration, idle_threshold: Duration, ) -> Result, AwaError> { if max_claimers <= 0 { return Ok(None); } let schema = self.schema(); let now = Utc::now(); let expires_at = now + TimeDelta::from_std(lease_ttl) .map_err(|err| AwaError::Validation(format!("invalid claimer lease ttl: {err}")))?; let idle_cutoff = now - TimeDelta::from_std(idle_threshold).map_err(|err| { AwaError::Validation(format!("invalid claimer idle threshold: {err}")) })?; let probe_start = if max_claimers > 1 { ((instance_id.as_u128() ^ (now.timestamp_millis() as u128)) % (max_claimers as u128)) as i16 } else { 0 }; if let Some(owned) = sqlx::query_as::<_, QueueClaimerLeaseRow>(audited_sql(format!( r#" SELECT claimer_slot, lease_epoch, last_claimed_at, expires_at FROM {schema}.queue_claimer_leases WHERE queue = $1 AND owner_instance_id = $2 AND expires_at > $3 ORDER BY claimer_slot LIMIT 1 "# ))) .bind(queue) .bind(instance_id) .bind(now) .fetch_optional(pool) .await .map_err(map_sqlx_error)? { return Ok(Some(owned)); } for offset in 0..max_claimers { let slot = (probe_start + offset) % max_claimers; if let Some(updated) = sqlx::query_as::<_, QueueClaimerLeaseRow>(audited_sql(format!( r#" UPDATE {schema}.queue_claimer_leases SET owner_instance_id = $3, lease_epoch = CASE WHEN owner_instance_id = $3 THEN lease_epoch ELSE lease_epoch + 1 END, leased_at = $4, last_claimed_at = $4, expires_at = $5 WHERE queue = $1 AND claimer_slot = $2 AND ( owner_instance_id = $3 OR expires_at <= $4 OR last_claimed_at <= $6 ) RETURNING claimer_slot, lease_epoch, last_claimed_at, expires_at "# ))) .bind(queue) .bind(slot) .bind(instance_id) .bind(now) .bind(expires_at) .bind(idle_cutoff) .fetch_optional(pool) .await .map_err(map_sqlx_error)? { return Ok(Some(updated)); } if let Some(inserted) = sqlx::query_as::<_, QueueClaimerLeaseRow>(audited_sql(format!( r#" INSERT INTO {schema}.queue_claimer_leases ( queue, claimer_slot, owner_instance_id, lease_epoch, leased_at, last_claimed_at, expires_at ) VALUES ($1, $2, $3, 0, $4, $4, $5) ON CONFLICT (queue, claimer_slot) DO NOTHING RETURNING claimer_slot, lease_epoch, last_claimed_at, expires_at "# ))) .bind(queue) .bind(slot) .bind(instance_id) .bind(now) .bind(expires_at) .fetch_optional(pool) .await .map_err(map_sqlx_error)? { return Ok(Some(inserted)); } } Ok(None) } #[doc(hidden)] pub fn cached_claimer_lease_count(&self) -> usize { self.claimer_gates .lock() .expect("claimer gate cache poisoned") .len() } fn cache_claimer_gate(&self, key: (String, Uuid), lease: QueueClaimerLeaseRow) { self.claimer_gates .lock() .expect("claimer gate cache poisoned") .insert( key, CachedClaimerGate { checked_at: Instant::now(), lease, }, ); } #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue, instance_id = %instance_id, claimer_slot = lease.claimer_slot), name = "queue_storage.mark_queue_claimer_active")] pub async fn mark_queue_claimer_active( &self, pool: &PgPool, queue: &str, instance_id: Uuid, lease: QueueClaimerLease, lease_ttl: Duration, ) -> Result { let schema = self.schema(); let now = Utc::now(); let expires_at = now + TimeDelta::from_std(lease_ttl) .map_err(|err| AwaError::Validation(format!("invalid claimer lease ttl: {err}")))?; let result = sqlx::query(audited_sql(format!( r#" UPDATE {schema}.queue_claimer_leases SET last_claimed_at = $5, expires_at = $6 WHERE queue = $1 AND claimer_slot = $2 AND owner_instance_id = $3 AND lease_epoch = $4 "# ))) .bind(queue) .bind(lease.claimer_slot) .bind(instance_id) .bind(lease.lease_epoch) .bind(now) .bind(expires_at) .execute(pool) .await .map_err(map_sqlx_error)?; Ok(result.rows_affected() == 1) } fn desired_queue_claimer_target( &self, current_target: Option, signal: &AvailableSignal, max_claimers: i16, ) -> i16 { // The signal source counts sequence positions reserved for enqueue // but not yet claimed. It can over-count deleted or uncommitted // positions, but already-claimed rows are excluded by the claim cursor // advance. // `backlog` is retained as a separate name in the threshold // table to keep room for shape tweaks that diverge from // `available` later. let available = signal.available.max(0) as u64; let backlog = available; let current = current_target.unwrap_or(1).clamp(1, max_claimers.max(1)); let max_four = 4.min(max_claimers.max(1)); let max_two = 2.min(max_claimers.max(1)); match current { 4.. => { if available >= 32 || backlog >= 16 { max_four } else if available >= 8 || backlog >= 4 { max_two } else { 1 } } 2..=3 => { if available >= 128 || backlog >= 64 { max_four } else if available >= 4 || backlog >= 2 { max_two } else { 1 } } _ => { if available >= 64 || backlog >= 32 { max_four } else if available >= 8 || backlog >= 4 { max_two } else { 1 } } } } async fn queue_claimer_target( &self, pool: &PgPool, queue: &str, max_claimers: i16, control_interval: Duration, ) -> Result { let schema = self.schema(); let now = Utc::now(); let stale_cutoff = now - TimeDelta::from_std(control_interval).map_err(|err| { AwaError::Validation(format!("invalid claimer control interval: {err}")) })?; if let Some(target) = sqlx::query_scalar::<_, i16>(audited_sql(format!( r#" SELECT target_claimers FROM {schema}.queue_claimer_state WHERE queue = $1 AND updated_at > $2 "# ))) .bind(queue) .bind(stale_cutoff) .fetch_optional(pool) .await .map_err(map_sqlx_error)? { return Ok(target.clamp(1, max_claimers.max(1))); } let current_target = sqlx::query_scalar::<_, i16>(audited_sql(format!( r#" SELECT target_claimers FROM {schema}.queue_claimer_state WHERE queue = $1 "# ))) .bind(queue) .fetch_optional(pool) .await .map_err(map_sqlx_error)?; let signal = self.queue_claimer_signal(pool, queue).await?; let desired = self.desired_queue_claimer_target(current_target, &signal, max_claimers); if let Some(updated) = sqlx::query_scalar::<_, i16>(audited_sql(format!( r#" INSERT INTO {schema}.queue_claimer_state (queue, target_claimers, updated_at) VALUES ($1, $2, $3) ON CONFLICT (queue) DO UPDATE SET target_claimers = EXCLUDED.target_claimers, updated_at = EXCLUDED.updated_at WHERE {schema}.queue_claimer_state.updated_at <= $4 RETURNING target_claimers "# ))) .bind(queue) .bind(desired) .bind(now) .bind(stale_cutoff) .fetch_optional(pool) .await .map_err(map_sqlx_error)? { return Ok(updated.clamp(1, max_claimers.max(1))); } Ok(current_target .unwrap_or(desired) .clamp(1, max_claimers.max(1))) } /// Cheap, dispatcher-grade available-count signal. /// /// Sums `enqueue_cursor - claim_cursor` across the queue's /// (queue, priority) lanes — one PK read into each of the two head tables per lane /// (typically ≤ 4 lanes per logical queue). The difference is an /// upper bound on the count of unclaimed ready rows: admin DELETEs /// of unclaimed jobs and in-flight enqueue reservations can leave a /// gap between `claim_cursor` and `enqueue_cursor`. The dispatcher /// tolerates this drift because the worst case is a wasted claim /// attempt that finds no committed rows. /// /// For an exact count, use [`Self::queue_counts_exact`], which /// scans `ready_entries`. async fn queue_claimer_signal( &self, pool: &PgPool, queue: &str, ) -> Result { let schema = self.schema(); let queues = self.physical_queues_for_logical(queue); let available: i64 = sqlx::query_scalar(audited_sql(format!( r#" SELECT COALESCE( sum(GREATEST( {schema}.sequence_next_value(qe.seq_name) - {schema}.sequence_next_value(qc.seq_name), 0 )), 0 )::bigint FROM {schema}.queue_enqueue_heads AS qe JOIN {schema}.queue_claim_heads AS qc ON qc.queue = qe.queue AND qc.priority = qe.priority AND qc.enqueue_shard = qe.enqueue_shard WHERE qe.queue = ANY($1) "# ))) .bind(&queues) .fetch_one(pool) .await .map_err(map_sqlx_error)?; Ok(AvailableSignal { available }) } #[allow(clippy::too_many_arguments)] #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue, instance_id = %instance_id), name = "queue_storage.claim_runtime_batch_with_aging_for_instance")] pub async fn claim_runtime_batch_with_aging_for_instance( &self, pool: &PgPool, queue: &str, max_batch: i64, deadline_duration: Duration, aging_interval: Duration, instance_id: Uuid, max_claimers: i16, lease_ttl: Duration, idle_threshold: Duration, ) -> Result, AwaError> { let gate_key = (queue.to_string(), instance_id); let cached_lease = self .claimer_gates .lock() .expect("claimer gate cache poisoned") .get(&gate_key) .filter(|gate| { gate.checked_at.elapsed() < CLAIMER_GATE_CACHE_TTL && gate.lease.expires_at > Utc::now() + TimeDelta::seconds(1) }) .map(|gate| gate.lease); let lease = match cached_lease { Some(lease) => lease, None => { let target_claimers = self .queue_claimer_target(pool, queue, max_claimers, Duration::from_millis(500)) .await?; let Some(lease) = self .acquire_queue_claimer_row( pool, queue, instance_id, target_claimers, lease_ttl, idle_threshold, ) .await? else { self.claimer_gates .lock() .expect("claimer gate cache poisoned") .remove(&gate_key); return Ok(Vec::new()); }; self.cache_claimer_gate(gate_key.clone(), lease); lease } }; let claimed = self .claim_runtime_batch_with_aging( pool, queue, max_batch, deadline_duration, aging_interval, ) .await?; let now = Utc::now(); if !claimed.is_empty() && lease.needs_refresh(now, lease_ttl, idle_threshold) { let still_owned = self .mark_queue_claimer_active(pool, queue, instance_id, lease.lease(), lease_ttl) .await?; if still_owned { let refreshed = QueueClaimerLeaseRow { last_claimed_at: now, expires_at: now + TimeDelta::from_std(lease_ttl).map_err(|err| { AwaError::Validation(format!("invalid claimer lease ttl: {err}")) })?, ..lease }; self.cache_claimer_gate(gate_key, refreshed); } else { self.claimer_gates .lock() .expect("claimer gate cache poisoned") .remove(&gate_key); } } Ok(claimed) } #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue), name = "queue_storage.claim_job_batch")] pub async fn claim_job_batch( &self, pool: &PgPool, queue: &str, max_batch: i64, deadline_duration: Duration, ) -> Result, AwaError> { self.claim_runtime_batch(pool, queue, max_batch, deadline_duration) .await .map(|claimed| claimed.into_iter().map(|row| row.job).collect()) } #[tracing::instrument(skip(self, pool, claimed), name = "queue_storage.complete_batch")] pub async fn complete_batch( &self, pool: &PgPool, claimed: &[ClaimedEntry], ) -> Result { self.complete_claimed_batch(pool, claimed) .await .map(|updated| updated.len()) } #[tracing::instrument( skip(self, pool, claimed), name = "queue_storage.complete_claimed_batch" )] pub async fn complete_claimed_batch( &self, pool: &PgPool, claimed: &[ClaimedEntry], ) -> Result, AwaError> { if claimed.is_empty() { return Ok(Vec::new()); } let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let lease_slots: Vec = claimed.iter().map(|entry| entry.lease_slot).collect(); let queues: Vec = claimed.iter().map(|entry| entry.queue.clone()).collect(); let priorities: Vec = claimed.iter().map(|entry| entry.priority).collect(); let enqueue_shards: Vec = claimed.iter().map(|entry| entry.enqueue_shard).collect(); let lane_seqs: Vec = claimed.iter().map(|entry| entry.lane_seq).collect(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" WITH completed(lease_slot, queue, priority, enqueue_shard, lane_seq) AS ( SELECT * FROM unnest($1::int[], $2::text[], $3::smallint[], $4::smallint[], $5::bigint[]) ) DELETE FROM {schema}.leases AS leases USING completed WHERE leases.lease_slot = completed.lease_slot AND leases.queue = completed.queue AND leases.priority = completed.priority AND leases.enqueue_shard = completed.enqueue_shard AND leases.lane_seq = completed.lane_seq RETURNING leases.ready_slot, leases.ready_generation, leases.job_id, leases.queue, leases.state, leases.priority, leases.attempt, leases.run_lease, leases.max_attempts, leases.lane_seq, leases.enqueue_shard, leases.heartbeat_at, leases.deadline_at, leases.attempted_at, leases.callback_id, leases.callback_timeout_at "# ))) .bind(&lease_slots) .bind(&queues) .bind(&priorities) .bind(&enqueue_shards) .bind(&lane_seqs) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(Vec::new()); } let completed_pairs: Vec<(i64, i64)> = deleted .iter() .map(|row| (row.job_id, row.run_lease)) .collect(); self.close_receipt_pairs_tx(&mut tx, &completed_pairs, "completed") .await?; let moved = self .hydrate_deleted_leases_tx(&mut tx, deleted, true) .await?; let finalized_at = Utc::now(); let mut done_rows = Vec::with_capacity(moved.len()); for entry in moved.iter().cloned() { let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state( entry.payload.clone(), entry.progress.clone(), )?)?; payload.set_progress(None); done_rows.push(entry.into_done_row( JobState::Completed, finalized_at, payload.into_json(), )); } self.insert_done_rows_tx(&mut tx, &done_rows, Some(JobState::Running)) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(moved .into_iter() .map(|entry| (entry.job_id, entry.run_lease)) .collect()) } fn receipt_fast_complete_candidate(entry: &ClaimedRuntimeJob) -> bool { entry.claim.lease_claim_receipt && entry.claim.receipt_id.is_some() && entry.job.unique_key.is_none() && is_compact_receipt_completion_metadata(&entry.job.metadata) && entry.job.tags.is_empty() && entry.job.errors.as_ref().is_none_or(Vec::is_empty) } async fn complete_receipt_runtime_batch_fast( &self, pool: &PgPool, claimed: &[ClaimedRuntimeJob], ) -> Result, AwaError> { if claimed.is_empty() { return Ok(Vec::new()); } let schema = self.schema(); let finalized_at = Utc::now(); let job_ids: Vec = claimed.iter().map(|entry| entry.job.id).collect(); let run_leases: Vec = claimed.iter().map(|entry| entry.job.run_lease).collect(); let mut by_partition: BTreeMap> = BTreeMap::new(); for entry in claimed { let claim_slot_index = ring_slot_index(entry.claim.claim_slot, self.claim_slot_count(), "claim")?; by_partition .entry(claim_slot_index) .or_default() .push(entry); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // Without the per-job closure-row insert, // lease_claim_closure_batches become the idempotence evidence. // Serialize duplicate completion attempts before the main statement // so a waiter takes its statement snapshot after the first // transaction commits and sees that evidence. let receipt_pairs: Vec<(i64, i64)> = job_ids .iter() .zip(run_leases.iter()) .map(|(job_id, run_lease)| (*job_id, *run_lease)) .collect(); if let Err(err) = self .lock_receipt_attempt_keys_tx(&mut tx, &receipt_pairs) .await { let _ = tx.rollback().await; return Err(err); } let mut rows = Vec::with_capacity(claimed.len()); // Outputs past the budget start another batch row in the same // transaction, so the flush still commits or fails as a whole. let mut groups = Vec::with_capacity(by_partition.len()); for (index, partition) in &by_partition { for group in split_by_output_budget(partition) { groups.push((*index, group)); } } for (claim_slot_index, group) in groups { let claim_rel = claim_child_name(schema, claim_slot_index); let claim_batch_rel = claim_batch_child_name(schema, claim_slot_index); let closure_rel = closure_child_name(schema, claim_slot_index); let closure_batch_rel = claim_closure_batch_child_name(schema, claim_slot_index); let receipt_batch_rel = format!("{schema}.receipt_completion_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let claim_slots: Vec = group.iter().map(|entry| entry.claim.claim_slot).collect(); let ready_slots: Vec = group.iter().map(|entry| entry.claim.ready_slot).collect(); let ready_generations: Vec = group .iter() .map(|entry| entry.claim.ready_generation) .collect(); let job_ids: Vec = group.iter().map(|entry| entry.job.id).collect(); let queues: Vec = group.iter().map(|entry| entry.job.queue.clone()).collect(); let priorities: Vec = group.iter().map(|entry| entry.claim.priority).collect(); let attempts: Vec = group.iter().map(|entry| entry.job.attempt).collect(); let run_leases: Vec = group.iter().map(|entry| entry.job.run_lease).collect(); let receipt_ids: Vec = group .iter() .map(|entry| { entry .claim .receipt_id .expect("receipt fast completion requires receipt_id") }) .collect(); let claim_batch_ids: Vec> = group .iter() .map(|entry| entry.claim.claim_batch_id) .collect(); let claim_batch_indices: Vec> = group .iter() .map(|entry| entry.claim.claim_batch_index) .collect(); let lane_seqs: Vec = group.iter().map(|entry| entry.claim.lane_seq).collect(); let enqueue_shards: Vec = group .iter() .map(|entry| entry.claim.enqueue_shard) .collect(); let attempted_ats: Vec>> = group.iter().map(|entry| entry.job.attempted_at).collect(); let finalized_ats: Vec> = vec![finalized_at; group.len()]; // ADR-048: outputs ride the batch row as an array aligned with // job_ids, and the columns are named only when some job has one. let outputs: Option>> = group .iter() .any(|entry| entry.job.output.is_some()) .then(|| { group .iter() .map(|entry| entry.job.output.as_ref().map(|output| output.to_string())) .collect() }); let (output_column, output_unnest, outputs_column, outputs_value) = if outputs.is_some() { ( ",\n output", ",\n $16::text[]::json[]", ",\n outputs", // NULL for a batch row whose jobs recorded no output. ",\n CASE WHEN bool_or(completed.output IS NOT NULL) \ THEN array_agg(completed.output ORDER BY completed.lane_seq, completed.job_id) END", ) } else { ("", "", "", "") }; let query = sqlx::query_as(audited_sql(format!( r#" WITH completed( claim_slot, ready_slot, ready_generation, job_id, queue, priority, attempt, run_lease, receipt_id, claim_batch_id, claim_batch_index, lane_seq, enqueue_shard, attempted_at, finalized_at{output_column} ) AS ( SELECT * FROM unnest( $1::int[], $2::int[], $3::bigint[], $4::bigint[], $5::text[], $6::smallint[], $7::smallint[], $8::bigint[], $9::bigint[], $10::bigint[], $11::int[], $12::bigint[], $13::smallint[], $14::timestamptz[], $15::timestamptz[]{output_unnest} ) ), row_claim_refs AS ( SELECT claims.claim_slot, claims.job_id, claims.run_lease, claims.receipt_id FROM {claim_rel} AS claims JOIN completed ON completed.claim_slot = claims.claim_slot AND completed.job_id = claims.job_id AND completed.run_lease = claims.run_lease AND completed.receipt_id = claims.receipt_id AND completed.claim_batch_id IS NULL WHERE NOT {closed_evidence} AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) ), batch_claim_refs AS ( SELECT claim_batches.claim_slot, items.job_id, items.run_lease, items.receipt_id FROM completed JOIN {claim_batch_rel} AS claim_batches ON claim_batches.claim_slot = completed.claim_slot AND claim_batches.batch_id = completed.claim_batch_id CROSS JOIN LATERAL ( SELECT claim_batches.job_ids[completed.claim_batch_index] AS job_id, claim_batches.run_leases[completed.claim_batch_index] AS run_lease, claim_batches.receipt_ids[completed.claim_batch_index] AS receipt_id ) AS items WHERE completed.claim_batch_id IS NOT NULL AND completed.claim_batch_index IS NOT NULL AND completed.claim_batch_index BETWEEN 1 AND claim_batches.claimed_count AND completed.job_id = items.job_id AND completed.run_lease = items.run_lease AND completed.receipt_id = items.receipt_id AND NOT EXISTS ( SELECT 1 FROM {closure_rel} AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) -- Compact successful completion treats receipt closure -- evidence as the same-attempt disposition fence. The -- non-success paths write explicit closure rows before -- moving the job to done/deferred/DLQ, so probing those -- terminal ledgers here only adds hot-path work. AND NOT EXISTS ( SELECT 1 FROM {closure_batch_rel} AS closure_batches WHERE closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = items.job_id AND lease.run_lease = items.run_lease ) ), claim_refs AS ( SELECT claim_slot, job_id, run_lease, receipt_id FROM row_claim_refs UNION ALL SELECT claim_slot, job_id, run_lease, receipt_id FROM batch_claim_refs ), deleted_attempts AS ( DELETE FROM {schema}.attempt_state AS attempt USING claim_refs WHERE attempt.job_id = claim_refs.job_id AND attempt.run_lease = claim_refs.run_lease RETURNING attempt.job_id ), completed_rows AS ( SELECT completed.* FROM completed JOIN claim_refs ON claim_refs.claim_slot = completed.claim_slot AND claim_refs.job_id = completed.job_id AND claim_refs.run_lease = completed.run_lease AND claim_refs.receipt_id = completed.receipt_id ), claim_closure_batches AS ( INSERT INTO {closure_batch_rel} ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT completed.claim_slot, completed.ready_slot, completed.ready_generation, 'completed', count(*)::int AS closed_count, array_agg(completed.receipt_id ORDER BY completed.lane_seq, completed.job_id), range_agg(int8range(completed.receipt_id, completed.receipt_id + 1, '[)') ORDER BY completed.receipt_id), max(completed.finalized_at) FROM completed_rows AS completed GROUP BY completed.claim_slot, completed.ready_slot, completed.ready_generation RETURNING claim_slot, ready_slot, ready_generation, receipt_ids ), terminal AS ( INSERT INTO {receipt_batch_rel} ( ready_slot, ready_generation, claim_slot, queue, priority, enqueue_shard, completed_count, job_ids, run_leases, lane_seqs, attempts, attempted_ats, finalized_at{outputs_column} ) SELECT completed.ready_slot, completed.ready_generation, completed.claim_slot, completed.queue, completed.priority, completed.enqueue_shard, count(*)::int AS completed_count, array_agg(completed.job_id ORDER BY completed.lane_seq, completed.job_id), array_agg(completed.run_lease ORDER BY completed.lane_seq, completed.job_id), array_agg(completed.lane_seq ORDER BY completed.lane_seq, completed.job_id), array_agg(completed.attempt ORDER BY completed.lane_seq, completed.job_id), array_agg(completed.attempted_at ORDER BY completed.lane_seq, completed.job_id), max(completed.finalized_at){outputs_value} FROM completed_rows AS completed GROUP BY completed.ready_slot, completed.ready_generation, completed.claim_slot, completed.queue, completed.priority, completed.enqueue_shard RETURNING ready_slot, ready_generation, claim_slot, queue, priority, enqueue_shard, job_ids, run_leases ) SELECT completed_rows.job_id, completed_rows.run_lease FROM completed_rows CROSS JOIN (SELECT count(*) FROM claim_closure_batches) AS closure_batch_write CROSS JOIN (SELECT count(*) FROM terminal) AS terminal_write "# ))); let query = query .bind(&claim_slots) .bind(&ready_slots) .bind(&ready_generations) .bind(&job_ids) .bind(&queues) .bind(&priorities) .bind(&attempts) .bind(&run_leases) .bind(&receipt_ids) .bind(&claim_batch_ids) .bind(&claim_batch_indices) .bind(&lane_seqs) .bind(&enqueue_shards) .bind(&attempted_ats) .bind(&finalized_ats); let query = match &outputs { Some(outputs) => query.bind(outputs), None => query, }; let completed: Vec<(i64, i64)> = match query.fetch_all(tx.as_mut()).await { Ok(rows) => rows, Err(err) => { let _ = tx.rollback().await; return Err(map_sqlx_error(err)); } }; rows.extend(completed); } tx.commit().await.map_err(map_sqlx_error)?; Ok(rows) } #[tracing::instrument( skip(self, pool, claimed), name = "queue_storage.complete_runtime_batch" )] pub async fn complete_runtime_batch( &self, pool: &PgPool, claimed: &[ClaimedRuntimeJob], ) -> Result, AwaError> { if claimed.is_empty() { return Ok(Vec::new()); } if self.lease_claim_receipts() && claimed.iter().all(Self::receipt_fast_complete_candidate) { let mut updated = self .complete_receipt_runtime_batch_fast(pool, claimed) .await?; if updated.len() == claimed.len() { return Ok(updated); } let updated_pairs: BTreeSet<(i64, i64)> = updated.iter().copied().collect(); let missed: Vec<_> = claimed .iter() .filter(|entry| !updated_pairs.contains(&(entry.job.id, entry.job.run_lease))) .cloned() .collect(); if !missed.is_empty() { updated.extend(self.complete_runtime_batch_slow(pool, &missed).await?); } return Ok(updated); } self.complete_runtime_batch_slow(pool, claimed).await } async fn complete_runtime_batch_slow( &self, pool: &PgPool, claimed: &[ClaimedRuntimeJob], ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .complete_runtime_batch_slow_in_tx(&mut tx, claimed) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Same as [`Self::complete_runtime_batch_slow`] but runs on the caller's /// transaction so additional writes — for example ADR-029 follow-up job /// inserts — can join the same commit. The caller is responsible for /// `commit()` / `rollback()`. Handles both receipt-claimed and /// materialised leases. pub async fn complete_runtime_batch_slow_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, claimed: &[ClaimedRuntimeJob], ) -> Result, AwaError> { if claimed.is_empty() { return Ok(Vec::new()); } let schema = self.schema(); let claimed_map: BTreeMap<(i64, i64), ClaimedRuntimeJob> = claimed .iter() .cloned() .map(|entry| ((entry.job.id, entry.job.run_lease), entry)) .collect(); if self.lease_claim_receipts() { let (mut receipt_claimed, mut materialized_claimed): (Vec<_>, Vec<_>) = claimed .iter() .cloned() .partition(|entry| entry.claim.lease_claim_receipt); let mut updated_all = Vec::new(); if !receipt_claimed.is_empty() { let receipt_pairs: Vec<(i64, i64)> = receipt_claimed .iter() .map(|entry| (entry.job.id, entry.job.run_lease)) .collect(); self.lock_receipt_attempts_tx(tx, &receipt_pairs).await?; // claim_slot rides along on `ClaimedEntry`, so receipt // completion can validate exact claim evidence and route the // explicit completed closure without an extra lookup. let receipt_claim_slots: Vec = receipt_claimed .iter() .map(|entry| entry.claim.claim_slot) .collect(); let receipt_job_ids: Vec = receipt_claimed.iter().map(|entry| entry.job.id).collect(); let receipt_run_leases: Vec = receipt_claimed .iter() .map(|entry| entry.job.run_lease) .collect(); let receipt_receipt_ids: Vec = receipt_claimed .iter() .map(|entry| { entry .claim .receipt_id .expect("receipt-backed slow completion requires receipt_id") }) .collect(); let receipt_claim_batch_ids: Vec> = receipt_claimed .iter() .map(|entry| entry.claim.claim_batch_id) .collect(); let receipt_claim_batch_indices: Vec> = receipt_claimed .iter() .map(|entry| entry.claim.claim_batch_index) .collect(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let updated: Vec<(i64, i64)> = sqlx::query_as(audited_sql(format!( r#" WITH completed(claim_slot, job_id, run_lease, receipt_id, claim_batch_id, claim_batch_index) AS ( SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[], $5::bigint[], $6::int[]) ), locked_row_claims AS ( SELECT claims.claim_slot, claims.job_id, claims.run_lease FROM {schema}.lease_claims AS claims JOIN completed ON completed.claim_slot = claims.claim_slot AND completed.job_id = claims.job_id AND completed.run_lease = claims.run_lease AND completed.receipt_id = claims.receipt_id WHERE NOT {closed_evidence} AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) FOR UPDATE OF claims ), locked_batch_claims AS ( SELECT claim_batches.claim_slot, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, items.run_lease, items.receipt_id FROM completed JOIN {schema}.lease_claim_batches AS claim_batches ON claim_batches.claim_slot = completed.claim_slot AND claim_batches.batch_id = completed.claim_batch_id CROSS JOIN LATERAL ( SELECT claim_batches.job_ids[completed.claim_batch_index] AS job_id, claim_batches.run_leases[completed.claim_batch_index] AS run_lease, claim_batches.receipt_ids[completed.claim_batch_index] AS receipt_id ) AS items WHERE completed.claim_batch_id IS NOT NULL AND completed.claim_batch_index IS NOT NULL AND completed.claim_batch_index BETWEEN 1 AND claim_batches.claimed_count AND completed.job_id = items.job_id AND completed.run_lease = items.run_lease AND completed.receipt_id = items.receipt_id AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = items.job_id AND lease.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) ), locked_claims AS ( SELECT claim_slot, job_id, run_lease FROM locked_row_claims UNION ALL SELECT claim_slot, job_id, run_lease FROM locked_batch_claims ), deleted_attempts AS ( DELETE FROM {schema}.attempt_state AS attempt USING locked_claims WHERE attempt.job_id = locked_claims.job_id AND attempt.run_lease = locked_claims.run_lease RETURNING attempt.job_id ), -- Row-sourced (deadline lease) claims close into the -- explicit ledger so the prune gates balance them -- against the lease_claims row via the closure JOIN. closed_claims AS ( INSERT INTO {schema}.lease_claim_closures ( claim_slot, job_id, run_lease, outcome, closed_at ) SELECT locked_row_claims.claim_slot, locked_row_claims.job_id, locked_row_claims.run_lease, 'completed', clock_timestamp() FROM locked_row_claims ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), -- Compact batch-sourced claims have no lease_claims row -- to JOIN, so close them into the batch ledger that the -- queue prune gate counts via compact_count. closed_batches AS ( INSERT INTO {schema}.lease_claim_closure_batches ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT locked_batch_claims.claim_slot, locked_batch_claims.ready_slot, locked_batch_claims.ready_generation, 'completed', count(*)::int, array_agg(locked_batch_claims.receipt_id ORDER BY locked_batch_claims.receipt_id), range_agg(int8range(locked_batch_claims.receipt_id, locked_batch_claims.receipt_id + 1, '[)') ORDER BY locked_batch_claims.receipt_id), clock_timestamp() FROM locked_batch_claims GROUP BY locked_batch_claims.claim_slot, locked_batch_claims.ready_slot, locked_batch_claims.ready_generation RETURNING claim_slot ), marked_claims AS ( UPDATE {schema}.lease_claims AS claims SET closed_at = COALESCE(claims.closed_at, closed_claims.closed_at) FROM closed_claims WHERE claims.claim_slot = closed_claims.claim_slot AND claims.job_id = closed_claims.job_id AND claims.run_lease = closed_claims.run_lease RETURNING claims.job_id ), -- Force the batch insert to run and report the compact -- claims it closed so the caller finalizes them too. closed_batch_pairs AS ( SELECT locked_batch_claims.job_id, locked_batch_claims.run_lease FROM locked_batch_claims WHERE EXISTS (SELECT 1 FROM closed_batches) ) SELECT job_id, run_lease FROM closed_claims UNION ALL SELECT job_id, run_lease FROM closed_batch_pairs "# ))) .bind(&receipt_claim_slots) .bind(&receipt_job_ids) .bind(&receipt_run_leases) .bind(&receipt_receipt_ids) .bind(&receipt_claim_batch_ids) .bind(&receipt_claim_batch_indices) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if !updated.is_empty() { let finalized_at = Utc::now(); let mut done_rows = Vec::with_capacity(updated.len()); for (job_id, run_lease) in &updated { if let Some(runtime_job) = claimed_map.get(&(*job_id, *run_lease)).cloned() { done_rows.push(runtime_job.into_done_row(finalized_at)?); } } self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running)) .await?; updated_all.extend(updated); } let updated_pairs: BTreeSet<(i64, i64)> = updated_all.iter().copied().collect(); let mut escalated_receipts = Vec::new(); for entry in receipt_claimed.drain(..) { if !updated_pairs.contains(&(entry.job.id, entry.job.run_lease)) { escalated_receipts.push(entry); } } materialized_claimed.extend(escalated_receipts); } if !materialized_claimed.is_empty() { let ready_slots: Vec = materialized_claimed .iter() .map(|entry| entry.claim.ready_slot) .collect(); let ready_generations: Vec = materialized_claimed .iter() .map(|entry| entry.claim.ready_generation) .collect(); let job_ids: Vec = materialized_claimed .iter() .map(|entry| entry.job.id) .collect(); let queues: Vec = materialized_claimed .iter() .map(|entry| entry.claim.queue.clone()) .collect(); let priorities: Vec = materialized_claimed .iter() .map(|entry| entry.claim.priority) .collect(); let enqueue_shards: Vec = materialized_claimed .iter() .map(|entry| entry.claim.enqueue_shard) .collect(); let lane_seqs: Vec = materialized_claimed .iter() .map(|entry| entry.claim.lane_seq) .collect(); let run_leases: Vec = materialized_claimed .iter() .map(|entry| entry.job.run_lease) .collect(); // CTE-as-DML: delete the leases and the matching attempt_state // rows in one round-trip. Receipt claims may materialize a // lease after the original claim, so the materialized lease can // live in a newer lease slot than the claim carried. Match on // the stable ready-lane and attempt identity instead. let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" WITH completed(ready_slot, ready_generation, job_id, queue, priority, enqueue_shard, lane_seq, run_lease) AS ( SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::text[], $5::smallint[], $6::smallint[], $7::bigint[], $8::bigint[]) ), deleted AS ( DELETE FROM {schema}.leases AS leases USING completed WHERE leases.ready_slot = completed.ready_slot AND leases.ready_generation = completed.ready_generation AND leases.job_id = completed.job_id AND leases.queue = completed.queue AND leases.priority = completed.priority AND leases.enqueue_shard = completed.enqueue_shard AND leases.lane_seq = completed.lane_seq AND leases.run_lease = completed.run_lease RETURNING leases.ready_slot, leases.ready_generation, leases.job_id, leases.queue, leases.state, leases.priority, leases.attempt, leases.run_lease, leases.max_attempts, leases.lane_seq, leases.enqueue_shard, leases.heartbeat_at, leases.deadline_at, leases.attempted_at, leases.callback_id, leases.callback_timeout_at ), del_attempts AS ( DELETE FROM {schema}.attempt_state AS attempt USING deleted WHERE attempt.job_id = deleted.job_id AND attempt.run_lease = deleted.run_lease RETURNING attempt.job_id ) SELECT ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at FROM deleted "# ))) .bind(&ready_slots) .bind(&ready_generations) .bind(&job_ids) .bind(&queues) .bind(&priorities) .bind(&enqueue_shards) .bind(&lane_seqs) .bind(&run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if !deleted.is_empty() { let completed_pairs: Vec<(i64, i64)> = deleted .iter() .map(|row| (row.job_id, row.run_lease)) .collect(); self.close_receipt_pairs_tx(tx, &completed_pairs, "completed") .await?; let finalized_at = Utc::now(); let mut done_rows = Vec::with_capacity(deleted.len()); for deleted_row in deleted { if let Some(runtime_job) = claimed_map .get(&(deleted_row.job_id, deleted_row.run_lease)) .cloned() { done_rows.push(runtime_job.into_done_row(finalized_at)?); updated_all.push((deleted_row.job_id, deleted_row.run_lease)); } } self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running)) .await?; } } return Ok(updated_all); } let lease_slots: Vec = claimed.iter().map(|entry| entry.claim.lease_slot).collect(); let queues: Vec = claimed .iter() .map(|entry| entry.claim.queue.clone()) .collect(); let priorities: Vec = claimed.iter().map(|entry| entry.claim.priority).collect(); let enqueue_shards: Vec = claimed .iter() .map(|entry| entry.claim.enqueue_shard) .collect(); let lane_seqs: Vec = claimed.iter().map(|entry| entry.claim.lane_seq).collect(); let run_leases: Vec = claimed.iter().map(|entry| entry.job.run_lease).collect(); // Single CTE-as-DML statement: delete the leases and the matching // attempt_state rows in one round-trip. The `deleted` CTE materialises // the lease deletion (so its RETURNING is observable to the // attempt-state delete and to the final SELECT), and `del_attempts` // hangs off it. Saves one round-trip per completion batch versus // issuing the attempt-state delete as a separate statement. let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" WITH completed(lease_slot, queue, priority, enqueue_shard, lane_seq, run_lease) AS ( SELECT * FROM unnest($1::int[], $2::text[], $3::smallint[], $4::smallint[], $5::bigint[], $6::bigint[]) ), deleted AS ( DELETE FROM {schema}.leases AS leases USING completed WHERE leases.lease_slot = completed.lease_slot AND leases.queue = completed.queue AND leases.priority = completed.priority AND leases.enqueue_shard = completed.enqueue_shard AND leases.lane_seq = completed.lane_seq AND leases.run_lease = completed.run_lease RETURNING leases.ready_slot, leases.ready_generation, leases.job_id, leases.queue, leases.state, leases.priority, leases.attempt, leases.run_lease, leases.max_attempts, leases.lane_seq, leases.enqueue_shard, leases.heartbeat_at, leases.deadline_at, leases.attempted_at, leases.callback_id, leases.callback_timeout_at ), del_attempts AS ( DELETE FROM {schema}.attempt_state AS attempt USING deleted WHERE attempt.job_id = deleted.job_id AND attempt.run_lease = deleted.run_lease RETURNING attempt.job_id ) SELECT ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at FROM deleted "# ))) .bind(&lease_slots) .bind(&queues) .bind(&priorities) .bind(&enqueue_shards) .bind(&lane_seqs) .bind(&run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { return Ok(Vec::new()); } let completed_pairs: Vec<(i64, i64)> = deleted .iter() .map(|row| (row.job_id, row.run_lease)) .collect(); self.close_receipt_pairs_tx(tx, &completed_pairs, "completed") .await?; let finalized_at = Utc::now(); let mut done_rows = Vec::with_capacity(deleted.len()); let mut updated = Vec::with_capacity(deleted.len()); for deleted_row in deleted { if let Some(runtime_job) = claimed_map .get(&(deleted_row.job_id, deleted_row.run_lease)) .cloned() { done_rows.push(runtime_job.into_done_row(finalized_at)?); updated.push((deleted_row.job_id, deleted_row.run_lease)); } } self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running)) .await?; Ok(updated) } #[tracing::instrument( skip(self, pool, completions), name = "queue_storage.complete_job_batch_by_id" )] pub async fn complete_job_batch_by_id( &self, pool: &PgPool, completions: &[(i64, i64)], ) -> Result, AwaError> { if completions.is_empty() { return Ok(Vec::new()); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .complete_job_batch_by_id_in_tx(&mut tx, completions) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Same as [`Self::complete_job_batch_by_id`] but runs on the caller's /// transaction so additional writes — for example ADR-029 follow-up job /// inserts — can join the same commit. The caller is responsible for /// `commit()` / `rollback()`. pub async fn complete_job_batch_by_id_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, completions: &[(i64, i64)], ) -> Result, AwaError> { if completions.is_empty() { return Ok(Vec::new()); } let schema = self.schema(); let job_ids: Vec = completions.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = completions .iter() .map(|(_, run_lease)| *run_lease) .collect(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" WITH completed(job_id, run_lease) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) ) DELETE FROM {schema}.leases AS leases USING completed WHERE leases.job_id = completed.job_id AND leases.run_lease = completed.run_lease RETURNING leases.ready_slot, leases.ready_generation, leases.job_id, leases.queue, leases.state, leases.priority, leases.attempt, leases.run_lease, leases.max_attempts, leases.lane_seq, leases.enqueue_shard, leases.heartbeat_at, leases.deadline_at, leases.attempted_at, leases.callback_id, leases.callback_timeout_at "# ))) .bind(&job_ids) .bind(&run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { return Ok(Vec::new()); } let completed_pairs: Vec<(i64, i64)> = deleted .iter() .map(|row| (row.job_id, row.run_lease)) .collect(); self.close_receipt_pairs_tx(tx, &completed_pairs, "completed") .await?; let moved = self.hydrate_deleted_leases_tx(tx, deleted, true).await?; let finalized_at = Utc::now(); let mut done_rows = Vec::with_capacity(moved.len()); for entry in moved.iter().cloned() { let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state( entry.payload.clone(), entry.progress.clone(), )?)?; payload.set_progress(None); done_rows.push(entry.into_done_row( JobState::Completed, finalized_at, payload.into_json(), )); } self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running)) .await?; Ok(moved .into_iter() .map(|entry| (entry.job_id, entry.run_lease)) .collect()) } /// Exact admin/UI-grade queue counts. Scans `ready_entries` rather /// than reading the head tables — slower, but unaffected by the /// transient gap between sequence reservations and committed, /// still-claimable ready rows. Use [`Self::queue_claimer_signal`] for the /// dispatcher hot path. /// /// The live-terminal portion reads retained compact receipt batches /// directly. Done-entry terminal rows read from folded /// `queue_terminal_live_counts` plus unrolled /// `queue_terminal_count_deltas` when [`Self::terminal_counter_trusted`] /// returns true, and fall back to a `count(*) FROM terminal_jobs` scan /// when not. The fallback exists for the rolling-upgrade window: older /// binaries may have written terminal rows without maintaining the /// counter/delta contract, so reads stay honest until the operator runs /// `awa storage rebuild-terminal-counters`. async fn queue_counts_exact( &self, pool: &PgPool, queue: &str, ) -> Result { let schema = self.schema(); let queues = self.physical_queues_for_logical(queue); // Canonical open-receipt-claim shape (row-local + compact batch), // reused by admin::state_counts. (The maintenance health gauge is // deliberately NOT a consumer — it stays leases-only so it can // never block on claim-ring ACCESS EXCLUSIVE locks.) let open_running = open_receipt_running_claims_sql(schema); let counter_trusted = self.terminal_counter_trusted(pool).await?; // The live-terminal CTE swaps between counter-fed and // scan-fed depending on trust. Build it as a string so the // outer query plan is otherwise identical between the two // paths. let live_terminal_cte = if counter_trusted { format!( "live_terminal AS ( SELECT GREATEST( 0, COALESCE(( SELECT SUM(live_terminal_count) FROM {schema}.queue_terminal_live_counts WHERE queue = ANY($1) ), 0) + COALESCE(( SELECT SUM(terminal_delta) FROM {schema}.queue_terminal_count_deltas WHERE queue = ANY($1) ), 0) + COALESCE(( SELECT SUM(completed_count) FROM {schema}.receipt_completion_batches WHERE queue = ANY($1) ), 0) - COALESCE(( SELECT count(*)::bigint FROM {schema}.receipt_completion_tombstones WHERE queue = ANY($1) ), 0) )::bigint AS terminal )" ) } else { format!( "live_terminal AS ( SELECT count(*)::bigint AS terminal FROM {schema}.terminal_jobs WHERE queue = ANY($1) )" ) }; let row: (i64, i64, i64, i64) = sqlx::query_as(audited_sql(format!( r#" WITH lane_counts AS ( -- Exact count: a ready row is available iff its -- lane_seq has not yet been passed by the lane's -- claim sequence cursor. Each shard within a (queue, -- priority) lane carries its own sequence, so the -- join matches on shard too — otherwise a ready row -- in shard A could be incorrectly compared against -- shard B's claim cursor. SELECT COALESCE(count(*)::bigint, 0) AS available FROM {schema}.ready_entries AS ready JOIN {schema}.queue_claim_heads AS claims ON claims.queue = ready.queue AND claims.priority = ready.priority AND claims.enqueue_shard = ready.enqueue_shard WHERE ready.queue = ANY($1) AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq AND tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation ) ), pruned_terminal AS ( -- The GREATEST legacy dedupe applies to the completed -- column only: queue_lanes never carried a failed -- column, so failed counts come from the rollups alone. -- Unfolded prune deltas (#371) are added on top of the -- folded rollups so counts stay exact between folds. SELECT COALESCE( sum( GREATEST( COALESCE(lanes.pruned_completed_count, 0), COALESCE(rollups.pruned_completed_count, 0) ) + COALESCE(pending.pruned_completed_delta, 0) ), 0 )::bigint AS completed, COALESCE( sum( COALESCE(rollups.pruned_failed_count, 0) + COALESCE(pending.pruned_failed_delta, 0) ), 0 )::bigint AS failed FROM ( SELECT queue, priority, pruned_completed_count FROM {schema}.queue_lanes WHERE queue = ANY($1) ) AS lanes FULL OUTER JOIN ( SELECT queue, priority, pruned_completed_count, pruned_failed_count FROM {schema}.queue_terminal_rollups WHERE queue = ANY($1) ) AS rollups USING (queue, priority) FULL OUTER JOIN ( SELECT queue, priority, sum(pruned_completed_delta)::bigint AS pruned_completed_delta, sum(pruned_failed_delta)::bigint AS pruned_failed_delta FROM {schema}.queue_terminal_rollup_deltas WHERE queue = ANY($1) GROUP BY queue, priority ) AS pending USING (queue, priority) ), live_running AS ( SELECT ( COALESCE(( SELECT count(*)::bigint FROM {schema}.leases WHERE queue = ANY($1) AND state = 'running' ), 0) + -- Receipt-backed running claims (row-local `lease_claims` -- plus compact `lease_claim_batches` members, #246), each -- anti-joined against every durable closure/terminal/ -- materialised-lease shape. Shared canonical shape reused -- by admin::state_counts (the health gauge stays -- leases-only by design). COALESCE(( SELECT count(*)::bigint FROM ({open_running}) AS open_running WHERE open_running.queue = ANY($1) ), 0) )::bigint AS running ), {live_terminal_cte} SELECT lane_counts.available, live_running.running, pruned_terminal.completed + pruned_terminal.failed + live_terminal.terminal AS terminal, pruned_terminal.failed AS pruned_failed FROM lane_counts CROSS JOIN pruned_terminal CROSS JOIN live_running CROSS JOIN live_terminal "# ))) .bind(&queues) .fetch_one(pool) .await .map_err(map_sqlx_error)?; let (available, running, terminal, pruned_failed) = row; Ok(QueueCounts { available, running, terminal, pruned_failed, }) } #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue), name = "queue_storage.queue_counts")] pub async fn queue_counts(&self, pool: &PgPool, queue: &str) -> Result { self.queue_counts_exact(pool, queue).await } /// Index-only queue depth probe — for observability / depth-target /// throttling. Returns the same shape as [`Self::queue_counts`] but /// skips the table scans that [`Self::queue_counts_exact`] needs for /// exact terminal accounting: /// /// - **available** is the same as the dispatcher's claim signal: /// `sum(GREATEST(enqueue_seq - claim_seq, 0))` over the shard /// head tables. No scan of `ready_entries`. This is an upper /// bound: admin DELETEs, committed gaps, and uncommitted enqueue /// reservations can leave the enqueue sequence ahead of the actual /// ready row count. Acceptable for depth-target throttling and /// dashboards; not suitable for exact billing-style counts. /// - **running** matches [`Self::queue_counts`]'s strict definition: /// `leases.state = 'running'` only. Receipt-plane claims that /// have not yet materialised a lease row are omitted (the /// exact path catches them via the `lease_claims` anti-join, but /// that anti-join is what this fast variant exists to avoid). /// `waiting_external` is *not* included — it's reported as part /// of admin's parked-callback view, not running. /// - **terminal** is read from the persisted /// `queue_terminal_rollups` denormaliser /// (`pruned_completed_count + pruned_failed_count`). /// Rows currently in `done_entries` or /// `receipt_completion_batches` that have not yet rolled up are /// not included. Strictly a lower bound; converges to the exact /// count when rotation prunes the live queue segment. (The name /// `terminal` is honest — this number counts `completed`, /// `failed`, and `cancelled` terminal facts with the same /// semantics as [`Self::queue_counts_exact`]; renamed from /// `completed` in #290.) /// /// All three counters are O(num shards) lookups against small head /// tables and `leases` index probes. Use this for high-cadence /// pollers (admin dashboards, depth-target producers, soak /// observability); use [`Self::queue_counts`] for admin tooling /// that needs the exact terminal count. #[tracing::instrument(skip(self, pool), fields(queue = %queue, messaging.destination.name = %queue), name = "queue_storage.queue_counts_fast")] pub async fn queue_counts_fast( &self, pool: &PgPool, queue: &str, ) -> Result { let schema = self.schema(); let queues = self.physical_queues_for_logical(queue); // available: dispatcher signal — already an index-only sum over // the (queue, priority, enqueue_shard) head tables. let available = self.queue_claimer_signal(pool, queue).await?.available; // running: leases.state = 'running' only, matching // queue_counts_exact's strict definition. Receipt-plane claims // that haven't materialised a lease row yet are documented as // omitted in the method-level doc. let running: i64 = sqlx::query_scalar(audited_sql(format!( r#" SELECT COALESCE(count(*)::bigint, 0) FROM {schema}.leases WHERE queue = ANY($1) AND state = 'running' "# ))) .bind(&queues) .fetch_one(pool) .await .map_err(map_sqlx_error)?; // terminal: denormalised rollup plus unfolded prune deltas // (#371). Live (un-rolled-up) done_entries rows are excluded — // see method-level docs. The GREATEST legacy dedupe applies to // the completed column only: queue_lanes never carried a failed // column. let (pruned_completed, pruned_failed): (i64, i64) = sqlx::query_as(audited_sql(format!( r#" SELECT COALESCE(sum( GREATEST( COALESCE(lanes.pruned_completed_count, 0), COALESCE(rollups.pruned_completed_count, 0) ) + COALESCE(pending.pruned_completed_delta, 0) ), 0)::bigint, COALESCE(sum( COALESCE(rollups.pruned_failed_count, 0) + COALESCE(pending.pruned_failed_delta, 0) ), 0)::bigint FROM ( SELECT queue, priority, pruned_completed_count FROM {schema}.queue_lanes WHERE queue = ANY($1) ) AS lanes FULL OUTER JOIN ( SELECT queue, priority, pruned_completed_count, pruned_failed_count FROM {schema}.queue_terminal_rollups WHERE queue = ANY($1) ) AS rollups USING (queue, priority) FULL OUTER JOIN ( SELECT queue, priority, sum(pruned_completed_delta)::bigint AS pruned_completed_delta, sum(pruned_failed_delta)::bigint AS pruned_failed_delta FROM {schema}.queue_terminal_rollup_deltas WHERE queue = ANY($1) GROUP BY queue, priority ) AS pending USING (queue, priority) "# ))) .bind(&queues) .fetch_one(pool) .await .map_err(map_sqlx_error)?; Ok(QueueCounts { available, running, terminal: pruned_completed + pruned_failed, pruned_failed, }) } /// Cumulative count of `failed` terminal rows pruned past the /// retention floor for a queue, summed over the queue's /// `queue_terminal_rollups` priorities (and stripes, for striped /// queues). Monotonic — rollups never decrease. These rows no /// longer exist in `done_entries` and cannot be retried. pub async fn pruned_failed_count_for_queue( &self, pool: &PgPool, queue: &str, ) -> Result { let schema = self.schema(); let queues = self.physical_queues_for_logical(queue); // Folded rollups plus unfolded prune deltas (#371), so the count // stays exact between maintenance folds. let pruned_failed: i64 = sqlx::query_scalar(audited_sql(format!( r#" SELECT COALESCE(( SELECT sum(pruned_failed_count) FROM {schema}.queue_terminal_rollups WHERE queue = ANY($1) ), 0)::bigint + COALESCE(( SELECT sum(pruned_failed_delta) FROM {schema}.queue_terminal_rollup_deltas WHERE queue = ANY($1) ), 0)::bigint "# ))) .bind(&queues) .fetch_one(pool) .await .map_err(map_sqlx_error)?; Ok(pruned_failed.max(0) as u64) } async fn retry_job_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, ) -> Result, AwaError> { let schema = self.schema(); let deleted_waiting: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE job_id = $1 AND state = 'waiting_external' RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .bind(job_id) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if !deleted_waiting.is_empty() { let waiting = self .hydrate_deleted_leases_tx(tx, deleted_waiting, true) .await? .into_iter() .next() .expect("deleted waiting lease"); let ready_payload = Self::payload_with_attempt_state( waiting.payload.clone(), waiting.progress.clone(), )?; let ready_row = ExistingReadyRow { attempt: 0, run_at: Utc::now(), attempted_at: None, ..waiting.clone().into_ready_row(Utc::now(), ready_payload) }; self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(waiting.state)) .await?; self.notify_queues_tx(tx, std::iter::once(waiting.queue.clone())) .await?; return Ok(Some( ReadyJobRow { job_id: ready_row.job_id, kind: ready_row.kind, queue: ready_row.queue, args: ready_row.args, priority: ready_row.priority, attempt: ready_row.attempt, run_lease: ready_row.run_lease, max_attempts: ready_row.max_attempts, run_at: ready_row.run_at, attempted_at: ready_row.attempted_at, created_at: ready_row.created_at, unique_key: ready_row.unique_key, payload: ready_row.payload, } .into_job_row()?, )); } let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); let terminal: Option = sqlx::query_as(audited_sql(format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries WHERE (job_id, finalized_at) IN ( SELECT job_id, finalized_at FROM {schema}.done_entries WHERE job_id = $1 AND state IN ('failed', 'cancelled') ORDER BY finalized_at DESC LIMIT 1 FOR UPDATE SKIP LOCKED ) RETURNING * ) SELECT {done_projection} FROM deleted AS done {ready_join} "# ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if let Some(terminal) = terminal { self.ensure_terminal_removed_receipt_closures_tx(tx, std::slice::from_ref(&terminal)) .await?; // The DELETE FROM done_entries above removes one terminal row; // append a negative delta so exact counts stay in lockstep. self.decrement_live_terminal_counters_tx( tx, &Self::done_rows_to_counter_keys(std::slice::from_ref(&terminal)), ) .await?; let ready_row = ExistingReadyRow { job_id: terminal.job_id, kind: terminal.kind, queue: terminal.queue.clone(), args: terminal.args, priority: terminal.priority, attempt: 0, run_lease: terminal.run_lease, max_attempts: terminal.max_attempts, run_at: Utc::now(), attempted_at: None, created_at: terminal.created_at, unique_key: terminal.unique_key, unique_states: terminal.unique_states, payload: terminal.payload, }; self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(terminal.state)) .await?; self.notify_queues_tx(tx, std::iter::once(terminal.queue.clone())) .await?; return Ok(Some( ReadyJobRow { job_id: ready_row.job_id, kind: ready_row.kind, queue: ready_row.queue, args: ready_row.args, priority: ready_row.priority, attempt: ready_row.attempt, run_lease: ready_row.run_lease, max_attempts: ready_row.max_attempts, run_at: ready_row.run_at, attempted_at: ready_row.attempted_at, created_at: ready_row.created_at, unique_key: ready_row.unique_key, payload: ready_row.payload, } .into_job_row()?, )); } Ok(None) } pub async fn retry_job(&self, pool: &PgPool, job_id: i64) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let row = self.retry_job_tx(&mut tx, job_id).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(row) } /// Retry jobs by id inside a single transaction. Duplicate ids are /// collapsed so each job is attempted at most once. Returns the /// retried rows together with the number of unique ids attempted: /// ids whose terminal row raced to another state or was pruned /// between the caller's scan and this call are skipped, so /// `attempted - retried.len()` is the count of unique ids that were /// requested but not retried. pub async fn retry_jobs_by_ids( &self, pool: &PgPool, ids: &[i64], ) -> Result<(Vec, u64), AwaError> { let unique_ids: BTreeSet = ids.iter().copied().collect(); if unique_ids.is_empty() { return Ok((Vec::new(), 0)); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let mut rows = Vec::with_capacity(unique_ids.len()); for job_id in &unique_ids { if let Some(row) = self.retry_job_tx(&mut tx, *job_id).await? { rows.push(row); } } tx.commit().await.map_err(map_sqlx_error)?; Ok((rows, unique_ids.len() as u64)) } /// Write a `` closure row for any matching open receipt. /// Idempotent: no-op if no row-local claim or compact claim-batch item /// exists for the `(job_id, run_lease)` pair, or if a closure already exists. Used /// by the admin cancel path to keep the receipt plane consistent /// with the job's new terminal state so rescue doesn't revive it. /// /// `FOR UPDATE` on the inner SELECT serialises the closure write /// against `ensure_running_leases_from_receipts_tx` /// (which also takes `FOR UPDATE` on the same claim evidence) and /// against concurrent rescue / re-close paths that might race the /// same `(job_id, run_lease)`. Without it, materialization could /// see the claim evidence, decide to materialize, and a concurrent /// admin cancel could write the closure between materialization's /// SELECT and the lease INSERT — leaving a `running` lease for a /// closed claim that admin cancel believes is fully shut down. async fn close_receipt_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: i64, outcome: &str, ) -> Result<(), AwaError> { self.close_receipt_pairs_tx(tx, &[(job_id, run_lease)], outcome) .await } async fn close_receipt_pairs_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, pairs: &[(i64, i64)], outcome: &str, ) -> Result<(), AwaError> { if pairs.is_empty() || !self.lease_claim_receipts() { return Ok(()); } let unique_pairs: BTreeSet<(i64, i64)> = pairs.iter().copied().collect(); let unique_jobs: Vec<(i64, i64)> = unique_pairs.iter().copied().collect(); self.lock_receipt_attempts_tx(tx, &unique_jobs).await?; let job_ids: Vec = unique_pairs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = unique_pairs .iter() .map(|(_, run_lease)| *run_lease) .collect(); let schema = self.schema(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); sqlx::query(audited_sql(format!( r#" WITH refs(job_id, run_lease) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) ), locked_claims AS ( SELECT claims.claim_slot, claims.job_id, claims.run_lease FROM {schema}.lease_claims AS claims JOIN refs ON refs.job_id = claims.job_id AND refs.run_lease = claims.run_lease WHERE NOT {closed_evidence} FOR UPDATE OF claims ), locked_batch_claims AS ( SELECT claim_batches.claim_slot, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, items.run_lease, items.receipt_id FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids ) AS items(job_id, run_lease, receipt_id) JOIN refs ON refs.job_id = items.job_id AND refs.run_lease = items.run_lease WHERE NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) FOR UPDATE OF claim_batches ), -- Row-sourced (deadline lease) claims keep their explicit -- closure: the queue/claim prune gates balance those against -- the lease_claims row via the closure JOIN. Compact -- batch-sourced claims have no lease_claims row to JOIN, so -- they are closed into lease_claim_closure_batches below. closing_claims AS ( SELECT locked_claims.claim_slot, locked_claims.job_id, locked_claims.run_lease, $3 AS outcome, clock_timestamp() AS closed_at FROM locked_claims ), inserted AS ( INSERT INTO {schema}.lease_claim_closures ( claim_slot, job_id, run_lease, outcome, closed_at ) SELECT claim_slot, job_id, run_lease, outcome, closed_at FROM closing_claims ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), inserted_batches AS ( INSERT INTO {schema}.lease_claim_closure_batches ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT locked_batch_claims.claim_slot, locked_batch_claims.ready_slot, locked_batch_claims.ready_generation, $3, count(*)::int, array_agg(locked_batch_claims.receipt_id ORDER BY locked_batch_claims.receipt_id), range_agg(int8range(locked_batch_claims.receipt_id, locked_batch_claims.receipt_id + 1, '[)') ORDER BY locked_batch_claims.receipt_id), clock_timestamp() FROM locked_batch_claims GROUP BY locked_batch_claims.claim_slot, locked_batch_claims.ready_slot, locked_batch_claims.ready_generation RETURNING claim_slot ), marked AS ( UPDATE {schema}.lease_claims AS claims SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) FROM inserted WHERE claims.claim_slot = inserted.claim_slot AND claims.job_id = inserted.job_id AND claims.run_lease = inserted.run_lease RETURNING claims.job_id ) SELECT (SELECT count(*) FROM marked) + (SELECT count(*) FROM inserted_batches) "# ))) .bind(&job_ids) .bind(&run_leases) .bind(outcome) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Emit a `pg_notify('awa:cancel', ...)` inside the cancel /// transaction so any worker runtime currently executing this /// `(job_id, run_lease)` learns about the cancellation on commit /// and fires its in-flight cancel flag. Notifications are /// automatically discarded on rollback. async fn notify_cancellation_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: i64, ) -> Result<(), AwaError> { let payload = serde_json::json!({ "job_id": job_id, "run_lease": run_lease }).to_string(); sqlx::query("SELECT pg_notify('awa:cancel', $1)") .bind(payload) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Delete and cancel the running lease of `job_id`, or of exactly /// `run_lease` when given. async fn cancel_materialized_lease_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: Option, ) -> Result, AwaError> { let schema = self.schema(); let deleted_lease: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE job_id = $1 AND state IN ('running', 'waiting_external') AND ($2::bigint IS NULL OR run_lease = $2) RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .bind(job_id) .bind(run_lease) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted_lease.is_empty() { return Ok(None); } let lease = self .hydrate_deleted_leases_tx(tx, deleted_lease, true) .await? .into_iter() .next() .expect("deleted running lease"); let done_payload = Self::payload_with_attempt_state(lease.payload.clone(), lease.progress.clone())?; let done = lease .clone() .into_done_row(JobState::Cancelled, Utc::now(), done_payload); self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(lease.state)) .await?; // Receipt-plane consistency: close any matching open // receipt so the ADR-023 anti-join no longer considers this // attempt live, and rescue doesn't try to revive it. self.close_receipt_tx(tx, lease.job_id, lease.run_lease, "cancelled") .await?; // Wake any worker currently executing this attempt. self.notify_cancellation_tx(tx, lease.job_id, lease.run_lease) .await?; Ok(Some(CancelJobTxResult { row: done.into_job_row()?, claim_cursor_advance: None, })) } /// Cancel the running attempt of `job_id` on a receipts-mode store. /// /// Closers disagree on lock order: receipt closers take the attempt's /// advisory lock before its claim rows, while lease-side closers delete /// the lease row first and then take the advisory lock. A cancel that /// waited in either order could close a cycle with the other, so it /// waits on nothing. Each try runs in a savepoint, takes the attempt /// lock and the row locks without waiting, and on contention rolls the /// savepoint back, releasing everything it took, before trying again /// for up to `CANCEL_BUSY_BUDGET`. A transaction that already holds an /// attempt lock, such as a later job of a bulk cancel, tries once, so it /// does not hold that lock against its closers while it retries. async fn cancel_running_attempt_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, ) -> Result { let holds_attempt_lock: bool = sqlx::query_scalar(audited_sql(format!( "SELECT COALESCE(current_setting('{ATTEMPT_LOCKS_HELD_SETTING}', true), '') = 'on'" ))) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; let deadline = Instant::now() + CANCEL_BUSY_BUDGET; let mut backoff = Duration::from_millis(5); loop { let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; let step = self .try_cancel_running_attempt_tx(&mut attempt, job_id) .await; let busy = match step { Ok(CancelStep::Done(outcome)) => { attempt.commit().await.map_err(map_sqlx_error)?; return Ok(outcome); } Ok(CancelStep::Again) => { attempt.rollback().await.map_err(map_sqlx_error)?; false } Ok(CancelStep::Busy) => { attempt.rollback().await.map_err(map_sqlx_error)?; true } Err(err) => { let _ = attempt.rollback().await; return Err(err); } }; if Instant::now() >= deadline || (busy && holds_attempt_lock) { tracing::info!( job_id, "cancel left a running attempt that another transaction kept busy" ); return Ok(RunningAttemptCancel::Left); } if busy { self.cancel_phase("attempt_busy").await; tokio::time::sleep(backoff).await; backoff = (backoff * 2).min(Duration::from_millis(100)); } } } async fn try_cancel_running_attempt_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, ) -> Result { let schema = self.schema(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); type ReceiptCancelRow = ( i32, i64, i32, i64, String, i16, i16, i16, i64, DateTime, i64, bool, ); let row_receipt_sql = format!( r#" SELECT claims.claim_slot, claims.run_lease, claims.ready_slot, claims.ready_generation, claims.queue, claims.priority, claims.attempt, claims.max_attempts, claims.lane_seq, claims.claimed_at, claims.receipt_id, false AS compact_batch FROM {schema}.lease_claims AS claims WHERE claims.job_id = $1 AND NOT {closed_evidence} ORDER BY claims.run_lease DESC LIMIT 1 FOR UPDATE OF claims NOWAIT "# ); let batch_receipt_sql = format!( r#" SELECT claim_batches.claim_slot, items.run_lease, claim_batches.ready_slot, claim_batches.ready_generation, claim_batches.queue, claim_batches.priority, items.attempt, items.max_attempts, items.lane_seq, claim_batches.claimed_at, items.receipt_id, true AS compact_batch FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids, claim_batches.lane_seqs, claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) WHERE items.job_id = $1 AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = items.job_id AND lease.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) ORDER BY items.run_lease DESC LIMIT 1 FOR UPDATE OF claim_batches NOWAIT "# ); let open_receipt_sql = format!( r#" WITH row_receipt AS ({row_receipt_sql}), batch_receipt AS ({batch_receipt_sql}) SELECT * FROM ( SELECT * FROM row_receipt UNION ALL SELECT * FROM batch_receipt ) AS receipt ORDER BY run_lease DESC LIMIT 1 "# ); // Find the attempt with plain reads; take nothing yet. let lease_run_lease: Option = sqlx::query_scalar(audited_sql(format!( "SELECT run_lease FROM {schema}.leases \ WHERE job_id = $1 AND state IN ('running', 'waiting_external') \ ORDER BY run_lease DESC LIMIT 1" ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; self.cancel_phase("lease_lookup").await; let found_receipt: Option = if lease_run_lease.is_none() { let find_receipt_sql = open_receipt_sql .replace("FOR UPDATE OF claims NOWAIT", "") .replace("FOR UPDATE OF claim_batches NOWAIT", ""); debug_assert!(!find_receipt_sql.contains("FOR UPDATE")); sqlx::query_as(audited_sql(find_receipt_sql)) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)? } else { None }; let Some(run_lease) = lease_run_lease.or(found_receipt.as_ref().map(|found| found.1)) else { // The two reads are separate statements: a lease materialized // between them hides the receipt from the second. Look again. let lease_appeared: bool = sqlx::query_scalar(audited_sql(format!( "SELECT EXISTS (SELECT 1 FROM {schema}.leases \ WHERE job_id = $1 AND state IN ('running', 'waiting_external'))" ))) .bind(job_id) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; return Ok(if lease_appeared { CancelStep::Again } else { CancelStep::Done(RunningAttemptCancel::NotRunning) }); }; if !self.try_lock_attempt_tx(tx, job_id, run_lease).await? { return Ok(CancelStep::Busy); } self.cancel_phase("attempt_locked").await; // Materializers take the attempt lock, so under it the attempt's // representation is settled. Resolve it with a fresh statement: a // lease that exists now, whatever the receipt kind, goes through the // lease path, whose row lock does not wait. let lease_now: bool = sqlx::query_scalar(audited_sql(format!( "SELECT EXISTS (SELECT 1 FROM {schema}.leases \ WHERE job_id = $1 AND run_lease = $2 \ AND state IN ('running', 'waiting_external'))" ))) .bind(job_id) .bind(run_lease) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; // Under the attempt lock, take the rows without waiting. A row lock // held elsewhere belongs to a closer that may be waiting for this // attempt lock, so contention means: let go and try again. if lease_now { let lease_lock = sqlx::query_scalar::<_, i32>(audited_sql(format!( "SELECT 1 FROM {schema}.leases \ WHERE job_id = $1 AND run_lease = $2 AND state IN ('running', 'waiting_external') \ FOR UPDATE NOWAIT" ))) .bind(job_id) .bind(run_lease) .fetch_optional(tx.as_mut()) .await; return match lease_lock { Err(err) if is_lock_contention_error(&err) => Ok(CancelStep::Busy), Err(err) => Err(map_sqlx_error(err)), Ok(None) => Ok(CancelStep::Again), Ok(Some(_)) => Ok( match self .cancel_materialized_lease_tx(tx, job_id, Some(run_lease)) .await? { Some(result) => { CancelStep::Done(RunningAttemptCancel::Cancelled(Box::new(result))) } None => CancelStep::Again, }, ), }; } // Each kind of receipt is locked by its own statement. A lookup that // locks a claim another transaction has just rewritten rechecks it, // and PostgreSQL 18 frees the statement's run-time partition pruning // state when that lock runs out of rows; a later pruning pass in the // same statement would read freed memory. let mut locked_receipt: Result, sqlx::Error> = Ok(None); for sql in [row_receipt_sql, batch_receipt_sql] { match sqlx::query_as::<_, ReceiptCancelRow>(audited_sql(sql)) .bind(job_id) .fetch_optional(tx.as_mut()) .await { Ok(Some(found)) => { let newer = match &locked_receipt { Ok(Some(current)) => found.1 > current.1, _ => true, }; if newer { locked_receipt = Ok(Some(found)); } } Ok(None) => {} Err(err) => { locked_receipt = Err(err); break; } } } let receipt = match locked_receipt { Err(err) if is_lock_contention_error(&err) => return Ok(CancelStep::Busy), Err(err) => return Err(map_sqlx_error(err)), Ok(Some(receipt)) if receipt.1 == run_lease => receipt, Ok(Some(_)) => return Ok(CancelStep::Again), Ok(None) => { // Closed by another closer, or materialized before this // cancel took the lock; the latter goes through the lease. let materialized: bool = sqlx::query_scalar(audited_sql(format!( "SELECT EXISTS (SELECT 1 FROM {schema}.leases \ WHERE job_id = $1 AND run_lease = $2 \ AND state IN ('running', 'waiting_external'))" ))) .bind(job_id) .bind(run_lease) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; return Ok(if materialized { CancelStep::Again } else { CancelStep::Done(RunningAttemptCancel::NotRunning) }); } }; let ( claim_slot, run_lease, ready_slot, ready_generation, queue, priority, attempt, max_attempts, lane_seq, claimed_at, receipt_id, compact_batch, ) = receipt; // Hydrate the ready row so we can synthesize the done // row with the original args/payload. let ready_match: Option = sqlx::query_as(audited_sql(format!( r#" SELECT ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM {schema}.ready_entries WHERE job_id = $1 AND ready_slot = $2 AND ready_generation = $3 AND queue = $4 AND lane_seq = $5 "# ))) .bind(job_id) .bind(ready_slot) .bind(ready_generation) .bind(&queue) .bind(lane_seq) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some(ready) = ready_match else { // Shouldn't happen: the claim references a ready row. return Ok(CancelStep::Done(RunningAttemptCancel::Left)); }; let done = DoneJobRow { ready_slot, ready_generation, job_id, kind: ready.kind, queue: queue.clone(), args: ready.args, state: JobState::Cancelled, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard: ready.enqueue_shard, run_at: ready.run_at, attempted_at: Some(claimed_at), finalized_at: Utc::now(), created_at: ready.created_at, unique_key: ready.unique_key, unique_states: ready.unique_states, payload: ready.payload, output: None, }; self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Running)) .await?; // Write the closure into the same claim partition. A // compact batch-sourced claim has no lease_claims row, so // it closes into the batch ledger the queue prune gate // counts via compact_count; a deadline lease claim keeps // its explicit closure so the gate balances it against the // lease_claims row. if compact_batch { sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.lease_claim_closure_batches ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) VALUES ( $1, $2, $3, 'cancelled', 1, ARRAY[$4::bigint], int8multirange(int8range($4::bigint, $4::bigint + 1, '[)')), clock_timestamp() ) "# ))) .bind(claim_slot) .bind(ready_slot) .bind(ready_generation) .bind(receipt_id) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } else { sqlx::query(audited_sql(format!( r#" WITH inserted AS ( INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at) VALUES ($1, $2, $3, 'cancelled', clock_timestamp()) ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), marked AS ( UPDATE {schema}.lease_claims AS claims SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) FROM inserted WHERE claims.claim_slot = inserted.claim_slot AND claims.job_id = inserted.job_id AND claims.run_lease = inserted.run_lease RETURNING claims.job_id ) SELECT count(*) FROM marked "# ))) .bind(claim_slot) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } self.notify_cancellation_tx(tx, job_id, run_lease).await?; Ok(CancelStep::Done(RunningAttemptCancel::Cancelled(Box::new( CancelJobTxResult { row: done.into_job_row()?, claim_cursor_advance: None, }, )))) } /// Try the attempt's advisory lock without waiting, marking the /// transaction as an attempt-lock holder when it is taken. async fn try_lock_attempt_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: i64, ) -> Result { let locked: bool = sqlx::query_scalar( "SELECT pg_catalog.pg_try_advisory_xact_lock(pg_catalog.hashtextextended(\ format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", ) .bind(job_id) .bind(run_lease) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; if locked { sqlx::query(audited_sql(format!( "SELECT pg_catalog.set_config('{ATTEMPT_LOCKS_HELD_SETTING}', 'on', true)" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } Ok(locked) } /// SQL true when the ready row `alias` has been claimed. A claim records /// its attempt in its own transaction but moves the lane's claim cursor /// past the row only after commit, so a claimed row can still sit at or /// after the cursor; claims skip such rows on the same evidence. Receipt /// claims leave an attempt batch, and materialized claims a later attempt /// of the job, which is checked whatever this store's mode, since admin /// callers build the store with default configuration. fn ready_row_claimed_sql(&self, alias: &str) -> String { let schema = self.schema(); let mut sql = format!( "EXISTS ( SELECT 1 FROM {schema}.ready_claim_attempt_batches AS attempts WHERE attempts.ready_slot = {alias}.ready_slot AND attempts.ready_generation = {alias}.ready_generation AND attempts.queue = {alias}.queue AND attempts.priority = {alias}.priority AND attempts.enqueue_shard = {alias}.enqueue_shard AND attempts.first_lane_seq <= {alias}.lane_seq AND attempts.next_lane_seq > {alias}.lane_seq AND attempts.lane_ranges @> int8range({alias}.lane_seq, {alias}.lane_seq + 1, '[)') )" ); for table in [ "leases", "done_entries", "deferred_jobs", "dlq_entries", "ready_entries", ] { sql.push_str(&format!( " OR EXISTS ( SELECT 1 FROM {schema}.{table} AS attempt WHERE attempt.job_id = {alias}.job_id AND attempt.run_lease = {alias}.run_lease + 1 )" )); } sql } async fn cancel_job_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, ) -> Result, AwaError> { let schema = self.schema(); self.lock_claim_heads_of_jobs_tx(tx, &[job_id]).await?; let claimed = self.ready_row_claimed_sql("ready"); let ready: Option = sqlx::query_as(audited_sql(format!( r#" WITH target AS ( SELECT ready.* FROM {schema}.ready_entries AS ready JOIN {schema}.queue_claim_heads AS claims ON claims.queue = ready.queue AND claims.priority = ready.priority AND claims.enqueue_shard = ready.enqueue_shard WHERE ready.job_id = $1 AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq AND tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation ) AND NOT ({claimed}) ORDER BY ready.lane_seq DESC LIMIT 1 FOR UPDATE OF ready SKIP LOCKED ), tombstone AS ( INSERT INTO {schema}.ready_tombstones ( ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id ) SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id FROM target ON CONFLICT DO NOTHING ) SELECT ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM target "# ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if let Some(ready) = ready { let done = ready .clone() .into_done_row(JobState::Cancelled, Utc::now(), ready.payload.clone()); self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Available)) .await?; // If the cancelled lane was *exactly* at the claim head, // advance the head past it so the derived cursor-difference count // drops by 1 immediately. // When other unclaimed lanes still sit between claim_seq // and the cancelled lane_seq, leave claim_seq alone — // advancing would skip past those still-claimable rows. // Non-head deletes can leave a bounded dispatcher over-count until // later committed rows on the lane are claimed. let claim_cursor_advance = ClaimCursorAdvance { queue: ready.queue.clone(), priority: ready.priority, enqueue_shard: ready.enqueue_shard, next_seq: ready.lane_seq + 1, only_if_current: Some(ready.lane_seq), evidence: None, }; return Ok(Some(CancelJobTxResult { row: done.into_job_row()?, claim_cursor_advance: Some(claim_cursor_advance), })); } if self.lease_claim_receipts() { match self.cancel_running_attempt_tx(tx, job_id).await? { RunningAttemptCancel::Cancelled(result) => return Ok(Some(*result)), RunningAttemptCancel::Left => return Ok(None), RunningAttemptCancel::NotRunning => {} } } else if let Some(result) = self.cancel_materialized_lease_tx(tx, job_id, None).await? { return Ok(Some(result)); } let deferred: Option = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.deferred_jobs WHERE job_id = $1 AND state IN ('scheduled', 'retryable') RETURNING job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload "# ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if let Some(deferred) = deferred { let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?; // A deferred-cancel never observed a claim, so it has no // shard assignment to inherit. The synthesized terminal row // is parked on shard 0 with a synthetic negative `lane_seq` // that keeps the `done_entries` PK // `(ready_slot, queue, priority, enqueue_shard, lane_seq)` // unique. `ensure_lane` registers shard 0 for the queue so // the lane row exists regardless of producer activity. self.ensure_lane(tx, &deferred.queue, deferred.priority, 0) .await?; let done = DoneJobRow { ready_slot, ready_generation, job_id: deferred.job_id, kind: deferred.kind, queue: deferred.queue.clone(), args: deferred.args, state: JobState::Cancelled, priority: deferred.priority, attempt: deferred.attempt, run_lease: deferred.run_lease, max_attempts: deferred.max_attempts, lane_seq: -deferred.job_id, enqueue_shard: 0, run_at: deferred.run_at, attempted_at: deferred.attempted_at, finalized_at: Utc::now(), created_at: deferred.created_at, unique_key: deferred.unique_key, unique_states: deferred.unique_states, payload: deferred.payload, output: None, }; self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(deferred.state)) .await?; return Ok(Some(CancelJobTxResult { row: done.into_job_row()?, claim_cursor_advance: None, })); } Ok(None) } pub async fn cancel_job(&self, pool: &PgPool, job_id: i64) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self.cancel_job_tx(&mut tx, job_id).await?; tx.commit().await.map_err(map_sqlx_error)?; if let Some(result) = result { if let Some(advance) = result.claim_cursor_advance.as_ref() { self.advance_claim_cursors(pool, std::slice::from_ref(advance)) .await; } Ok(Some(result.row)) } else { Ok(None) } } /// Transaction-scoped cancel used by [`crate::admin::cancel_tx`]. Runs the /// full cancellation on the caller's transaction and returns the cancelled /// row. Unlike [`Self::cancel_job`] this does NOT perform the post-commit /// claim-cursor advance, so the derived queue depth may briefly over-count /// by one until later committed rows on the lane are claimed. pub(crate) async fn cancel_job_in_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, ) -> Result, AwaError> { Ok(self .cancel_job_tx(tx, job_id) .await? .map(|result| result.row)) } pub async fn cancel_jobs_by_ids( &self, pool: &PgPool, ids: &[i64], ) -> Result, AwaError> { if ids.is_empty() { return Ok(Vec::new()); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; self.lock_claim_heads_of_jobs_tx(&mut tx, ids).await?; let mut rows = Vec::with_capacity(ids.len()); let mut claim_cursor_advances = Vec::new(); for job_id in ids { if let Some(result) = self.cancel_job_tx(&mut tx, *job_id).await? { if let Some(advance) = result.claim_cursor_advance { claim_cursor_advances.push(advance); } rows.push(result.row); } } tx.commit().await.map_err(map_sqlx_error)?; self.advance_claim_cursors(pool, &claim_cursor_advances) .await; Ok(rows) } /// Cancel every available, scheduled, retryable and waiting job in /// `queue` and its stripes (`queue#N`), returning how many were cancelled. /// /// Jobs enqueued after the drain starts are left alone, and so are running /// jobs. Scheduled, retryable and available jobs are cancelled a chunk at /// a time, each chunk in its own transaction, so a large backlog costs a /// few statements per chunk rather than per job and holds row locks for /// one chunk at a time. Rows another transaction has locked are skipped. /// Waiting jobs go through [`Self::cancel_jobs_by_ids`]. pub(crate) async fn drain_queue(&self, pool: &PgPool, queue: &str) -> Result { let schema = self.schema(); let job_id_bound: i64 = sqlx::query_scalar(audited_sql(format!( "SELECT {schema}.sequence_next_value('job_id_seq')" ))) .fetch_one(pool) .await .map_err(map_sqlx_error)?; let mut drained = 0u64; let mut after_job_id = i64::MIN; loop { let mut tx = pool .begin_with(DRAIN_TRANSACTION_BEGIN) .await .map_err(map_sqlx_error)?; let (cancelled, last_job_id) = self .cancel_deferred_chunk_tx(&mut tx, queue, after_job_id, job_id_bound) .await?; tx.commit().await.map_err(map_sqlx_error)?; drained += cancelled as u64; if (cancelled as i64) < DRAIN_BATCH_SIZE { break; } after_job_id = last_job_id; } // Read after the deferred pass, so jobs promoted while it ran are // drained from the ready lanes. let mut lanes: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT claims.queue, claims.priority, claims.enqueue_shard, {schema}.sequence_next_value(claims.seq_name) AS from_seq, {schema}.sequence_next_value(enqueues.seq_name) AS to_seq FROM {schema}.queue_claim_heads AS claims JOIN {schema}.queue_enqueue_heads AS enqueues ON enqueues.queue = claims.queue AND enqueues.priority = claims.priority AND enqueues.enqueue_shard = claims.enqueue_shard WHERE {queues} "#, queues = queue_and_stripes_sql("claims.queue", "$1"), ))) .bind(queue) .fetch_all(pool) .await .map_err(map_sqlx_error)?; lanes.retain(|lane| lane.from_seq < lane.to_seq); let drained_lanes = lanes.clone(); while !lanes.is_empty() { let mut tx = pool .begin_with(DRAIN_TRANSACTION_BEGIN) .await .map_err(map_sqlx_error)?; self.lock_claim_heads_tx(&mut tx, &lanes).await?; let cancelled = self.cancel_available_chunk_tx(&mut tx, &mut lanes).await?; tx.commit().await.map_err(map_sqlx_error)?; self.skip_spent_claim_prefixes(pool, &lanes).await?; drained += cancelled as u64; if (cancelled as i64) < DRAIN_BATCH_SIZE { break; } lanes.retain(|lane| lane.from_seq < lane.to_seq); } self.skip_spent_claim_prefixes(pool, &drained_lanes).await?; let waiting: Vec = sqlx::query_scalar(audited_sql(format!( r#" SELECT job_id FROM {schema}.leases WHERE {queues} AND state = 'waiting_external' AND job_id < $2 ORDER BY job_id "#, queues = queue_and_stripes_sql("queue", "$1"), ))) .bind(queue) .bind(job_id_bound) .fetch_all(pool) .await .map_err(map_sqlx_error)?; drained += self.cancel_jobs_by_ids(pool, &waiting).await?.len() as u64; Ok(drained) } /// Cancel up to [`DRAIN_BATCH_SIZE`] scheduled and retryable jobs of /// `queue` with ids in `(after_job_id, job_id_bound)`, returning how many /// were cancelled and the largest cancelled id. async fn cancel_deferred_chunk_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, queue: &str, after_job_id: i64, job_id_bound: i64, ) -> Result<(usize, i64), AwaError> { let schema = self.schema(); let deferred: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.deferred_jobs WHERE job_id IN ( SELECT job_id FROM {schema}.deferred_jobs WHERE {queues} AND state IN ('scheduled', 'retryable') AND job_id > $2 AND job_id < $3 ORDER BY job_id LIMIT $4 FOR UPDATE SKIP LOCKED ) RETURNING job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload "#, queues = queue_and_stripes_sql("queue", "$1"), ))) .bind(queue) .bind(after_job_id) .bind(job_id_bound) .bind(DRAIN_BATCH_SIZE) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some(last_job_id) = deferred.iter().map(|row| row.job_id).max() else { return Ok((0, after_job_id)); }; let cancelled = deferred.len(); let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?; let lanes: BTreeSet<(String, i16)> = deferred .iter() .map(|row| (row.queue.clone(), row.priority)) .collect(); for (queue, priority) in &lanes { self.ensure_lane(tx, queue, *priority, 0).await?; } let finalized_at = Utc::now(); let mut by_state: HashMap> = HashMap::new(); for deferred in deferred { // Same shape as a single deferred cancel: shard 0 and a negative // lane_seq keep the done_entries key unique. let done = DoneJobRow { ready_slot, ready_generation, job_id: deferred.job_id, kind: deferred.kind, queue: deferred.queue, args: deferred.args, state: JobState::Cancelled, priority: deferred.priority, attempt: deferred.attempt, run_lease: deferred.run_lease, max_attempts: deferred.max_attempts, lane_seq: -deferred.job_id, enqueue_shard: 0, run_at: deferred.run_at, attempted_at: deferred.attempted_at, finalized_at, created_at: deferred.created_at, unique_key: deferred.unique_key, unique_states: deferred.unique_states, payload: deferred.payload, output: None, }; by_state.entry(deferred.state).or_default().push(done); } for (old_state, rows) in by_state { self.insert_done_rows_tx(tx, &rows, Some(old_state)).await?; } Ok((cancelled, last_job_id)) } /// Cancel up to [`DRAIN_BATCH_SIZE`] available jobs from `lanes`, /// returning how many were cancelled and the claim-cursor advances that /// skip them. async fn cancel_available_chunk_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, lanes: &mut [DrainLane], ) -> Result { let schema = self.schema(); let claimed = self.ready_row_claimed_sql("ready"); // Unclaimed rows are found the way a claim finds them: from each // lane's claim cursor through its segments, then along the ready key // within each segment's ring slot. The ready key leads with the ring // slot, so a lane range without it scans every retained ready row. // Rows are locked before the limit applies, so a row another // transaction holds does not shorten the chunk. The tombstone check // shares one NOT with the claim check: on its own, NOT EXISTS is // planned as an anti join that reads the lane's tombstones again for // every segment, and a lane enqueued one job at a time has a segment // per row. Only rows whose // tombstone this statement inserts are cancelled, so a concurrent // cancel of the same row cannot produce a second done row. let ready: Vec = sqlx::query_as(audited_sql(format!( r#" WITH heads AS ( SELECT lanes.queue, lanes.priority, lanes.enqueue_shard, GREATEST(lanes.from_seq, {schema}.sequence_next_value(claims.seq_name)) AS claim_seq, lanes.to_seq AS next_seq FROM unnest($1::text[], $2::smallint[], $3::smallint[], $4::bigint[], $5::bigint[]) AS lanes(queue, priority, enqueue_shard, from_seq, to_seq) JOIN {schema}.queue_claim_heads AS claims ON claims.queue = lanes.queue AND claims.priority = lanes.priority AND claims.enqueue_shard = lanes.enqueue_shard ), target AS ( SELECT ready.* FROM heads CROSS JOIN LATERAL ( SELECT segment.ready_slot, segment.ready_generation, GREATEST(segment.first_lane_seq, heads.claim_seq) AS from_seq, LEAST(segment.next_lane_seq, heads.next_seq) AS to_seq FROM {schema}.ready_segments AS segment WHERE segment.queue = heads.queue AND segment.priority = heads.priority AND segment.enqueue_shard = heads.enqueue_shard AND segment.next_lane_seq > heads.claim_seq AND segment.first_lane_seq < heads.next_seq ORDER BY segment.next_lane_seq ) AS segment CROSS JOIN LATERAL ( SELECT ready.* FROM {schema}.ready_entries AS ready WHERE ready.ready_slot = segment.ready_slot AND ready.ready_generation = segment.ready_generation AND ready.queue = heads.queue AND ready.priority = heads.priority AND ready.enqueue_shard = heads.enqueue_shard AND ready.lane_seq >= segment.from_seq AND ready.lane_seq < segment.to_seq AND NOT ( EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation AND tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq ) OR {claimed} ) ORDER BY ready.lane_seq FOR UPDATE OF ready SKIP LOCKED ) AS ready LIMIT $6 ), tombstone AS ( INSERT INTO {schema}.ready_tombstones ( ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id ) SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id FROM target ON CONFLICT DO NOTHING RETURNING ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq ) SELECT target.ready_slot, target.ready_generation, target.job_id, target.kind, target.queue, target.args, target.priority, target.attempt, target.run_lease, target.max_attempts, target.lane_seq, target.enqueue_shard, target.run_at, target.attempted_at, target.created_at, target.unique_key, target.unique_states, COALESCE(target.payload, '{{}}'::jsonb) AS payload FROM target JOIN tombstone ON tombstone.ready_slot = target.ready_slot AND tombstone.ready_generation = target.ready_generation AND tombstone.queue = target.queue AND tombstone.priority = target.priority AND tombstone.enqueue_shard = target.enqueue_shard AND tombstone.lane_seq = target.lane_seq "# ))) .bind(lanes.iter().map(|lane| lane.queue.as_str()).collect::>()) .bind(lanes.iter().map(|lane| lane.priority).collect::>()) .bind(lanes.iter().map(|lane| lane.enqueue_shard).collect::>()) .bind(lanes.iter().map(|lane| lane.from_seq).collect::>()) .bind(lanes.iter().map(|lane| lane.to_seq).collect::>()) .bind(DRAIN_BATCH_SIZE) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if ready.is_empty() { return Ok(0); } // The next chunk resumes after the rows this one cancelled, so each // lane is read once even where a claimed or locked row stays behind. for row in &ready { if let Some(lane) = lanes.iter_mut().find(|lane| { lane.queue == row.queue && lane.priority == row.priority && lane.enqueue_shard == row.enqueue_shard }) { lane.from_seq = lane.from_seq.max(row.lane_seq + 1); } } let finalized_at = Utc::now(); let done_rows: Vec = ready .into_iter() .map(|row| { let payload = row.payload.clone(); row.into_done_row(JobState::Cancelled, finalized_at, payload) }) .collect(); self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Available)) .await?; Ok(done_rows.len()) } /// Lock the claim heads of the lanes holding ready rows of `job_ids`, so /// the transaction's next statement sees any claim of those rows that was /// in progress. See [`Self::lock_claim_heads_tx`]. async fn lock_claim_heads_of_jobs_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_ids: &[i64], ) -> Result<(), AwaError> { let schema = self.schema(); sqlx::query(audited_sql(format!( r#" SELECT 1 FROM {schema}.queue_claim_heads AS claims WHERE (claims.queue, claims.priority, claims.enqueue_shard) IN ( SELECT ready.queue, ready.priority, ready.enqueue_shard FROM {schema}.ready_entries AS ready WHERE ready.job_id = ANY($1) ) ORDER BY claims.queue COLLATE "C", claims.priority, claims.enqueue_shard FOR UPDATE OF claims "# ))) .bind(job_ids) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Lock the claim heads of `lanes`. A claim holds its lane's head until /// it commits, so once this returns the transaction's next statement sees /// every claim on these lanes. Heads are locked in the byte order that /// cursor advances use. async fn lock_claim_heads_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, lanes: &[DrainLane], ) -> Result<(), AwaError> { let schema = self.schema(); sqlx::query(audited_sql(format!( r#" SELECT 1 FROM {schema}.queue_claim_heads AS claims JOIN unnest($1::text[], $2::smallint[], $3::smallint[]) AS lanes(queue, priority, enqueue_shard) ON claims.queue = lanes.queue AND claims.priority = lanes.priority AND claims.enqueue_shard = lanes.enqueue_shard ORDER BY claims.queue COLLATE "C", claims.priority, claims.enqueue_shard FOR UPDATE OF claims "# ))) .bind( lanes .iter() .map(|lane| lane.queue.as_str()) .collect::>(), ) .bind(lanes.iter().map(|lane| lane.priority).collect::>()) .bind( lanes .iter() .map(|lane| lane.enqueue_shard) .collect::>(), ) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Move each lane's claim cursor over cancelled and claimed rows, up to /// the lane's `to_seq`. A claim crosses cancelled rows only one claim /// batch per poll, so without this the jobs behind a drained range wait /// for many polls before a worker reaches them. async fn skip_spent_claim_prefixes( &self, pool: &PgPool, lanes: &[DrainLane], ) -> Result<(), AwaError> { let schema = self.schema(); let claimed = self.ready_row_claimed_sql("ready"); let mut lanes = lanes.to_vec(); while !lanes.is_empty() { let mut tx = pool .begin_with(DRAIN_TRANSACTION_BEGIN) .await .map_err(map_sqlx_error)?; self.lock_claim_heads_tx(&mut tx, &lanes).await?; // Each window is a range of positions, so a live row inside it is // always seen. Untombstoned rows are checked for a claim only from // the first, a few at a time: claims past the cursor are the few // whose claimer has not yet moved it. The tombstone check is a // column behind OFFSET 0, so it is not pushed down as a NOT EXISTS, // for the reason given in `cancel_available_chunk_tx`. A gap before a visible row // is an enqueue that rolled back: enqueues on a lane commit in // position order. let unfinished: Vec<(String, i16, i16)> = sqlx::query_as(audited_sql(format!( r#" WITH windows AS ( SELECT lanes.queue, lanes.priority, lanes.enqueue_shard, claims.seq_name, {schema}.sequence_next_value(claims.seq_name) AS claim_seq, LEAST({schema}.sequence_next_value(claims.seq_name) + $5, lanes.to_seq) AS window_end, lanes.to_seq FROM unnest($1::text[], $2::smallint[], $3::smallint[], $4::bigint[]) AS lanes(queue, priority, enqueue_shard, to_seq) JOIN {schema}.queue_claim_heads AS claims ON claims.queue = lanes.queue AND claims.priority = lanes.priority AND claims.enqueue_shard = lanes.enqueue_shard ), scanned AS ( SELECT windows.*, unclaimed.first_live, COALESCE( unclaimed.first_live, CASE WHEN unclaimed.checked = $6 THEN unclaimed.last_checked + 1 ELSE visible.last_seen + 1 END ) AS next_seq FROM windows CROSS JOIN LATERAL ( SELECT max(ready.lane_seq) AS last_seen FROM {schema}.ready_segments AS segment CROSS JOIN LATERAL ( SELECT ready.lane_seq FROM {schema}.ready_entries AS ready WHERE ready.ready_slot = segment.ready_slot AND ready.ready_generation = segment.ready_generation AND ready.queue = windows.queue AND ready.priority = windows.priority AND ready.enqueue_shard = windows.enqueue_shard AND ready.lane_seq >= GREATEST(segment.first_lane_seq, windows.claim_seq) AND ready.lane_seq < LEAST(segment.next_lane_seq, windows.window_end) ORDER BY ready.lane_seq ) AS ready WHERE segment.queue = windows.queue AND segment.priority = windows.priority AND segment.enqueue_shard = windows.enqueue_shard AND segment.next_lane_seq > windows.claim_seq AND segment.first_lane_seq < windows.window_end ) AS visible CROSS JOIN LATERAL ( SELECT min(ready.lane_seq) FILTER (WHERE NOT ({claimed})) AS first_live, max(ready.lane_seq) AS last_checked, count(*) AS checked FROM ( SELECT ready.* FROM {schema}.ready_segments AS segment CROSS JOIN LATERAL ( SELECT ready.*, EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation AND tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq ) AS tombstoned FROM {schema}.ready_entries AS ready WHERE ready.ready_slot = segment.ready_slot AND ready.ready_generation = segment.ready_generation AND ready.queue = windows.queue AND ready.priority = windows.priority AND ready.enqueue_shard = windows.enqueue_shard AND ready.lane_seq >= GREATEST(segment.first_lane_seq, windows.claim_seq) AND ready.lane_seq < LEAST(segment.next_lane_seq, windows.window_end) ORDER BY ready.lane_seq OFFSET 0 ) AS ready WHERE segment.queue = windows.queue AND segment.priority = windows.priority AND segment.enqueue_shard = windows.enqueue_shard AND segment.next_lane_seq > windows.claim_seq AND segment.first_lane_seq < windows.window_end AND NOT ready.tombstoned ORDER BY ready.lane_seq LIMIT $6 ) AS ready ) AS unclaimed WHERE windows.claim_seq < windows.to_seq ), advanced AS MATERIALIZED ( SELECT scanned.*, {schema}.set_sequence_next(scanned.seq_name, scanned.next_seq) AS moved FROM scanned WHERE scanned.next_seq > scanned.claim_seq ) SELECT queue, priority, enqueue_shard FROM advanced WHERE first_live IS NULL AND next_seq < to_seq "# ))) .bind(lanes.iter().map(|lane| lane.queue.as_str()).collect::>()) .bind(lanes.iter().map(|lane| lane.priority).collect::>()) .bind(lanes.iter().map(|lane| lane.enqueue_shard).collect::>()) .bind(lanes.iter().map(|lane| lane.to_seq).collect::>()) .bind(DRAIN_BATCH_SIZE) .bind(DRAIN_CLAIM_CHECK_ROWS) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; lanes.retain(|lane| { unfinished.iter().any(|(queue, priority, enqueue_shard)| { &lane.queue == queue && lane.priority == *priority && lane.enqueue_shard == *enqueue_shard }) }); } Ok(()) } pub async fn set_priority( &self, pool: &PgPool, job_id: i64, priority: i16, ) -> Result { if !(1..=4).contains(&priority) { return Err(AwaError::Validation( "priority must be between 1 and 4".to_string(), )); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self.set_priority_tx(&mut tx, job_id, priority).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } pub async fn set_priority_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, priority: i16, ) -> Result { if !(1..=4).contains(&priority) { return Err(AwaError::Validation( "priority must be between 1 and 4".to_string(), )); } if self .update_deferred_batch_fields_tx(tx, job_id, None, Some(priority)) .await? { return Ok(true); } let result = self .move_ready_batch_fields_tx(tx, job_id, None, Some(priority)) .await?; Ok(result.moved) } pub async fn move_queue( &self, pool: &PgPool, job_id: i64, queue: &str, priority: Option, ) -> Result { if queue.is_empty() || queue.len() > 200 { return Err(AwaError::Validation( "destination queue must be 1..=200 characters".to_string(), )); } if let Some(priority) = priority { if !(1..=4).contains(&priority) { return Err(AwaError::Validation( "priority must be between 1 and 4".to_string(), )); } } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self.move_queue_tx(&mut tx, job_id, queue, priority).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } pub async fn move_queue_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, queue: &str, priority: Option, ) -> Result { if queue.is_empty() || queue.len() > 200 { return Err(AwaError::Validation( "destination queue must be 1..=200 characters".to_string(), )); } if let Some(priority) = priority { if !(1..=4).contains(&priority) { return Err(AwaError::Validation( "priority must be between 1 and 4".to_string(), )); } } if self .update_deferred_batch_fields_tx(tx, job_id, Some(queue), priority) .await? { return Ok(true); } let result = self .move_ready_batch_fields_tx(tx, job_id, Some(queue), priority) .await?; Ok(result.moved) } async fn update_deferred_batch_fields_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, queue: Option<&str>, priority: Option, ) -> Result { let schema = self.schema(); let row: Option = sqlx::query_as(audited_sql(format!( r#" SELECT job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM {schema}.deferred_jobs WHERE job_id = $1 AND state = 'scheduled' FOR UPDATE SKIP LOCKED "# ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some(row) = row else { return Ok(false); }; let old_queue = row.queue.clone(); let old_priority = row.priority; let requested_queue = queue.unwrap_or(&old_queue); let old_logical_queue = self.logical_queue_name(&old_queue).to_string(); let new_queue = if queue.is_some() && requested_queue != old_queue && requested_queue != old_logical_queue { self.queue_stripe_for_enqueue(requested_queue, &row.unique_key, row.job_id) } else { old_queue.clone() }; let new_priority = priority.unwrap_or(old_priority); if new_queue == old_queue && new_priority == old_priority { return Ok(false); } let mut payload = RuntimePayload::from_json(row.payload)?; let metadata = payload.metadata.as_object_mut().ok_or_else(|| { AwaError::Validation("queue storage payload metadata must be a JSON object".to_string()) })?; if queue.is_some() { metadata .entry("_awa_original_queue".to_string()) .or_insert_with(|| serde_json::Value::from(old_logical_queue)); } if priority.is_some() { metadata .entry("_awa_original_priority".to_string()) .or_insert_with(|| serde_json::Value::from(i64::from(old_priority))); } sqlx::query(audited_sql(format!( r#" UPDATE {schema}.deferred_jobs SET queue = $2, priority = $3, payload = $4 WHERE job_id = $1 "# ))) .bind(job_id) .bind(new_queue) .bind(new_priority) .bind(storage_payload(&payload.into_json())) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(true) } async fn move_ready_batch_fields_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, queue: Option<&str>, priority: Option, ) -> Result { let schema = self.schema(); self.lock_claim_heads_of_jobs_tx(tx, &[job_id]).await?; let claimed = self.ready_row_claimed_sql("ready"); let ready: Option = sqlx::query_as(audited_sql(format!( r#" WITH target AS ( SELECT ready.* FROM {schema}.ready_entries AS ready JOIN {schema}.queue_claim_heads AS claims ON claims.queue = ready.queue AND claims.priority = ready.priority AND claims.enqueue_shard = ready.enqueue_shard WHERE ready.job_id = $1 AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq AND tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation ) AND NOT ({claimed}) ORDER BY ready.lane_seq DESC LIMIT 1 FOR UPDATE OF ready SKIP LOCKED ) SELECT ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM target "# ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some(ready) = ready else { return Ok(ReadyBatchMoveResult { moved: false }); }; let old_queue = ready.queue.clone(); let old_priority = ready.priority; let requested_queue = queue.unwrap_or(&old_queue); let old_logical_queue = self.logical_queue_name(&old_queue).to_string(); let new_queue = if queue.is_some() && requested_queue != old_queue && requested_queue != old_logical_queue { self.queue_stripe_for_enqueue(requested_queue, &ready.unique_key, ready.job_id) } else { old_queue.clone() }; let new_priority = priority.unwrap_or(old_priority); if new_queue == old_queue && new_priority == old_priority { return Ok(ReadyBatchMoveResult { moved: false }); } sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.ready_tombstones ( ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id ) VALUES ($1, $2, $3, $4, $5, $6, $7) ON CONFLICT DO NOTHING "# ))) .bind(ready.ready_slot) .bind(ready.ready_generation) .bind(&ready.queue) .bind(ready.priority) .bind(ready.enqueue_shard) .bind(ready.lane_seq) .bind(ready.job_id) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; let mut payload = RuntimePayload::from_json(ready.payload.clone())?; let metadata = payload.metadata.as_object_mut().ok_or_else(|| { AwaError::Validation("queue storage payload metadata must be a JSON object".to_string()) })?; if queue.is_some() { metadata .entry("_awa_original_queue".to_string()) .or_insert_with(|| serde_json::Value::from(old_logical_queue)); } if priority.is_some() { metadata .entry("_awa_original_priority".to_string()) .or_insert_with(|| serde_json::Value::from(i64::from(old_priority))); } let notify_queue = new_queue.clone(); let ready_row = ready.into_existing_ready_row(new_queue, new_priority, payload.into_json()); self.insert_existing_ready_rows_tx(tx, vec![ready_row], Some(JobState::Available)) .await?; self.notify_queues_tx(tx, std::iter::once(notify_queue)) .await?; Ok(ReadyBatchMoveResult { moved: true }) } pub async fn age_waiting_priorities( &self, pool: &PgPool, aging_interval: Duration, limit: i64, ) -> Result, AwaError> { if limit <= 0 { return Ok(Vec::new()); } let cutoff = Utc::now() - TimeDelta::from_std(aging_interval) .map_err(|err| AwaError::Validation(format!("invalid aging interval: {err}")))?; let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let moved: Vec = sqlx::query_as(audited_sql(format!( r#" WITH target AS ( SELECT ready.* FROM {schema}.ready_entries AS ready JOIN {schema}.queue_claim_heads AS claims ON claims.queue = ready.queue AND claims.priority = ready.priority AND claims.enqueue_shard = ready.enqueue_shard WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) AND ready.priority > 1 AND ready.run_at <= $1 AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.queue = ready.queue AND tomb.priority = ready.priority AND tomb.enqueue_shard = ready.enqueue_shard AND tomb.lane_seq = ready.lane_seq AND tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation ) ORDER BY ready.run_at ASC, ready.lane_seq ASC LIMIT $2 FOR UPDATE OF ready SKIP LOCKED ), tombstones AS ( INSERT INTO {schema}.ready_tombstones ( ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id ) SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id FROM target ON CONFLICT DO NOTHING ) SELECT ready_slot, ready_generation, job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM target "# ))) .bind(cutoff) .bind(limit) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if moved.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(Vec::new()); } let mut ids = Vec::with_capacity(moved.len()); let mut queues = BTreeSet::new(); let mut ready_rows = Vec::with_capacity(moved.len()); for row in moved { ids.push(row.job_id); queues.insert(row.queue.clone()); let mut payload = RuntimePayload::from_json(row.payload)?; let metadata = payload.metadata.as_object_mut().ok_or_else(|| { AwaError::Validation( "queue storage payload metadata must be a JSON object".to_string(), ) })?; metadata .entry("_awa_original_priority".to_string()) .or_insert_with(|| serde_json::Value::from(i64::from(row.priority))); ready_rows.push(ExistingReadyRow { job_id: row.job_id, kind: row.kind, queue: row.queue, args: row.args, priority: row.priority - 1, attempt: row.attempt, run_lease: row.run_lease, max_attempts: row.max_attempts, run_at: row.run_at, attempted_at: row.attempted_at, created_at: row.created_at, unique_key: row.unique_key, unique_states: row.unique_states, payload: payload.into_json(), }); } // Aging tombstones the source lane without moving its claim cursor, // then re-inserts on the target lane. The source lane can temporarily // over-count by `moved`; later claims advance over the tombstones. // The drift is bounded by aging rate × poll interval. self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(JobState::Available)) .await?; self.notify_queues_tx(&mut tx, queues).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(ids) } fn with_progress( payload: serde_json::Value, progress: Option, ) -> Result { let mut payload = RuntimePayload::from_json(payload)?; payload.set_progress(progress); Ok(payload.into_json()) } async fn take_callback_result( &self, pool: &PgPool, job_id: i64, run_lease: i64, ) -> Result { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let mut row: Option = sqlx::query_as(audited_sql(format!( r#" SELECT job_id, run_lease, progress, callback_filter, callback_on_complete, callback_on_fail, callback_transform, callback_result FROM {} WHERE job_id = $1 AND run_lease = $2 FOR UPDATE "#, self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some(mut row) = row.take() else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(serde_json::Value::Null); }; let result = row .callback_result .take() .unwrap_or(serde_json::Value::Null); if row.progress.is_none() && row.callback_filter.is_none() && row.callback_on_complete.is_none() && row.callback_on_fail.is_none() && row.callback_transform.is_none() { sqlx::query(audited_sql(format!( "DELETE FROM {} WHERE job_id = $1 AND run_lease = $2", self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } else { sqlx::query(audited_sql(format!( "UPDATE {} SET callback_result = NULL, updated_at = clock_timestamp() WHERE job_id = $1 AND run_lease = $2", self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } async fn backoff_at_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, attempt: i16, max_attempts: i16, ) -> Result, AwaError> { sqlx::query_scalar("SELECT clock_timestamp() + awa.backoff_duration($1, $2)") .bind(attempt) .bind(max_attempts) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error) } async fn notify_queues_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, queues: impl IntoIterator, ) -> Result<(), AwaError> { // BTreeSet dedups so we never emit the same NOTIFY twice per // transaction. Multi-queue enqueues fold into a single round-trip via // `unnest($1::text[])` rather than one round-trip per channel. let channels: Vec = queues .into_iter() .map(|queue| format!("awa:{}", self.logical_queue_name(&queue))) .collect::>() .into_iter() .collect(); if channels.is_empty() { return Ok(()); } sqlx::query("SELECT pg_notify(channel, '') FROM unnest($1::text[]) AS channel") .bind(&channels) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Take the attempts' completion locks, then lock their row-local claims. /// /// The statements that follow lock these claims again. A statement that /// locks a claim another transaction has just rewritten rechecks it, and /// PostgreSQL 18 frees that statement's run-time partition pruning state /// once the lock runs out of rows, so a later pruning pass in it reads /// freed memory. Locked here, any recheck happens in a statement that /// prunes nothing, and the later statements find the claims already held. async fn lock_receipt_attempts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, jobs: &[(i64, i64)], ) -> Result<(), AwaError> { if jobs.is_empty() { return Ok(()); } self.lock_receipt_attempt_keys_tx(tx, jobs).await?; let unique_jobs: BTreeSet<(i64, i64)> = jobs.iter().copied().collect(); let job_ids: Vec = unique_jobs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = unique_jobs .iter() .map(|(_, run_lease)| *run_lease) .collect(); let schema = self.schema(); sqlx::query(audited_sql(format!( r#" SELECT 1 FROM {schema}.lease_claims AS claims JOIN unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease) ON input.job_id = claims.job_id AND input.run_lease = claims.run_lease ORDER BY claims.job_id, claims.run_lease FOR UPDATE OF claims "# ))) .bind(&job_ids) .bind(&run_leases) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } /// Serialize the receipt paths that complete or close the same attempts. async fn lock_receipt_attempt_keys_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, jobs: &[(i64, i64)], ) -> Result<(), AwaError> { if jobs.is_empty() { return Ok(()); } let unique_jobs: BTreeSet<(i64, i64)> = jobs.iter().copied().collect(); let job_ids: Vec = unique_jobs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = unique_jobs .iter() .map(|(_, run_lease)| *run_lease) .collect(); sqlx::query(audited_sql(format!( r#" WITH locks AS MATERIALIZED ( SELECT DISTINCT pg_catalog.hashtextextended( format('awa.receipt.complete:%s:%s', job_id, run_lease), 0 ) AS lock_key FROM unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease) ORDER BY lock_key ), acquired AS MATERIALIZED ( SELECT pg_catalog.pg_advisory_xact_lock(lock_key) FROM locks ) SELECT count(*) FROM acquired CROSS JOIN ( SELECT pg_catalog.set_config('{ATTEMPT_LOCKS_HELD_SETTING}', 'on', true) ) AS marked "# ))) .bind(&job_ids) .bind(&run_leases) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(()) } async fn ensure_running_leases_from_receipts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, jobs: &[(i64, i64)], ) -> Result { if jobs.is_empty() { return Ok(0); } self.lock_receipt_attempts_tx(tx, jobs).await?; let schema = self.schema(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); let inserted: i64 = sqlx::query_scalar(audited_sql(format!( r#" WITH inflight(job_id, run_lease) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) ), -- #371: MATERIALIZED pins the max-generation ledger read to a -- single evaluation. Without it PG18 inlines this one-reference -- CTE into the `CROSS JOIN lease_ring` below, and the planner -- re-runs the `ORDER BY generation DESC LIMIT 1` index descent -- once per materialized claim row (loops=N). The pre-ledger -- `lease_ring_state` singleton was a single cached heap read; the -- ledger descent cost scales with unfolded ledger depth, so the -- per-row re-execution regressed the completion hot path. lease_ring AS MATERIALIZED ( SELECT slot AS lease_slot, generation AS lease_generation FROM {schema}.lease_ring_rotations ORDER BY generation DESC LIMIT 1 ), row_claim_refs AS ( -- Source claim metadata directly from the partitioned -- lease_claims table anti-joined against every durable -- closure evidence shape. SELECT claims.claim_slot, claims.job_id, claims.run_lease, claims.ready_slot, claims.ready_generation, claims.queue, claims.priority, claims.attempt, claims.max_attempts, claims.lane_seq, claims.enqueue_shard, claims.claimed_at, claims.deadline_at FROM {schema}.lease_claims AS claims JOIN inflight ON inflight.job_id = claims.job_id AND inflight.run_lease = claims.run_lease WHERE NOT {closed_evidence} FOR UPDATE OF claims ), batch_claim_refs AS ( SELECT claim_batches.claim_slot, items.job_id, items.run_lease, claim_batches.ready_slot, claim_batches.ready_generation, claim_batches.queue, claim_batches.priority, items.attempt, items.max_attempts, items.lane_seq, claim_batches.enqueue_shard, claim_batches.claimed_at, claim_batches.deadline_at FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids, claim_batches.lane_seqs, claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) JOIN inflight ON inflight.job_id = items.job_id AND inflight.run_lease = items.run_lease WHERE NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) FOR UPDATE OF claim_batches ), claim_refs AS ( SELECT * FROM row_claim_refs UNION ALL SELECT * FROM batch_claim_refs ), already_live AS ( SELECT claim_refs.job_id, claim_refs.run_lease FROM claim_refs WHERE EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claim_refs.job_id AND lease.run_lease = claim_refs.run_lease ) ), inserted AS ( INSERT INTO {schema}.leases ( lease_slot, lease_generation, ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at ) SELECT lease_ring.lease_slot, lease_ring.lease_generation, claim_refs.ready_slot, claim_refs.ready_generation, claim_refs.job_id, claim_refs.queue, 'running'::awa.job_state, claim_refs.priority, claim_refs.attempt, claim_refs.run_lease, claim_refs.max_attempts, claim_refs.lane_seq, claim_refs.enqueue_shard, clock_timestamp(), -- Preserve the per-claim deadline so the lease-side -- deadline rescue path picks up materialized claims -- without an extra hop. NULL when receipts mode is -- on with `deadline_duration = 0` (the short-job -- shape that needs no deadline at all). claim_refs.deadline_at, claim_refs.claimed_at FROM claim_refs CROSS JOIN lease_ring WHERE NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claim_refs.job_id AND lease.run_lease = claim_refs.run_lease ) RETURNING job_id, run_lease ), marked AS ( UPDATE {schema}.lease_claims AS claims SET materialized_at = clock_timestamp() FROM ( SELECT job_id, run_lease FROM inserted UNION SELECT job_id, run_lease FROM already_live ) AS moved WHERE claims.job_id = moved.job_id AND claims.run_lease = moved.run_lease RETURNING claims.job_id ) SELECT count(*)::bigint FROM ( SELECT job_id, run_lease FROM inserted UNION SELECT job_id, run_lease FROM already_live ) AS moved "# ))) .bind(&job_ids) .bind(&run_leases) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(inserted as usize) } async fn ensure_mutable_running_attempt_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: i64, ) -> Result<(), AwaError> { if self.lease_claim_receipts() { self.ensure_running_leases_from_receipts_tx(tx, &[(job_id, run_lease)]) .await?; } Ok(()) } async fn upsert_attempt_state_from_receipts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, jobs: &[(i64, i64)], ) -> Result { if jobs.is_empty() { return Ok(0); } self.lock_receipt_attempts_tx(tx, jobs).await?; let schema = self.schema(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); let updated: i64 = sqlx::query_scalar(audited_sql(format!( r#" WITH inflight(job_id, run_lease) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) ), row_claim_refs AS ( -- Source open-claim identity from lease_claims -- anti-joined against every durable closure evidence -- shape. SELECT claims.job_id, claims.run_lease FROM {schema}.lease_claims AS claims JOIN inflight ON inflight.job_id = claims.job_id AND inflight.run_lease = claims.run_lease WHERE NOT {closed_evidence} FOR UPDATE OF claims ), batch_claim_refs AS ( SELECT items.job_id, items.run_lease FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids ) AS items(job_id, run_lease, receipt_id) JOIN inflight ON inflight.job_id = items.job_id AND inflight.run_lease = items.run_lease WHERE NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) FOR UPDATE OF claim_batches ), claim_refs AS ( SELECT * FROM row_claim_refs UNION ALL SELECT * FROM batch_claim_refs ), upserted AS ( INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at) SELECT claim_refs.job_id, claim_refs.run_lease, clock_timestamp(), clock_timestamp() FROM claim_refs ON CONFLICT (job_id, run_lease) DO UPDATE SET heartbeat_at = clock_timestamp(), updated_at = clock_timestamp() RETURNING job_id, run_lease ), marked AS ( UPDATE {schema}.lease_claims AS claims SET materialized_at = COALESCE(claims.materialized_at, clock_timestamp()) FROM row_claim_refs WHERE claims.job_id = row_claim_refs.job_id AND claims.run_lease = row_claim_refs.run_lease RETURNING claims.job_id ) SELECT count(*)::bigint FROM upserted "# ))) .bind(&job_ids) .bind(&run_leases) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(updated as usize) } async fn upsert_attempt_state_progress_from_receipts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, jobs: &[(i64, i64, serde_json::Value)], ) -> Result { if jobs.is_empty() { return Ok(0); } let receipt_jobs: Vec<(i64, i64)> = jobs .iter() .map(|(job_id, run_lease, _)| (*job_id, *run_lease)) .collect(); self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?; let schema = self.schema(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let job_ids: Vec = jobs.iter().map(|(job_id, _, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect(); let progress: Vec = jobs .iter() .map(|(_, _, progress)| progress.clone()) .collect(); let updated: i64 = sqlx::query_scalar(audited_sql(format!( r#" WITH inflight(job_id, run_lease, progress) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::jsonb[]) ), row_claim_refs AS ( -- Same anti-join pattern as the heartbeat-only path -- above. SELECT claims.job_id, claims.run_lease, inflight.progress FROM {schema}.lease_claims AS claims JOIN inflight ON inflight.job_id = claims.job_id AND inflight.run_lease = claims.run_lease WHERE NOT {closed_evidence} FOR UPDATE OF claims ), batch_claim_refs AS ( SELECT items.job_id, items.run_lease, inflight.progress FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids ) AS items(job_id, run_lease, receipt_id) JOIN inflight ON inflight.job_id = items.job_id AND inflight.run_lease = items.run_lease WHERE NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) FOR UPDATE OF claim_batches ), claim_refs AS ( SELECT * FROM row_claim_refs UNION ALL SELECT * FROM batch_claim_refs ), upserted AS ( INSERT INTO {schema}.attempt_state ( job_id, run_lease, heartbeat_at, progress, updated_at ) SELECT claim_refs.job_id, claim_refs.run_lease, clock_timestamp(), claim_refs.progress, clock_timestamp() FROM claim_refs ON CONFLICT (job_id, run_lease) DO UPDATE SET heartbeat_at = clock_timestamp(), progress = EXCLUDED.progress, updated_at = clock_timestamp() RETURNING job_id, run_lease ), marked AS ( UPDATE {schema}.lease_claims AS claims SET materialized_at = COALESCE(claims.materialized_at, clock_timestamp()) FROM row_claim_refs WHERE claims.job_id = row_claim_refs.job_id AND claims.run_lease = row_claim_refs.run_lease RETURNING claims.job_id ) SELECT count(*)::bigint FROM upserted "# ))) .bind(&job_ids) .bind(&run_leases) .bind(&progress) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(updated as usize) } /// Skip children with only bounded empty-index metadata. Callers already /// hold ACCESS EXCLUSIVE and have proved closure. DELETE/VACUUM can leave /// expanded indexes or TOAST storage even after the main heap reaches zero; /// these children must truncate to preserve physical reclamation. async fn truncate_used_children_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, children: &[&str], ) -> Result { let mut targets: Vec<&str> = Vec::with_capacity(children.len()); for child in children { let has_reclaimable_storage: bool = sqlx::query_scalar(RECLAIMABLE_CHILD_STORAGE_SQL) .bind(child) .fetch_one(tx.as_mut()) .await?; if has_reclaimable_storage { targets.push(child); } } if targets.is_empty() { return Ok(0); } sqlx::query(audited_sql(format!( "TRUNCATE TABLE {}", targets.join(", ") ))) .execute(tx.as_mut()) .await?; Ok(targets.len()) } /// `write_receipt_closures` is false when the caller has already closed /// the receipt for every row in `deleted` in this transaction. async fn hydrate_deleted_leases_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, deleted: Vec, write_receipt_closures: bool, ) -> Result, AwaError> { if deleted.is_empty() { return Ok(Vec::new()); } let schema = self.schema(); let ready_slots: Vec = deleted.iter().map(|row| row.ready_slot).collect(); let ready_generations: Vec = deleted.iter().map(|row| row.ready_generation).collect(); let queues: Vec = deleted.iter().map(|row| row.queue.clone()).collect(); let enqueue_shards: Vec = deleted.iter().map(|row| row.enqueue_shard).collect(); let lane_seqs: Vec = deleted.iter().map(|row| row.lane_seq).collect(); let job_ids: Vec = deleted.iter().map(|row| row.job_id).collect(); let run_leases: Vec = deleted.iter().map(|row| row.run_lease).collect(); let ready_rows: Vec = sqlx::query_as(audited_sql(format!( r#" WITH refs(ready_slot, ready_generation, queue, enqueue_shard, lane_seq, job_id) AS ( SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::bigint[], $6::bigint[]) ) SELECT ready.ready_slot, ready.ready_generation, ready.job_id, ready.kind, ready.queue, ready.args, ready.lane_seq, ready.enqueue_shard, ready.run_at, ready.created_at, ready.unique_key, ready.unique_states, COALESCE(ready.payload, '{{}}'::jsonb) AS payload FROM refs JOIN {schema}.ready_entries AS ready ON ready.ready_slot = refs.ready_slot AND ready.ready_generation = refs.ready_generation AND ready.queue = refs.queue AND ready.enqueue_shard = refs.enqueue_shard AND ready.lane_seq = refs.lane_seq AND ready.job_id = refs.job_id "# ))) .bind(&ready_slots) .bind(&ready_generations) .bind(&queues) .bind(&enqueue_shards) .bind(&lane_seqs) .bind(&job_ids) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; let attempt_rows: Vec = sqlx::query_as(audited_sql(format!( r#" WITH refs(job_id, run_lease) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) ) DELETE FROM {schema}.attempt_state AS attempt USING refs WHERE attempt.job_id = refs.job_id AND attempt.run_lease = refs.run_lease RETURNING attempt.job_id, attempt.run_lease, attempt.progress, attempt.callback_filter, attempt.callback_on_complete, attempt.callback_on_fail, attempt.callback_transform, attempt.callback_result "# ))) .bind(&job_ids) .bind(&run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; // Hydrate runs as part of every rescue path that DELETE'd // a leases row (heartbeat / deadline / callback timeout // rescue, plus admin cancel of running attempts). For // receipt-backed attempts those leases came from // `ensure_running_leases_from_receipts_tx`, which leaves a // `lease_claims` row behind with `materialized_at` set. The // rescue itself closes the lease but never wrote a closure // for the original receipt, so the claim sat "open" until // partition prune — `load_job` and any // `lease_claims`-aware count then double-counted the // attempt as `running` even after it had moved to // retryable / failed / completed. Write the closure here so // the receipt plane mirrors the lease plane: when the lease // is gone, the receipt is gone too. if write_receipt_closures { sqlx::query(audited_sql(format!( r#" WITH refs(job_id, run_lease) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) ), row_claims AS ( SELECT claims.claim_slot, claims.job_id, claims.run_lease, 'rescue', clock_timestamp() FROM {schema}.lease_claims AS claims JOIN refs ON refs.job_id = claims.job_id AND refs.run_lease = claims.run_lease ), -- Compact batch-sourced claims (including those materialized -- into `leases` without a lease_claims row) have no row to -- JOIN, so they close into the batch ledger the queue prune -- gate counts via compact_count. The double-close guards -- below keep a re-hydrated rescued receipt from writing a -- second batch closure. batch_claims AS ( SELECT claim_batches.claim_slot, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, items.run_lease, items.receipt_id FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids ) AS items(job_id, run_lease, receipt_id) JOIN refs ON refs.job_id = items.job_id AND refs.run_lease = items.run_lease WHERE NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) ), inserted AS ( INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at) SELECT * FROM row_claims ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), inserted_batches AS ( INSERT INTO {schema}.lease_claim_closure_batches ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT batch_claims.claim_slot, batch_claims.ready_slot, batch_claims.ready_generation, 'rescue', count(*)::int, array_agg(batch_claims.receipt_id ORDER BY batch_claims.receipt_id), range_agg(int8range(batch_claims.receipt_id, batch_claims.receipt_id + 1, '[)') ORDER BY batch_claims.receipt_id), clock_timestamp() FROM batch_claims GROUP BY batch_claims.claim_slot, batch_claims.ready_slot, batch_claims.ready_generation RETURNING claim_slot ), marked AS ( UPDATE {schema}.lease_claims AS claims SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) FROM inserted WHERE claims.claim_slot = inserted.claim_slot AND claims.job_id = inserted.job_id AND claims.run_lease = inserted.run_lease RETURNING claims.job_id ) SELECT (SELECT count(*) FROM marked) + (SELECT count(*) FROM inserted_batches) "# ))) .bind(&job_ids) .bind(&run_leases) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } let ready_map: BTreeMap<(i32, i64, String, i16, i64, i64), ReadySnapshotRow> = ready_rows .into_iter() .map(|row| { ( ( row.ready_slot, row.ready_generation, row.queue.clone(), row.enqueue_shard, row.lane_seq, row.job_id, ), row, ) }) .collect(); let attempt_map: BTreeMap<(i64, i64), AttemptStateRow> = attempt_rows .into_iter() .map(|row| ((row.job_id, row.run_lease), row)) .collect(); let mut hydrated = Vec::with_capacity(deleted.len()); for deleted_row in deleted { let ready = ready_map .get(&( deleted_row.ready_slot, deleted_row.ready_generation, deleted_row.queue.clone(), deleted_row.enqueue_shard, deleted_row.lane_seq, deleted_row.job_id, )) .ok_or_else(|| { AwaError::Validation(format!( "queue storage ready row missing for deleted lease job {} run_lease {}", deleted_row.job_id, deleted_row.run_lease )) })?; let attempt = attempt_map.get(&(deleted_row.job_id, deleted_row.run_lease)); hydrated.push(LeaseTransitionRow { ready_slot: deleted_row.ready_slot, ready_generation: deleted_row.ready_generation, job_id: deleted_row.job_id, kind: ready.kind.clone(), queue: ready.queue.clone(), args: ready.args.clone(), state: deleted_row.state, priority: deleted_row.priority, attempt: deleted_row.attempt, run_lease: deleted_row.run_lease, max_attempts: deleted_row.max_attempts, lane_seq: deleted_row.lane_seq, enqueue_shard: deleted_row.enqueue_shard, run_at: ready.run_at, attempted_at: deleted_row.attempted_at, created_at: ready.created_at, unique_key: ready.unique_key.clone(), unique_states: ready.unique_states.clone(), payload: ready.payload.clone(), progress: attempt.and_then(|row| row.progress.clone()), }); } Ok(hydrated) } /// Close one compact receipt claim by its batch identity. The claim /// batch row is addressed by primary key and the item by index, and the /// closure evidence probes are all index lookups, so the cost does not /// grow with the number of open claim batches. The per-attempt advisory /// lock taken by the caller serialises closers; the batch row itself is /// not locked, matching the compact completion path. #[allow(clippy::too_many_arguments)] async fn close_compact_receipt_claim_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: i64, outcome: &str, claim_slot: i32, batch_id: i64, batch_index: i32, receipt_id: i64, ) -> Result, AwaError> { let schema = self.schema(); let slot = usize::try_from(claim_slot) .map_err(|_| AwaError::Validation(format!("invalid claim slot {claim_slot}")))?; let claim_batch_rel = claim_batch_child_name(schema, slot); let closure_rel = closure_child_name(schema, slot); let closure_batch_rel = claim_closure_batch_child_name(schema, slot); sqlx::query_as(audited_sql(format!( r#" WITH target AS ( SELECT claim_batches.claim_slot, claim_batches.ready_slot, claim_batches.ready_generation, claim_batches.job_ids[$3] AS job_id, claim_batches.queue, 'running'::awa.job_state AS state, claim_batches.priority, claim_batches.attempts[$3] AS attempt, claim_batches.run_leases[$3] AS run_lease, claim_batches.max_attempts[$3] AS max_attempts, claim_batches.lane_seqs[$3] AS lane_seq, claim_batches.enqueue_shard, claim_batches.claimed_at AS attempted_at FROM {claim_batch_rel} AS claim_batches WHERE claim_batches.claim_slot = $1 AND claim_batches.batch_id = $2 AND $3 BETWEEN 1 AND claim_batches.claimed_count AND claim_batches.job_ids[$3] = $4 AND claim_batches.run_leases[$3] = $5 AND claim_batches.receipt_ids[$3] = $6 AND NOT EXISTS ( SELECT 1 FROM {closure_rel} AS closures WHERE closures.claim_slot = $1 AND closures.job_id = $4 AND closures.run_lease = $5 ) AND NOT EXISTS ( SELECT 1 FROM {closure_batch_rel} AS closure_batches WHERE closure_batches.claim_slot = $1 AND closure_batches.receipt_ranges @> $6::bigint ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = $4 AND done.run_lease = $5 ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = $4 AND deferred.run_lease = $5 ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = $4 AND dlq.run_lease = $5 ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = $4 AND lease.run_lease = $5 ) ), inserted_batches AS ( INSERT INTO {schema}.lease_claim_closure_batches ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT target.claim_slot, target.ready_slot, target.ready_generation, $7, 1, ARRAY[$6::bigint], int8multirange(int8range($6::bigint, $6::bigint + 1, '[)')), clock_timestamp() FROM target RETURNING claim_slot ) SELECT target.ready_slot, target.ready_generation, target.job_id, target.queue, target.state, target.priority, target.attempt, target.run_lease, target.max_attempts, target.lane_seq, target.enqueue_shard, target.attempted_at FROM target WHERE EXISTS (SELECT 1 FROM inserted_batches) "# ))) .bind(claim_slot) .bind(batch_id) .bind(batch_index) .bind(job_id) .bind(run_lease) .bind(receipt_id) .bind(outcome) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error) } async fn close_open_receipt_claim_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: i64, outcome: &str, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { if !self.lease_claim_receipts() { return Ok(None); } self.lock_receipt_attempts_tx(tx, &[(job_id, run_lease)]) .await?; if let Some(claim) = claim.filter(|claim| claim.lease_claim_receipt) { if let (Some(batch_id), Some(batch_index), Some(receipt_id)) = ( claim.claim_batch_id, claim.claim_batch_index, claim.receipt_id, ) { let deleted = self .close_compact_receipt_claim_tx( tx, job_id, run_lease, outcome, claim.claim_slot, batch_id, batch_index, receipt_id, ) .await?; if deleted.is_empty() { return Ok(None); } let moved = self.hydrate_deleted_leases_tx(tx, deleted, false).await?; return Ok(moved.into_iter().next()); } } let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" WITH row_target AS ( -- Target is the open claim identified from the -- partitioned lease_claims table anti-joined against -- durable closure evidence. SELECT claims.claim_slot, NULL::bigint AS batch_id, claims.receipt_id, claims.ready_slot, claims.ready_generation, claims.job_id, claims.queue, 'running'::awa.job_state AS state, claims.priority, claims.attempt, claims.run_lease, claims.max_attempts, claims.lane_seq, claims.enqueue_shard, claims.claimed_at AS attempted_at FROM {schema}.lease_claims AS claims WHERE claims.job_id = $1 AND claims.run_lease = $2 AND claims.closed_at IS NULL AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claims.claim_slot AND closures.job_id = claims.job_id AND closures.run_lease = claims.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.receipt_ranges @> claims.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = claims.job_id AND done.run_lease = claims.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = claims.job_id AND deferred.run_lease = claims.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = claims.job_id AND dlq.run_lease = claims.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) FOR UPDATE OF claims ), batch_target AS ( SELECT claim_batches.claim_slot, claim_batches.batch_id, items.receipt_id, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, claim_batches.queue, 'running'::awa.job_state AS state, claim_batches.priority, items.attempt, items.run_lease, items.max_attempts, items.lane_seq, claim_batches.enqueue_shard, claim_batches.claimed_at AS attempted_at FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids, claim_batches.lane_seqs, claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) WHERE items.job_id = $1 AND items.run_lease = $2 AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = items.job_id AND lease.run_lease = items.run_lease ) FOR UPDATE OF claim_batches ), target AS ( SELECT * FROM row_target UNION ALL SELECT * FROM batch_target LIMIT 1 ), -- A row target (batch_id IS NULL) is a deadline lease claim: -- close it explicitly so the prune gates balance it against -- the lease_claims row. A batch target has no lease_claims -- row to JOIN, so it closes into the batch ledger below. inserted AS ( INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at) SELECT target.claim_slot, target.job_id, target.run_lease, $3, clock_timestamp() FROM target WHERE target.batch_id IS NULL ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), inserted_batches AS ( INSERT INTO {schema}.lease_claim_closure_batches ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT target.claim_slot, target.ready_slot, target.ready_generation, $3, count(*)::int, array_agg(target.receipt_id ORDER BY target.receipt_id), range_agg(int8range(target.receipt_id, target.receipt_id + 1, '[)') ORDER BY target.receipt_id), clock_timestamp() FROM target WHERE target.batch_id IS NOT NULL GROUP BY target.claim_slot, target.ready_slot, target.ready_generation RETURNING claim_slot ), marked AS ( UPDATE {schema}.lease_claims AS claims SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) FROM inserted WHERE claims.claim_slot = inserted.claim_slot AND claims.job_id = inserted.job_id AND claims.run_lease = inserted.run_lease RETURNING claims.job_id ), -- The batch ledger has no job_id/run_lease column to RETURN, -- so derive the closed batch target from `target` gated on -- the batch insert having fired (target is LIMIT 1). closed_target AS ( SELECT claim_slot, job_id, run_lease FROM inserted UNION ALL SELECT target.claim_slot, target.job_id, target.run_lease FROM target WHERE target.batch_id IS NOT NULL AND EXISTS (SELECT 1 FROM inserted_batches) ) SELECT target.ready_slot, target.ready_generation, target.job_id, target.queue, target.state, target.priority, target.attempt, target.run_lease, target.max_attempts, target.lane_seq, target.enqueue_shard, NULL::timestamptz AS heartbeat_at, NULL::timestamptz AS deadline_at, target.attempted_at, NULL::uuid AS callback_id, NULL::timestamptz AS callback_timeout_at FROM target JOIN closed_target ON closed_target.claim_slot = target.claim_slot AND closed_target.job_id = target.job_id AND closed_target.run_lease = target.run_lease "# ))) .bind(job_id) .bind(run_lease) .bind(outcome) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { return Ok(None); } let moved = self.hydrate_deleted_leases_tx(tx, deleted, true).await?; Ok(moved.into_iter().next()) } async fn take_running_attempt_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_id: i64, run_lease: i64, receipt_outcome: &str, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { if let Some(moved) = self .close_open_receipt_claim_tx(tx, job_id, run_lease, receipt_outcome, claim) .await? { return Ok(Some(moved)); } let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE job_id = $1 AND run_lease = $2 AND state = 'running' RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .bind(job_id) .bind(run_lease) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { return Ok(None); } let moved = self.hydrate_deleted_leases_tx(tx, deleted, true).await?; Ok(moved.into_iter().next()) } async fn rescue_stale_receipt_claims_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, cutoff: DateTime, ) -> Result, AwaError> { let mut rescued = Vec::new(); let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?; if let Some(slot) = preferred_slot { let mut slot_rescued = self .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); if remaining == 0 { return Ok(rescued); } } for slot in 0..self.claim_slot_count() { let slot = slot as i32; if Some(slot) == preferred_slot { continue; } let mut slot_rescued = self .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); if remaining == 0 { break; } } Ok(rescued) } async fn rescue_stale_receipt_claims_for_slot_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, cutoff: DateTime, rescue_limit: i64, ) -> Result, AwaError> { let schema = self.schema(); let claim_child = claim_child_name(schema, slot as usize); let claim_batch_child = claim_batch_child_name(schema, slot as usize); let closure_child = closure_child_name(schema, slot as usize); let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize); let candidate_ctes = format!( r#" cursor_row AS ( SELECT rescue_cursor_claimed_at, rescue_cursor_job_id, rescue_cursor_run_lease FROM {schema}.claim_ring_slots WHERE slot = $1 FOR UPDATE ), claim_source AS MATERIALIZED ( SELECT claims.claim_slot, NULL::bigint AS batch_id, claims.ready_slot, claims.ready_generation, claims.job_id, claims.queue, claims.priority, claims.attempt, claims.run_lease, claims.max_attempts, claims.lane_seq, claims.enqueue_shard, claims.receipt_id, claims.claimed_at, claims.closed_at, false AS compact_batch FROM {claim_child} AS claims WHERE claims.claim_slot = $1 UNION ALL SELECT claim_batches.claim_slot, claim_batches.batch_id, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, claim_batches.queue, claim_batches.priority, items.attempt, items.run_lease, items.max_attempts, items.lane_seq, claim_batches.enqueue_shard, items.receipt_id, claim_batches.claimed_at, NULL::timestamptz AS closed_at, true AS compact_batch FROM {claim_batch_child} AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids, claim_batches.lane_seqs, claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) WHERE claim_batches.claim_slot = $1 ), after_cursor AS MATERIALIZED ( SELECT claim_source.claim_slot, claim_source.job_id, claim_source.run_lease, row_number() OVER ( ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease ) AS rn FROM claim_source CROSS JOIN cursor_row WHERE claim_source.claimed_at < $2 AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease) > ( cursor_row.rescue_cursor_claimed_at, cursor_row.rescue_cursor_job_id, cursor_row.rescue_cursor_run_lease ) ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease LIMIT $3 ), after_count AS ( SELECT count(*)::bigint AS count FROM after_cursor ), before_cursor AS MATERIALIZED ( SELECT claim_source.claim_slot, claim_source.job_id, claim_source.run_lease, after_count.count + row_number() OVER ( ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease ) AS rn FROM claim_source CROSS JOIN cursor_row CROSS JOIN after_count WHERE after_count.count < $3 AND claim_source.claimed_at < $2 AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease) <= ( cursor_row.rescue_cursor_claimed_at, cursor_row.rescue_cursor_job_id, cursor_row.rescue_cursor_run_lease ) ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease LIMIT (SELECT GREATEST($3 - count, 0) FROM after_count) ), candidate_keys AS MATERIALIZED ( SELECT claim_slot, job_id, run_lease, rn FROM after_cursor UNION ALL SELECT claim_slot, job_id, run_lease, rn FROM before_cursor ), candidates AS MATERIALIZED ( SELECT claim_source.claim_slot, claim_source.batch_id, claim_source.ready_slot, claim_source.ready_generation, claim_source.job_id, claim_source.queue, claim_source.priority, claim_source.attempt, claim_source.run_lease, claim_source.max_attempts, claim_source.lane_seq, claim_source.enqueue_shard, claim_source.receipt_id, claim_source.claimed_at, claim_source.closed_at, claim_source.compact_batch, COALESCE(attempt.heartbeat_at, claim_source.claimed_at) < $2 AS is_stale, ( claim_source.closed_at IS NOT NULL OR EXISTS ( SELECT 1 FROM {closure_child} AS closures WHERE closures.claim_slot = claim_source.claim_slot AND closures.job_id = claim_source.job_id AND closures.run_lease = claim_source.run_lease ) OR EXISTS ( SELECT 1 FROM {closure_batch_child} AS closure_batches WHERE closure_batches.receipt_ranges @> claim_source.receipt_id ) OR EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = claim_source.job_id AND done.run_lease = claim_source.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = claim_source.job_id AND deferred.run_lease = claim_source.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = claim_source.job_id AND dlq.run_lease = claim_source.run_lease ) ) AS is_closed, EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claim_source.job_id AND lease.run_lease = claim_source.run_lease ) AS is_lease_managed, candidate_keys.rn FROM candidate_keys JOIN claim_source ON claim_source.claim_slot = candidate_keys.claim_slot AND claim_source.job_id = candidate_keys.job_id AND claim_source.run_lease = candidate_keys.run_lease LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = claim_source.job_id AND attempt.run_lease = claim_source.run_lease ), stale_candidates AS ( SELECT candidates.* FROM candidates WHERE NOT is_closed AND NOT is_lease_managed AND is_stale ORDER BY rn LIMIT $4 )"# ); // Closers take the attempt's advisory lock before their statement // snapshot. Rescue takes the same locks in a statement of its own, // so the statement that decides and writes the closure starts after // any closer that held a lock first has committed and sees its // evidence. Only attempts locked here may be closed below. let (locked_job_ids, locked_run_leases): (Vec, Vec) = sqlx::query_as::<_, (i64, i64)>(audited_sql(format!( r#" WITH {candidate_ctes} SELECT job_id, run_lease FROM stale_candidates WHERE pg_catalog.pg_try_advisory_xact_lock( pg_catalog.hashtextextended( format('awa.receipt.complete:%s:%s', job_id, run_lease), 0 ) ) "# ))) .bind(slot) .bind(cutoff) .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? .into_iter() .unzip(); let rescued: Vec = sqlx::query_as(audited_sql(format!( r#" WITH {candidate_ctes}, stale_row_locked AS ( SELECT stale_candidates.* FROM stale_candidates JOIN {claim_child} AS claims ON claims.claim_slot = stale_candidates.claim_slot AND claims.job_id = stale_candidates.job_id AND claims.run_lease = stale_candidates.run_lease WHERE NOT stale_candidates.compact_batch AND claims.closed_at IS NULL AND NOT EXISTS ( SELECT 1 FROM {closure_child} AS closures WHERE closures.claim_slot = claims.claim_slot AND closures.job_id = claims.job_id AND closures.run_lease = claims.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {closure_batch_child} AS closure_batches WHERE closure_batches.receipt_ranges @> claims.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = claims.job_id AND done.run_lease = claims.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = claims.job_id AND deferred.run_lease = claims.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = claims.job_id AND dlq.run_lease = claims.run_lease ) AND (claims.job_id, claims.run_lease) IN ( SELECT locked.job_id, locked.run_lease FROM unnest($5::bigint[], $6::bigint[]) AS locked(job_id, run_lease) ) -- A claim that already materialized into `leases` is -- on the lease-side heartbeat-rescue path (see -- `rescue_stale_heartbeats`). Rescuing it again here -- would write a second closure for an attempt the -- runtime is still tracking via its lease row. AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) FOR UPDATE OF claims SKIP LOCKED ), stale_batch_locked AS ( SELECT stale_candidates.* FROM stale_candidates JOIN {claim_batch_child} AS claim_batches ON claim_batches.claim_slot = stale_candidates.claim_slot AND claim_batches.batch_id = stale_candidates.batch_id WHERE stale_candidates.compact_batch AND NOT EXISTS ( SELECT 1 FROM {closure_child} AS closures WHERE closures.claim_slot = stale_candidates.claim_slot AND closures.job_id = stale_candidates.job_id AND closures.run_lease = stale_candidates.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {closure_batch_child} AS closure_batches WHERE closure_batches.receipt_ranges @> stale_candidates.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = stale_candidates.job_id AND done.run_lease = stale_candidates.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = stale_candidates.job_id AND deferred.run_lease = stale_candidates.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = stale_candidates.job_id AND dlq.run_lease = stale_candidates.run_lease ) AND (stale_candidates.job_id, stale_candidates.run_lease) IN ( SELECT locked.job_id, locked.run_lease FROM unnest($5::bigint[], $6::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = stale_candidates.job_id AND lease.run_lease = stale_candidates.run_lease ) FOR UPDATE OF claim_batches SKIP LOCKED ), stale_locked AS ( SELECT * FROM stale_row_locked UNION ALL SELECT * FROM stale_batch_locked ), -- Row-sourced (deadline lease) claims close explicitly so the -- prune gates balance them against the lease_claims row. -- Compact batch-sourced claims have no lease_claims row to -- JOIN, so they close into the batch ledger the queue prune -- gate counts via compact_count. inserted AS ( INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at) SELECT stale_locked.claim_slot, stale_locked.job_id, stale_locked.run_lease, 'rescued', clock_timestamp() FROM stale_locked WHERE NOT stale_locked.compact_batch ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), inserted_batches AS ( INSERT INTO {schema}.lease_claim_closure_batches ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT stale_locked.claim_slot, stale_locked.ready_slot, stale_locked.ready_generation, 'rescued', count(*)::int, array_agg(stale_locked.receipt_id ORDER BY stale_locked.receipt_id), range_agg(int8range(stale_locked.receipt_id, stale_locked.receipt_id + 1, '[)') ORDER BY stale_locked.receipt_id), clock_timestamp() FROM stale_locked WHERE stale_locked.compact_batch GROUP BY stale_locked.claim_slot, stale_locked.ready_slot, stale_locked.ready_generation RETURNING claim_slot ), closed_locked AS ( SELECT claim_slot, job_id, run_lease FROM inserted UNION ALL SELECT stale_locked.claim_slot, stale_locked.job_id, stale_locked.run_lease FROM stale_locked WHERE stale_locked.compact_batch AND EXISTS (SELECT 1 FROM inserted_batches) ), marked AS ( UPDATE {claim_child} AS claims SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) FROM inserted WHERE claims.job_id = inserted.job_id AND claims.run_lease = inserted.run_lease RETURNING claims.job_id ), annotated AS ( SELECT candidates.*, ( candidates.is_closed OR candidates.is_lease_managed OR NOT candidates.is_stale OR EXISTS ( SELECT 1 FROM closed_locked WHERE closed_locked.claim_slot = candidates.claim_slot AND closed_locked.job_id = candidates.job_id AND closed_locked.run_lease = candidates.run_lease ) ) AS advanceable FROM candidates ), bounded AS ( SELECT annotated.*, min(CASE WHEN NOT annotated.advanceable THEN annotated.rn END) OVER () AS first_blocked_rn FROM annotated ), advance_target AS ( SELECT claimed_at, job_id, run_lease FROM bounded WHERE first_blocked_rn IS NULL OR rn < first_blocked_rn ORDER BY rn DESC LIMIT 1 ), advance_cursor AS ( UPDATE {schema}.claim_ring_slots AS slots SET rescue_cursor_claimed_at = advance_target.claimed_at, rescue_cursor_job_id = advance_target.job_id, rescue_cursor_run_lease = advance_target.run_lease FROM advance_target WHERE slots.slot = $1 RETURNING slots.slot ), cursor_advance AS ( SELECT count(*) FROM advance_cursor ) SELECT stale_locked.ready_slot, stale_locked.ready_generation, stale_locked.job_id, stale_locked.queue, 'running'::awa.job_state AS state, stale_locked.priority, stale_locked.attempt, stale_locked.run_lease, stale_locked.max_attempts, stale_locked.lane_seq, stale_locked.enqueue_shard, stale_locked.claimed_at AS attempted_at FROM stale_locked JOIN closed_locked ON closed_locked.claim_slot = stale_locked.claim_slot AND closed_locked.job_id = stale_locked.job_id AND closed_locked.run_lease = stale_locked.run_lease CROSS JOIN cursor_advance "# ))) .bind(slot) .bind(cutoff) .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(rescued) } /// Receipt-side counterpart to `rescue_expired_deadlines`: scans both /// receipt-claim representations for attempts whose per-claim /// `deadline_at` has passed but which still don't have a closure or a /// materialized lease row: /// /// * compact `lease_claim_batches` rows — the shape the claim path /// writes for deadline-backed claims since #246; the batch carries /// one shared `deadline_at`, and expired members are force-closed /// into `lease_claim_closure_batches`; /// * row-local `lease_claims` rows — legacy in-flight claims written /// before the compact-deadline upgrade; these keep the original /// `lease_claim_closures` closure shape. /// /// Each match is returned for the maintenance caller to convert into /// a deferred / DLQ row, exactly as the lease-side path does. /// /// The anti-joins mirror `rescue_stale_receipt_claims_tx`'s /// disambiguation: a claim that has already materialized into /// `leases` is on the lease-side deadline-rescue path, and /// rescuing it here would double-close it. /// /// Unlike the lease-side scan, receipt deadline rescue walks each /// claim partition through a tiny cursor ordered by deadline (one /// cursor per representation). That keeps a long MVCC horizon from /// making every maintenance tick re-prove old successful receipt /// completions until claim prune can truncate the partition. async fn rescue_expired_receipt_deadlines_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ) -> Result, AwaError> { let mut rescued = Vec::new(); let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?; let mut slots: Vec = Vec::with_capacity(self.claim_slot_count() + 1); if let Some(slot) = preferred_slot { slots.push(slot); } for slot in 0..self.claim_slot_count() { let slot = slot as i32; if Some(slot) != preferred_slot { slots.push(slot); } } for slot in slots { let mut slot_rescued = self .rescue_expired_batch_deadlines_for_slot_tx(tx, slot, remaining) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); if remaining == 0 { break; } let mut slot_rescued = self .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); if remaining == 0 { break; } } Ok(rescued) } async fn oldest_initialized_claim_slot_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ) -> Result, AwaError> { let schema = self.schema(); let preferred_slot = sqlx::query_as::<_, (i32, i64, i32)>(audited_sql(format!( r#" SELECT ledger.slot, ledger.generation, state.slot_count FROM {schema}.claim_ring_state AS state CROSS JOIN LATERAL ( SELECT slot, generation FROM {schema}.claim_ring_rotations ORDER BY generation DESC LIMIT 1 ) AS ledger WHERE state.singleton = TRUE "# ))) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)? .and_then(|(current_slot, generation, slot_count)| { oldest_initialized_ring_slot(current_slot, generation, slot_count) .map(|(slot, _generation)| slot) .filter(|slot| *slot >= 0 && (*slot as usize) < self.claim_slot_count()) }); Ok(preferred_slot) } async fn rescue_expired_receipt_deadlines_for_slot_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, rescue_limit: i64, ) -> Result, AwaError> { if rescue_limit <= 0 { return Ok(Vec::new()); } let schema = self.schema(); let claim_child = claim_child_name(schema, slot as usize); let closure_child = closure_child_name(schema, slot as usize); let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_child, &closure_batch_child, "claims"); let candidate_ctes = format!( r#" cursor_row AS ( SELECT deadline_cursor_deadline_at, deadline_cursor_job_id, deadline_cursor_run_lease FROM {schema}.claim_ring_slots WHERE slot = $1 FOR UPDATE ), after_cursor AS MATERIALIZED ( SELECT claims.claim_slot, claims.job_id, claims.run_lease, row_number() OVER ( ORDER BY claims.deadline_at, claims.job_id, claims.run_lease ) AS rn FROM {claim_child} AS claims CROSS JOIN cursor_row WHERE claims.claim_slot = $1 AND claims.deadline_at IS NOT NULL AND claims.deadline_at < clock_timestamp() AND (claims.deadline_at, claims.job_id, claims.run_lease) > ( cursor_row.deadline_cursor_deadline_at, cursor_row.deadline_cursor_job_id, cursor_row.deadline_cursor_run_lease ) ORDER BY claims.deadline_at, claims.job_id, claims.run_lease LIMIT $2 ), after_count AS ( SELECT count(*)::bigint AS count FROM after_cursor ), before_cursor AS MATERIALIZED ( SELECT claims.claim_slot, claims.job_id, claims.run_lease, after_count.count + row_number() OVER ( ORDER BY claims.deadline_at, claims.job_id, claims.run_lease ) AS rn FROM {claim_child} AS claims CROSS JOIN cursor_row CROSS JOIN after_count WHERE after_count.count < $2 AND claims.claim_slot = $1 AND claims.deadline_at IS NOT NULL AND claims.deadline_at < clock_timestamp() AND (claims.deadline_at, claims.job_id, claims.run_lease) <= ( cursor_row.deadline_cursor_deadline_at, cursor_row.deadline_cursor_job_id, cursor_row.deadline_cursor_run_lease ) ORDER BY claims.deadline_at, claims.job_id, claims.run_lease LIMIT (SELECT GREATEST($2 - count, 0) FROM after_count) ), candidate_keys AS MATERIALIZED ( SELECT claim_slot, job_id, run_lease, rn FROM after_cursor UNION ALL SELECT claim_slot, job_id, run_lease, rn FROM before_cursor ), candidates AS MATERIALIZED ( SELECT claims.claim_slot, claims.ready_slot, claims.ready_generation, claims.job_id, claims.queue, 'running'::awa.job_state AS state, claims.priority, claims.attempt, claims.run_lease, claims.max_attempts, claims.lane_seq, claims.enqueue_shard, claims.claimed_at, claims.deadline_at, claims.deadline_at < clock_timestamp() AS is_expired, {closed_evidence} AS is_closed, EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) AS is_lease_managed, candidate_keys.rn FROM candidate_keys JOIN {claim_child} AS claims ON claims.claim_slot = candidate_keys.claim_slot AND claims.job_id = candidate_keys.job_id AND claims.run_lease = candidate_keys.run_lease ), expired_candidates AS ( SELECT candidates.* FROM candidates WHERE NOT is_closed AND NOT is_lease_managed AND is_expired ORDER BY rn LIMIT $3 )"# ); // Lock before deciding; see `rescue_stale_receipt_claims_for_slot_tx`. let (locked_job_ids, locked_run_leases): (Vec, Vec) = sqlx::query_as::<_, (i64, i64)>(audited_sql(format!( r#" WITH {candidate_ctes} SELECT job_id, run_lease FROM expired_candidates WHERE pg_catalog.pg_try_advisory_xact_lock( pg_catalog.hashtextextended( format('awa.receipt.complete:%s:%s', job_id, run_lease), 0 ) ) "# ))) .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? .into_iter() .unzip(); let rescued: Vec = sqlx::query_as(audited_sql(format!( r#" WITH {candidate_ctes}, expired_locked AS ( SELECT expired_candidates.* FROM expired_candidates JOIN {claim_child} AS claims ON claims.claim_slot = expired_candidates.claim_slot AND claims.job_id = expired_candidates.job_id AND claims.run_lease = expired_candidates.run_lease WHERE NOT {closed_evidence} AND claims.deadline_at < clock_timestamp() AND (claims.job_id, claims.run_lease) IN ( SELECT locked.job_id, locked.run_lease FROM unnest($4::bigint[], $5::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) FOR UPDATE OF claims SKIP LOCKED ), inserted AS ( INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at) SELECT expired_locked.claim_slot, expired_locked.job_id, expired_locked.run_lease, 'deadline_expired', clock_timestamp() FROM expired_locked ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING RETURNING claim_slot, job_id, run_lease, closed_at ), marked AS ( UPDATE {claim_child} AS claims SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) FROM inserted WHERE claims.job_id = inserted.job_id AND claims.run_lease = inserted.run_lease RETURNING claims.job_id ), annotated AS ( SELECT candidates.*, ( candidates.is_closed OR candidates.is_lease_managed OR EXISTS ( SELECT 1 FROM inserted WHERE inserted.claim_slot = candidates.claim_slot AND inserted.job_id = candidates.job_id AND inserted.run_lease = candidates.run_lease ) ) AS advanceable FROM candidates ), bounded AS ( SELECT annotated.*, min(CASE WHEN NOT annotated.advanceable THEN annotated.rn END) OVER () AS first_blocked_rn FROM annotated ), advance_target AS ( SELECT deadline_at, job_id, run_lease FROM bounded WHERE first_blocked_rn IS NULL OR rn < first_blocked_rn ORDER BY rn DESC LIMIT 1 ), advance_cursor AS ( UPDATE {schema}.claim_ring_slots AS slots SET deadline_cursor_deadline_at = advance_target.deadline_at, deadline_cursor_job_id = advance_target.job_id, deadline_cursor_run_lease = advance_target.run_lease FROM advance_target WHERE slots.slot = $1 RETURNING slots.slot ), cursor_advance AS ( SELECT count(*) FROM advance_cursor ) SELECT expired_locked.ready_slot, expired_locked.ready_generation, expired_locked.job_id, expired_locked.queue, expired_locked.state, expired_locked.priority, expired_locked.attempt, expired_locked.run_lease, expired_locked.max_attempts, expired_locked.lane_seq, expired_locked.enqueue_shard, expired_locked.claimed_at AS attempted_at FROM expired_locked JOIN inserted ON inserted.claim_slot = expired_locked.claim_slot AND inserted.job_id = expired_locked.job_id AND inserted.run_lease = expired_locked.run_lease CROSS JOIN cursor_advance "# ))) .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(rescued) } /// Compact-batch counterpart of /// `rescue_expired_receipt_deadlines_for_slot_tx` (#246). Deadline-backed /// claims are written as one `lease_claim_batches` row per claimed batch /// with a single shared `deadline_at`, so the sweep cursor is /// `(deadline_at, batch_id)` on `claim_ring_slots.batch_deadline_cursor_*` /// and the candidate window includes only batches whose deadline has /// already passed (index-backed by the partial per-child /// `(deadline_at, batch_id) WHERE deadline_at IS NOT NULL` index). /// /// Expired batches expand into members; members WITHOUT durable closure /// evidence and WITHOUT a materialized lease are force-closed by /// appending a `lease_claim_closure_batches` row with outcome /// `'deadline_expired'` — the same closure shape compact claims use on /// every other exit path, so the claim/queue prune count proofs see the /// closure. Members that completed, materialized into `leases` /// (lease-side deadline rescue owns those), or already closed are /// skipped. The cursor advances past a batch only when every member is /// proven closed, lease-managed, or was closed by this transaction, so a /// partially-rescued batch (rescue-budget cutoff or a lost advisory-lock /// race) is revisited next tick. async fn rescue_expired_batch_deadlines_for_slot_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, rescue_limit: i64, ) -> Result, AwaError> { if rescue_limit <= 0 { return Ok(Vec::new()); } let schema = self.schema(); let claim_batch_child = claim_batch_child_name(schema, slot as usize); let closure_child = closure_child_name(schema, slot as usize); let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize); let candidate_ctes = format!( r#" cursor_row AS ( SELECT batch_deadline_cursor_deadline_at, batch_deadline_cursor_batch_id FROM {schema}.claim_ring_slots WHERE slot = $1 FOR UPDATE ), after_cursor AS MATERIALIZED ( SELECT claim_batches.claim_slot, claim_batches.batch_id, row_number() OVER ( ORDER BY claim_batches.deadline_at, claim_batches.batch_id ) AS rn FROM {claim_batch_child} AS claim_batches CROSS JOIN cursor_row WHERE claim_batches.claim_slot = $1 AND claim_batches.deadline_at IS NOT NULL AND claim_batches.deadline_at < clock_timestamp() AND (claim_batches.deadline_at, claim_batches.batch_id) > ( cursor_row.batch_deadline_cursor_deadline_at, cursor_row.batch_deadline_cursor_batch_id ) ORDER BY claim_batches.deadline_at, claim_batches.batch_id LIMIT $2 ), after_count AS ( SELECT count(*)::bigint AS count FROM after_cursor ), before_cursor AS MATERIALIZED ( SELECT claim_batches.claim_slot, claim_batches.batch_id, after_count.count + row_number() OVER ( ORDER BY claim_batches.deadline_at, claim_batches.batch_id ) AS rn FROM {claim_batch_child} AS claim_batches CROSS JOIN cursor_row CROSS JOIN after_count WHERE after_count.count < $2 AND claim_batches.claim_slot = $1 AND claim_batches.deadline_at IS NOT NULL AND claim_batches.deadline_at < clock_timestamp() AND (claim_batches.deadline_at, claim_batches.batch_id) <= ( cursor_row.batch_deadline_cursor_deadline_at, cursor_row.batch_deadline_cursor_batch_id ) ORDER BY claim_batches.deadline_at, claim_batches.batch_id LIMIT (SELECT GREATEST($2 - count, 0) FROM after_count) ), candidate_keys AS MATERIALIZED ( SELECT claim_slot, batch_id, rn FROM after_cursor UNION ALL SELECT claim_slot, batch_id, rn FROM before_cursor ), -- Every member of a candidate batch, annotated with the same -- durable closure evidence and materialized-lease anti-joins the -- row-local deadline rescue uses. The batch deadline is shared, -- so every member of a candidate batch is expired by -- construction. members AS MATERIALIZED ( SELECT claim_batches.claim_slot, claim_batches.batch_id, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, claim_batches.queue, 'running'::awa.job_state AS state, claim_batches.priority, items.attempt, items.run_lease, items.max_attempts, items.lane_seq, claim_batches.enqueue_shard, items.receipt_id, claim_batches.claimed_at, claim_batches.deadline_at, ( EXISTS ( SELECT 1 FROM {closure_child} AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) OR EXISTS ( SELECT 1 FROM {closure_batch_child} AS closure_batches WHERE closure_batches.receipt_ranges @> items.receipt_id ) OR EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) ) AS is_closed, EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = items.job_id AND lease.run_lease = items.run_lease ) AS is_lease_managed, candidate_keys.rn FROM candidate_keys JOIN {claim_batch_child} AS claim_batches ON claim_batches.claim_slot = candidate_keys.claim_slot AND claim_batches.batch_id = candidate_keys.batch_id CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids, claim_batches.lane_seqs, claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) ), expired_members AS ( SELECT members.* FROM members WHERE NOT is_closed AND NOT is_lease_managed ORDER BY rn, receipt_id LIMIT $3 )"# ); // Lock before deciding; see `rescue_stale_receipt_claims_for_slot_tx`. let (locked_job_ids, locked_run_leases): (Vec, Vec) = sqlx::query_as::<_, (i64, i64)>(audited_sql(format!( r#" WITH {candidate_ctes} SELECT job_id, run_lease FROM expired_members WHERE pg_catalog.pg_try_advisory_xact_lock( pg_catalog.hashtextextended( format('awa.receipt.complete:%s:%s', job_id, run_lease), 0 ) ) "# ))) .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? .into_iter() .unzip(); let rescued: Vec = sqlx::query_as(audited_sql(format!( r#" WITH {candidate_ctes}, expired_locked AS ( SELECT expired_members.* FROM expired_members JOIN {claim_batch_child} AS claim_batches ON claim_batches.claim_slot = expired_members.claim_slot AND claim_batches.batch_id = expired_members.batch_id WHERE claim_batches.deadline_at < clock_timestamp() AND NOT EXISTS ( SELECT 1 FROM {closure_child} AS closures WHERE closures.claim_slot = expired_members.claim_slot AND closures.job_id = expired_members.job_id AND closures.run_lease = expired_members.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {closure_batch_child} AS closure_batches WHERE closure_batches.receipt_ranges @> expired_members.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = expired_members.job_id AND done.run_lease = expired_members.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = expired_members.job_id AND deferred.run_lease = expired_members.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = expired_members.job_id AND dlq.run_lease = expired_members.run_lease ) AND (expired_members.job_id, expired_members.run_lease) IN ( SELECT locked.job_id, locked.run_lease FROM unnest($4::bigint[], $5::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = expired_members.job_id AND lease.run_lease = expired_members.run_lease ) FOR UPDATE OF claim_batches SKIP LOCKED ), -- Compact batch-sourced claims have no lease_claims row to mark, -- so they close into the batch closure ledger both prune count -- proofs read via compact closed_count totals. inserted_batches AS ( INSERT INTO {closure_batch_child} ( claim_slot, ready_slot, ready_generation, outcome, closed_count, receipt_ids, receipt_ranges, closed_at ) SELECT expired_locked.claim_slot, expired_locked.ready_slot, expired_locked.ready_generation, 'deadline_expired', count(*)::int, array_agg(expired_locked.receipt_id ORDER BY expired_locked.receipt_id), range_agg(int8range(expired_locked.receipt_id, expired_locked.receipt_id + 1, '[)') ORDER BY expired_locked.receipt_id), clock_timestamp() FROM expired_locked GROUP BY expired_locked.claim_slot, expired_locked.ready_slot, expired_locked.ready_generation RETURNING claim_slot ), closed_locked AS ( SELECT expired_locked.claim_slot, expired_locked.job_id, expired_locked.run_lease FROM expired_locked WHERE EXISTS (SELECT 1 FROM inserted_batches) ), annotated AS ( SELECT members.*, ( members.is_closed OR members.is_lease_managed OR EXISTS ( SELECT 1 FROM closed_locked WHERE closed_locked.claim_slot = members.claim_slot AND closed_locked.job_id = members.job_id AND closed_locked.run_lease = members.run_lease ) ) AS advanceable FROM members ), -- The cursor may only advance past a batch when every one of its -- members is advanceable; a single open member (rescue-budget -- cutoff, lost advisory-lock race) blocks there so the batch is -- revisited. bounded AS ( SELECT annotated.*, min(CASE WHEN NOT annotated.advanceable THEN annotated.rn END) OVER () AS first_blocked_rn FROM annotated ), advance_target AS ( SELECT deadline_at, batch_id FROM bounded WHERE first_blocked_rn IS NULL OR rn < first_blocked_rn ORDER BY rn DESC LIMIT 1 ), advance_cursor AS ( UPDATE {schema}.claim_ring_slots AS slots SET batch_deadline_cursor_deadline_at = advance_target.deadline_at, batch_deadline_cursor_batch_id = advance_target.batch_id FROM advance_target WHERE slots.slot = $1 RETURNING slots.slot ), cursor_advance AS ( SELECT count(*) FROM advance_cursor ) SELECT expired_locked.ready_slot, expired_locked.ready_generation, expired_locked.job_id, expired_locked.queue, expired_locked.state, expired_locked.priority, expired_locked.attempt, expired_locked.run_lease, expired_locked.max_attempts, expired_locked.lane_seq, expired_locked.enqueue_shard, expired_locked.claimed_at AS attempted_at FROM expired_locked JOIN closed_locked ON closed_locked.claim_slot = expired_locked.claim_slot AND closed_locked.job_id = expired_locked.job_id AND closed_locked.run_lease = expired_locked.run_lease CROSS JOIN cursor_advance "# ))) .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(rescued) } pub async fn load_job(&self, pool: &PgPool, job_id: i64) -> Result, AwaError> { match self.load_job_inner(pool, job_id, true).await { Err(err) if err.is_undefined_column() => { self.forget_job_output_columns(); self.load_job_inner(pool, job_id, false).await } other => other, } } /// [`Self::load_job`] without the job's output (ADR-048), for list /// reads: a compact-completed job's output is read by detoasting its /// batch's whole output array. #[doc(hidden)] pub async fn load_job_without_output( &self, pool: &PgPool, job_id: i64, ) -> Result, AwaError> { self.load_job_inner(pool, job_id, false).await } async fn load_job_inner( &self, pool: &PgPool, job_id: i64, include_output: bool, ) -> Result, AwaError> { let schema = self.schema(); let output_column = if include_output && self.supports_job_output(pool).await? { ",\n output" } else { "" }; let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let mut candidates = Vec::new(); let ready_rows: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT job_id, kind, queue, args, priority, attempt, run_lease, max_attempts, run_at, attempted_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM {schema}.ready_entries WHERE job_id = $1 ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC "#, ))) .bind(job_id) .fetch_all(pool) .await .map_err(map_sqlx_error)?; for row in ready_rows { candidates.push(row.into_job_row()?); } let deferred_rows: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload FROM {schema}.deferred_jobs WHERE job_id = $1 "#, ))) .bind(job_id) .fetch_all(pool) .await .map_err(map_sqlx_error)?; for row in deferred_rows { candidates.push(row.into_job_row()?); } let lease_rows: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT lease.ready_slot, lease.ready_generation, lease.job_id, ready.kind, ready.queue, ready.args, lease.state, lease.priority, lease.attempt, lease.run_lease, lease.max_attempts, lease.lane_seq, ready.run_at, COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at, lease.deadline_at, lease.attempted_at, NULL::timestamptz AS finalized_at, ready.created_at, ready.unique_key, ready.unique_states, lease.callback_id, lease.callback_timeout_at, attempt.callback_filter, attempt.callback_on_complete, attempt.callback_on_fail, attempt.callback_transform, COALESCE(ready.payload, '{{}}'::jsonb) AS payload, attempt.progress, attempt.callback_result FROM {schema}.leases AS lease JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue AND ready.priority = lease.priority AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease WHERE lease.job_id = $1 ORDER BY lease.run_lease DESC "#, ))) .bind(job_id) .fetch_all(pool) .await .map_err(map_sqlx_error)?; for row in lease_rows { candidates.push(row.into_job_row()?); } // Report receipt-backed attempts as running by anti-joining // lease_claims against every durable closure evidence shape. let lease_claim_rows: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT claims.ready_slot, claims.ready_generation, claims.job_id, ready.kind, ready.queue, ready.args, 'running'::awa.job_state AS state, claims.priority, claims.attempt, claims.run_lease, claims.max_attempts, claims.lane_seq, ready.run_at, attempt.heartbeat_at, claims.deadline_at, claims.claimed_at AS attempted_at, NULL::timestamptz AS finalized_at, ready.created_at, ready.unique_key, ready.unique_states, NULL::uuid AS callback_id, NULL::timestamptz AS callback_timeout_at, attempt.callback_filter, attempt.callback_on_complete, attempt.callback_on_fail, attempt.callback_transform, COALESCE(ready.payload, '{{}}'::jsonb) AS payload, attempt.progress, attempt.callback_result FROM {schema}.lease_claims AS claims JOIN {schema}.ready_entries AS ready ON ready.ready_slot = claims.ready_slot AND ready.ready_generation = claims.ready_generation AND ready.queue = claims.queue AND ready.priority = claims.priority AND ready.enqueue_shard = claims.enqueue_shard AND ready.lane_seq = claims.lane_seq LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = claims.job_id AND attempt.run_lease = claims.run_lease WHERE claims.job_id = $1 AND NOT {closed_evidence} -- Exclude claims that have already been materialized into -- leases — the lease-backed branch above already reports -- those. AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id AND lease.run_lease = claims.run_lease ) ORDER BY claims.run_lease DESC "#, ))) .bind(job_id) .fetch_all(pool) .await .map_err(map_sqlx_error)?; for row in lease_claim_rows { candidates.push(row.into_job_row()?); } // Receipt claims are stored as compact batches. Expand // them only for this admin read, and report still-open items as // running until durable closure, terminal, or materialized-lease // evidence supersedes the claim. let lease_claim_batch_rows: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, ready.kind, ready.queue, ready.args, 'running'::awa.job_state AS state, claim_batches.priority, items.attempt, items.run_lease, items.max_attempts, items.lane_seq, ready.run_at, attempt.heartbeat_at, claim_batches.deadline_at, claim_batches.claimed_at AS attempted_at, NULL::timestamptz AS finalized_at, ready.created_at, ready.unique_key, ready.unique_states, NULL::uuid AS callback_id, NULL::timestamptz AS callback_timeout_at, attempt.callback_filter, attempt.callback_on_complete, attempt.callback_on_fail, attempt.callback_transform, COALESCE(ready.payload, '{{}}'::jsonb) AS payload, attempt.progress, attempt.callback_result FROM {schema}.lease_claim_batches AS claim_batches CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, claim_batches.receipt_ids, claim_batches.lane_seqs, claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) JOIN {schema}.ready_entries AS ready ON ready.ready_slot = claim_batches.ready_slot AND ready.ready_generation = claim_batches.ready_generation AND ready.queue = claim_batches.queue AND ready.enqueue_shard = claim_batches.enqueue_shard AND ready.lane_seq = items.lane_seq AND ready.job_id = items.job_id LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = items.job_id AND attempt.run_lease = items.run_lease WHERE items.job_id = $1 AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = claim_batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = items.job_id AND lease.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = items.job_id AND done.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = items.job_id AND deferred.run_lease = items.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) ORDER BY items.run_lease DESC "#, ))) .bind(job_id) .fetch_all(pool) .await .map_err(map_sqlx_error)?; for row in lease_claim_batch_rows { candidates.push(row.into_job_row()?); } let done_rows: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT ready_slot, ready_generation, job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload{output_column} FROM {schema}.terminal_jobs AS done WHERE done.job_id = $1 ORDER BY done.run_lease DESC, done.finalized_at DESC "#, ))) .bind(job_id) .fetch_all(pool) .await .map_err(map_sqlx_error)?; for row in done_rows { candidates.push(row.into_job_row()?); } let dlq_rows: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload, dlq_reason, dlq_at, original_run_lease FROM {schema}.dlq_entries WHERE job_id = $1 ORDER BY dlq_at DESC "#, ))) .bind(job_id) .fetch_all(pool) .await .map_err(map_sqlx_error)?; for row in dlq_rows { candidates.push(row.into_job_row()?); } Ok(candidates.into_iter().max_by_key(|job| { ( job.run_lease, transition_timestamp(job), state_rank(job.state), ) })) } pub async fn register_callback( &self, pool: &PgPool, job_id: i64, run_lease: i64, timeout: Duration, ) -> Result { let callback_id = Uuid::new_v4(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; self.ensure_mutable_running_attempt_tx(&mut tx, job_id, run_lease) .await?; let updated = sqlx::query(audited_sql(format!( r#" UPDATE {} SET callback_id = $2, callback_timeout_at = clock_timestamp() + make_interval(secs => $3) WHERE job_id = $1 AND state = 'running' AND run_lease = $4 "#, self.leases_table() ))) .bind(job_id) .bind(callback_id) .bind(timeout.as_secs_f64()) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; if updated.rows_affected() == 0 { tx.rollback().await.map_err(map_sqlx_error)?; return Err(AwaError::Validation("job is not in running state".into())); } sqlx::query(audited_sql(format!( r#" UPDATE {} SET callback_filter = NULL, callback_on_complete = NULL, callback_on_fail = NULL, callback_transform = NULL, updated_at = clock_timestamp() WHERE job_id = $1 AND run_lease = $2 "#, self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( r#" DELETE FROM {} WHERE job_id = $1 AND run_lease = $2 AND progress IS NULL AND callback_result IS NULL AND callback_filter IS NULL AND callback_on_complete IS NULL AND callback_on_fail IS NULL AND callback_transform IS NULL "#, self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; Ok(callback_id) } pub async fn register_callback_with_config( &self, pool: &PgPool, job_id: i64, run_lease: i64, timeout: Duration, config: &CallbackConfig, ) -> Result { if config.is_empty() { return self .register_callback(pool, job_id, run_lease, timeout) .await; } #[cfg(feature = "cel")] { for (name, expr) in [ ("filter", &config.filter), ("on_complete", &config.on_complete), ("on_fail", &config.on_fail), ("transform", &config.transform), ] { if let Some(src) = expr { let program = cel::Program::compile(src).map_err(|e| { AwaError::Validation(format!("invalid CEL expression for {name}: {e}")) })?; let references = program.references(); let bad_vars: Vec = references .variables() .into_iter() .filter(|v| *v != "payload") .map(str::to_string) .collect(); if !bad_vars.is_empty() { return Err(AwaError::Validation(format!( "CEL expression for {name} references undeclared variable(s): {}; only 'payload' is available", bad_vars.join(", ") ))); } } } } #[cfg(not(feature = "cel"))] { if !config.is_empty() { return Err(AwaError::Validation( "CEL expressions require the 'cel' feature".into(), )); } } let callback_id = Uuid::new_v4(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; self.ensure_mutable_running_attempt_tx(&mut tx, job_id, run_lease) .await?; let updated = sqlx::query(audited_sql(format!( r#" UPDATE {} SET callback_id = $2, callback_timeout_at = clock_timestamp() + make_interval(secs => $3) WHERE job_id = $1 AND state = 'running' AND run_lease = $4 "#, self.leases_table() ))) .bind(job_id) .bind(callback_id) .bind(timeout.as_secs_f64()) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; if updated.rows_affected() == 0 { tx.rollback().await.map_err(map_sqlx_error)?; return Err(AwaError::Validation("job is not in running state".into())); } sqlx::query(audited_sql(format!( r#" INSERT INTO {} ( job_id, run_lease, callback_filter, callback_on_complete, callback_on_fail, callback_transform, updated_at ) VALUES ($1, $2, $3, $4, $5, $6, clock_timestamp()) ON CONFLICT (job_id, run_lease) DO UPDATE SET callback_filter = EXCLUDED.callback_filter, callback_on_complete = EXCLUDED.callback_on_complete, callback_on_fail = EXCLUDED.callback_on_fail, callback_transform = EXCLUDED.callback_transform, updated_at = clock_timestamp() "#, self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .bind(&config.filter) .bind(&config.on_complete) .bind(&config.on_fail) .bind(&config.transform) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; Ok(callback_id) } pub async fn cancel_callback( &self, pool: &PgPool, job_id: i64, run_lease: i64, ) -> Result { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = sqlx::query(audited_sql(format!( r#" UPDATE {} SET callback_id = NULL, callback_timeout_at = NULL WHERE job_id = $1 AND callback_id IS NOT NULL AND state = 'running' AND run_lease = $2 "#, self.leases_table() ))) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; if result.rows_affected() == 0 { tx.rollback().await.map_err(map_sqlx_error)?; return Ok(false); } sqlx::query(audited_sql(format!( r#" UPDATE {} SET callback_filter = NULL, callback_on_complete = NULL, callback_on_fail = NULL, callback_transform = NULL, updated_at = clock_timestamp() WHERE job_id = $1 AND run_lease = $2 "#, self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( r#" DELETE FROM {} WHERE job_id = $1 AND run_lease = $2 AND progress IS NULL AND callback_result IS NULL AND callback_filter IS NULL AND callback_on_complete IS NULL AND callback_on_fail IS NULL AND callback_transform IS NULL "#, self.attempt_state_table() ))) .bind(job_id) .bind(run_lease) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; Ok(true) } /// Load the currently-active lease row for `job_id` (running or /// waiting_external) inside a caller-owned transaction. Used by ADR-029 /// helpers that need the post-park snapshot — including the /// `callback_id` and `callback_timeout_at` written by /// `register_callback()` — without leaving the transaction that just /// performed the parking UPDATE. pub async fn load_active_lease_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, ) -> Result, AwaError> { let schema = self.schema(); let row: Option = sqlx::query_as(audited_sql(format!( r#" SELECT lease.ready_slot, lease.ready_generation, lease.job_id, ready.kind, ready.queue, ready.args, lease.state, lease.priority, lease.attempt, lease.run_lease, lease.max_attempts, lease.lane_seq, ready.run_at, COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at, lease.deadline_at, lease.attempted_at, NULL::timestamptz AS finalized_at, ready.created_at, ready.unique_key, ready.unique_states, lease.callback_id, lease.callback_timeout_at, attempt.callback_filter, attempt.callback_on_complete, attempt.callback_on_fail, attempt.callback_transform, COALESCE(ready.payload, '{{}}'::jsonb) AS payload, attempt.progress, attempt.callback_result FROM {schema}.leases AS lease JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue AND ready.priority = lease.priority AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease WHERE lease.job_id = $1 AND lease.run_lease = $2 "#, ))) .bind(job_id) .bind(run_lease) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; row.map(LeaseJobRow::into_job_row).transpose() } pub async fn enter_callback_wait( &self, pool: &PgPool, job_id: i64, run_lease: i64, callback_id: Uuid, ) -> Result { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let entered = self .enter_callback_wait_in_tx(&mut tx, job_id, run_lease, callback_id) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(entered) } /// Transaction-aware variant of [`Self::enter_callback_wait`] (ADR-029). /// Returns whether the row transitioned to `waiting_external`. pub async fn enter_callback_wait_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, callback_id: Uuid, ) -> Result { let result = sqlx::query(audited_sql(format!( r#" UPDATE {} SET state = 'waiting_external', heartbeat_at = NULL, deadline_at = NULL WHERE job_id = $1 AND state = 'running' AND run_lease = $2 AND callback_id = $3 "#, self.leases_table() ))) .bind(job_id) .bind(run_lease) .bind(callback_id) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(result.rows_affected() > 0) } pub async fn check_callback_state( &self, pool: &PgPool, job_id: i64, callback_id: Uuid, ) -> Result { let row: Option<(JobState, Option, i64, Option)> = sqlx::query_as(audited_sql(format!( r#" SELECT lease.state, lease.callback_id, lease.run_lease, attempt.callback_result FROM {} AS lease LEFT JOIN {} AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease WHERE lease.job_id = $1 ORDER BY lease.run_lease DESC LIMIT 1 "#, self.leases_table(), self.attempt_state_table() ))) .bind(job_id) .fetch_optional(pool) .await .map_err(map_sqlx_error)?; match row { Some((JobState::Running, None, run_lease, Some(_))) => { let result = self.take_callback_result(pool, job_id, run_lease).await?; Ok(CallbackPollResult::Resolved(result)) } Some((state, Some(current_callback_id), _, _)) if current_callback_id != callback_id => { Ok(CallbackPollResult::Stale { token: callback_id, current: current_callback_id, state, }) } Some((JobState::WaitingExternal, Some(current), _, _)) if current == callback_id => { Ok(CallbackPollResult::Pending) } Some((state, _, _, _)) => Ok(CallbackPollResult::UnexpectedState { token: callback_id, state, }), None => { if let Some(job) = self.load_job(pool, job_id).await? { Ok(CallbackPollResult::UnexpectedState { token: callback_id, state: job.state, }) } else { Ok(CallbackPollResult::NotFound) } } } } pub async fn callback_job( &self, pool: &PgPool, callback_id: Uuid, run_lease: Option, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .callback_job_in_tx(&mut tx, callback_id, run_lease, false) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::callback_job`] (ADR-029). /// When `for_update` is `true` the join's lease row is locked with /// `FOR UPDATE OF lease`, mirroring the canonical `resolve_callback` /// lookup that takes a row lock on `awa.jobs_hot` before evaluating /// the callback policy and committing the resulting transition. pub async fn callback_job_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, callback_id: Uuid, run_lease: Option, for_update: bool, ) -> Result, AwaError> { let lock_clause = if for_update { "FOR UPDATE OF lease" } else { "" }; let row: Option = sqlx::query_as(audited_sql(format!( r#" SELECT lease.ready_slot, lease.ready_generation, lease.job_id, ready.kind, ready.queue, ready.args, lease.state, lease.priority, lease.attempt, lease.run_lease, lease.max_attempts, lease.lane_seq, ready.run_at, COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at, lease.deadline_at, lease.attempted_at, NULL::timestamptz AS finalized_at, ready.created_at, ready.unique_key, ready.unique_states, lease.callback_id, lease.callback_timeout_at, attempt.callback_filter, attempt.callback_on_complete, attempt.callback_on_fail, attempt.callback_transform, COALESCE(ready.payload, '{{}}'::jsonb) AS payload, attempt.progress, attempt.callback_result FROM {} AS lease JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue AND ready.priority = lease.priority AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease WHERE lease.callback_id = $1 AND lease.state IN ('waiting_external', 'running') AND ($2::bigint IS NULL OR lease.run_lease = $2) ORDER BY lease.run_lease DESC LIMIT 1 {lock_clause} "#, self.leases_table(), schema = self.schema(), ))) .bind(callback_id) .bind(run_lease) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; row.map(LeaseJobRow::into_job_row).transpose() } #[tracing::instrument(skip(self, pool, payload), name = "queue_storage.complete_external")] pub async fn complete_external( &self, pool: &PgPool, callback_id: Uuid, payload: Option, run_lease: Option, resume: bool, ) -> Result { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .complete_external_in_tx(&mut tx, callback_id, payload, run_lease, resume) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::complete_external`] (ADR-029). /// The non-resume path returns the post-completion `JobRow` directly /// from the `done_row` insert. The resume path returns the parked-row /// snapshot reloaded inside the same transaction via /// [`Self::load_active_lease_in_tx`] — i.e. it does not leave the /// caller's transaction to refresh state. pub async fn complete_external_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, callback_id: Uuid, payload: Option, run_lease: Option, resume: bool, ) -> Result { if resume { let resumed: Option<(i64, i64)> = sqlx::query_as(audited_sql(format!( r#" UPDATE {} SET state = 'running', callback_id = NULL, callback_timeout_at = NULL, heartbeat_at = clock_timestamp() WHERE callback_id = $1 AND state IN ('waiting_external', 'running') AND ($2::bigint IS NULL OR run_lease = $2) RETURNING job_id, run_lease "#, self.leases_table() ))) .bind(callback_id) .bind(run_lease) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some((job_id, resumed_run_lease)) = resumed else { return Err(AwaError::CallbackNotFound { callback_id: callback_id.to_string(), }); }; sqlx::query(audited_sql(format!( r#" INSERT INTO {} ( job_id, run_lease, callback_filter, callback_on_complete, callback_on_fail, callback_transform, callback_result, updated_at ) VALUES ($1, $2, NULL, NULL, NULL, NULL, $3, clock_timestamp()) ON CONFLICT (job_id, run_lease) DO UPDATE SET callback_filter = NULL, callback_on_complete = NULL, callback_on_fail = NULL, callback_transform = NULL, callback_result = EXCLUDED.callback_result, updated_at = clock_timestamp() "#, self.attempt_state_table() ))) .bind(job_id) .bind(resumed_run_lease) .bind(payload.unwrap_or(serde_json::Value::Null)) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; return self .load_active_lease_in_tx(tx, job_id, resumed_run_lease) .await? .ok_or(AwaError::CallbackNotFound { callback_id: callback_id.to_string(), }); } self.complete_external_with_output_in_tx(tx, callback_id, run_lease, None) .await } /// Complete a job waiting on `callback_id` and store `output` as its /// output (ADR-048), if this schema has the output columns. The caller /// validates the output. Returns the completed row, whose `output` is /// `None` when it could not be stored. #[doc(hidden)] pub async fn complete_external_with_output_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, callback_id: Uuid, run_lease: Option, output: Option, ) -> Result { let output = match output { Some(output) if self.supports_job_output(tx.as_mut()).await? => Some(output), _ => None, }; let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE callback_id = $1 AND state IN ('waiting_external', 'running') AND ($2::bigint IS NULL OR run_lease = $2) RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .bind(callback_id) .bind(run_lease) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { return Err(AwaError::CallbackNotFound { callback_id: callback_id.to_string(), }); } let completed_pairs: Vec<(i64, i64)> = deleted .iter() .map(|row| (row.job_id, row.run_lease)) .collect(); self.close_receipt_pairs_tx(tx, &completed_pairs, "completed") .await?; let moved = self.hydrate_deleted_leases_tx(tx, deleted, true).await?; let moved = moved.into_iter().next().expect("deleted callback lease"); let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state( moved.payload.clone(), moved.progress.clone(), )?)?; payload.set_progress(None); let mut done_row = moved .clone() .into_done_row(JobState::Completed, Utc::now(), payload.into_json()); done_row.output = output; self.insert_done_rows_tx(tx, std::slice::from_ref(&done_row), Some(moved.state)) .await?; done_row.into_job_row() } pub async fn fail_external( &self, pool: &PgPool, callback_id: Uuid, error: &str, run_lease: Option, ) -> Result { self.fail_external_with_error_entry( pool, callback_id, serde_json::json!({ "error": error }), run_lease, ) .await } pub async fn fail_external_with_error_entry( &self, pool: &PgPool, callback_id: Uuid, error_entry: serde_json::Value, run_lease: Option, ) -> Result { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .fail_external_with_error_entry_in_tx(&mut tx, callback_id, error_entry, run_lease) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::fail_external_with_error_entry`] /// (ADR-029). pub async fn fail_external_with_error_entry_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, callback_id: Uuid, error_entry: serde_json::Value, run_lease: Option, ) -> Result { let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE callback_id = $1 AND state IN ('waiting_external', 'running') AND ($2::bigint IS NULL OR run_lease = $2) RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .bind(callback_id) .bind(run_lease) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { return Err(AwaError::CallbackNotFound { callback_id: callback_id.to_string(), }); } let failed_pairs: Vec<(i64, i64)> = deleted .iter() .map(|row| (row.job_id, row.run_lease)) .collect(); self.close_receipt_pairs_tx(tx, &failed_pairs, "failed") .await?; let moved = self.hydrate_deleted_leases_tx(tx, deleted, true).await?; let moved = moved.into_iter().next().expect("deleted callback lease"); let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state( moved.payload.clone(), moved.progress.clone(), )?)?; let mut error_entry = match error_entry { serde_json::Value::Object(map) => serde_json::Value::Object(map), other => serde_json::json!({ "error": other }), }; let error_obj = error_entry .as_object_mut() .ok_or_else(|| AwaError::Validation("callback error entry must be an object".into()))?; error_obj .entry("attempt".to_string()) .or_insert_with(|| serde_json::Value::from(i64::from(moved.attempt))); error_obj .entry("at".to_string()) .or_insert_with(|| serde_json::Value::String(Utc::now().to_rfc3339())); error_obj .entry("terminal".to_string()) .or_insert(serde_json::Value::Bool(true)); payload.push_error(error_entry); let done_row = moved .clone() .into_done_row(JobState::Failed, Utc::now(), payload.into_json()); self.insert_done_rows_tx(tx, std::slice::from_ref(&done_row), Some(moved.state)) .await?; done_row.into_job_row() } pub async fn retry_external( &self, pool: &PgPool, callback_id: Uuid, run_lease: Option, ) -> Result { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .retry_external_in_tx(&mut tx, callback_id, run_lease) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::retry_external`] (ADR-029). pub async fn retry_external_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, callback_id: Uuid, run_lease: Option, ) -> Result { let schema = self.schema(); let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE callback_id = $1 AND state = 'waiting_external' AND ($2::bigint IS NULL OR run_lease = $2) RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .bind(callback_id) .bind(run_lease) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { return Err(AwaError::CallbackNotFound { callback_id: callback_id.to_string(), }); } let retryable_pairs: Vec<(i64, i64)> = deleted .iter() .map(|row| (row.job_id, row.run_lease)) .collect(); self.close_receipt_pairs_tx(tx, &retryable_pairs, "retryable") .await?; let moved = self.hydrate_deleted_leases_tx(tx, deleted, true).await?; let moved = moved.into_iter().next().expect("deleted callback lease"); let ready_payload = Self::payload_with_attempt_state(moved.payload.clone(), moved.progress.clone())?; let ready_row = ExistingReadyRow { attempt: 0, run_at: Utc::now(), ..moved.clone().into_ready_row(Utc::now(), ready_payload) }; self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(moved.state)) .await?; self.notify_queues_tx(tx, std::iter::once(moved.queue.clone())) .await?; ReadyJobRow { job_id: ready_row.job_id, kind: ready_row.kind, queue: ready_row.queue, args: ready_row.args, priority: ready_row.priority, attempt: ready_row.attempt, run_lease: ready_row.run_lease, max_attempts: ready_row.max_attempts, run_at: ready_row.run_at, attempted_at: ready_row.attempted_at, created_at: ready_row.created_at, unique_key: ready_row.unique_key, payload: ready_row.payload, } .into_job_row() } pub async fn heartbeat_callback( &self, pool: &PgPool, callback_id: Uuid, timeout: Duration, ) -> Result { let updated: Option<(i64, i64)> = sqlx::query_as(audited_sql(format!( r#" UPDATE {} SET callback_timeout_at = clock_timestamp() + make_interval(secs => $2) WHERE callback_id = $1 AND state = 'waiting_external' RETURNING job_id, run_lease "#, self.leases_table() ))) .bind(callback_id) .bind(timeout.as_secs_f64()) .fetch_optional(pool) .await .map_err(map_sqlx_error)?; let Some((job_id, _run_lease)) = updated else { return Err(AwaError::CallbackNotFound { callback_id: callback_id.to_string(), }); }; self.load_job(pool, job_id) .await? .ok_or(AwaError::CallbackNotFound { callback_id: callback_id.to_string(), }) } pub async fn flush_progress( &self, pool: &PgPool, job_id: i64, run_lease: i64, progress: serde_json::Value, ) -> Result<(), AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; if self.lease_claim_receipts() { self.upsert_attempt_state_progress_from_receipts_tx( &mut tx, &[(job_id, run_lease, progress.clone())], ) .await?; } sqlx::query(audited_sql(format!( r#" INSERT INTO {} (job_id, run_lease, progress, updated_at) SELECT lease.job_id, lease.run_lease, $3, clock_timestamp() FROM {} AS lease WHERE lease.job_id = $1 AND lease.run_lease = $2 AND lease.state IN ('running', 'waiting_external') ON CONFLICT (job_id, run_lease) DO UPDATE SET progress = EXCLUDED.progress, updated_at = clock_timestamp() "#, self.attempt_state_table(), self.leases_table() ))) .bind(job_id) .bind(run_lease) .bind(progress) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; Ok(()) } /// Heartbeat compact receipt claims by their batch identity. Each claim /// batch is addressed by primary key and the attempt by item index, with /// index-only closure probes and no batch row lock, so the cost is /// O(in-flight attempts) rather than O(open claim batches). pub async fn heartbeat_claims_batch( &self, pool: &PgPool, claims: &[ClaimHeartbeatRef], ) -> Result { if claims.is_empty() { return Ok(0); } let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let pairs: Vec<(i64, i64)> = claims .iter() .map(|claim| (claim.job_id, claim.run_lease)) .collect(); self.lock_receipt_attempt_keys_tx(&mut tx, &pairs).await?; let job_ids: Vec = claims.iter().map(|claim| claim.job_id).collect(); let run_leases: Vec = claims.iter().map(|claim| claim.run_lease).collect(); let claim_slots: Vec = claims.iter().map(|claim| claim.claim_slot).collect(); let batch_ids: Vec = claims.iter().map(|claim| claim.claim_batch_id).collect(); let batch_indices: Vec = claims.iter().map(|claim| claim.claim_batch_index).collect(); let receipt_ids: Vec = claims.iter().map(|claim| claim.receipt_id).collect(); let updated: i64 = sqlx::query_scalar(audited_sql(format!( r#" WITH inflight(job_id, run_lease, claim_slot, batch_id, batch_index, receipt_id) AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::int[], $4::bigint[], $5::int[], $6::bigint[]) ), open_claims AS ( SELECT inflight.job_id, inflight.run_lease FROM inflight JOIN {schema}.lease_claim_batches AS claim_batches ON claim_batches.claim_slot = inflight.claim_slot AND claim_batches.batch_id = inflight.batch_id WHERE inflight.batch_index BETWEEN 1 AND claim_batches.claimed_count AND claim_batches.job_ids[inflight.batch_index] = inflight.job_id AND claim_batches.run_leases[inflight.batch_index] = inflight.run_lease AND claim_batches.receipt_ids[inflight.batch_index] = inflight.receipt_id AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = inflight.claim_slot AND closures.job_id = inflight.job_id AND closures.run_lease = inflight.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = inflight.claim_slot AND closure_batches.receipt_ranges @> inflight.receipt_id ) AND NOT EXISTS ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = inflight.job_id AND done.run_lease = inflight.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = inflight.job_id AND deferred.run_lease = inflight.run_lease ) AND NOT EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = inflight.job_id AND dlq.run_lease = inflight.run_lease ) ), upserted AS ( INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at) SELECT open_claims.job_id, open_claims.run_lease, clock_timestamp(), clock_timestamp() FROM open_claims ON CONFLICT (job_id, run_lease) DO UPDATE SET heartbeat_at = clock_timestamp(), updated_at = clock_timestamp() RETURNING job_id ) SELECT count(*)::bigint FROM upserted "# ))) .bind(&job_ids) .bind(&run_leases) .bind(&claim_slots) .bind(&batch_ids) .bind(&batch_indices) .bind(&receipt_ids) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; Ok(updated as usize) } pub async fn heartbeat_batch( &self, pool: &PgPool, jobs: &[(i64, i64)], ) -> Result { if jobs.is_empty() { return Ok(0); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let mut updated = 0_usize; if self.lease_claim_receipts() { // #169 B1: in receipts mode, attempt_state is the // authoritative heartbeat home. Skip the `UPDATE leases SET // heartbeat_at = ...` entirely — that write was the // dominant per-heartbeat non-HOT update on the partitioned // `leases` table (every state-indexed partial index pays 2 // dead entries per write under sustained churn). Compat // reads in `LeaseJobRow` SELECTs and the `awa.jobs` view // COALESCE attempt_state.heartbeat_at first. updated += self .upsert_attempt_state_from_receipts_tx(&mut tx, jobs) .await?; } else { // Legacy non-receipts mode (custom schemas with // `lease_claim_receipts=FALSE`): the `leases.heartbeat_at` // write is still the heartbeat home, since there is no // upsert_attempt_state path firing. let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); let result = sqlx::query(audited_sql(format!( r#" WITH inflight AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[]) AS v(job_id, run_lease) ) UPDATE {table} SET heartbeat_at = clock_timestamp() FROM inflight WHERE {table}.job_id = inflight.job_id AND {table}.run_lease = inflight.run_lease AND {table}.state = 'running' "#, table = self.leases_table(), ))) .bind(&job_ids) .bind(&run_leases) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; updated += result.rows_affected() as usize; } tx.commit().await.map_err(map_sqlx_error)?; Ok(updated) } pub async fn heartbeat_progress_batch( &self, pool: &PgPool, jobs: &[(i64, i64, serde_json::Value)], ) -> Result { if jobs.is_empty() { return Ok(0); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let updated = if self.lease_claim_receipts() { // #169 B1: receipts mode is the only supported shape for // the default `awa` schema. attempt_state already carries // heartbeat_at + progress, so the `UPDATE leases SET // heartbeat_at` + nested `INSERT INTO attempt_state` CTE // is collapsed into the single attempt_state upsert. The // upsert sources open-claim identity from `lease_claims` // anti-joined against durable closure evidence so it picks // up every open claim regardless of whether // `materialize_claims` has fanned it out to `leases` yet. self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs) .await? } else { // Legacy non-receipts mode keeps the old CTE shape that // updates leases.heartbeat_at + leases.progress // (via attempt_state upsert) in a single round-trip. let schema = self.schema(); let job_ids: Vec = jobs.iter().map(|(job_id, _, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect(); let progress: Vec = jobs.iter().map(|(_, _, value)| value.clone()).collect(); let lease_updated: i64 = sqlx::query_scalar(audited_sql(format!( r#" WITH inflight AS ( SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::jsonb[]) AS v(job_id, run_lease, progress) ), updated AS ( UPDATE {table} AS lease SET heartbeat_at = clock_timestamp() FROM inflight WHERE lease.job_id = inflight.job_id AND lease.run_lease = inflight.run_lease AND lease.state = 'running' RETURNING lease.job_id, lease.run_lease, inflight.progress ), upsert_attempt AS ( INSERT INTO {schema}.attempt_state (job_id, run_lease, progress, updated_at) SELECT job_id, run_lease, progress, clock_timestamp() FROM updated ON CONFLICT (job_id, run_lease) DO UPDATE SET progress = EXCLUDED.progress, updated_at = clock_timestamp() ) SELECT count(*)::bigint FROM updated "#, table = self.leases_table(), ))) .bind(&job_ids) .bind(&run_leases) .bind(&progress) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; lease_updated as usize }; tx.commit().await.map_err(map_sqlx_error)?; Ok(updated) } pub async fn retry_after( &self, pool: &PgPool, job_id: i64, run_lease: i64, retry_after: Duration, progress: Option, ) -> Result, AwaError> { self.retry_after_with_claim(pool, job_id, run_lease, retry_after, progress, None) .await } pub async fn retry_after_with_claim( &self, pool: &PgPool, job_id: i64, run_lease: i64, retry_after: Duration, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .retry_after_in_tx_with_claim(&mut tx, job_id, run_lease, retry_after, progress, claim) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::retry_after`]. Caller owns the /// transaction lifecycle so the move can commit atomically alongside /// follow-up `INSERT`s (ADR-029). pub async fn retry_after_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, retry_after: Duration, progress: Option, ) -> Result, AwaError> { self.retry_after_in_tx_with_claim(tx, job_id, run_lease, retry_after, progress, None) .await } pub async fn retry_after_in_tx_with_claim( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, retry_after: Duration, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { self.retry_after_with_output_in_tx( tx, job_id, run_lease, retry_after, progress, claim, None, ) .await } /// [`Self::retry_after_with_claim`] for an attempt that recorded an /// output (ADR-048). #[doc(hidden)] #[allow(clippy::too_many_arguments)] pub async fn retry_after_with_output( &self, pool: &PgPool, job_id: i64, run_lease: i64, retry_after: Duration, progress: Option, claim: Option<&ClaimedEntry>, output: Option, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .retry_after_with_output_in_tx( &mut tx, job_id, run_lease, retry_after, progress, claim, output, ) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// [`Self::retry_after_in_tx_with_claim`] for an attempt that recorded an /// output. `RetryAfter` writes no error entry of its own, so one is /// written to hold the output, and only then (ADR-048). #[doc(hidden)] #[allow(clippy::too_many_arguments)] pub async fn retry_after_with_output_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, retry_after: Duration, progress: Option, claim: Option<&ClaimedEntry>, output: Option, ) -> Result, AwaError> { let Some(moved) = self .take_running_attempt_tx(tx, job_id, run_lease, "retryable", claim) .await? else { return Ok(None); }; let now = self.current_timestamp_tx(tx).await?; let mut payload = Self::with_progress(moved.payload.clone(), progress.or(moved.progress.clone()))?; if output.is_some() { let mut runtime_payload = RuntimePayload::from_json(payload)?; runtime_payload.push_error(lifecycle_entry( AttemptError::new(format!( "retry requested after {}s", retry_after.as_secs_f64() )) .with_output(output), moved.attempt, false, )); payload = runtime_payload.into_json(); } let deferred = moved.clone().into_deferred_row( JobState::Retryable, now + TimeDelta::from_std(retry_after).map_err(|err| { AwaError::Validation(format!("invalid retry_after duration: {err}")) })?, Some(now), payload, ); self.insert_deferred_rows_tx(tx, vec![deferred.clone()], Some(moved.state)) .await?; Ok(Some(deferred.into_job_row()?)) } pub async fn snooze( &self, pool: &PgPool, job_id: i64, run_lease: i64, snooze_for: Duration, progress: Option, ) -> Result, AwaError> { self.snooze_with_claim(pool, job_id, run_lease, snooze_for, progress, None) .await } pub async fn snooze_with_claim( &self, pool: &PgPool, job_id: i64, run_lease: i64, snooze_for: Duration, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let Some(moved) = self .take_running_attempt_tx(&mut tx, job_id, run_lease, "scheduled", claim) .await? else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(None); }; let now = self.current_timestamp_tx(&mut tx).await?; let payload = Self::with_progress(moved.payload.clone(), progress.or(moved.progress.clone()))?; let mut deferred = moved.clone().into_deferred_row( JobState::Scheduled, now + TimeDelta::from_std(snooze_for) .map_err(|err| AwaError::Validation(format!("invalid snooze duration: {err}")))?, None, payload, ); deferred.attempt = deferred.attempt.saturating_sub(1); self.insert_deferred_rows_tx(&mut tx, vec![deferred.clone()], Some(moved.state)) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(Some(deferred.into_job_row()?)) } /// Insert a canonical attempt's transition-time deferred successor, /// degrading to a cancelled terminal row if a newer duplicate already /// owns the unique claim. Canonical deletion releases any claim before /// this call, so the deferred insert must acquire from `None` even when /// the old `running` state was inside the mask. pub(crate) async fn insert_migrated_deferred_or_cancel_duplicate_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: JobRow, unique_states: Option, duplicate_error: &str, ) -> Result { let payload = Self::payload_from_parts(row.metadata, row.tags, row.errors, row.progress)?; let deferred = DeferredJobRow { job_id: row.id, kind: row.kind, queue: row.queue, args: row.args, state: row.state, priority: row.priority, attempt: row.attempt, run_lease: row.run_lease, max_attempts: row.max_attempts, run_at: row.run_at, attempted_at: row.attempted_at, finalized_at: row.finalized_at, created_at: row.created_at, unique_key: row.unique_key, unique_states, payload, }; { let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; match self .insert_deferred_rows_tx(&mut savepoint, vec![deferred.clone()], None) .await { Ok(_) => { savepoint.commit().await.map_err(map_sqlx_error)?; return deferred.into_job_row(); } Err(AwaError::UniqueConflict { .. }) => { savepoint.rollback().await.map_err(map_sqlx_error)?; } Err(err) => return Err(err), } } let mut cancelled_payload = RuntimePayload::from_json(deferred.payload.clone())?; cancelled_payload.push_error(lifecycle_error(duplicate_error, deferred.attempt, true)); let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?; self.ensure_lane(tx, &deferred.queue, deferred.priority, 0) .await?; let mut done = DoneJobRow { ready_slot, ready_generation, job_id: deferred.job_id, kind: deferred.kind, queue: deferred.queue, args: deferred.args, state: JobState::Cancelled, priority: deferred.priority, attempt: deferred.attempt, run_lease: deferred.run_lease, max_attempts: deferred.max_attempts, // The successor never entered a ready lane. Mirror deferred-job // cancellation with a synthetic negative sequence on shard 0. lane_seq: -deferred.job_id, enqueue_shard: 0, run_at: deferred.run_at, attempted_at: deferred.attempted_at, finalized_at: self.current_timestamp_tx(tx).await?, created_at: deferred.created_at, unique_key: deferred.unique_key, unique_states: deferred.unique_states, payload: cancelled_payload.into_json(), output: None, }; { let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; match self .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), None) .await { Ok(_) => { savepoint.commit().await.map_err(map_sqlx_error)?; return done.into_job_row(); } Err(AwaError::UniqueConflict { .. }) => { savepoint.rollback().await.map_err(map_sqlx_error)?; } Err(err) => return Err(err), } } // If `cancelled` itself claims uniqueness, the newer holder still // wins. The cancelled evidence carries no claim to release, so strip // the uniqueness fields exactly like rescue duplicate cancellation. done.unique_key = None; done.unique_states = None; self.insert_done_rows_tx(tx, std::slice::from_ref(&done), None) .await?; done.into_job_row() } /// Rescue-path fallback for unique-claim conflicts: insert the rescued /// attempt's deferred successor, degrading to a cancelled terminal row /// when the job's unique claim is held by a newer duplicate — the claim /// holder wins, and one poisoned row must not abort the whole batched /// rescue transaction. The insert attempt runs inside a savepoint so the /// conflict leaves the outer transaction usable. async fn insert_rescued_deferred_or_cancel_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: LeaseTransitionRow, deferred: DeferredJobRow, duplicate_error: &str, ) -> Result { let old_state = row.state; { let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; match self .insert_deferred_rows_tx(&mut savepoint, vec![deferred.clone()], Some(old_state)) .await { Ok(_) => { savepoint.commit().await.map_err(map_sqlx_error)?; return deferred.into_job_row(); } Err(AwaError::UniqueConflict { .. }) => { savepoint.rollback().await.map_err(map_sqlx_error)?; } Err(err) => return Err(err), } } self.cancel_rescued_duplicate_tx(tx, row, deferred.payload, duplicate_error) .await } /// Settle one attempt that maintenance rescued (stale heartbeat, expired /// deadline or callback timeout): back to `retryable` with backoff while /// attempts remain, `failed` once the rescued attempt was the last. A job /// that kills or hangs its worker every time therefore stops after /// `max_attempts`, like one whose handler returns an error. /// /// An exhausted attempt on a queue whose DLQ is enabled (`dlq_reason` /// is `Some`) goes straight to `dlq_entries` in this transaction, as /// `fail_to_dlq` does for a handler error, so there is no window in /// which it sits `failed` outside the DLQ. async fn settle_rescued_attempt_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: LeaseTransitionRow, error: &str, dlq_reason: Option<&str>, ) -> Result { let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state( row.payload.clone(), row.progress.clone(), )?)?; let exhausted = row.attempt >= row.max_attempts; payload.push_error(lifecycle_error(error, row.attempt, exhausted)); let duplicate_error = "rescued as duplicate: unique claim held by a newer job"; if exhausted { let finalized_at = Utc::now(); if let Some(dlq_reason) = dlq_reason { let dlq = row.clone().into_dlq_row( finalized_at, payload.into_json(), dlq_reason.to_string(), finalized_at, ); return self .insert_rescued_dlq_or_cancel_tx(tx, row, dlq, duplicate_error) .await; } let done = row.clone() .into_done_row(JobState::Failed, finalized_at, payload.into_json()); self.insert_rescued_done_or_cancel_tx(tx, row, done, duplicate_error) .await } else { let deferred = row.clone().into_deferred_row( JobState::Retryable, self.backoff_at_tx(tx, row.attempt, row.max_attempts) .await?, Some(Utc::now()), payload.into_json(), ); self.insert_rescued_deferred_or_cancel_tx(tx, row, deferred, duplicate_error) .await } } /// Same fallback for rescue arms that write a terminal row directly /// (an exhausted rescued attempt): a `failed` insert that conflicts /// with a newer duplicate's claim becomes a cancellation instead. async fn insert_rescued_done_or_cancel_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: LeaseTransitionRow, done: DoneJobRow, duplicate_error: &str, ) -> Result { let old_state = row.state; { let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; match self .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state)) .await { Ok(_) => { savepoint.commit().await.map_err(map_sqlx_error)?; return done.into_job_row(); } Err(AwaError::UniqueConflict { .. }) => { savepoint.rollback().await.map_err(map_sqlx_error)?; } Err(err) => return Err(err), } } self.cancel_rescued_duplicate_tx(tx, row, done.payload, duplicate_error) .await } /// The DLQ counterpart of [`Self::insert_rescued_done_or_cancel_tx`]: a /// DLQ row holds its unique claim as `failed`, so it conflicts with a /// newer duplicate exactly when the `failed` terminal row would. async fn insert_rescued_dlq_or_cancel_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: LeaseTransitionRow, dlq: DlqJobRow, duplicate_error: &str, ) -> Result { let old_state = row.state; { let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; match self .insert_dlq_rows_tx(&mut savepoint, std::slice::from_ref(&dlq), Some(old_state)) .await { Ok(_) => { savepoint.commit().await.map_err(map_sqlx_error)?; return dlq.into_job_row(); } Err(AwaError::UniqueConflict { .. }) => { savepoint.rollback().await.map_err(map_sqlx_error)?; } Err(err) => return Err(err), } } self.cancel_rescued_duplicate_tx(tx, row, dlq.payload, duplicate_error) .await } /// Write the cancelled terminal row for a rescue that lost its unique /// claim to a newer duplicate. `cancelled` normally sits outside the /// claiming mask so the insert is conflict-free; a mask that claims /// `cancelled` as well conflicts again, and then the row is written with /// its unique key stripped — the attempt held no claim (that is what /// made it conflict), so there is nothing to release, and losing the key /// on this terminal row beats aborting rescue for the whole cluster. async fn cancel_rescued_duplicate_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, row: LeaseTransitionRow, base_payload: serde_json::Value, duplicate_error: &str, ) -> Result { let old_state = row.state; let mut payload = RuntimePayload::from_json(base_payload)?; payload.push_error(lifecycle_error(duplicate_error, row.attempt, true)); let done = row.into_done_row(JobState::Cancelled, Utc::now(), payload.into_json()); { let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; match self .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state)) .await { Ok(_) => { savepoint.commit().await.map_err(map_sqlx_error)?; return done.into_job_row(); } Err(AwaError::UniqueConflict { .. }) => { savepoint.rollback().await.map_err(map_sqlx_error)?; } Err(err) => return Err(err), } } let mut stripped = done; stripped.unique_key = None; stripped.unique_states = None; self.insert_done_rows_tx(tx, std::slice::from_ref(&stripped), Some(old_state)) .await?; stripped.into_job_row() } pub async fn cancel_running( &self, pool: &PgPool, job_id: i64, run_lease: i64, reason: impl Into, progress: Option, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .cancel_running_in_tx(&mut tx, job_id, run_lease, reason, progress) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::cancel_running`] (ADR-029). pub async fn cancel_running_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, reason: impl Into, progress: Option, ) -> Result, AwaError> { self.cancel_running_in_tx_with_claim(tx, job_id, run_lease, reason, progress, None) .await } pub async fn cancel_running_in_tx_with_claim( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, reason: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let Some(moved) = self .take_running_attempt_tx(tx, job_id, run_lease, "cancelled", claim) .await? else { return Ok(None); }; let mut payload = RuntimePayload::from_json(Self::with_progress( moved.payload.clone(), progress.or(moved.progress.clone()), )?)?; let reason = reason.into(); payload.push_error(lifecycle_entry( AttemptError::new(format!("cancelled: {}", reason.message)).with_output(reason.output), moved.attempt, false, )); let done = moved .clone() .into_done_row(JobState::Cancelled, Utc::now(), payload.into_json()); self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state)) .await?; Ok(Some(done.into_job_row()?)) } pub async fn fail_terminal( &self, pool: &PgPool, job_id: i64, run_lease: i64, error: impl Into, progress: Option, ) -> Result, AwaError> { self.fail_terminal_with_claim(pool, job_id, run_lease, error, progress, None) .await } pub async fn fail_terminal_with_claim( &self, pool: &PgPool, job_id: i64, run_lease: i64, error: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .fail_terminal_in_tx_with_claim(&mut tx, job_id, run_lease, error, progress, claim) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::fail_terminal`] (ADR-029). pub async fn fail_terminal_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, error: impl Into, progress: Option, ) -> Result, AwaError> { self.fail_terminal_in_tx_with_claim(tx, job_id, run_lease, error, progress, None) .await } pub async fn fail_terminal_in_tx_with_claim( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, error: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let Some(moved) = self .take_running_attempt_tx(tx, job_id, run_lease, "failed", claim) .await? else { return Ok(None); }; let mut payload = RuntimePayload::from_json(Self::with_progress( moved.payload.clone(), progress.or(moved.progress.clone()), )?)?; payload.push_error(lifecycle_entry(error.into(), moved.attempt, true)); let done = moved .clone() .into_done_row(JobState::Failed, Utc::now(), payload.into_json()); self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state)) .await?; Ok(Some(done.into_job_row()?)) } pub async fn fail_to_dlq( &self, pool: &PgPool, job_id: i64, run_lease: i64, dlq_reason: &str, error: impl Into, progress: Option, ) -> Result, AwaError> { self.fail_to_dlq_with_claim(pool, job_id, run_lease, dlq_reason, error, progress, None) .await } #[allow(clippy::too_many_arguments)] pub async fn fail_to_dlq_with_claim( &self, pool: &PgPool, job_id: i64, run_lease: i64, dlq_reason: &str, error: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .fail_to_dlq_in_tx_with_claim( &mut tx, job_id, run_lease, dlq_reason, error, progress, claim, ) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::fail_to_dlq`] (ADR-029). pub async fn fail_to_dlq_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, dlq_reason: &str, error: impl Into, progress: Option, ) -> Result, AwaError> { self.fail_to_dlq_in_tx_with_claim(tx, job_id, run_lease, dlq_reason, error, progress, None) .await } #[allow(clippy::too_many_arguments)] pub async fn fail_to_dlq_in_tx_with_claim( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, dlq_reason: &str, error: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let Some(moved) = self .take_running_attempt_tx(tx, job_id, run_lease, "dlq", claim) .await? else { return Ok(None); }; let finalized_at = Utc::now(); let dlq_at = finalized_at; let mut payload = RuntimePayload::from_json(Self::with_progress( moved.payload.clone(), progress.or(moved.progress.clone()), )?)?; payload.push_error(lifecycle_entry(error.into(), moved.attempt, true)); let dlq_row = moved.clone().into_dlq_row( finalized_at, payload.into_json(), dlq_reason.to_string(), dlq_at, ); self.insert_dlq_rows_tx(tx, std::slice::from_ref(&dlq_row), Some(moved.state)) .await?; Ok(Some(dlq_row.into_job_row()?)) } pub async fn move_failed_to_dlq( &self, pool: &PgPool, job_id: i64, dlq_reason: &str, ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); let moved: Option = sqlx::query_as(audited_sql(format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries WHERE (job_id, finalized_at) IN ( SELECT job_id, finalized_at FROM {schema}.done_entries WHERE job_id = $1 AND state = 'failed' ORDER BY finalized_at DESC LIMIT 1 FOR UPDATE SKIP LOCKED ) RETURNING * ) SELECT {done_projection} FROM deleted AS done {ready_join} "# ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some(moved) = moved else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(None); }; self.ensure_terminal_removed_receipt_closures_tx(&mut tx, std::slice::from_ref(&moved)) .await?; // DLQ inserts are outside done_entries, so the source terminal DELETE // appends a negative delta. self.decrement_live_terminal_counters_tx( &mut tx, &Self::done_rows_to_counter_keys(std::slice::from_ref(&moved)), ) .await?; let dlq_row = moved .clone() .into_dlq_row(dlq_reason.to_string(), Utc::now()); self.insert_dlq_rows_tx(&mut tx, std::slice::from_ref(&dlq_row), Some(moved.state)) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(Some(dlq_row.into_job_row()?)) } #[tracing::instrument( skip(self, pool, dlq_reason), fields(kind = ?kind, queue = ?queue), name = "queue_storage.bulk_move_failed_to_dlq" )] pub async fn bulk_move_failed_to_dlq( &self, pool: &PgPool, kind: Option<&str>, queue: Option<&str>, dlq_reason: &str, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); let moved: Vec = sqlx::query_as(audited_sql(format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries WHERE state = 'failed' AND ($1::text IS NULL OR kind = $1) AND ($2::text IS NULL OR queue = $2) RETURNING * ) SELECT {done_projection} FROM deleted AS done {ready_join} "# ))) .bind(kind) .bind(queue) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if moved.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(0); } self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &moved) .await?; // Bulk DLQ move deletes from done_entries. Append negative deltas by // the per-group magnitudes of the moved rows. self.decrement_live_terminal_counters_tx(&mut tx, &Self::done_rows_to_counter_keys(&moved)) .await?; let dlq_at = Utc::now(); let rows: Vec = moved .into_iter() .map(|row| row.into_dlq_row(dlq_reason.to_string(), dlq_at)) .collect(); self.insert_dlq_rows_tx(&mut tx, &rows, Some(JobState::Failed)) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(rows.len() as u64) } pub async fn retry_from_dlq( &self, pool: &PgPool, job_id: i64, opts: &RetryFromDlqOpts, ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let moved: Option = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.dlq_entries WHERE job_id = $1 RETURNING job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload, dlq_reason, dlq_at, original_run_lease "# ))) .bind(job_id) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; let Some(moved) = moved else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(None); }; let queue = opts.queue.clone().unwrap_or_else(|| moved.queue.clone()); let priority = opts.priority.unwrap_or(moved.priority); let mut payload = RuntimePayload::from_json(moved.payload.clone())?; payload.set_progress(None); let payload = payload.into_json(); if let Some(run_at) = opts.run_at.filter(|run_at| *run_at > Utc::now()) { let deferred = moved.into_retry_deferred_row(queue, priority, run_at, payload); self.insert_deferred_rows_tx(&mut tx, vec![deferred.clone()], Some(JobState::Failed)) .await?; tx.commit().await.map_err(map_sqlx_error)?; return Ok(Some(deferred.into_job_row()?)); } let ready = moved.into_retry_ready_row(queue.clone(), priority, Utc::now(), payload); self.insert_existing_ready_rows_tx(&mut tx, vec![ready.clone()], Some(JobState::Failed)) .await?; self.notify_queues_tx(&mut tx, std::iter::once(queue)) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(Some( ReadyJobRow { job_id: ready.job_id, kind: ready.kind, queue: ready.queue, args: ready.args, priority: ready.priority, attempt: ready.attempt, run_lease: ready.run_lease, max_attempts: ready.max_attempts, run_at: ready.run_at, attempted_at: ready.attempted_at, created_at: ready.created_at, unique_key: ready.unique_key, payload: ready.payload, } .into_job_row()?, )) } #[tracing::instrument( skip(self, pool, filter), fields(kind = ?filter.kind, queue = ?filter.queue, tag = ?filter.tag), name = "queue_storage.bulk_retry_from_dlq" )] pub async fn bulk_retry_from_dlq( &self, pool: &PgPool, filter: &ListDlqFilter, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let moved: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.dlq_entries WHERE ($1::text IS NULL OR kind = $1) AND ($2::text IS NULL OR queue = $2) AND ($3::text IS NULL OR payload -> 'tags' ? $3) AND ( ($4::bigint IS NULL AND $5::timestamptz IS NULL) OR ($4::bigint IS NOT NULL AND $5::timestamptz IS NULL AND job_id < $4) OR ($4::bigint IS NULL AND $5::timestamptz IS NOT NULL AND dlq_at < $5) OR ( $4::bigint IS NOT NULL AND $5::timestamptz IS NOT NULL AND (dlq_at, job_id) < ($5, $4) ) ) RETURNING job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload, dlq_reason, dlq_at, original_run_lease "# ))) .bind(&filter.kind) .bind(&filter.queue) .bind(&filter.tag) .bind(filter.before_id) .bind(filter.before_dlq_at) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if moved.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(0); } let run_at = Utc::now(); let mut queues = BTreeSet::new(); let mut ready_rows = Vec::with_capacity(moved.len()); for moved_row in moved { let queue = moved_row.queue.clone(); let priority = moved_row.priority; queues.insert(queue.clone()); let mut payload = RuntimePayload::from_json(moved_row.payload.clone())?; payload.set_progress(None); ready_rows.push(moved_row.into_retry_ready_row( queue, priority, run_at, payload.into_json(), )); } let revived = ready_rows.len() as u64; self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(JobState::Failed)) .await?; self.notify_queues_tx(&mut tx, queues).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(revived) } pub async fn discard_failed_by_kind(&self, pool: &PgPool, kind: &str) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); let deleted_done: Vec = sqlx::query_as(audited_sql(format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries WHERE kind = $1 AND state = 'failed' RETURNING * ) SELECT {done_projection} FROM deleted AS done {ready_join} "# ))) .bind(kind) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; let deleted_dlq: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.dlq_entries WHERE kind = $1 RETURNING job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload, dlq_reason, dlq_at, original_run_lease "# ))) .bind(kind) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; // Discard removes terminal `failed` rows from done_entries (and // `failed` DLQ rows from dlq_entries — exact terminal counts are keyed // on done_entries only). Append negative deltas by the per-group // magnitudes of `deleted_done`. self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &deleted_done) .await?; self.decrement_live_terminal_counters_tx( &mut tx, &Self::done_rows_to_counter_keys(&deleted_done), ) .await?; let mut unique_releases: Vec<_> = deleted_done .iter() .map(|row| { ( row.job_id, &row.unique_key, row.unique_states.as_deref(), row.state, ) }) .chain(deleted_dlq.iter().map(|row| { ( row.job_id, &row.unique_key, row.unique_states.as_deref(), row.state, ) })) .collect(); unique_releases .sort_unstable_by(|left, right| left.1.cmp(right.1).then(left.0.cmp(&right.0))); for (job_id, unique_key, unique_states, state) in unique_releases { self.sync_unique_claim( &mut tx, job_id, unique_key, unique_states, Some(state), None, ) .await?; } tx.commit().await.map_err(map_sqlx_error)?; Ok((deleted_done.len() + deleted_dlq.len()) as u64) } pub async fn fail_retryable( &self, pool: &PgPool, job_id: i64, run_lease: i64, error: impl Into, progress: Option, ) -> Result, AwaError> { self.fail_retryable_with_claim(pool, job_id, run_lease, error, progress, None) .await } pub async fn fail_retryable_with_claim( &self, pool: &PgPool, job_id: i64, run_lease: i64, error: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .fail_retryable_in_tx_with_claim(&mut tx, job_id, run_lease, error, progress, claim) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } /// Transaction-aware variant of [`Self::fail_retryable`] (ADR-029). pub async fn fail_retryable_in_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, error: impl Into, progress: Option, ) -> Result, AwaError> { self.fail_retryable_in_tx_with_claim(tx, job_id, run_lease, error, progress, None) .await } pub async fn fail_retryable_in_tx_with_claim( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, error: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { self.fail_retryable_at_tx(tx, job_id, run_lease, error.into(), progress, claim, None) .await } /// [`Self::fail_retryable_with_claim`] with an explicit retry delay in /// place of the backoff schedule, for test harnesses that must not see /// the job become claimable again. #[doc(hidden)] #[allow(clippy::too_many_arguments)] pub async fn fail_retryable_after_with_claim( &self, pool: &PgPool, job_id: i64, run_lease: i64, delay: Duration, error: impl Into, progress: Option, claim: Option<&ClaimedEntry>, ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let result = self .fail_retryable_at_tx( &mut tx, job_id, run_lease, error.into(), progress, claim, Some(delay), ) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(result) } #[allow(clippy::too_many_arguments)] async fn fail_retryable_at_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, job_id: i64, run_lease: i64, error: AttemptError, progress: Option, claim: Option<&ClaimedEntry>, delay: Option, ) -> Result, AwaError> { let Some(moved) = self .take_running_attempt_tx(tx, job_id, run_lease, "retryable", claim) .await? else { return Ok(None); }; let mut payload = RuntimePayload::from_json(Self::with_progress( moved.payload.clone(), progress.or(moved.progress.clone()), )?)?; let exhausted = moved.attempt >= moved.max_attempts; payload.push_error(lifecycle_entry(error, moved.attempt, exhausted)); if exhausted { let done = moved .clone() .into_done_row(JobState::Failed, Utc::now(), payload.into_json()); self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state)) .await?; return Ok(Some(done.into_job_row()?)); } let run_at = match delay { Some(delay) => { Utc::now() + TimeDelta::from_std(delay).map_err(|err| { AwaError::Validation(format!("invalid retry delay: {err}")) })? } None => { self.backoff_at_tx(tx, moved.attempt, moved.max_attempts) .await? } }; let deferred = moved.clone().into_deferred_row( JobState::Retryable, run_at, Some(Utc::now()), payload.into_json(), ); self.insert_deferred_rows_tx(tx, vec![deferred.clone()], Some(moved.state)) .await?; Ok(Some(deferred.into_job_row()?)) } /// Rescue attempts whose heartbeat went stale. Exhausted attempts fail /// into `done_entries`; see [`Self::rescue_stale_heartbeats_with_dlq`] /// to dead-letter them instead. pub async fn rescue_stale_heartbeats( &self, pool: &PgPool, staleness: Duration, ) -> Result, AwaError> { self.rescue_stale_heartbeats_with_dlq(pool, staleness, &|_| false) .await } /// [`Self::rescue_stale_heartbeats`], moving an exhausted attempt into /// `dlq_entries` (reason `heartbeat_stale`) in the rescue transaction /// when `dlq_enabled` returns true for its queue. #[tracing::instrument( skip(self, pool, dlq_enabled), name = "queue_storage.rescue_stale_heartbeats" )] pub async fn rescue_stale_heartbeats_with_dlq( &self, pool: &PgPool, staleness: Duration, dlq_enabled: &(dyn Fn(&str) -> bool + Sync), ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let cutoff = Utc::now() - TimeDelta::from_std(staleness) .map_err(|err| AwaError::Validation(format!("invalid staleness: {err}")))?; // #169 B1: the staleness predicate prefers // `attempt_state.heartbeat_at` (receipts-mode source of truth, // where heartbeat_batch writes go) and falls back to // `leases.heartbeat_at` (legacy non-receipts source). Two // cases this covers: // // * Receipts mode, claim materialized into `leases` via // callback registration / progress upsert / equivalent. // The leases row was written once at materialize time and // its `heartbeat_at` is stale by definition. The fresh // value lives on `attempt_state.heartbeat_at`. COALESCE // picks `attempt` so a healthy worker isn't falsely // rescued — and a dead worker IS rescued, which the // receipt-side path can't do (its anti-join below // excludes materialized leases to avoid double-closure). // * Legacy non-receipts mode: attempt_state.heartbeat_at is // never written; COALESCE falls back to // leases.heartbeat_at — same shape as pre-B1. // // The dropped `idx_state_hb` (v025) doesn't hurt this scan: // the planner can satisfy the `state='running'` prefix via the // surviving `(state, deadline_at)` or // `(state, callback_timeout_at)` indexes, followed by a heap // recheck of the COALESCE. Bounded by running-lease count and // called at 30s cadence — cheap. let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases AS target WHERE (target.job_id, target.run_lease) IN ( SELECT lease.job_id, lease.run_lease FROM {schema}.leases AS lease LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease WHERE lease.state = 'running' AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1 ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC LIMIT 500 FOR UPDATE OF lease SKIP LOCKED ) RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .bind(cutoff) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; let rescued_receipts = if self.lease_claim_receipts() { self.rescue_stale_receipt_claims_tx(&mut tx, cutoff).await? } else { Vec::new() }; if deleted.is_empty() && rescued_receipts.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(Vec::new()); } let moved_leases = self .hydrate_deleted_leases_tx(&mut tx, deleted, true) .await?; let moved_receipts = self .hydrate_deleted_leases_tx(&mut tx, rescued_receipts, true) .await?; let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len()); for row in moved_leases { let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); rescued.push( self.settle_rescued_attempt_tx( &mut tx, row, "heartbeat stale: worker presumed dead", dlq_reason, ) .await?, ); } for row in moved_receipts { let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); rescued.push( self.settle_rescued_attempt_tx( &mut tx, row, "receipt claim stale: worker presumed dead", dlq_reason, ) .await?, ); } tx.commit().await.map_err(map_sqlx_error)?; Ok(rescued) } /// Rescue attempts past their hard deadline. Exhausted attempts fail /// into `done_entries`; see [`Self::rescue_expired_deadlines_with_dlq`] /// to dead-letter them instead. pub async fn rescue_expired_deadlines(&self, pool: &PgPool) -> Result, AwaError> { self.rescue_expired_deadlines_with_dlq(pool, &|_| false) .await } /// [`Self::rescue_expired_deadlines`], moving an exhausted attempt into /// `dlq_entries` (reason `deadline_exceeded`) in the rescue transaction /// when `dlq_enabled` returns true for its queue. #[tracing::instrument( skip(self, pool, dlq_enabled), name = "queue_storage.rescue_expired_deadlines" )] pub async fn rescue_expired_deadlines_with_dlq( &self, pool: &PgPool, dlq_enabled: &(dyn Fn(&str) -> bool + Sync), ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE job_id IN ( SELECT job_id FROM {schema}.leases WHERE state = 'running' AND deadline_at IS NOT NULL AND deadline_at < clock_timestamp() ORDER BY deadline_at ASC LIMIT 500 FOR UPDATE SKIP LOCKED ) RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; // Receipts-mode short-path claims hold their deadline on // `lease_claim_batches.deadline_at` (compact batches, #246) or on // legacy row-local `lease_claims.deadline_at` rather than on a // `leases` row, so the receipt-plane needs its own scan; merge // all populations into one `moved` set so the maintenance caller // observes a single rescue batch per tick. let receipt_deleted = if self.lease_claim_receipts() { self.rescue_expired_receipt_deadlines_tx(&mut tx).await? } else { Vec::new() }; if deleted.is_empty() && receipt_deleted.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(Vec::new()); } let mut moved = self .hydrate_deleted_leases_tx(&mut tx, deleted, true) .await?; moved.extend( self.hydrate_deleted_leases_tx(&mut tx, receipt_deleted, true) .await?, ); let mut rescued = Vec::with_capacity(moved.len()); for row in moved { let dlq_reason = dlq_enabled(&row.queue).then_some("deadline_exceeded"); rescued.push( self.settle_rescued_attempt_tx(&mut tx, row, "hard deadline exceeded", dlq_reason) .await?, ); } tx.commit().await.map_err(map_sqlx_error)?; Ok(rescued) } /// Rescue attempts whose callback timed out. Exhausted attempts fail /// into `done_entries`; see [`Self::rescue_expired_callbacks_with_dlq`] /// to dead-letter them instead. pub async fn rescue_expired_callbacks(&self, pool: &PgPool) -> Result, AwaError> { self.rescue_expired_callbacks_with_dlq(pool, &|_| false) .await } /// [`Self::rescue_expired_callbacks`], moving an exhausted attempt into /// `dlq_entries` (reason `callback_timeout`) in the rescue transaction /// when `dlq_enabled` returns true for its queue. #[tracing::instrument( skip(self, pool, dlq_enabled), name = "queue_storage.rescue_expired_callbacks" )] pub async fn rescue_expired_callbacks_with_dlq( &self, pool: &PgPool, dlq_enabled: &(dyn Fn(&str) -> bool + Sync), ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases WHERE job_id IN ( SELECT job_id FROM {schema}.leases WHERE state = 'waiting_external' AND callback_timeout_at IS NOT NULL AND callback_timeout_at < clock_timestamp() ORDER BY callback_timeout_at ASC LIMIT 500 FOR UPDATE SKIP LOCKED ) RETURNING ready_slot, ready_generation, job_id, queue, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, heartbeat_at, deadline_at, attempted_at, callback_id, callback_timeout_at "# ))) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deleted.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(Vec::new()); } let moved = self .hydrate_deleted_leases_tx(&mut tx, deleted, true) .await?; let mut rescued = Vec::with_capacity(moved.len()); for row in moved { let dlq_reason = dlq_enabled(&row.queue).then_some("callback_timeout"); rescued.push( self.settle_rescued_attempt_tx(&mut tx, row, "callback timed out", dlq_reason) .await?, ); } tx.commit().await.map_err(map_sqlx_error)?; Ok(rescued) } pub async fn promote_due( &self, pool: &PgPool, state: JobState, batch_size: i64, ) -> Result { if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { return Ok(0); } let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let moved: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.deferred_jobs WHERE job_id IN ( SELECT job_id FROM {schema}.deferred_jobs WHERE state = $1 AND run_at <= clock_timestamp() AND NOT EXISTS ( SELECT 1 FROM awa.queue_meta WHERE queue = {schema}.deferred_jobs.queue AND paused = TRUE ) ORDER BY run_at ASC, priority ASC, job_id ASC LIMIT $2 FOR UPDATE SKIP LOCKED ) RETURNING job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload "# ))) .bind(state) .bind(batch_size) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if moved.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(0); } let ready_rows: Vec = moved .iter() .cloned() .map(|row| ExistingReadyRow { job_id: row.job_id, kind: row.kind, queue: row.queue, args: row.args, priority: row.priority, attempt: row.attempt, run_lease: row.run_lease, max_attempts: row.max_attempts, run_at: Utc::now(), attempted_at: row.attempted_at, created_at: row.created_at, unique_key: row.unique_key, unique_states: row.unique_states, payload: row.payload, }) .collect(); let queues = ready_rows .iter() .map(|row| row.queue.clone()) .collect::>(); self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(state)) .await?; self.notify_queues_tx(&mut tx, queues).await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(moved.len()) } async fn relation_has_rows_tx( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, relation: &str, ) -> Result { sqlx::query_scalar(audited_sql(format!( "SELECT EXISTS (SELECT 1 FROM {relation} LIMIT 1)" ))) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error) } /// Idle-gate probe (#371): true when every relation in `relations` is /// empty. Short-circuits on the first non-empty relation, so under /// active traffic the extra cost over the existing busy gate is a /// single `EXISTS ... LIMIT 1` probe; when the ring really is idle all /// probes hit empty heaps. Same SQL shape (and lock strength — /// AccessShare only) as the per-table busy checks above. async fn relations_all_empty_tx( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, relations: &[String], ) -> Result { for relation in relations { if Self::relation_has_rows_tx(tx, relation).await? { return Ok(false); } } Ok(true) } async fn slot_still_filling_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, family: RingFamily, cursor: (i32, i64), cadence: RingRotationCadence, children: &[String], ) -> Result { if cadence.min_slot_age.is_zero() { return Ok(false); } // Sample at most `min_slot_rows` rows per child so the probe stays // bounded on a full slot; one row is enough to know it is non-empty. let sample_limit = cadence.min_slot_rows.clamp(1, i64::MAX as u64); let row_sources = children .iter() .map(|relation| format!("(SELECT 1 FROM {relation} LIMIT {sample_limit})")) .collect::>() .join(" UNION ALL "); // The ledger row of the current generation records when the slot // opened on the database clock, so the age ceiling survives a change // of maintenance leader; it is absent only while compat authority's // shadow ledger lags, when the local observation stands alone. let schema = self.schema(); let ledger = family.ledger_relname(); let (sampled_rows, bytes, open_secs): (i64, i64, Option) = sqlx::query_as(audited_sql(format!( "WITH occupancy AS (SELECT count(*) AS sampled FROM ({row_sources}) AS rows) \ SELECT sampled, CASE WHEN sampled > 0 THEN (\ SELECT COALESCE(SUM(pg_total_relation_size(relation::regclass)), 0)::bigint \ FROM unnest($1::text[]) AS relation\ ) ELSE 0::bigint END, \ (SELECT EXTRACT(EPOCH FROM clock_timestamp() - rotated_at)::float8 \ FROM {schema}.{ledger} WHERE generation = $2) \ FROM occupancy" ))) .bind(children) .bind(cursor.1) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; let has_rows = sampled_rows > 0; let mut observations = self .slot_fill_observations .lock() .expect("slot fill observations poisoned"); if !has_rows { observations.remove(&family); return Ok(false); } let now = Instant::now(); let observation = observations.entry(family).or_insert(SlotFillObservation { cursor, first_nonempty_at: now, }); if observation.cursor != cursor { *observation = SlotFillObservation { cursor, first_nonempty_at: now, }; } let bytes = u64::try_from(bytes).unwrap_or(0); if cadence.target_slot_bytes > 0 && bytes >= cadence.target_slot_bytes { return Ok(false); } let age = now.duration_since(observation.first_nonempty_at); let open_for = open_secs .filter(|secs| secs.is_finite() && *secs > 0.0) .map(Duration::from_secs_f64) .unwrap_or_default(); if !cadence.max_slot_age.is_zero() && age.max(open_for) >= cadence.max_slot_age { return Ok(false); } if age < cadence.min_slot_age { tracing::debug!(?family, ?age, "slot left open: younger than min_slot_age"); return Ok(true); } let sampled_rows = u64::try_from(sampled_rows).unwrap_or(0); if sampled_rows < cadence.min_slot_rows { tracing::debug!( ?family, sampled_rows, min_slot_rows = cadence.min_slot_rows, "slot left open: below min_slot_rows" ); return Ok(true); } Ok(false) } /// Seal the current queue-ring slot on every call that passes the busy /// and idle gates. Equivalent to [`Self::rotate_with_cadence`] with /// [`RingRotationCadence::EAGER`]. pub async fn rotate(&self, pool: &PgPool) -> Result { self.rotate_with_cadence(pool, RingRotationCadence::EAGER) .await } #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate")] pub async fn rotate_with_cadence( &self, pool: &PgPool, cadence: RingRotationCadence, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // #371 staged upgrade: pick the cursor authority once, then thread // it through the lock/cursor/append seams so this whole rotation // runs in one mode. In ledger mode the per-ring advisory lock // serializes rotate ↔ prune (non-blocking, skip on contention); in // compat mode the ring-state singleton FOR UPDATE serializes against // live 0.6 rotators too. let (authority, acquired) = self .acquire_ring_serializer_tx(&mut tx, RingFamily::Queue) .await?; if !acquired { let (current_slot, _) = self .ring_cursor_tx(&mut tx, RingFamily::Queue, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Queue).await?; tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: (current_slot + 1).rem_euclid(slot_count), busy: BusyCounts::default(), }); } let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Queue, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Queue).await?; let state = (current_slot, generation, slot_count); let current_children = [ ready_child_name(schema, state.0 as usize), ready_claim_attempt_batch_child_name(schema, state.0 as usize), done_child_name(schema, state.0 as usize), ready_tombstone_child_name(schema, state.0 as usize), ready_segment_child_name(schema, state.0 as usize), receipt_completion_batch_child_name(schema, state.0 as usize), receipt_completion_tombstone_child_name(schema, state.0 as usize), terminal_delta_child_name(schema, state.0 as usize), ]; if self .slot_still_filling_tx( &mut tx, RingFamily::Queue, (state.0, state.1), cadence, ¤t_children, ) .await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedFilling { slot: state.0, generation: state.1, }); } let next_slot = (state.0 + 1).rem_euclid(state.2); let ready_busy = Self::relation_has_rows_tx(&mut tx, &ready_child_name(schema, next_slot as usize)) .await?; let claim_attempt_batch_busy = Self::relation_has_rows_tx( &mut tx, &ready_claim_attempt_batch_child_name(schema, next_slot as usize), ) .await?; let done_busy = Self::relation_has_rows_tx(&mut tx, &done_child_name(schema, next_slot as usize)) .await?; let tombstone_busy = Self::relation_has_rows_tx( &mut tx, &ready_tombstone_child_name(schema, next_slot as usize), ) .await?; let segment_busy = Self::relation_has_rows_tx( &mut tx, &ready_segment_child_name(schema, next_slot as usize), ) .await?; let receipt_batch_busy = Self::relation_has_rows_tx( &mut tx, &receipt_completion_batch_child_name(schema, next_slot as usize), ) .await?; let receipt_tombstone_busy = Self::relation_has_rows_tx( &mut tx, &receipt_completion_tombstone_child_name(schema, next_slot as usize), ) .await?; let terminal_delta_busy = Self::relation_has_rows_tx( &mut tx, &terminal_delta_child_name(schema, next_slot as usize), ) .await?; if ready_busy || claim_attempt_batch_busy || done_busy || tombstone_busy || segment_busy || receipt_batch_busy || receipt_tombstone_busy || terminal_delta_busy { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: next_slot, busy: BusyCounts { queue_ready: busy_indicator(ready_busy), queue_claim_attempt_batches: busy_indicator(claim_attempt_batch_busy), queue_done: busy_indicator(done_busy), queue_tombstones: busy_indicator(tombstone_busy), queue_ready_segments: busy_indicator(segment_busy), queue_receipt_completion_batches: busy_indicator(receipt_batch_busy), queue_receipt_completion_tombstones: busy_indicator(receipt_tombstone_busy), queue_terminal_deltas: busy_indicator(terminal_delta_busy), ..Default::default() }, }); } // Idle gate (#371): skip the rotation only when the ENTIRE ring // is empty — every partition of every ring family, not just the // current slot. The ledger append below would otherwise only grow // the `queue_ring_rotations` ledger with a row that carries no // sealed data (and the folder cannot trim it while an MVCC horizon // is pinned). // // The whole-ring scope is load-bearing for reclamation: prune // never advances a slot's generation (only rotation does, by // reusing slots), so `prune_oldest` keeps selecting the same // already-truncated slot while rotation is frozen. Gating on // current-slot emptiness alone would freeze a drained ring that // still holds rows in later sealed slots and strand them until // new traffic resumed rotation. Probing the parents keeps // rotation cycling until prune has reclaimed everything, and // freezes only when there is genuinely nothing left to seal OR // prune. (A queue holding failed rows inside the retention // floor keeps rotating until retention expires — those rows // need live reclamation machinery by definition.) // // Race note: a concurrent enqueue that commits right after this // probe is benign. Rotation is enablement, not obligation (the // TLA+ rotate actions are guarded and carry no fairness), so // skipping only delays sealing that write by one rotate tick — // the next tick observes the row and rotates. The skip never // loses data and never leaves a non-empty slot unsealed while // traffic flows. let ring_parents = [ format!("{schema}.ready_entries"), format!("{schema}.ready_claim_attempt_batches"), format!("{schema}.done_entries"), format!("{schema}.ready_tombstones"), format!("{schema}.ready_segments"), format!("{schema}.receipt_completion_batches"), format!("{schema}.receipt_completion_tombstones"), format!("{schema}.queue_terminal_count_deltas"), ]; if Self::relations_all_empty_tx(&mut tx, &ring_parents).await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedIdle { slot: state.0, generation: state.1, }); } let next_generation = state.1 + 1; // Advance the cursor (ledger CAS-append, or compat column CAS + // ledger shadow). A `false` return is a lost race with a concurrent // rotator — treat it as busy, never retry with the same generation. if !self .append_ring_rotation_tx( &mut tx, RingFamily::Queue, next_generation, next_slot, authority, ) .await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: next_slot, busy: BusyCounts::default(), }); } tx.commit().await.map_err(map_sqlx_error)?; Ok(RotateOutcome::Rotated { slot: next_slot, generation: next_generation, }) } /// Seal the current lease-ring slot on every call that passes the busy /// and idle gates. Equivalent to [`Self::rotate_leases_with_cadence`] /// with [`RingRotationCadence::EAGER`]. pub async fn rotate_leases(&self, pool: &PgPool) -> Result { self.rotate_leases_with_cadence(pool, RingRotationCadence::EAGER) .await } #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_leases")] pub async fn rotate_leases_with_cadence( &self, pool: &PgPool, cadence: RingRotationCadence, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // The per-ring serialization lock guards against prune_oldest_leases // and parallel rotators (#371). In ledger mode it is the per-ring // advisory lock; in compat mode (#371 staged upgrade) it is the // lease_ring_state singleton FOR UPDATE, which also serializes // against live 0.6 rotators. Without it two rotators can both pass // the busy-check, both compute the same next_slot, and the loser's // CAS wastes work. `RotateLeasesPlan` in // `correctness/storage/AwaStorageLockOrder.tla` requires this lock // as the first acquired resource for the rotation tx. let (authority, acquired) = self .acquire_ring_serializer_tx(&mut tx, RingFamily::Lease) .await?; if !acquired { let (current_slot, _) = self .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Lease).await?; tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: (current_slot + 1).rem_euclid(slot_count), busy: BusyCounts::default(), }); } let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Lease).await?; let state = (current_slot, generation, slot_count); let current_children = [lease_child_name(schema, state.0 as usize)]; if self .slot_still_filling_tx( &mut tx, RingFamily::Lease, (state.0, state.1), cadence, ¤t_children, ) .await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedFilling { slot: state.0, generation: state.1, }); } let next_slot = (state.0 + 1).rem_euclid(state.2); let lease_busy = Self::relation_has_rows_tx(&mut tx, &lease_child_name(schema, next_slot as usize)) .await?; if lease_busy { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: next_slot, busy: BusyCounts { leases: busy_indicator(lease_busy), ..Default::default() }, }); } // Idle gate (#371): skip only when the ENTIRE lease ring is // empty (parent probe, all partitions). Otherwise the ledger // append below only grows the `lease_ring_rotations` ledger (the // lease ring was the top residual pinned-horizon dead-tuple // accumulator at the historical 250ms singleton-UPDATE cadence — // ~14.4k dead tuples/hour). See `rotate` for both the race // analysis and why the whole-ring scope is load-bearing (prune // never advances slot generations; freezing rotation with sealed // rows behind would strand them until traffic resumed). // Completion and VACUUM can empty heaps while retaining indexes/TOAST. // Seal these children so prune can reclaim physical storage, while // genuinely idle rings with only empty-index overhead remain frozen. let has_reclaimable_storage: bool = sqlx::query_scalar(RECLAIMABLE_CHILD_STORAGE_SQL) .bind(format!("{schema}.leases")) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; if !has_reclaimable_storage { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedIdle { slot: state.0, generation: state.1, }); } let next_generation = state.1 + 1; if !self .append_ring_rotation_tx( &mut tx, RingFamily::Lease, next_generation, next_slot, authority, ) .await? { // Another rotator beat us to the CAS. Report the bounded // presence evidence we sampled before giving up. tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: next_slot, busy: BusyCounts { leases: busy_indicator(lease_busy), ..Default::default() }, }); } tx.commit().await.map_err(map_sqlx_error)?; Ok(RotateOutcome::Rotated { slot: next_slot, generation: next_generation, }) } /// Rebuild terminal counters from scratch by truncating folded live counts /// and pending deltas, then re-aggregating terminal history. Run this when: /// /// - upgrading from a pre-#290 fleet, where in-flight binaries wrote /// terminal history without maintaining the counter, /// - after any incident that may have left the counter inconsistent /// with terminal history, or /// - as a routine drift-recovery step before relying on counter-fed /// reads for billing-grade accuracy. /// /// The rebuild takes a schema-scoped rebuild lock so two rebuilds do not /// overlap, then waits for the queue ring's rotation lock, so it starts /// only between maintenance ticks; ticks that arrive meanwhile skip /// rather than wait. It then locks the terminal tables exclusively, so /// completions, failures and retries that write terminal rows wait for /// it to commit. Other schemas are unaffected. /// /// **Operator note:** this is best run on a quiesced fleet (workers /// paused or fully drained). Terminal writes during the rebuild wait on /// its table locks, and a long rebuild stalls them for its duration. The /// rebuild itself is O(rows in `{schema}.done_entries`). Compact receipt /// completions remain counted directly from retained /// `receipt_completion_batches` until queue prune folds them into /// permanent rollups. #[tracing::instrument(skip(self, pool), name = "queue_storage.rebuild_terminal_counters")] pub async fn rebuild_terminal_counters(&self, pool: &PgPool) -> Result { let schema = self.schema(); let install_lock_name = format!("awa.queue_storage.install:{schema}"); let rebuild_lock_name = format!("awa.queue_storage.rebuild_terminal_counters:{schema}"); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") .bind(&install_lock_name) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") .bind(&rebuild_lock_name) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; self.lock_ring_serializers_tx(&mut tx, &[RingFamily::Queue]) .await?; sqlx::query(audited_sql(format!( "SELECT 1 FROM {schema}.queue_ring_state WHERE singleton = TRUE FOR UPDATE" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( "LOCK TABLE {schema}.done_entries, \ {schema}.receipt_completion_batches, \ {schema}.receipt_completion_tombstones, \ {schema}.queue_terminal_count_deltas, \ {schema}.queue_terminal_live_counts \ IN ACCESS EXCLUSIVE MODE" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; sqlx::query(audited_sql(format!( "TRUNCATE TABLE {schema}.queue_terminal_live_counts, {schema}.queue_terminal_count_deltas" ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; // One row per lane once the ring authority flip has fenced N-1 // maintenance; per job-id bucket before it, which is the shape N-1 // folds apply their deltas to (ADR-041). let counter_bucket = match self.ring_authority_tx(&mut tx).await? { RingAuthority::Ledger => "0::smallint".to_string(), RingAuthority::Columns => format!( "mod(mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) \ + {TERMINAL_COUNTER_BUCKETS}::bigint, {TERMINAL_COUNTER_BUCKETS}::bigint)::smallint" ), }; let inserted: i64 = sqlx::query_scalar(audited_sql(format!( r#" WITH inserted AS ( INSERT INTO {schema}.queue_terminal_live_counts AS counts ( ready_slot, queue, priority, enqueue_shard, counter_bucket, live_terminal_count ) SELECT ready_slot, queue, priority, enqueue_shard, {counter_bucket} AS counter_bucket, count(*)::bigint FROM {schema}.done_entries GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket RETURNING 1 ) SELECT COALESCE(count(*), 0)::bigint FROM inserted "# ))) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; // Flip the trust marker. From this point the read path // (queue_counts_exact) uses the counter for done-entry terminal rows; // before this call, it falls back to scanning terminal_jobs. sqlx::query(audited_sql(format!( r#" UPDATE {schema}.queue_ring_state SET terminal_counter_trusted_at = now() WHERE singleton = TRUE "# ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; Ok(inserted) } /// Check whether the live terminal counter has been marked trusted /// for exact reads. Returns `true` for fresh installs (the trust /// marker is auto-set by `prepare_schema` when `done_entries` is /// empty) and after a successful /// [`Self::rebuild_terminal_counters`]; returns `false` on an /// existing install that has not yet been rebuilt under the new /// counter-aware code path. /// /// `queue_counts_exact` consults this to decide between the /// counter-fed fast path and the scan-based fallback. Single-row /// PK fetch; negligible cost per call. pub async fn terminal_counter_trusted(&self, pool: &PgPool) -> Result { let schema = self.schema(); let trusted: Option = sqlx::query_scalar(audited_sql(format!( "SELECT terminal_counter_trusted_at IS NOT NULL \ FROM {schema}.queue_ring_state WHERE singleton = TRUE" ))) .fetch_optional(pool) .await .map_err(map_sqlx_error)?; // Missing row = pre-#290 schema that hasn't run the new // prepare_schema yet. Treat as untrusted; the scan path is // still correct. Ok(trusted.unwrap_or(false)) } /// Fold append-only `done_entries` terminal-count deltas into /// `queue_terminal_live_counts` for sealed queue slots. /// /// `done_entries` terminal inserts and deletes append signed rows into /// `queue_terminal_count_deltas` instead of updating the live counter on /// the user-facing hot path. Compact receipt completions are counted from /// retained batch rows directly. Exact reads sum folded live counts plus /// pending deltas and retained compact batches, so this rollup can run /// asynchronously. It intentionally /// skips the current queue slot and any slot with active leases or open /// receipt claims; that keeps rollup away from the segment receiving hot /// completions and avoids racing future terminal deltas for the same /// ready generation. Rollup also stands down while another backend pins the /// MVCC horizon, leaving the append-only deltas visible to exact reads /// until the horizon clears. #[tracing::instrument(skip(self, pool), name = "queue_storage.rollup_terminal_count_deltas")] pub async fn rollup_terminal_count_deltas( &self, pool: &PgPool, max_slots: usize, ) -> Result { if max_slots == 0 { return Ok(TerminalDeltaRollupOutcome::default()); } let slots = self.terminal_delta_rollup_candidates(pool).await?; let mut outcome = TerminalDeltaRollupOutcome::default(); for (slot, generation) in slots { if outcome.rolled_slots >= max_slots { break; } match self .rollup_terminal_count_delta_slot(pool, slot, generation) .await? { TerminalDeltaSlotRollup::Empty => {} TerminalDeltaSlotRollup::Rolled { delta_rows, grouped_keys, } => { outcome.rolled_slots += 1; outcome.delta_rows += delta_rows; outcome.grouped_keys += grouped_keys; } TerminalDeltaSlotRollup::SkippedActive => { outcome.skipped_active_slots += 1; } TerminalDeltaSlotRollup::SkippedMvccPinned => { outcome.skipped_mvcc_pinned = true; break; } TerminalDeltaSlotRollup::Blocked => { outcome.blocked_slots += 1; } } } Ok(outcome) } /// Does another backend hold a snapshot old enough to strand the versions /// the horizon-gated folds delete? See [`MVCC_HORIZON_PIN_MIN_AGE`] for why /// this is age-gated rather than tripping on any live `backend_xmin`. /// /// A backend counts as pinning when it is either: /// - **idle in transaction with a live xid** — it holds its snapshot /// indefinitely with no in-flight query to bound the hold, so any age is /// dangerous; or /// - **actively holding a snapshot** (`backend_xmin` set) whose transaction /// has been open longer than [`MVCC_HORIZON_PIN_MIN_AGE`] — a genuinely /// long-lived reader, not a transient hot-path claim/enqueue. /// /// `xact_start` is NULL for a backend not in a transaction; the age /// comparison naturally excludes those. It can lag `backend_xmin` by a /// statement boundary, but a backend with a snapshot and no transaction /// start is a momentary state we are content to let the next tick catch. async fn terminal_delta_rollup_mvcc_horizon_pinned_tx( tx: &mut sqlx::Transaction<'_, Postgres>, ) -> Result { let pinned: bool = sqlx::query_scalar( r#" SELECT EXISTS ( SELECT 1 FROM pg_stat_activity WHERE datname = current_database() AND pid <> pg_backend_pid() AND backend_type = 'client backend' AND ( ( backend_xmin IS NOT NULL AND xact_start IS NOT NULL AND xact_start <= now() - make_interval(secs => $1) ) OR ( backend_xid IS NOT NULL AND state LIKE 'idle in transaction%' ) ) ) "#, ) .bind(MVCC_HORIZON_PIN_MIN_AGE.as_secs_f64()) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; Ok(pinned) } async fn terminal_delta_rollup_candidates( &self, pool: &PgPool, ) -> Result, AwaError> { let schema = self.schema(); // Sealed slots and their last-open generations are derived from // the rotation-ledger cursor (#371); no per-slot generation rows // exist anymore. let (current_slot, generation, slot_count): (i32, i64, i32) = sqlx::query_as(audited_sql(format!( r#" SELECT ledger.slot, ledger.generation, state.slot_count FROM {schema}.queue_ring_state AS state CROSS JOIN LATERAL ( SELECT slot, generation FROM {schema}.queue_ring_rotations ORDER BY generation DESC LIMIT 1 ) AS ledger WHERE state.singleton = TRUE "# ))) .fetch_one(pool) .await .map_err(map_sqlx_error)?; let sealed_slots = initialized_sealed_ring_slots(current_slot, generation, slot_count); let mut pending_slots = Vec::new(); for (slot, generation) in sealed_slots { let Ok(slot_index) = usize::try_from(slot) else { continue; }; let delta_child = terminal_delta_child_name(schema, slot_index); let has_pending: bool = sqlx::query_scalar(audited_sql(format!( r#" SELECT EXISTS ( SELECT 1 FROM {delta_child} WHERE ready_generation = $1 LIMIT 1 ) "# ))) .bind(generation) .fetch_one(pool) .await .map_err(map_sqlx_error)?; if has_pending { pending_slots.push((slot, generation)); } } Ok(pending_slots) } async fn rollup_terminal_count_delta_slot( &self, pool: &PgPool, slot: i32, generation: i64, ) -> Result { let schema = self.schema(); let delta_child = terminal_delta_child_name(schema, slot as usize); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // Ring rotation serialization first (stabilises the cursor against // rotate and prune — advisory lock in ledger mode, singleton FOR // UPDATE in compat mode), then the slot row lock — the same order // every queue-ring maintenance tx uses (`RollupTerminalDeltasPlan` // in `correctness/storage/AwaStorageLockOrder.tla`). let (authority, acquired) = self .acquire_ring_serializer_tx(&mut tx, RingFamily::Queue) .await?; if !acquired { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::Blocked); } let (current_slot, current_generation) = self .ring_cursor_tx(&mut tx, RingFamily::Queue, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Queue).await?; let slot_locked: Option = sqlx::query_scalar(audited_sql(format!( r#" SELECT slot FROM {schema}.queue_ring_slots WHERE slot = $1 FOR UPDATE "# ))) .bind(slot) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if slot_locked.is_none() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::Empty); } if current_slot == slot { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::SkippedActive); } let Some(slot_generation) = ring_slot_generation(current_slot, current_generation, slot_count, slot) else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::Empty); }; if slot_generation != generation { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::Empty); } if Self::terminal_delta_rollup_mvcc_horizon_pinned_tx(&mut tx).await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::SkippedMvccPinned); } let pending_ready = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; if pending_ready { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::SkippedActive); } set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_delta = sqlx::query(audited_sql(format!( "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" ))) .execute(tx.as_mut()) .await; match lock_delta { Ok(_) => {} Err(err) if is_lock_contention_error(&err) => { let _ = tx.rollback().await; return Ok(TerminalDeltaSlotRollup::Blocked); } Err(err) => { let _ = tx.rollback().await; return Err(map_sqlx_error(err)); } } let active_leases = queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; if active_leases { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::SkippedActive); } let unclosed_claim_refs = queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; if unclosed_claim_refs { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::SkippedActive); } let delta_rows: i64 = sqlx::query_scalar(audited_sql(format!( r#" SELECT count(*)::bigint FROM {delta_child} WHERE ready_generation = $1 "# ))) .bind(generation) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; if delta_rows == 0 { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalDeltaSlotRollup::Empty); } // In ledger authority each touched lane collapses to one bucket-0 row // holding its old total (summed over every bucket) plus the signed // delta, which keeps a negative delta exact when the matching positive // count sits in another bucket. Before the flip, N-1 maintenance can // still win the election and folds per job-id bucket, so the fold // keeps that shape until the flip fences it (ADR-041). let grouped_keys: i64 = match authority { RingAuthority::Ledger => sqlx::query_scalar(audited_sql(format!( r#" WITH grouped AS MATERIALIZED ( SELECT ready_slot, queue, priority, enqueue_shard, SUM(terminal_delta)::bigint AS delta FROM {delta_child} WHERE ready_generation = $1 GROUP BY ready_slot, queue, priority, enqueue_shard HAVING SUM(terminal_delta) <> 0 ), collapsed AS ( DELETE FROM {schema}.queue_terminal_live_counts AS counts USING grouped WHERE counts.ready_slot = grouped.ready_slot AND counts.queue = grouped.queue AND counts.priority = grouped.priority AND counts.enqueue_shard = grouped.enqueue_shard RETURNING counts.ready_slot, counts.queue, counts.priority, counts.enqueue_shard, counts.live_terminal_count ), folded AS ( SELECT grouped.ready_slot, grouped.queue, grouped.priority, grouped.enqueue_shard, (grouped.delta + COALESCE(SUM(collapsed.live_terminal_count), 0))::bigint AS live_terminal_count FROM grouped LEFT JOIN collapsed ON collapsed.ready_slot = grouped.ready_slot AND collapsed.queue = grouped.queue AND collapsed.priority = grouped.priority AND collapsed.enqueue_shard = grouped.enqueue_shard GROUP BY grouped.ready_slot, grouped.queue, grouped.priority, grouped.enqueue_shard, grouped.delta ), inserted AS ( INSERT INTO {schema}.queue_terminal_live_counts ( ready_slot, queue, priority, enqueue_shard, counter_bucket, live_terminal_count ) SELECT ready_slot, queue, priority, enqueue_shard, 0::smallint, live_terminal_count FROM folded WHERE live_terminal_count <> 0 ORDER BY ready_slot, queue, priority, enqueue_shard RETURNING 1 ) SELECT count(*)::bigint FROM grouped "# ))) .bind(generation) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?, RingAuthority::Columns => sqlx::query_scalar(audited_sql(format!( r#" WITH grouped AS MATERIALIZED ( SELECT ready_slot, queue, priority, enqueue_shard, counter_bucket, SUM(terminal_delta)::bigint AS delta FROM {delta_child} WHERE ready_generation = $1 GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket HAVING SUM(terminal_delta) <> 0 ), updated AS ( UPDATE {schema}.queue_terminal_live_counts AS counts SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta) FROM grouped WHERE counts.ready_slot = grouped.ready_slot AND counts.queue = grouped.queue AND counts.priority = grouped.priority AND counts.enqueue_shard = grouped.enqueue_shard AND counts.counter_bucket = grouped.counter_bucket RETURNING counts.ready_slot, counts.queue, counts.priority, counts.enqueue_shard, counts.counter_bucket ), inserted AS ( INSERT INTO {schema}.queue_terminal_live_counts AS counts ( ready_slot, queue, priority, enqueue_shard, counter_bucket, live_terminal_count ) SELECT grouped.ready_slot, grouped.queue, grouped.priority, grouped.enqueue_shard, grouped.counter_bucket, grouped.delta FROM grouped WHERE grouped.delta > 0 AND NOT EXISTS ( SELECT 1 FROM updated WHERE updated.ready_slot = grouped.ready_slot AND updated.queue = grouped.queue AND updated.priority = grouped.priority AND updated.enqueue_shard = grouped.enqueue_shard AND updated.counter_bucket = grouped.counter_bucket ) ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE SET live_terminal_count = GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count) RETURNING 1 ) SELECT count(*)::bigint FROM grouped "# ))) .bind(generation) .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?, }; let truncate_delta = sqlx::query(audited_sql(format!("TRUNCATE TABLE {delta_child}"))) .execute(tx.as_mut()) .await; match truncate_delta { Ok(_) => { tx.commit().await.map_err(map_sqlx_error)?; Ok(TerminalDeltaSlotRollup::Rolled { delta_rows, grouped_keys, }) } Err(err) if is_lock_contention_error(&err) => { let _ = tx.rollback().await; Ok(TerminalDeltaSlotRollup::Blocked) } Err(err) => { let _ = tx.rollback().await; Err(map_sqlx_error(err)) } } } /// Fold pending `queue_terminal_rollup_deltas` rows into the /// permanent `queue_terminal_rollups` denormaliser (#371). /// /// Queue prune appends per-(queue, priority) delta rows instead of /// upserting the rollups inside the prune transaction, so a pinned /// MVCC horizon cannot pile up dead rollup versions. This fold runs /// from the maintenance leader (piggybacking the 30s terminal-rollup /// tick), stands down while any other backend pins the MVCC horizon /// (the same guard the terminal-count delta rollup uses), and /// otherwise consumes the delta rows with `DELETE ... RETURNING`, /// aggregates them, and upserts the rollups in deterministic /// (queue, priority) order with the historical `GREATEST(0, ...)` /// clamp applied at fold time. /// /// Exact readers ([`Self::queue_counts`], [`Self::queue_counts_fast`], /// [`Self::pruned_failed_count_for_queue`]) add unfolded delta sums to /// their rollup reads, so counts stay exact regardless of fold timing. #[tracing::instrument(skip(self, pool), name = "queue_storage.fold_terminal_rollup_deltas")] pub async fn fold_terminal_rollup_deltas( &self, pool: &PgPool, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; if Self::terminal_delta_rollup_mvcc_horizon_pinned_tx(&mut tx).await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalRollupFoldOutcome { skipped_mvcc_pinned: true, ..Default::default() }); } let deltas: Vec<(String, i16, i64, i64)> = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.queue_terminal_rollup_deltas RETURNING queue, priority, pruned_completed_delta, pruned_failed_delta "# ))) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; if deltas.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(TerminalRollupFoldOutcome::default()); } let folded_delta_rows = deltas.len(); let folded_keys = deltas .iter() .map(|(queue, priority, _, _)| (queue.clone(), *priority)) .collect::>() .len(); self.adjust_terminal_rollups_batch(&mut tx, deltas.into_iter()) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(TerminalRollupFoldOutcome { folded_delta_rows, folded_keys, skipped_mvcc_pinned: false, }) } /// Trim the ring-rotation ledgers to one full wrap (#371). /// /// Rotation appends one ledger row per advance; this fold deletes /// rows older than `max_generation - (slot_count - 1)`, retaining one /// row per slot so every sealed slot's last-open generation stays /// directly readable from the ledger. Like the terminal-delta /// rollup, the fold stands down while another backend pins the MVCC /// horizon, so the deleted versions are immediately reclaimable when /// they are created. #[tracing::instrument(skip(self, pool), name = "queue_storage.fold_ring_rotation_ledgers")] pub async fn fold_ring_rotation_ledgers( &self, pool: &PgPool, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; if Self::terminal_delta_rollup_mvcc_horizon_pinned_tx(&mut tx).await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RingLedgerFoldOutcome { skipped_mvcc_pinned: true, ..Default::default() }); } let mut trimmed_rows = 0_u64; for family in [RingFamily::Queue, RingFamily::Lease, RingFamily::Claim] { let ledger = family.ledger_relname(); let state = family.state_relname(); let deleted = sqlx::query(audited_sql(format!( r#" DELETE FROM {schema}.{ledger} WHERE generation < ( SELECT max(generation) FROM {schema}.{ledger} ) - ( SELECT slot_count - 1 FROM {schema}.{state} WHERE singleton = TRUE )::bigint "# ))) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; trimmed_rows += deleted.rows_affected(); } tx.commit().await.map_err(map_sqlx_error)?; Ok(RingLedgerFoldOutcome { trimmed_rows, skipped_mvcc_pinned: false, }) } /// Prune the oldest sealed queue ring slot. /// /// `failed_retention` is the floor for `failed` terminal rows: rows /// with `finalized_at` inside the floor are re-homed into the live /// `done_entries` segment (still retryable) instead of being /// truncated with the slot. `Duration::ZERO` disables the floor and /// prunes every terminal row. Granularity is whole seconds, /// matching DLQ retention. The floor applies to `failed` only — /// `cancelled` is an explicit operator action and is always pruned. #[tracing::instrument( skip(self, pool), name = "queue_storage.prune_oldest", fields(slot = tracing::field::Empty, generation = tracing::field::Empty) )] pub async fn prune_oldest( &self, pool: &PgPool, failed_retention: Duration, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // Ring rotation serialization keeps the cursor stable for the whole // prune so rotate cannot hand the slot being truncated out to fresh // writes. Ledger mode -> advisory lock; compat mode (#371 staged // upgrade) -> queue_ring_state singleton FOR UPDATE, matching the // pre-#371 discipline and serializing against live 0.6 rotators. let (authority, acquired) = self .acquire_ring_serializer_tx(&mut tx, RingFamily::Queue) .await?; if !acquired { let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Queue, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Queue).await?; tx.commit().await.map_err(map_sqlx_error)?; return Ok( match oldest_initialized_ring_slot(current_slot, generation, slot_count) { Some((slot, _)) => PruneOutcome::Blocked { slot }, None => PruneOutcome::Noop, }, ); } let (current_slot, current_generation) = self .ring_cursor_tx(&mut tx, RingFamily::Queue, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Queue).await?; let state = (current_slot,); let Some((slot, generation)) = oldest_initialized_ring_slot(current_slot, current_generation, slot_count) else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::Noop); }; tracing::Span::current().record("slot", slot); tracing::Span::current().record("generation", generation); // Slot row lock: second resource in the queue-ring lock order // (ring advisory lock → slot row → child ACCESS EXCLUSIVE). let slot_locked: Option = sqlx::query_scalar(audited_sql(format!( r#" SELECT slot FROM {schema}.queue_ring_slots WHERE slot = $1 FOR UPDATE "# ))) .bind(slot) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if slot_locked.is_none() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::Noop); }; let ready_child = ready_child_name(schema, slot as usize); let claim_attempt_child = ready_claim_attempt_batch_child_name(schema, slot as usize); let done_child = done_child_name(schema, slot as usize); let tomb_child = ready_tombstone_child_name(schema, slot as usize); let segment_child = ready_segment_child_name(schema, slot as usize); let receipt_batch_child = receipt_completion_batch_child_name(schema, slot as usize); let receipt_tomb_child = receipt_completion_tombstone_child_name(schema, slot as usize); let delta_child = terminal_delta_child_name(schema, slot as usize); // Queue prune has three common skip gates. Check them before // taking ACCESS EXCLUSIVE on the ready/terminal child tables so // a known-busy slot does not block hot claim reads behind a // doomed truncate attempt. These are only MVCC snapshots; the // same gates are checked again after the exclusive locks. let active_leases = queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; if active_leases { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::QueueActiveLeases, count: busy_indicator(active_leases), }); } let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; if pending { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::QueuePendingReady, count: busy_indicator(pending), }); } let unclosed_claim_refs = queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; if unclosed_claim_refs { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::QueueUnclosedClaimRefs, count: busy_indicator(unclosed_claim_refs), }); } set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_started = Instant::now(); // Claims read segments before ready rows, while producers write in // the opposite order. Queue behind segment readers, then abandon // the transaction rather than wait with an exclusive child held. let lock_tables = async { sqlx::query(audited_sql(format!( "LOCK TABLE {segment_child} IN ACCESS EXCLUSIVE MODE" ))) .execute(tx.as_mut()) .await?; sqlx::query(audited_sql(format!( "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE NOWAIT" ))) .execute(tx.as_mut()) .await } .await; let lock_duration = lock_started.elapsed(); if let Err(err) = lock_tables { let _ = tx.rollback().await; if is_lock_contention_error(&err) { return Ok(PruneOutcome::Blocked { slot }); } return Err(map_sqlx_error(err)); } // The pre-lock gates are only an MVCC snapshot. If an in-flight // producer or claimer committed while the exclusive partition // lock waited, these post-lock gates see it before TRUNCATE. let active_leases = queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; if active_leases { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::QueueActiveLeases, count: busy_indicator(active_leases), }); } let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; if pending { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::QueuePendingReady, count: busy_indicator(pending), }); } let unclosed_claim_refs = queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; if unclosed_claim_refs { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::QueueUnclosedClaimRefs, count: busy_indicator(unclosed_claim_refs), }); } let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); let retention_floor = failed_retention_secs > 0; // #337: failed terminal rows still inside the retention floor // are re-homed into the live done segment before the TRUNCATE // so they stay retryable. The ring rotation lock is held above, // which serializes against rotate — the current slot (and its // generation, read from the ledger cursor) cannot move while // carried rows are re-inserted into it. let carried_failed_rows = if retention_floor { // Carried rows are widened to the self-contained synthetic // done-row shape (the COALESCE chain mirrors // `done_row_projection`) because the ready rows they would // otherwise hydrate from are truncated with this slot. // Deliberately no ON CONFLICT: lane_seq comes from // never-resetting sequences, so a PK collision means // corrupted terminal state and must abort the prune loudly // rather than silently drop a terminal fact. let carried = sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.done_entries ( ready_slot, ready_generation, job_id, kind, queue, args, state, priority, attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, attempted_at, finalized_at, created_at, unique_key, unique_states, payload ) SELECT $2, $3, done.job_id, done.kind, done.queue, COALESCE(done.args, ready.args, '{{}}'::jsonb), done.state, done.priority, done.attempt, done.run_lease, COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), done.lane_seq, done.enqueue_shard, COALESCE(done.run_at, ready.run_at, done.finalized_at), COALESCE(done.attempted_at, ready.attempted_at), done.finalized_at, COALESCE(done.created_at, ready.created_at, done.finalized_at), COALESCE(done.unique_key, ready.unique_key), COALESCE(done.unique_states, ready.unique_states), COALESCE(done.payload, ready.payload, '{{}}'::jsonb) FROM {done_child} AS done LEFT JOIN {ready_child} AS ready ON ready.ready_slot = done.ready_slot AND ready.ready_generation = done.ready_generation AND ready.queue = done.queue AND ready.priority = done.priority AND ready.enqueue_shard = done.enqueue_shard AND ready.lane_seq = done.lane_seq WHERE done.state = 'failed' AND done.finalized_at >= now() - make_interval(secs => $1::bigint) "# ))) .bind(failed_retention_secs) .bind(state.0) .bind(current_generation) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)? .rows_affected(); if carried > 0 { // This slot's live counter rows and delta segment are // wiped below, so the carried rows' exact-count // evidence moves with them: re-append positive deltas // under the current slot, grouped exactly like the // completion path's delta append. sqlx::query(audited_sql(format!( r#" INSERT INTO {schema}.queue_terminal_count_deltas ( ready_slot, ready_generation, queue, priority, enqueue_shard, counter_bucket, terminal_delta ) SELECT $2, $3, queue, priority, enqueue_shard, mod( mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) + {TERMINAL_COUNTER_BUCKETS}::bigint, {TERMINAL_COUNTER_BUCKETS}::bigint )::smallint AS counter_bucket, count(*)::bigint FROM {done_child} WHERE state = 'failed' AND finalized_at >= now() - make_interval(secs => $1::bigint) GROUP BY queue, priority, enqueue_shard, counter_bucket ORDER BY queue, priority, enqueue_shard, counter_bucket "# ))) .bind(failed_retention_secs) .bind(state.0) .bind(current_generation) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; } carried } else { 0 }; // #290: scan the about-to-be-truncated partition for the rollup // fold. The rollup column is *permanent* state, so we MUST fold // from ground truth (the `{done_child}` partition itself), not // from `queue_terminal_live_counts` — the counter can drift // briefly during a rolling upgrade from a pre-counter binary // and folding drift into the rollup would bake it in forever. // The counter rows for this slot are still cleaned up after the // truncate; reads from the counter (queue_counts_exact in #305) // may transiently disagree with the rollup until the operator // runs `awa storage rebuild-terminal-counters`, but the rollup // itself stays authoritative. See PR #304 reviewer finding // "High: A1 can persist stale counter state before the // read-switch PR" for the trade-off. // // Carried failed rows are excluded from both columns: they are // still live in `done_entries`, so folding them into the // permanent rollup would double-count them. let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = sqlx::query_as(audited_sql(format!( r#" WITH done_counts AS ( SELECT queue, priority, (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed, (count(*) FILTER ( WHERE state = 'failed' AND ( $2::boolean IS FALSE OR finalized_at < now() - make_interval(secs => $1::bigint) ) ))::bigint AS failed FROM {done_child} GROUP BY queue, priority ), -- Scan the entire partition being truncated (all generations), -- mirroring done_counts' whole-child scan above. The rotate guard -- keeps a partition to one generation at prune time, so this is the -- same set today; counting the whole child keeps the rollup fold -- conservative even if that invariant is ever weakened (a stray -- generation's batches would otherwise be truncated without being -- folded into the rollup -> terminal-count undercount). batch_rows AS ( SELECT batch.ready_generation, batch.queue, batch.priority, item.job_id, item.run_lease FROM {receipt_batch_child} AS batch CROSS JOIN LATERAL unnest( batch.job_ids, batch.run_leases ) AS item(job_id, run_lease) ), batch_counts AS ( SELECT batch_rows.queue, batch_rows.priority, count(*)::bigint AS completed FROM batch_rows WHERE NOT EXISTS ( SELECT 1 FROM {receipt_tomb_child} AS tomb WHERE tomb.ready_generation = batch_rows.ready_generation AND tomb.job_id = batch_rows.job_id AND tomb.run_lease = batch_rows.run_lease ) GROUP BY batch_rows.queue, batch_rows.priority ), keys AS ( SELECT queue, priority FROM done_counts UNION SELECT queue, priority FROM batch_counts ) SELECT keys.queue, keys.priority, ( COALESCE(done_counts.completed, 0) + COALESCE(batch_counts.completed, 0) )::bigint AS pruned_completed_count, COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count FROM keys LEFT JOIN done_counts USING (queue, priority) LEFT JOIN batch_counts USING (queue, priority) "# ))) .bind(failed_retention_secs) .bind(retention_floor) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; let truncate_started = Instant::now(); let truncate = self .truncate_used_children_tx( &mut tx, &[ &ready_child, &claim_attempt_child, &done_child, &tomb_child, &segment_child, &receipt_batch_child, &receipt_tomb_child, &delta_child, ], ) .await; let truncate_duration = truncate_started.elapsed(); match truncate { Ok(_) => { if !pruned_terminal_counts.is_empty() { // #371: append the pruned counts as delta rows instead // of upserting `queue_terminal_rollups` here — the // upsert under a pinned MVCC horizon left one // unreclaimable dead rollup version per prune. The // maintenance leader folds the deltas // (`fold_terminal_rollup_deltas`) when the horizon is // clear; exact readers add unfolded delta sums to // their rollup reads. self.append_terminal_rollup_deltas_tx( &mut tx, pruned_terminal_counts.into_iter(), ) .await?; } // #290: the live counter rows for this slot are about to // be orphans (their underlying done_entries rows just got // truncated). Delete them in the same transaction. The // rollup fold above already captured ground-truth from // the partition scan; this just cleans up the counter // index entries so a future insert into a re-rotated // slot starts from zero. sqlx::query(audited_sql(format!( "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" ))) .bind(slot) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; let commit_started = Instant::now(); tx.commit().await.map_err(map_sqlx_error)?; let commit_duration = commit_started.elapsed(); if carried_failed_rows > 0 { tracing::info!( slot, carried_failed_rows, "Carried failed terminal rows inside the retention floor to the live done segment" ); } Ok(PruneOutcome::Pruned { slot, generation, carried_failed_rows, durations: PruneDurations { lock: lock_duration, truncate: truncate_duration, commit: commit_duration, }, }) } Err(err) if is_lock_contention_error(&err) => { let _ = tx.rollback().await; Ok(PruneOutcome::Blocked { slot }) } Err(err) => { let _ = tx.rollback().await; Err(map_sqlx_error(err)) } } } #[tracing::instrument( skip(self, pool), name = "queue_storage.prune_oldest_leases", fields(slot = tracing::field::Empty, generation = tracing::field::Empty) )] pub async fn prune_oldest_leases(&self, pool: &PgPool) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // `PruneLeasesPlan` in // `correctness/storage/AwaStorageLockOrder.tla` requires the // sequence lease-ring rotation serialization → // `lease_ring_slots[slot] FOR UPDATE` → `ACCESS EXCLUSIVE` on // the child. The child lock is bounded by a short // transaction-local lock_timeout because pure NOWAIT can starve // under continuous parent-partition readers, while an unbounded // wait would put maintenance at the head of the relation-lock // queue indefinitely. Without the rotation serialization a // concurrent rotator can flip the cursor under the prune's liveness // check and prune what should be the active partition. Ledger mode // -> advisory lock; compat mode (#371 staged upgrade) -> singleton // FOR UPDATE. let (authority, acquired) = self .acquire_ring_serializer_tx(&mut tx, RingFamily::Lease) .await?; if !acquired { let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Lease).await?; tx.commit().await.map_err(map_sqlx_error)?; return Ok( match oldest_initialized_ring_slot(current_slot, generation, slot_count) { Some((slot, _)) => PruneOutcome::Blocked { slot }, None => PruneOutcome::Noop, }, ); } let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Lease).await?; let state = (current_slot, generation, slot_count); let Some((slot, generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2) else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::Noop); }; tracing::Span::current().record("slot", slot); tracing::Span::current().record("generation", generation); let slot_locked: Option = sqlx::query_scalar(audited_sql(format!( r#" SELECT slot FROM {schema}.lease_ring_slots WHERE slot = $1 FOR UPDATE "# ))) .bind(slot) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if slot_locked.is_none() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::Noop); } let lease_child = lease_child_name(schema, slot as usize); set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_started = Instant::now(); let lock_table = sqlx::query(audited_sql(format!( "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" ))) .execute(tx.as_mut()) .await; let lock_duration = lock_started.elapsed(); if let Err(err) = lock_table { let _ = tx.rollback().await; if is_lock_contention_error(&err) { return Ok(PruneOutcome::Blocked { slot }); } return Err(map_sqlx_error(err)); } // The ring rotation serialization held above keeps the cursor // stable, but keep the explicit recheck to document the truncate // precondition. let (current_slot, _) = self .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) .await?; if current_slot == slot { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::LeaseCurrent, count: 0, }); } let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?; if active_leases { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::LeaseActive, count: busy_indicator(active_leases), }); } let truncate_started = Instant::now(); let truncate = self .truncate_used_children_tx(&mut tx, &[&lease_child]) .await; let truncate_duration = truncate_started.elapsed(); match truncate { Ok(_) => { let commit_started = Instant::now(); tx.commit().await.map_err(map_sqlx_error)?; let commit_duration = commit_started.elapsed(); Ok(PruneOutcome::Pruned { slot, generation, carried_failed_rows: 0, durations: PruneDurations { lock: lock_duration, truncate: truncate_duration, commit: commit_duration, }, }) } Err(err) if is_lock_contention_error(&err) => { let _ = tx.rollback().await; Ok(PruneOutcome::Blocked { slot }) } Err(err) => { let _ = tx.rollback().await; Err(map_sqlx_error(err)) } } } pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> { sqlx::query(audited_sql(format!("VACUUM {}", self.leases_table()))) .execute(pool) .await .map_err(map_sqlx_error)?; Ok(()) } /// ADR-023 claim-ring rotation. Parallel of `rotate_leases`. /// /// Advances the claim-ring cursor by appending to the /// `claim_ring_rotations` ledger (the generation PK insert is the /// compare-and-swap). Before flipping the cursor the target partition /// must be drained: the /// `lease_claims_`, `lease_claim_batches_`, /// `lease_claim_closures_`, and /// `lease_claim_closure_batches_` child tables /// must be empty. This is what the `rotate → prune → rotate` ring /// invariant requires — we only hand out a slot to new claims when a /// prior prune has truncated it. /// Seal the current claim-ring slot on every call that passes the busy /// and idle gates. Equivalent to [`Self::rotate_claims_with_cadence`] /// with [`RingRotationCadence::EAGER`]. pub async fn rotate_claims(&self, pool: &PgPool) -> Result { self.rotate_claims_with_cadence(pool, RingRotationCadence::EAGER) .await } #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_claims")] pub async fn rotate_claims_with_cadence( &self, pool: &PgPool, cadence: RingRotationCadence, ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // Per-ring serialization lock: guards prune_oldest_claims and // parallel rotators (#371). Ledger mode -> advisory lock; compat // mode -> claim_ring_state singleton FOR UPDATE (also serializes vs // live 0.6 rotators). let (authority, acquired) = self .acquire_ring_serializer_tx(&mut tx, RingFamily::Claim) .await?; if !acquired { let (current_slot, _) = self .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Claim).await?; tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: (current_slot + 1).rem_euclid(slot_count), busy: BusyCounts::default(), }); } let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Claim).await?; let state = (current_slot, generation, slot_count); let current_children = [ claim_child_name(schema, state.0 as usize), claim_batch_child_name(schema, state.0 as usize), closure_child_name(schema, state.0 as usize), claim_closure_batch_child_name(schema, state.0 as usize), ]; if self .slot_still_filling_tx( &mut tx, RingFamily::Claim, (state.0, state.1), cadence, ¤t_children, ) .await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedFilling { slot: state.0, generation: state.1, }); } let next_slot = (state.0 + 1).rem_euclid(state.2); // Busy check: all children of the incoming slot must be empty. // A non-empty `lease_claims_` means the previous lap's // prune hasn't run (or didn't complete); rotating anyway would // mix fresh claims with legacy rows and defeat the point of // partitioning. Non-empty closure children mean prune fell behind // on receipt-closure evidence specifically. let claim_busy = Self::relation_has_rows_tx(&mut tx, &claim_child_name(schema, next_slot as usize)) .await?; let claim_batch_busy = Self::relation_has_rows_tx( &mut tx, &claim_batch_child_name(schema, next_slot as usize), ) .await?; let closure_busy = Self::relation_has_rows_tx(&mut tx, &closure_child_name(schema, next_slot as usize)) .await?; let closure_batch_busy = Self::relation_has_rows_tx( &mut tx, &claim_closure_batch_child_name(schema, next_slot as usize), ) .await?; if claim_busy || claim_batch_busy || closure_busy || closure_batch_busy { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: next_slot, busy: BusyCounts { claims: busy_indicator(claim_busy || claim_batch_busy), closures: busy_indicator(closure_busy), closure_batches: busy_indicator(closure_batch_busy), ..Default::default() }, }); } // Idle gate (#371): skip only when the ENTIRE claim ring is // empty (parent probes, all partitions). Otherwise the ledger // append below only grows the `claim_ring_rotations` ledger (the // pre-ledger singleton UPDATE accumulated ~3.6k dead tuples/hour // under a pinned MVCC horizon at the historical cadence). See // `rotate` for both the race analysis and why the whole-ring scope // is load-bearing (prune never advances slot generations; freezing // rotation with sealed receipt evidence behind would strand it // until traffic resumed). let ring_parents = [ format!("{schema}.lease_claims"), format!("{schema}.lease_claim_batches"), format!("{schema}.lease_claim_closures"), format!("{schema}.lease_claim_closure_batches"), ]; if Self::relations_all_empty_tx(&mut tx, &ring_parents).await? { tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedIdle { slot: state.0, generation: state.1, }); } let next_generation = state.1 + 1; if !self .append_ring_rotation_tx( &mut tx, RingFamily::Claim, next_generation, next_slot, authority, ) .await? { // Lost the CAS race. Report the bounded presence // evidence we sampled before giving up. tx.commit().await.map_err(map_sqlx_error)?; return Ok(RotateOutcome::SkippedBusy { slot: next_slot, busy: BusyCounts { claims: busy_indicator(claim_busy || claim_batch_busy), closures: busy_indicator(closure_busy), closure_batches: busy_indicator(closure_batch_busy), ..Default::default() }, }); } tx.commit().await.map_err(map_sqlx_error)?; Ok(RotateOutcome::Rotated { slot: next_slot, generation: next_generation, }) } /// ADR-023 claim-ring prune. Parallel of `prune_oldest_leases`. /// /// Reclaims the oldest initialized (sealed) claim-ring slot by /// `TRUNCATE`-ing its `lease_claims_`, /// `lease_claim_batches_`, `lease_claim_closures_`, and /// `lease_claim_closure_batches_` children. Takes the full ADR-023 /// lock sequence: /// /// 1. The claim-ring rotation advisory lock (serialises with rotate; /// #371 — replaces the pre-ledger `FOR UPDATE` on the /// `claim_ring_state` singleton). /// 2. `FOR UPDATE` on the target `claim_ring_slots` row. /// 3. Proves every claim in the sealed partition has closure evidence. /// 4. `LOCK TABLE ACCESS EXCLUSIVE` on all claim-ring children with a /// short transaction-local `lock_timeout` (serialises with in-flight /// claim/complete/rescue writers; gives up under sustained /// contention). /// 5. Rechecks the slot is not the current one and that every /// claim still has closure evidence. /// 6. `TRUNCATE` all claim-ring children in a single statement. /// /// The "every claim has closure evidence" precondition is what ADR-023 /// calls `PartitionTruncateSafety`. If an open claim remains in the /// partition, prune returns `SkippedActive` and the claim has to /// drain by normal completion or be rescued by /// `rescue_stale_receipt_claims_tx` before prune will try again. #[tracing::instrument( skip(self, pool), name = "queue_storage.prune_oldest_claims", fields(slot = tracing::field::Empty, generation = tracing::field::Empty) )] pub async fn prune_oldest_claims(&self, pool: &PgPool) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; // Claim-ring rotation serialization: first resource in // `PruneClaimsPlan` (`correctness/storage/AwaStorageLockOrder.tla`); // serialises with rotate_claims and parallel pruners. Ledger mode -> // advisory lock; compat mode -> claim_ring_state singleton FOR // UPDATE (also serializes vs live 0.6 rotators). let (authority, acquired) = self .acquire_ring_serializer_tx(&mut tx, RingFamily::Claim) .await?; if !acquired { let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Claim).await?; tx.commit().await.map_err(map_sqlx_error)?; return Ok( match oldest_initialized_ring_slot(current_slot, generation, slot_count) { Some((slot, _)) => PruneOutcome::Blocked { slot }, None => PruneOutcome::Noop, }, ); } let (current_slot, generation) = self .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) .await?; let slot_count = self.ring_slot_count_tx(&mut tx, RingFamily::Claim).await?; let state = (current_slot, generation, slot_count); let Some((slot, generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2) else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::Noop); }; tracing::Span::current().record("slot", slot); tracing::Span::current().record("generation", generation); // Lock the slot row so concurrent rotate/prune observe the same // state machine transition. let slot_locked: Option = sqlx::query_scalar(audited_sql(format!( r#" SELECT slot FROM {schema}.claim_ring_slots WHERE slot = $1 FOR UPDATE "# ))) .bind(slot) .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; if slot_locked.is_none() { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::Noop); } let claim_child = claim_child_name(schema, slot as usize); let claim_batch_child = claim_batch_child_name(schema, slot as usize); let closure_child = closure_child_name(schema, slot as usize); let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize); // Before taking ACCESS EXCLUSIVE on the child partitions, prove // the sealed slot is actually reclaimable. This optimistic proof // avoids blocking claim/complete traffic behind a prune attempt // that is already known to be doomed. A claimer that read the // old current slot before rotation may still commit while the // exclusive lock waits, so the open-claim proof is repeated once // the child locks are held. let open_claims = claim_prune_has_open_claims_tx( &mut tx, schema, &claim_child, &claim_batch_child, &closure_child, &closure_batch_child, ) .await?; if open_claims { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::ClaimOpen, count: busy_indicator(open_claims), }); } set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_started = Instant::now(); let lock_tables = sqlx::query(audited_sql(format!( "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" ))) .execute(tx.as_mut()) .await; let lock_duration = lock_started.elapsed(); if let Err(err) = lock_tables { let _ = tx.rollback().await; if is_lock_contention_error(&err) { return Ok(PruneOutcome::Blocked { slot }); } return Err(map_sqlx_error(err)); } // After taking ACCESS EXCLUSIVE, recheck that the slot is still // not the current one. The ring rotation serialization should // already make this stable, but keeping the explicit gate documents // the truncate precondition. let (current_slot, _) = self .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) .await?; if current_slot == slot { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::ClaimCurrent, count: 0, }); } let open_claims = claim_prune_has_open_claims_tx( &mut tx, schema, &claim_child, &claim_batch_child, &closure_child, &closure_batch_child, ) .await?; if open_claims { tx.commit().await.map_err(map_sqlx_error)?; return Ok(PruneOutcome::SkippedActive { slot, reason: SkipReason::ClaimOpen, count: busy_indicator(open_claims), }); } let truncate_started = Instant::now(); let truncate = self .truncate_used_children_tx( &mut tx, &[ &claim_child, &claim_batch_child, &closure_child, &closure_batch_child, ], ) .await; let truncate_duration = truncate_started.elapsed(); match truncate { Ok(_) => { sqlx::query(audited_sql(format!( r#" UPDATE {schema}.claim_ring_slots SET rescue_cursor_claimed_at = '-infinity'::timestamptz, rescue_cursor_job_id = 0, rescue_cursor_run_lease = 0, deadline_cursor_deadline_at = '-infinity'::timestamptz, deadline_cursor_job_id = 0, deadline_cursor_run_lease = 0, batch_deadline_cursor_deadline_at = '-infinity'::timestamptz, batch_deadline_cursor_batch_id = 0 WHERE slot = $1 "# ))) .bind(slot) .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; let commit_started = Instant::now(); tx.commit().await.map_err(map_sqlx_error)?; let commit_duration = commit_started.elapsed(); Ok(PruneOutcome::Pruned { slot, generation, carried_failed_rows: 0, durations: PruneDurations { lock: lock_duration, truncate: truncate_duration, commit: commit_duration, }, }) } Err(err) if is_lock_contention_error(&err) => { let _ = tx.rollback().await; Ok(PruneOutcome::Blocked { slot }) } Err(err) => { let _ = tx.rollback().await; Err(map_sqlx_error(err)) } } } fn job_id_sequence(&self) -> String { format!("{}.job_id_seq", self.schema()) } fn leases_table(&self) -> String { format!("{}.{}", self.schema(), self.leases_relname()) } fn attempt_state_table(&self) -> String { format!("{}.{}", self.schema(), self.attempt_state_relname()) } } #[cfg(test)] mod writer_bind_tests { /// Multi-row `VALUES` writers bind one parameter per column per row, so a /// large batch exceeds PostgreSQL's 65,535 bind-parameter limit. Writers /// send one array per column through `unnest` instead. #[test] fn queue_storage_writers_do_not_bind_per_row() { let needle = concat!(".push", "_values("); let offenders: Vec = include_str!("queue_storage.rs") .lines() .enumerate() .filter(|(_, line)| line.contains(needle)) .map(|(index, _)| index + 1) .collect(); assert!( offenders.is_empty(), "queue_storage.rs builds a per-row VALUES insert at lines {offenders:?}; \ use INSERT ... SELECT FROM unnest(...) with one array bind per column" ); } }