Skip to main content

awa_model/
queue_storage.rs

1use crate::admin::{CallbackConfig, CallbackPollResult};
2use crate::dlq::{ListDlqFilter, RetryFromDlqOpts};
3use crate::error::AwaError;
4use crate::insert::prepare_row_raw;
5use crate::{InsertParams, JobRow, JobState};
6use chrono::TimeDelta;
7use chrono::{DateTime, Utc};
8use sqlx::Acquire;
9use sqlx::{PgPool, Postgres, QueryBuilder};
10use std::collections::hash_map::DefaultHasher;
11use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
12use std::hash::{Hash, Hasher};
13use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering};
14use std::sync::Mutex;
15use std::time::Duration;
16use uuid::Uuid;
17
18const DEFAULT_SCHEMA: &str = "awa";
19const DEFAULT_QUEUE_SLOT_COUNT: usize = 16;
20const DEFAULT_LEASE_SLOT_COUNT: usize = 8;
21const DEFAULT_CLAIM_SLOT_COUNT: usize = 8;
22const DEFAULT_QUEUE_STRIPE_COUNT: usize = 1;
23const QUEUE_STRIPE_DELIMITER: &str = "#";
24const COPY_NULL_SENTINEL: &str = "__AWA_NULL__";
25const COPY_CHUNK_TARGET_BYTES: usize = 256 * 1024;
26const TERMINAL_COUNTER_BUCKETS: i16 = 256;
27const RECEIPT_RESCUE_BATCH_LIMIT: i64 = 500;
28const RECEIPT_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000;
29const RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000;
30const DEFAULT_PRUNE_LOCK_TIMEOUT: Duration = Duration::from_millis(10);
31
32/// Portable 64-bit hash over raw ordering-key bytes.
33///
34/// This is intentionally simple enough to implement byte-for-byte in
35/// `awa.insert_job_compat`, so SQL, Rust, and Python producers that
36/// pass the same ordering-key bytes can derive the same routing facts.
37pub fn ordering_key_hash64(ordering_key: &[u8]) -> u64 {
38    let mut hash: u128 = 14_695_981_039_346_656_037;
39    const PRIME: u128 = 1_099_511_628_211;
40    const MASK: u128 = u64::MAX as u128;
41    for byte in ordering_key {
42        hash = hash.wrapping_mul(PRIME).wrapping_add(*byte as u128) & MASK;
43    }
44    hash as u64
45}
46
47/// Deterministically map an ordering key to a shard in `[0, shards)`.
48///
49/// Inputs sharing a key always produce the same shard, which is what
50/// lets producers preserve FIFO within a key when the destination
51/// queue is sharded. `shards <= 1` returns shard 0 unconditionally.
52///
53/// This deliberately remains the raw `ordering_key_hash64(key) % shards`
54/// mapping used by SQL compatibility producers.
55pub fn shard_for_ordering_key(ordering_key: &[u8], shards: i16) -> i16 {
56    if shards <= 1 {
57        return 0;
58    }
59    (ordering_key_hash64(ordering_key) % shards as u64) as i16
60}
61
62fn terminal_counter_bucket(job_id: i64) -> i16 {
63    job_id.rem_euclid(TERMINAL_COUNTER_BUCKETS as i64) as i16
64}
65
66#[derive(Debug, Clone)]
67pub struct QueueStorageConfig {
68    pub schema: String,
69    pub queue_slot_count: usize,
70    pub lease_slot_count: usize,
71    /// Number of child partitions the receipt ring splits
72    /// `lease_claims`, `lease_claim_batches`, and receipt-closure evidence
73    /// across (ADR-023).
74    /// Mirrors `lease_slot_count`: a small fixed set of slots
75    /// reclaimed by rotation + TRUNCATE rather than by row-level
76    /// DELETE.
77    pub claim_slot_count: usize,
78    pub queue_stripe_count: usize,
79    /// Use the receipt-plane short path for zero-deadline jobs:
80    /// claim writes compact batches into `lease_claim_batches` for the
81    /// zero-deadline hot path, while deadline-backed attempts keep row-local
82    /// evidence in `lease_claims` for indexed deadline rescue. Compact claims
83    /// can still materialize into `leases` when the worker later needs mutable
84    /// attempt state. Successful compact
85    /// completion writes durable terminal history through
86    /// `receipt_completion_batches` and compact claim-local closure evidence
87    /// in `lease_claim_closure_batches`, or falls back to `done_entries`, while
88    /// non-success exits and cold terminal-delete paths materialize explicit
89    /// closures in `lease_claim_closures`.
90    /// Receipt claim evidence is reclaimed by claim-ring rotation +
91    /// TRUNCATE.
92    /// Default `true`.
93    /// Set to `false` to force every claim through the legacy
94    /// `leases` materialization path.
95    pub lease_claim_receipts: bool,
96}
97
98impl Default for QueueStorageConfig {
99    fn default() -> Self {
100        Self {
101            schema: DEFAULT_SCHEMA.to_string(),
102            queue_slot_count: DEFAULT_QUEUE_SLOT_COUNT,
103            lease_slot_count: DEFAULT_LEASE_SLOT_COUNT,
104            claim_slot_count: DEFAULT_CLAIM_SLOT_COUNT,
105            queue_stripe_count: DEFAULT_QUEUE_STRIPE_COUNT,
106            lease_claim_receipts: true,
107        }
108    }
109}
110
111#[derive(Debug, Clone, PartialEq, Eq, sqlx::FromRow)]
112pub struct ClaimedEntry {
113    pub queue: String,
114    pub priority: i16,
115    pub lane_seq: i64,
116    pub ready_slot: i32,
117    pub ready_generation: i64,
118    pub lease_slot: i32,
119    pub lease_generation: i64,
120    /// ADR-023: the `claim_slot` partition this attempt's receipt
121    /// evidence landed in. Receipt-closure paths use this to co-locate
122    /// explicit `lease_claim_closures`; compact successful receipt
123    /// completions keep claim-local evidence in
124    /// `lease_claim_closure_batches` instead.
125    pub claim_slot: i32,
126    /// Stable identity for immutable receipt claim evidence. Present for
127    /// receipt-backed claims and used by the compact completion path to
128    /// validate the exact claim attempt without row-locking it.
129    pub receipt_id: Option<i64>,
130    /// Compact claim batch row identity for zero-deadline receipt claims.
131    /// Row-local receipt claims leave this unset.
132    pub claim_batch_id: Option<i64>,
133    /// One-based item position inside `lease_claim_batches.*_ids` arrays.
134    /// Present with `claim_batch_id` and used as an O(1) compact proof.
135    pub claim_batch_index: Option<i32>,
136    pub lease_claim_receipt: bool,
137    /// The enqueue shard the row was claimed from. Routes the
138    /// terminal `done_entries` write onto the correct shard's
139    /// `(ready_slot, queue, priority, enqueue_shard, lane_seq)` key
140    /// and is the join predicate for receipt and admin lookups that
141    /// touch `queue_claim_heads` / `ready_entries` / `leases`.
142    pub enqueue_shard: i16,
143}
144
145#[derive(Debug, Clone)]
146pub struct ClaimedRuntimeJob {
147    pub claim: ClaimedEntry,
148    pub job: JobRow,
149    pub unique_states: Option<String>,
150}
151
152#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)]
153pub struct QueueClaimerLease {
154    pub claimer_slot: i16,
155    pub lease_epoch: i64,
156}
157
158#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)]
159struct QueueClaimerLeaseRow {
160    claimer_slot: i16,
161    lease_epoch: i64,
162    last_claimed_at: DateTime<Utc>,
163    expires_at: DateTime<Utc>,
164}
165
166impl QueueClaimerLeaseRow {
167    fn lease(self) -> QueueClaimerLease {
168        QueueClaimerLease {
169            claimer_slot: self.claimer_slot,
170            lease_epoch: self.lease_epoch,
171        }
172    }
173
174    fn needs_refresh(
175        self,
176        now: DateTime<Utc>,
177        lease_ttl: Duration,
178        idle_threshold: Duration,
179    ) -> bool {
180        let Ok(idle_refresh_delta) = TimeDelta::from_std(idle_threshold / 2) else {
181            return true;
182        };
183        let Ok(expiry_refresh_delta) = TimeDelta::from_std(lease_ttl / 2) else {
184            return true;
185        };
186
187        self.last_claimed_at <= now - idle_refresh_delta
188            || self.expires_at <= now + expiry_refresh_delta
189    }
190}
191
192#[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)]
193pub struct QueueClaimerState {
194    pub target_claimers: i16,
195}
196
197impl ClaimedRuntimeJob {
198    fn into_done_row(self, finalized_at: DateTime<Utc>) -> Result<DoneJobRow, AwaError> {
199        let payload = QueueStorage::payload_from_parts(
200            self.job.metadata,
201            self.job.tags,
202            self.job.errors,
203            None,
204        )?;
205
206        Ok(DoneJobRow {
207            ready_slot: self.claim.ready_slot,
208            ready_generation: self.claim.ready_generation,
209            job_id: self.job.id,
210            kind: self.job.kind,
211            queue: self.job.queue,
212            args: self.job.args,
213            state: JobState::Completed,
214            priority: self.claim.priority,
215            attempt: self.job.attempt,
216            run_lease: self.job.run_lease,
217            max_attempts: self.job.max_attempts,
218            lane_seq: self.claim.lane_seq,
219            enqueue_shard: self.claim.enqueue_shard,
220            run_at: self.job.run_at,
221            attempted_at: self.job.attempted_at,
222            finalized_at,
223            created_at: self.job.created_at,
224            unique_key: self.job.unique_key,
225            unique_states: self.unique_states,
226            payload,
227        })
228    }
229}
230
231#[derive(Debug, Clone, Copy, PartialEq, Eq)]
232pub struct QueueCounts {
233    pub available: i64,
234    pub running: i64,
235    /// Count of rows in *any* terminal state (`completed`, `failed`, or
236    /// `cancelled`) for the queue. The name reflects what the field
237    /// actually counts: it is `{schema}.terminal_jobs` semantics, not
238    /// `count(*) WHERE state = 'completed'`. Physical terminal facts may
239    /// live in `done_entries` or compact receipt completion batches.
240    /// The historical name `completed` was a misnomer —
241    /// `queue_counts_exact` has always included failed and cancelled
242    /// terminals; renamed in #290 along with the counter-backed read
243    /// path.
244    pub terminal: i64,
245    /// Cumulative count of `failed` terminal rows pruned past the
246    /// failed-retention floor, summed from
247    /// `queue_terminal_rollups.pruned_failed_count`. These rows no
248    /// longer exist in `done_entries` and cannot be retried.
249    /// Monotonically non-decreasing — rollups never shrink.
250    pub pruned_failed: i64,
251}
252
253/// Cheap available-only signal used by the dispatcher's claimer-sizing
254/// control loop. Derives the count from enqueue and claim sequence cursors
255/// summed over the queue's physical stripes — two PK reads per lane, O(few
256/// rows) regardless of backlog size.
257///
258/// This is intentionally a separate type from [`QueueCounts`]: the
259/// dispatcher claim hot path only consumes the available count, and
260/// returning a `QueueCounts` with two perpetually-zero fields would
261/// invite future code to read `.running` or `.terminal` and silently
262/// get wrong answers. Code that legitimately needs the full counts
263/// should call [`QueueStorage::queue_counts`].
264#[derive(Debug, Clone, Copy, PartialEq, Eq)]
265pub(crate) struct AvailableSignal {
266    pub available: i64,
267}
268
269#[derive(Debug, Clone, Copy, PartialEq, Eq)]
270pub enum RotateOutcome {
271    Rotated {
272        slot: i32,
273        generation: i64,
274    },
275    /// Target slot has live state; rotation deferred. `busy` carries a
276    /// bounded per-table presence indicator observed at the gate (only fields
277    /// relevant to the ring being rotated are populated).
278    SkippedBusy {
279        slot: i32,
280        busy: BusyCounts,
281    },
282}
283
284/// Per-table presence observed at a rotation gate. Each non-zero value means
285/// "this table was non-empty"; it is not an exact row count. Each ring populates
286/// only the fields meaningful for it; unused fields stay zero. The
287/// maintenance loop emits one OTel metric label per non-zero field so
288/// dashboards can attribute "rotation pinned" to the responsible side.
289#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
290pub struct BusyCounts {
291    /// Queue ring: rows in the next `ready_entries` child.
292    pub queue_ready: i64,
293    /// Queue ring: rows in the next `ready_claim_attempt_batches` child.
294    pub queue_claim_attempt_batches: i64,
295    /// Queue ring: rows in the next `done_entries` child.
296    pub queue_done: i64,
297    /// Queue ring: rows in the next `ready_tombstones` child.
298    pub queue_tombstones: i64,
299    /// Queue ring: rows in the next `ready_segments` child.
300    pub queue_ready_segments: i64,
301    /// Queue ring: rows in the next `receipt_completion_batches` child.
302    pub queue_receipt_completion_batches: i64,
303    /// Queue ring: rows in the next `receipt_completion_tombstones` child.
304    pub queue_receipt_completion_tombstones: i64,
305    /// Queue ring: rows in the next `queue_terminal_count_deltas` child.
306    pub queue_terminal_deltas: i64,
307    /// Lease ring: rows in the next `leases` child.
308    pub leases: i64,
309    /// Claim ring: rows in the next `lease_claims` child.
310    pub claims: i64,
311    /// Claim ring: rows in the next `lease_claim_closures` child.
312    pub closures: i64,
313    /// Claim ring: rows in the next compact closure-batch child.
314    pub closure_batches: i64,
315}
316
317#[derive(Debug, Clone, Copy, PartialEq, Eq)]
318pub enum PruneOutcome {
319    Noop,
320    Pruned {
321        slot: i32,
322        /// Failed terminal rows inside the failed-retention floor that
323        /// the queue prune re-homed into the live `done_entries`
324        /// segment instead of dropping. Always zero for the lease and
325        /// claim rings.
326        carried_failed_rows: u64,
327    },
328    /// Lock acquisition timed out (held-tx, lock contention).
329    Blocked {
330        slot: i32,
331    },
332    /// Target slot still has live state. `reason` discriminates which gate
333    /// fired and `count` gives its magnitude.
334    SkippedActive {
335        slot: i32,
336        reason: SkipReason,
337        count: i64,
338    },
339}
340
341#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
342pub struct TerminalDeltaRollupOutcome {
343    pub rolled_slots: usize,
344    pub delta_rows: i64,
345    pub grouped_keys: i64,
346    pub skipped_active_slots: usize,
347    pub blocked_slots: usize,
348    pub skipped_mvcc_pinned: bool,
349}
350
351#[derive(Debug, Clone, Copy, PartialEq, Eq)]
352enum TerminalDeltaSlotRollup {
353    Empty,
354    Rolled { delta_rows: i64, grouped_keys: i64 },
355    SkippedActive,
356    SkippedMvccPinned,
357    Blocked,
358}
359
360/// Discriminator for [`PruneOutcome::SkippedActive`].
361///
362/// Multiple gates can fire `SkippedActive` for the same ring (e.g. queue
363/// prune checks both `active_leases` and `pending_ready`). Carrying the
364/// reason separately from `count` lets dashboards split out "ring saturated
365/// because backlog never drained" from "leases lingering on prior
366/// generation" without re-parsing log lines.
367#[derive(Debug, Clone, Copy, PartialEq, Eq)]
368pub enum SkipReason {
369    /// Queue prune: leases on the prior generation persist.
370    QueueActiveLeases,
371    /// Queue prune: receipt claims still need same-slot terminal evidence.
372    QueueUnclosedClaimRefs,
373    /// Queue prune: ready rows without matching done or tombstone evidence.
374    QueuePendingReady,
375    /// Lease prune: target slot equals the current slot (rotator race).
376    LeaseCurrent,
377    /// Lease prune: pending leases on target slot.
378    LeaseActive,
379    /// Claim prune: target slot equals the current slot (rotator race).
380    ClaimCurrent,
381    /// Claim prune: open claims on target slot (no closure evidence).
382    ClaimOpen,
383}
384
385impl SkipReason {
386    /// Stable, low-cardinality label suitable for OTel metric attributes.
387    pub fn as_str(self) -> &'static str {
388        match self {
389            Self::QueueActiveLeases => "queue.active_leases",
390            Self::QueueUnclosedClaimRefs => "queue.unclosed_claim_refs",
391            Self::QueuePendingReady => "queue.pending_ready",
392            Self::LeaseCurrent => "lease.current",
393            Self::LeaseActive => "lease.active",
394            Self::ClaimCurrent => "claim.current",
395            Self::ClaimOpen => "claim.open",
396        }
397    }
398}
399
400fn map_sqlx_error(err: sqlx::Error) -> AwaError {
401    if let sqlx::Error::Database(ref db_err) = err {
402        if db_err.code().as_deref() == Some("23505") {
403            return AwaError::UniqueConflict {
404                constraint: db_err.constraint().map(|c| c.to_string()),
405            };
406        }
407    }
408    AwaError::Database(err)
409}
410
411fn is_lock_contention_error(err: &sqlx::Error) -> bool {
412    matches!(
413        err,
414        sqlx::Error::Database(db_err) if db_err.code().as_deref() == Some("55P03")
415    )
416}
417
418async fn set_prune_lock_timeout_tx(
419    tx: &mut sqlx::Transaction<'_, Postgres>,
420    timeout: Duration,
421) -> Result<(), AwaError> {
422    let millis = timeout.as_millis().max(1).min(i64::MAX as u128);
423    let timeout = format!("{millis}ms");
424    sqlx::query("SELECT set_config('lock_timeout', $1, true)")
425        .bind(timeout)
426        .execute(tx.as_mut())
427        .await
428        .map_err(map_sqlx_error)?;
429    Ok(())
430}
431
432fn validate_ident(ident: &str) -> Result<(), AwaError> {
433    let mut chars = ident.chars();
434    match chars.next() {
435        Some(first) if first.is_ascii_lowercase() || first == '_' => {}
436        _ => {
437            return Err(AwaError::Validation(format!(
438                "invalid SQL identifier: {ident}"
439            )));
440        }
441    }
442
443    if chars.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_') {
444        Ok(())
445    } else {
446        Err(AwaError::Validation(format!(
447            "invalid SQL identifier: {ident}"
448        )))
449    }
450}
451
452fn ready_child_name(schema: &str, slot: usize) -> String {
453    format!("{schema}.ready_entries_{slot}")
454}
455
456fn ready_claim_attempt_batch_child_name(schema: &str, slot: usize) -> String {
457    format!("{schema}.ready_claim_attempt_batches_{slot}")
458}
459
460fn done_child_name(schema: &str, slot: usize) -> String {
461    format!("{schema}.done_entries_{slot}")
462}
463
464fn ready_tombstone_child_name(schema: &str, slot: usize) -> String {
465    format!("{schema}.ready_tombstones_{slot}")
466}
467
468fn ready_segment_child_name(schema: &str, slot: usize) -> String {
469    format!("{schema}.ready_segments_{slot}")
470}
471
472fn receipt_completion_batch_child_name(schema: &str, slot: usize) -> String {
473    format!("{schema}.receipt_completion_batches_{slot}")
474}
475
476fn receipt_completion_tombstone_child_name(schema: &str, slot: usize) -> String {
477    format!("{schema}.receipt_completion_tombstones_{slot}")
478}
479
480fn terminal_delta_child_name(schema: &str, slot: usize) -> String {
481    format!("{schema}.queue_terminal_count_deltas_{slot}")
482}
483
484fn done_ready_join(schema: &str, done_alias: &str, ready_alias: &str) -> String {
485    format!(
486        r#"
487            LEFT JOIN {schema}.ready_entries AS {ready_alias}
488              ON {ready_alias}.ready_slot = {done_alias}.ready_slot
489             AND {ready_alias}.ready_generation = {done_alias}.ready_generation
490             AND {ready_alias}.queue = {done_alias}.queue
491             AND {ready_alias}.priority = {done_alias}.priority
492             AND {ready_alias}.enqueue_shard = {done_alias}.enqueue_shard
493             AND {ready_alias}.lane_seq = {done_alias}.lane_seq
494        "#
495    )
496}
497
498fn done_row_projection(done_alias: &str, ready_alias: &str) -> String {
499    format!(
500        r#"
501                {done_alias}.ready_slot,
502                {done_alias}.ready_generation,
503                {done_alias}.job_id,
504                {done_alias}.kind,
505                {done_alias}.queue,
506                COALESCE({done_alias}.args, {ready_alias}.args, '{{}}'::jsonb) AS args,
507                {done_alias}.state,
508                {done_alias}.priority,
509                {done_alias}.attempt,
510                {done_alias}.run_lease,
511                COALESCE({done_alias}.max_attempts, {ready_alias}.max_attempts, 25::smallint) AS max_attempts,
512                {done_alias}.lane_seq,
513                {done_alias}.enqueue_shard,
514                COALESCE({done_alias}.run_at, {ready_alias}.run_at, {done_alias}.finalized_at) AS run_at,
515                COALESCE({done_alias}.attempted_at, {ready_alias}.attempted_at) AS attempted_at,
516                {done_alias}.finalized_at,
517                COALESCE({done_alias}.created_at, {ready_alias}.created_at, {done_alias}.finalized_at) AS created_at,
518                COALESCE({done_alias}.unique_key, {ready_alias}.unique_key) AS unique_key,
519                COALESCE({done_alias}.unique_states, {ready_alias}.unique_states) AS unique_states,
520                COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) AS payload
521        "#
522    )
523}
524
525fn lease_child_name(schema: &str, slot: usize) -> String {
526    format!("{schema}.leases_{slot}")
527}
528
529fn claim_child_name(schema: &str, slot: usize) -> String {
530    format!("{schema}.lease_claims_{slot}")
531}
532
533fn claim_batch_child_name(schema: &str, slot: usize) -> String {
534    format!("{schema}.lease_claim_batches_{slot}")
535}
536
537fn closure_child_name(schema: &str, slot: usize) -> String {
538    format!("{schema}.lease_claim_closures_{slot}")
539}
540
541fn claim_closure_batch_child_name(schema: &str, slot: usize) -> String {
542    format!("{schema}.lease_claim_closure_batches_{slot}")
543}
544
545fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result<usize, AwaError> {
546    let index = usize::try_from(slot).map_err(|_| {
547        AwaError::Validation(format!(
548            "invalid {ring} ring slot {slot}: slot must be non-negative"
549        ))
550    })?;
551    if index >= slot_count {
552        return Err(AwaError::Validation(format!(
553            "invalid {ring} ring slot {slot}: configured slot count is {slot_count}"
554        )));
555    }
556    Ok(index)
557}
558
559fn receipt_closed_evidence_sql(
560    schema: &str,
561    closure_rel: &str,
562    closure_batch_rel: &str,
563    claims_alias: &str,
564) -> String {
565    // `receipt_id` is allocated from a global sequence, so it identifies a
566    // claim without a `claim_slot` predicate. Compact closure batches retain
567    // the raw receipt array for audit/debugging, but hot membership checks use
568    // the per-child GiST index on the derived multirange.
569    format!(
570        r#"
571        (
572            {claims_alias}.closed_at IS NOT NULL
573            OR EXISTS (
574                SELECT 1 FROM {closure_rel} AS closures
575                WHERE closures.claim_slot = {claims_alias}.claim_slot
576                  AND closures.job_id = {claims_alias}.job_id
577                  AND closures.run_lease = {claims_alias}.run_lease
578            )
579            OR EXISTS (
580                SELECT 1
581                FROM {closure_batch_rel} AS closure_batches
582                WHERE closure_batches.receipt_ranges @> {claims_alias}.receipt_id
583            )
584            OR EXISTS (
585                SELECT 1 FROM {schema}.done_entries AS done
586                WHERE done.job_id = {claims_alias}.job_id
587                  AND done.run_lease = {claims_alias}.run_lease
588            )
589            OR EXISTS (
590                SELECT 1 FROM {schema}.deferred_jobs AS deferred
591                WHERE deferred.job_id = {claims_alias}.job_id
592                  AND deferred.run_lease = {claims_alias}.run_lease
593            )
594            OR EXISTS (
595                SELECT 1 FROM {schema}.dlq_entries AS dlq
596                WHERE dlq.job_id = {claims_alias}.job_id
597                  AND dlq.run_lease = {claims_alias}.run_lease
598            )
599        )
600        "#
601    )
602}
603
604fn busy_indicator(has_rows: bool) -> i64 {
605    if has_rows {
606        1
607    } else {
608        0
609    }
610}
611
612async fn queue_prune_has_active_leases_tx(
613    tx: &mut sqlx::Transaction<'_, Postgres>,
614    schema: &str,
615    slot: i32,
616    generation: i64,
617) -> Result<bool, AwaError> {
618    sqlx::query_scalar(&format!(
619        r#"
620        SELECT EXISTS (
621            SELECT 1
622            FROM {schema}.leases
623            WHERE ready_slot = $1
624              AND ready_generation = $2
625            LIMIT 1
626        )
627        "#
628    ))
629    .bind(slot)
630    .bind(generation)
631    .fetch_one(tx.as_mut())
632    .await
633    .map_err(map_sqlx_error)
634}
635
636async fn queue_prune_has_pending_ready_tx(
637    tx: &mut sqlx::Transaction<'_, Postgres>,
638    schema: &str,
639    ready_child: &str,
640    generation: i64,
641) -> Result<bool, AwaError> {
642    sqlx::query_scalar(&format!(
643        r#"
644        WITH claim_cursors AS MATERIALIZED (
645            SELECT
646                queue,
647                priority,
648                enqueue_shard,
649                {schema}.sequence_next_value(seq_name) AS claim_seq
650            FROM {schema}.queue_claim_heads
651        )
652        SELECT EXISTS (
653            SELECT 1
654            FROM claim_cursors AS claims
655            CROSS JOIN LATERAL (
656                SELECT 1
657                FROM {ready_child} AS ready
658                WHERE ready.ready_generation = $1
659                  AND ready.queue = claims.queue
660                  AND ready.priority = claims.priority
661                  AND ready.enqueue_shard = claims.enqueue_shard
662                  AND ready.lane_seq >= claims.claim_seq
663                LIMIT 1
664            ) AS pending_ready
665            LIMIT 1
666        )
667        "#
668    ))
669    .bind(generation)
670    .fetch_one(tx.as_mut())
671    .await
672    .map_err(map_sqlx_error)
673}
674
675async fn queue_prune_has_unclosed_claim_refs_tx(
676    tx: &mut sqlx::Transaction<'_, Postgres>,
677    schema: &str,
678    slot: i32,
679    generation: i64,
680) -> Result<bool, AwaError> {
681    let count_proves_claim_refs_closed: bool = sqlx::query_scalar(&format!(
682        r#"
683        WITH claim_count AS (
684            SELECT count(*)::bigint AS total
685            FROM {schema}.lease_claims AS claims
686            WHERE claims.ready_slot = $1
687              AND claims.ready_generation = $2
688        ),
689        compact_claim_count AS (
690            SELECT COALESCE(sum(claimed_count), 0)::bigint AS total
691            FROM {schema}.lease_claim_batches AS batches
692            WHERE batches.ready_slot = $1
693              AND batches.ready_generation = $2
694        ),
695        explicit_count AS (
696            SELECT count(*)::bigint AS total
697            FROM {schema}.lease_claims AS claims
698            JOIN {schema}.lease_claim_closures AS closures
699              ON closures.claim_slot = claims.claim_slot
700             AND closures.job_id = claims.job_id
701             AND closures.run_lease = claims.run_lease
702            WHERE claims.ready_slot = $1
703              AND claims.ready_generation = $2
704        ),
705        compact_count AS (
706            SELECT COALESCE(sum(closed_count), 0)::bigint AS total
707            FROM {schema}.lease_claim_closure_batches AS batches
708            WHERE batches.ready_slot = $1
709              AND batches.ready_generation = $2
710        )
711        SELECT claim_count.total + compact_claim_count.total =
712               explicit_count.total + compact_count.total
713        FROM claim_count, compact_claim_count, explicit_count, compact_count
714        "#
715    ))
716    .bind(slot)
717    .bind(generation)
718    .fetch_one(tx.as_mut())
719    .await
720    .map_err(map_sqlx_error)?;
721    if count_proves_claim_refs_closed {
722        return Ok(false);
723    }
724
725    // The exact anti-join over compact claim batches requires unnesting
726    // retained claim history. Under a hot MVCC horizon that proof can become
727    // more expensive than the prune it guards. A count mismatch is enough to
728    // make truncate unsafe, so skip and try again after closures catch up.
729    Ok(true)
730}
731
732async fn claim_prune_has_open_claims_tx(
733    tx: &mut sqlx::Transaction<'_, Postgres>,
734    _schema: &str,
735    claim_child: &str,
736    claim_batch_child: &str,
737    closure_child: &str,
738    closure_batch_child: &str,
739) -> Result<bool, AwaError> {
740    let count_proves_claims_closed: bool = sqlx::query_scalar(&format!(
741        r#"
742        WITH claim_count AS (
743            SELECT count(*)::bigint AS total FROM {claim_child}
744        ),
745        compact_claim_count AS (
746            SELECT COALESCE(sum(claimed_count), 0)::bigint AS total
747            FROM {claim_batch_child}
748        ),
749        explicit_count AS (
750            SELECT count(*)::bigint AS total FROM {closure_child}
751        ),
752        compact_count AS (
753            SELECT COALESCE(sum(closed_count), 0)::bigint AS total
754            FROM {closure_batch_child}
755        )
756        SELECT claim_count.total + compact_claim_count.total =
757               explicit_count.total + compact_count.total
758        FROM claim_count, compact_claim_count, explicit_count, compact_count
759        "#
760    ))
761    .fetch_one(tx.as_mut())
762    .await
763    .map_err(map_sqlx_error)?;
764    if count_proves_claims_closed {
765        return Ok(false);
766    }
767
768    // A count mismatch means at least one claim is not proven closed by the
769    // append-only closure ledgers. Returning SkippedActive is conservative and
770    // avoids an unbounded anti-join over retained compact claim batches.
771    Ok(true)
772}
773
774fn oldest_initialized_ring_slot(
775    current_slot: i32,
776    generation: i64,
777    slot_count: i32,
778) -> Option<(i32, i64)> {
779    if slot_count <= 1 {
780        return None;
781    }
782
783    let initialized_slots = (generation + 1).min(slot_count as i64) as i32;
784    if initialized_slots <= 1 {
785        return None;
786    }
787
788    let offset = initialized_slots - 1;
789    let oldest_slot = (current_slot - offset).rem_euclid(slot_count);
790    let oldest_generation = generation - offset as i64;
791    if oldest_generation < 0 {
792        return None;
793    }
794
795    Some((oldest_slot, oldest_generation))
796}
797
798#[cfg(test)]
799mod identifier_tests {
800    use super::{validate_ident, QueueStorage, QueueStorageConfig};
801
802    #[test]
803    fn queue_storage_schema_identifiers_are_lowercase_unquoted_names() {
804        for ident in ["awa", "awa_queue_storage", "_awa123"] {
805            validate_ident(ident).expect("identifier should be accepted");
806        }
807
808        for ident in ["Awa", "awa-queue", "123awa", "awa.queue"] {
809            assert!(
810                validate_ident(ident).is_err(),
811                "identifier should be rejected: {ident}"
812            );
813        }
814    }
815
816    #[test]
817    fn default_queue_storage_schema_requires_default_physical_shape() {
818        for config in [
819            QueueStorageConfig {
820                queue_slot_count: 32,
821                ..Default::default()
822            },
823            QueueStorageConfig {
824                lease_slot_count: 4,
825                ..Default::default()
826            },
827            QueueStorageConfig {
828                claim_slot_count: 4,
829                ..Default::default()
830            },
831            QueueStorageConfig {
832                lease_claim_receipts: false,
833                ..Default::default()
834            },
835        ] {
836            let err = QueueStorage::new(config).expect_err("default awa schema shape must reject");
837            assert!(
838                err.to_string()
839                    .contains("default `awa` queue-storage schema"),
840                "unexpected error: {err}"
841            );
842        }
843
844        QueueStorage::new(QueueStorageConfig {
845            schema: "awa_custom".to_string(),
846            queue_slot_count: 4,
847            lease_slot_count: 2,
848            claim_slot_count: 2,
849            lease_claim_receipts: false,
850            ..Default::default()
851        })
852        .expect("custom schema should allow custom physical shape");
853    }
854}
855
856#[cfg(test)]
857mod shard_routing_tests {
858    use super::shard_for_ordering_key;
859    use std::collections::HashSet;
860
861    #[test]
862    fn shards_le_one_collapse_to_zero() {
863        assert_eq!(shard_for_ordering_key(b"customer-42", 1), 0);
864        assert_eq!(shard_for_ordering_key(b"", 1), 0);
865        assert_eq!(shard_for_ordering_key(b"customer-42", 0), 0);
866    }
867
868    #[test]
869    fn same_key_lands_on_same_shard() {
870        let key = b"customer-42";
871        let first = shard_for_ordering_key(key, 8);
872        for _ in 0..100 {
873            assert_eq!(shard_for_ordering_key(key, 8), first);
874        }
875    }
876
877    #[test]
878    fn shard_is_within_range() {
879        for n in 0..256u32 {
880            let key = format!("order-{n}");
881            let shard = shard_for_ordering_key(key.as_bytes(), 8);
882            assert!((0..8).contains(&shard));
883        }
884    }
885
886    #[test]
887    fn distinct_keys_spread_across_shards() {
888        let mut hit: HashSet<i16> = HashSet::new();
889        for n in 0..1024u32 {
890            let key = format!("order-{n}");
891            hit.insert(shard_for_ordering_key(key.as_bytes(), 8));
892        }
893        assert_eq!(hit.len(), 8, "1024 distinct keys should cover all 8 shards");
894    }
895}
896
897#[cfg(test)]
898mod ring_slot_tests {
899    use super::oldest_initialized_ring_slot;
900
901    #[test]
902    fn oldest_initialized_ring_slot_is_none_until_second_slot_exists() {
903        assert_eq!(oldest_initialized_ring_slot(0, 0, 8), None);
904    }
905
906    #[test]
907    fn oldest_initialized_ring_slot_tracks_partial_ring_startup() {
908        assert_eq!(oldest_initialized_ring_slot(1, 1, 8), Some((0, 0)));
909        assert_eq!(oldest_initialized_ring_slot(2, 2, 8), Some((0, 0)));
910        assert_eq!(oldest_initialized_ring_slot(3, 3, 8), Some((0, 0)));
911    }
912
913    #[test]
914    fn oldest_initialized_ring_slot_wraps_after_full_rotation() {
915        assert_eq!(oldest_initialized_ring_slot(7, 7, 8), Some((0, 0)));
916        assert_eq!(oldest_initialized_ring_slot(0, 8, 8), Some((1, 1)));
917        assert_eq!(oldest_initialized_ring_slot(1, 9, 8), Some((2, 2)));
918    }
919}
920
921#[cfg(test)]
922mod claim_cursor_advance_tests {
923    use super::{ClaimCursorAdvance, QueueStorage};
924
925    fn advance(next_seq: i64, only_if_current: Option<i64>) -> ClaimCursorAdvance {
926        ClaimCursorAdvance {
927            queue: "queue".to_string(),
928            priority: 2,
929            enqueue_shard: 0,
930            next_seq,
931            only_if_current,
932        }
933    }
934
935    #[test]
936    fn normalize_claim_cursor_advances_sorts_conditional_lane_updates() {
937        let normalized = QueueStorage::normalize_claim_cursor_advances(&[
938            advance(7, Some(6)),
939            advance(6, Some(5)),
940            advance(8, Some(7)),
941        ]);
942
943        let ordered: Vec<(i64, i64)> = normalized
944            .iter()
945            .map(|advance| (advance.only_if_current.unwrap(), advance.next_seq))
946            .collect();
947        assert_eq!(ordered, vec![(5, 6), (6, 7), (7, 8)]);
948    }
949
950    #[test]
951    fn normalize_claim_cursor_advances_coalesces_unconditional_lane_updates() {
952        let normalized = QueueStorage::normalize_claim_cursor_advances(&[
953            advance(3, None),
954            advance(5, None),
955            advance(4, Some(3)),
956        ]);
957
958        assert_eq!(normalized.len(), 1);
959        assert_eq!(normalized[0].next_seq, 5);
960        assert_eq!(normalized[0].only_if_current, None);
961    }
962}
963
964fn default_payload_metadata() -> serde_json::Value {
965    serde_json::json!({})
966}
967
968fn is_empty_json_object(value: &serde_json::Value) -> bool {
969    value.as_object().is_some_and(serde_json::Map::is_empty)
970}
971
972fn is_compact_receipt_completion_metadata(value: &serde_json::Value) -> bool {
973    let Some(metadata) = value.as_object() else {
974        return false;
975    };
976
977    metadata
978        .keys()
979        .all(|key| key == "_awa_original_priority" || key == "_awa_original_queue")
980}
981
982#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
983struct RuntimePayload {
984    #[serde(
985        default = "default_payload_metadata",
986        skip_serializing_if = "is_empty_json_object"
987    )]
988    metadata: serde_json::Value,
989    #[serde(default, skip_serializing_if = "Vec::is_empty")]
990    tags: Vec<String>,
991    #[serde(default, skip_serializing_if = "Vec::is_empty")]
992    errors: Vec<serde_json::Value>,
993    #[serde(default, skip_serializing_if = "Option::is_none")]
994    progress: Option<serde_json::Value>,
995}
996
997impl Default for RuntimePayload {
998    fn default() -> Self {
999        Self {
1000            metadata: default_payload_metadata(),
1001            tags: Vec::new(),
1002            errors: Vec::new(),
1003            progress: None,
1004        }
1005    }
1006}
1007
1008impl RuntimePayload {
1009    fn from_json(value: serde_json::Value) -> Result<Self, AwaError> {
1010        if value.is_null() {
1011            return Ok(Self::default());
1012        }
1013        let payload: Self = serde_json::from_value(value)?;
1014        if !payload.metadata.is_object() {
1015            return Err(AwaError::Validation(
1016                "queue storage payload metadata must be a JSON object".to_string(),
1017            ));
1018        }
1019        Ok(payload)
1020    }
1021
1022    fn into_json(self) -> serde_json::Value {
1023        serde_json::to_value(self).expect("runtime payload serializes")
1024    }
1025
1026    fn errors_option(&self) -> Option<Vec<serde_json::Value>> {
1027        (!self.errors.is_empty()).then(|| self.errors.clone())
1028    }
1029
1030    fn push_error(&mut self, error: serde_json::Value) {
1031        self.errors.push(error);
1032    }
1033
1034    fn set_progress(&mut self, progress: Option<serde_json::Value>) {
1035        self.progress = progress;
1036    }
1037
1038    fn insert_callback_result(&mut self, payload: Option<serde_json::Value>) {
1039        let metadata = self
1040            .metadata
1041            .as_object_mut()
1042            .expect("runtime payload metadata object");
1043        metadata.insert(
1044            "_awa_callback_result".to_string(),
1045            payload.unwrap_or(serde_json::Value::Null),
1046        );
1047    }
1048}
1049
1050#[cfg(test)]
1051mod runtime_payload_tests {
1052    use super::{
1053        is_compact_receipt_completion_metadata, storage_payload, terminal_storage_payload,
1054        RuntimePayload,
1055    };
1056
1057    #[test]
1058    fn default_runtime_payload_serializes_compactly() {
1059        assert_eq!(
1060            RuntimePayload::default().into_json(),
1061            serde_json::json!({}),
1062            "default payloads should not write empty metadata/tags/errors/progress"
1063        );
1064        assert_eq!(
1065            storage_payload(&RuntimePayload::default().into_json()),
1066            None
1067        );
1068    }
1069
1070    #[test]
1071    fn missing_runtime_payload_fields_round_trip_with_defaults() {
1072        let payload = RuntimePayload::from_json(serde_json::json!({})).unwrap();
1073
1074        assert_eq!(payload.metadata, serde_json::json!({}));
1075        assert!(payload.tags.is_empty());
1076        assert!(payload.errors.is_empty());
1077        assert_eq!(payload.progress, None);
1078        assert_eq!(payload.into_json(), serde_json::json!({}));
1079    }
1080
1081    #[test]
1082    fn null_runtime_payload_round_trips_with_defaults() {
1083        let payload = RuntimePayload::from_json(serde_json::Value::Null).unwrap();
1084
1085        assert_eq!(payload.metadata, serde_json::json!({}));
1086        assert!(payload.tags.is_empty());
1087        assert!(payload.errors.is_empty());
1088        assert_eq!(payload.progress, None);
1089        assert_eq!(storage_payload(&payload.into_json()), None);
1090    }
1091
1092    #[test]
1093    fn compact_receipt_completion_metadata_only_allows_awa_provenance() {
1094        assert!(is_compact_receipt_completion_metadata(&serde_json::json!(
1095            {}
1096        )));
1097        assert!(is_compact_receipt_completion_metadata(
1098            &serde_json::json!({ "_awa_original_priority": 4 })
1099        ));
1100        assert!(is_compact_receipt_completion_metadata(&serde_json::json!({
1101            "_awa_original_queue": "default",
1102            "_awa_original_priority": 4
1103        })));
1104        assert!(!is_compact_receipt_completion_metadata(
1105            &serde_json::json!({ "tenant": "acme" })
1106        ));
1107        assert!(!is_compact_receipt_completion_metadata(&serde_json::json!(
1108            null
1109        )));
1110    }
1111
1112    #[test]
1113    fn legacy_expanded_runtime_payload_round_trips_to_compact_form() {
1114        let payload = RuntimePayload::from_json(serde_json::json!({
1115            "metadata": {},
1116            "tags": [],
1117            "errors": [],
1118            "progress": null
1119        }))
1120        .unwrap();
1121
1122        assert_eq!(payload.metadata, serde_json::json!({}));
1123        assert!(payload.tags.is_empty());
1124        assert!(payload.errors.is_empty());
1125        assert_eq!(payload.progress, None);
1126        assert_eq!(payload.into_json(), serde_json::json!({}));
1127    }
1128
1129    #[test]
1130    fn non_default_runtime_payload_fields_are_preserved() {
1131        let payload = RuntimePayload::from_json(serde_json::json!({
1132            "metadata": { "source": "test" },
1133            "tags": ["fast"],
1134            "errors": [{ "message": "boom" }],
1135            "progress": { "step": 1 }
1136        }))
1137        .unwrap();
1138
1139        assert_eq!(
1140            payload.into_json(),
1141            serde_json::json!({
1142                "metadata": { "source": "test" },
1143                "tags": ["fast"],
1144                "errors": [{ "message": "boom" }],
1145                "progress": { "step": 1 }
1146            })
1147        );
1148    }
1149
1150    #[test]
1151    fn unchanged_terminal_payload_elides_storage_copy() {
1152        let payload = serde_json::json!({
1153            "metadata": { "source": "test" },
1154            "tags": ["fast"]
1155        });
1156
1157        assert_eq!(terminal_storage_payload(&payload, Some(&payload)), None);
1158
1159        let changed = serde_json::json!({
1160            "metadata": { "source": "test" },
1161            "tags": ["fast"],
1162            "errors": [{ "message": "boom" }]
1163        });
1164        assert_eq!(
1165            terminal_storage_payload(&changed, Some(&payload)),
1166            Some(&changed)
1167        );
1168    }
1169}
1170
1171fn unique_state_claims(unique_states: Option<&str>, state: JobState) -> bool {
1172    let Some(bitmask) = unique_states else {
1173        return false;
1174    };
1175    let idx = state.bit_position() as usize;
1176    bitmask.as_bytes().get(idx).is_some_and(|bit| *bit == b'1')
1177}
1178
1179fn write_copy_field(buf: &mut Vec<u8>, value: &str) {
1180    if value.contains(',')
1181        || value.contains('"')
1182        || value.contains('\n')
1183        || value.contains('\r')
1184        || value.contains('\\')
1185        || value == COPY_NULL_SENTINEL
1186    {
1187        buf.push(b'"');
1188        for byte in value.bytes() {
1189            if byte == b'"' {
1190                buf.push(b'"');
1191            }
1192            buf.push(byte);
1193        }
1194        buf.push(b'"');
1195    } else {
1196        buf.extend_from_slice(value.as_bytes());
1197    }
1198}
1199
1200fn write_copy_json(buf: &mut Vec<u8>, value: &serde_json::Value) {
1201    let json = serde_json::to_string(value).expect("JSON serialization should not fail");
1202    write_copy_field(buf, &json);
1203}
1204
1205fn storage_payload(value: &serde_json::Value) -> Option<&serde_json::Value> {
1206    (!is_storage_payload_empty(value)).then_some(value)
1207}
1208
1209fn terminal_storage_payload<'a>(
1210    value: &'a serde_json::Value,
1211    ready_payload: Option<&serde_json::Value>,
1212) -> Option<&'a serde_json::Value> {
1213    if is_storage_payload_empty(value) || ready_payload.is_some_and(|ready| ready == value) {
1214        None
1215    } else {
1216        Some(value)
1217    }
1218}
1219
1220fn is_storage_payload_empty(value: &serde_json::Value) -> bool {
1221    value.is_null() || is_empty_json_object(value)
1222}
1223
1224fn write_copy_storage_payload(buf: &mut Vec<u8>, value: &serde_json::Value) {
1225    match storage_payload(value) {
1226        Some(value) => write_copy_json(buf, value),
1227        None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1228    }
1229}
1230
1231fn write_copy_datetime(buf: &mut Vec<u8>, value: DateTime<Utc>) {
1232    write_copy_field(buf, &value.to_rfc3339());
1233}
1234
1235fn write_copy_optional_datetime(buf: &mut Vec<u8>, value: Option<DateTime<Utc>>) {
1236    match value {
1237        Some(value) => write_copy_datetime(buf, value),
1238        None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1239    }
1240}
1241
1242fn write_copy_optional_bytes(buf: &mut Vec<u8>, value: &Option<Vec<u8>>) {
1243    match value {
1244        Some(bytes) => {
1245            let bytea_hex = format!("\\x{}", hex::encode(bytes));
1246            write_copy_field(buf, &bytea_hex);
1247        }
1248        None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1249    }
1250}
1251
1252fn write_copy_optional_string(buf: &mut Vec<u8>, value: Option<&str>) {
1253    match value {
1254        Some(value) => write_copy_field(buf, value),
1255        None => buf.extend_from_slice(COPY_NULL_SENTINEL.as_bytes()),
1256    }
1257}
1258
1259fn write_ready_copy_row(
1260    buf: &mut Vec<u8>,
1261    ready_slot: i32,
1262    ready_generation: i64,
1263    row: &RuntimeReadyInsert,
1264) {
1265    buf.extend_from_slice(ready_slot.to_string().as_bytes());
1266    buf.push(b',');
1267    buf.extend_from_slice(ready_generation.to_string().as_bytes());
1268    buf.push(b',');
1269    buf.extend_from_slice(row.job_id.to_string().as_bytes());
1270    buf.push(b',');
1271    write_copy_field(buf, &row.kind);
1272    buf.push(b',');
1273    write_copy_field(buf, &row.queue);
1274    buf.push(b',');
1275    write_copy_json(buf, &row.args);
1276    buf.push(b',');
1277    buf.extend_from_slice(row.priority.to_string().as_bytes());
1278    buf.push(b',');
1279    buf.extend_from_slice(row.attempt.to_string().as_bytes());
1280    buf.push(b',');
1281    buf.extend_from_slice(row.run_lease.to_string().as_bytes());
1282    buf.push(b',');
1283    buf.extend_from_slice(row.max_attempts.to_string().as_bytes());
1284    buf.push(b',');
1285    buf.extend_from_slice(row.lane_seq.to_string().as_bytes());
1286    buf.push(b',');
1287    buf.extend_from_slice(row.enqueue_shard.to_string().as_bytes());
1288    buf.push(b',');
1289    write_copy_datetime(buf, row.run_at);
1290    buf.push(b',');
1291    write_copy_optional_datetime(buf, row.attempted_at);
1292    buf.push(b',');
1293    write_copy_datetime(buf, row.created_at);
1294    buf.push(b',');
1295    write_copy_optional_bytes(buf, &row.unique_key);
1296    buf.push(b',');
1297    write_copy_optional_string(buf, row.unique_states.as_deref());
1298    buf.push(b',');
1299    write_copy_storage_payload(buf, &row.payload);
1300    buf.push(b'\n');
1301}
1302
1303fn write_deferred_copy_row(buf: &mut Vec<u8>, row: &DeferredJobRow) {
1304    buf.extend_from_slice(row.job_id.to_string().as_bytes());
1305    buf.push(b',');
1306    write_copy_field(buf, &row.kind);
1307    buf.push(b',');
1308    write_copy_field(buf, &row.queue);
1309    buf.push(b',');
1310    write_copy_json(buf, &row.args);
1311    buf.push(b',');
1312    write_copy_field(buf, &row.state.to_string());
1313    buf.push(b',');
1314    buf.extend_from_slice(row.priority.to_string().as_bytes());
1315    buf.push(b',');
1316    buf.extend_from_slice(row.attempt.to_string().as_bytes());
1317    buf.push(b',');
1318    buf.extend_from_slice(row.run_lease.to_string().as_bytes());
1319    buf.push(b',');
1320    buf.extend_from_slice(row.max_attempts.to_string().as_bytes());
1321    buf.push(b',');
1322    write_copy_datetime(buf, row.run_at);
1323    buf.push(b',');
1324    write_copy_optional_datetime(buf, row.attempted_at);
1325    buf.push(b',');
1326    write_copy_optional_datetime(buf, row.finalized_at);
1327    buf.push(b',');
1328    write_copy_datetime(buf, row.created_at);
1329    buf.push(b',');
1330    write_copy_optional_bytes(buf, &row.unique_key);
1331    buf.push(b',');
1332    write_copy_optional_string(buf, row.unique_states.as_deref());
1333    buf.push(b',');
1334    write_copy_storage_payload(buf, &row.payload);
1335    buf.push(b'\n');
1336}
1337
1338fn lifecycle_error(error: impl Into<String>, attempt: i16, terminal: bool) -> serde_json::Value {
1339    let mut value = serde_json::json!({
1340        "error": error.into(),
1341        "attempt": attempt,
1342        "at": Utc::now().to_rfc3339(),
1343    });
1344    if terminal {
1345        value["terminal"] = serde_json::Value::Bool(true);
1346    }
1347    value
1348}
1349
1350fn transition_timestamp(job: &JobRow) -> DateTime<Utc> {
1351    job.finalized_at
1352        .or(job.heartbeat_at)
1353        .or(job.deadline_at)
1354        .or(job.attempted_at)
1355        .unwrap_or(job.run_at)
1356}
1357
1358fn state_rank(state: JobState) -> u8 {
1359    match state {
1360        JobState::Running | JobState::WaitingExternal => 4,
1361        JobState::Retryable | JobState::Scheduled => 3,
1362        JobState::Available => 2,
1363        JobState::Completed | JobState::Failed | JobState::Cancelled => 1,
1364    }
1365}
1366
1367#[derive(Debug, Clone, sqlx::FromRow)]
1368struct ReadyJobRow {
1369    job_id: i64,
1370    kind: String,
1371    queue: String,
1372    args: serde_json::Value,
1373    priority: i16,
1374    attempt: i16,
1375    run_lease: i64,
1376    max_attempts: i16,
1377    run_at: DateTime<Utc>,
1378    attempted_at: Option<DateTime<Utc>>,
1379    created_at: DateTime<Utc>,
1380    unique_key: Option<Vec<u8>>,
1381    payload: serde_json::Value,
1382}
1383
1384impl ReadyJobRow {
1385    fn into_job_row(self) -> Result<JobRow, AwaError> {
1386        let payload = RuntimePayload::from_json(self.payload)?;
1387        Ok(JobRow {
1388            id: self.job_id,
1389            kind: self.kind,
1390            queue: self.queue,
1391            args: self.args,
1392            state: JobState::Available,
1393            priority: self.priority,
1394            attempt: self.attempt,
1395            run_lease: self.run_lease,
1396            max_attempts: self.max_attempts,
1397            run_at: self.run_at,
1398            heartbeat_at: None,
1399            deadline_at: None,
1400            attempted_at: self.attempted_at,
1401            finalized_at: None,
1402            created_at: self.created_at,
1403            errors: payload.errors_option(),
1404            metadata: payload.metadata,
1405            tags: payload.tags,
1406            unique_key: self.unique_key,
1407            unique_states: None,
1408            callback_id: None,
1409            callback_timeout_at: None,
1410            callback_filter: None,
1411            callback_on_complete: None,
1412            callback_on_fail: None,
1413            callback_transform: None,
1414            progress: payload.progress,
1415        })
1416    }
1417}
1418
1419#[derive(Debug, Clone, sqlx::FromRow)]
1420struct ReadyTransitionRow {
1421    ready_slot: i32,
1422    ready_generation: i64,
1423    job_id: i64,
1424    kind: String,
1425    queue: String,
1426    args: serde_json::Value,
1427    priority: i16,
1428    attempt: i16,
1429    run_lease: i64,
1430    max_attempts: i16,
1431    lane_seq: i64,
1432    enqueue_shard: i16,
1433    run_at: DateTime<Utc>,
1434    attempted_at: Option<DateTime<Utc>>,
1435    created_at: DateTime<Utc>,
1436    unique_key: Option<Vec<u8>>,
1437    unique_states: Option<String>,
1438    payload: serde_json::Value,
1439}
1440
1441#[derive(Debug, Clone)]
1442struct ClaimCursorAdvance {
1443    queue: String,
1444    priority: i16,
1445    enqueue_shard: i16,
1446    next_seq: i64,
1447    only_if_current: Option<i64>,
1448}
1449
1450type ClaimCursorLaneKey = (String, i16, i16);
1451type ConditionalClaimCursorAdvances = BTreeMap<i64, i64>;
1452type GroupedClaimCursorAdvances =
1453    BTreeMap<ClaimCursorLaneKey, (Option<i64>, ConditionalClaimCursorAdvances)>;
1454type TerminalCounterKey = (i32, i64, String, i16, i16, i16);
1455
1456struct CancelJobTxResult {
1457    row: JobRow,
1458    claim_cursor_advance: Option<ClaimCursorAdvance>,
1459}
1460
1461struct ReadyBatchMoveResult {
1462    moved: bool,
1463}
1464
1465impl ReadyTransitionRow {
1466    fn into_existing_ready_row(
1467        self,
1468        queue: String,
1469        priority: i16,
1470        payload: serde_json::Value,
1471    ) -> ExistingReadyRow {
1472        ExistingReadyRow {
1473            job_id: self.job_id,
1474            kind: self.kind,
1475            queue,
1476            args: self.args,
1477            priority,
1478            attempt: self.attempt,
1479            run_lease: self.run_lease,
1480            max_attempts: self.max_attempts,
1481            run_at: self.run_at,
1482            attempted_at: self.attempted_at,
1483            created_at: self.created_at,
1484            unique_key: self.unique_key,
1485            unique_states: self.unique_states,
1486            payload,
1487        }
1488    }
1489
1490    fn into_done_row(
1491        self,
1492        state: JobState,
1493        finalized_at: DateTime<Utc>,
1494        payload: serde_json::Value,
1495    ) -> DoneJobRow {
1496        DoneJobRow {
1497            ready_slot: self.ready_slot,
1498            ready_generation: self.ready_generation,
1499            job_id: self.job_id,
1500            kind: self.kind,
1501            queue: self.queue,
1502            args: self.args,
1503            state,
1504            priority: self.priority,
1505            attempt: self.attempt,
1506            run_lease: self.run_lease,
1507            max_attempts: self.max_attempts,
1508            lane_seq: self.lane_seq,
1509            enqueue_shard: self.enqueue_shard,
1510            run_at: self.run_at,
1511            attempted_at: self.attempted_at,
1512            finalized_at,
1513            created_at: self.created_at,
1514            unique_key: self.unique_key,
1515            unique_states: self.unique_states,
1516            payload,
1517        }
1518    }
1519}
1520
1521#[derive(Debug, Clone, sqlx::FromRow)]
1522struct ReadyJobLeaseRow {
1523    ready_slot: i32,
1524    ready_generation: i64,
1525    lane_seq: i64,
1526    enqueue_shard: i16,
1527    lease_slot: i32,
1528    lease_generation: i64,
1529    claim_slot: i32,
1530    receipt_id: Option<i64>,
1531    claim_batch_id: Option<i64>,
1532    claim_batch_index: Option<i32>,
1533    job_id: i64,
1534    kind: String,
1535    queue: String,
1536    args: serde_json::Value,
1537    lane_priority: i16,
1538    priority: i16,
1539    attempt: i16,
1540    run_lease: i64,
1541    max_attempts: i16,
1542    run_at: DateTime<Utc>,
1543    heartbeat_at: Option<DateTime<Utc>>,
1544    deadline_at: Option<DateTime<Utc>>,
1545    attempted_at: Option<DateTime<Utc>>,
1546    created_at: DateTime<Utc>,
1547    unique_key: Option<Vec<u8>>,
1548    unique_states: Option<String>,
1549    payload: serde_json::Value,
1550}
1551
1552impl ReadyJobLeaseRow {
1553    fn claim_ref(&self, lease_claim_receipt: bool) -> ClaimedEntry {
1554        ClaimedEntry {
1555            queue: self.queue.clone(),
1556            priority: self.lane_priority,
1557            lane_seq: self.lane_seq,
1558            ready_slot: self.ready_slot,
1559            ready_generation: self.ready_generation,
1560            lease_slot: self.lease_slot,
1561            lease_generation: self.lease_generation,
1562            claim_slot: self.claim_slot,
1563            receipt_id: self.receipt_id,
1564            claim_batch_id: self.claim_batch_id,
1565            claim_batch_index: self.claim_batch_index,
1566            lease_claim_receipt,
1567            enqueue_shard: self.enqueue_shard,
1568        }
1569    }
1570
1571    fn into_job_row(self) -> Result<JobRow, AwaError> {
1572        let mut payload = RuntimePayload::from_json(self.payload)?;
1573        if self.priority < self.lane_priority {
1574            let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
1575                AwaError::Validation(
1576                    "queue storage payload metadata must be a JSON object".to_string(),
1577                )
1578            })?;
1579            metadata
1580                .entry("_awa_original_priority".to_string())
1581                .or_insert_with(|| serde_json::Value::from(i64::from(self.lane_priority)));
1582        }
1583
1584        Ok(JobRow {
1585            id: self.job_id,
1586            kind: self.kind,
1587            queue: self.queue,
1588            args: self.args,
1589            state: JobState::Running,
1590            priority: self.priority,
1591            attempt: self.attempt,
1592            run_lease: self.run_lease,
1593            max_attempts: self.max_attempts,
1594            run_at: self.run_at,
1595            heartbeat_at: self.heartbeat_at,
1596            deadline_at: self.deadline_at,
1597            attempted_at: self.attempted_at,
1598            finalized_at: None,
1599            created_at: self.created_at,
1600            errors: payload.errors_option(),
1601            metadata: payload.metadata,
1602            tags: payload.tags,
1603            unique_key: self.unique_key,
1604            unique_states: None,
1605            callback_id: None,
1606            callback_timeout_at: None,
1607            callback_filter: None,
1608            callback_on_complete: None,
1609            callback_on_fail: None,
1610            callback_transform: None,
1611            progress: payload.progress,
1612        })
1613    }
1614
1615    fn into_claimed_runtime_job(
1616        self,
1617        lease_claim_receipt: bool,
1618    ) -> Result<ClaimedRuntimeJob, AwaError> {
1619        let claim = self.claim_ref(lease_claim_receipt);
1620        let unique_states = self.unique_states.clone();
1621        let job = self.into_job_row()?;
1622        Ok(ClaimedRuntimeJob {
1623            claim,
1624            job,
1625            unique_states,
1626        })
1627    }
1628}
1629
1630#[derive(Debug, Clone)]
1631struct RuntimeReadyRow {
1632    kind: String,
1633    queue: String,
1634    args: serde_json::Value,
1635    priority: i16,
1636    attempt: i16,
1637    run_lease: i64,
1638    max_attempts: i16,
1639    run_at: DateTime<Utc>,
1640    attempted_at: Option<DateTime<Utc>>,
1641    created_at: DateTime<Utc>,
1642    unique_key: Option<Vec<u8>>,
1643    unique_states: Option<String>,
1644    payload: serde_json::Value,
1645    /// Optional caller-supplied key. When the destination queue has
1646    /// `enqueue_shards > 1`, all rows sharing the same key route to
1647    /// the same shard so FIFO within the key is preserved. `None`
1648    /// falls back to the per-store rotor.
1649    ordering_key: Option<Vec<u8>>,
1650}
1651
1652#[derive(Debug, Clone)]
1653struct RuntimeReadyInsert {
1654    job_id: i64,
1655    kind: String,
1656    queue: String,
1657    args: serde_json::Value,
1658    priority: i16,
1659    attempt: i16,
1660    run_lease: i64,
1661    max_attempts: i16,
1662    run_at: DateTime<Utc>,
1663    attempted_at: Option<DateTime<Utc>>,
1664    lane_seq: i64,
1665    /// Enqueue shard the row belongs to. Selected per row by
1666    /// `shard_for_enqueue`; defaults to 0 when the queue's
1667    /// `enqueue_shards` is 1 (the default).
1668    enqueue_shard: i16,
1669    created_at: DateTime<Utc>,
1670    unique_key: Option<Vec<u8>>,
1671    unique_states: Option<String>,
1672    payload: serde_json::Value,
1673}
1674
1675#[derive(Debug)]
1676struct ReadySegmentInsert {
1677    queue: String,
1678    priority: i16,
1679    enqueue_shard: i16,
1680    first_lane_seq: i64,
1681    next_lane_seq: i64,
1682    /// `run_at` for every row in this segment. Segment construction splits on
1683    /// `run_at` so claim-time priority aging stays exact after the claim cursor
1684    /// advances inside a multi-row segment.
1685    first_run_at: DateTime<Utc>,
1686}
1687
1688#[cfg(test)]
1689mod ready_segment_tests {
1690    use super::{QueueStorage, RuntimeReadyInsert};
1691    use chrono::{Duration, TimeZone, Utc};
1692
1693    fn ready_row(lane_seq: i64, run_at: chrono::DateTime<Utc>) -> RuntimeReadyInsert {
1694        RuntimeReadyInsert {
1695            job_id: lane_seq,
1696            kind: "segment_test".to_string(),
1697            queue: "segment-q".to_string(),
1698            args: serde_json::json!({}),
1699            priority: 2,
1700            attempt: 0,
1701            run_lease: 0,
1702            max_attempts: 25,
1703            run_at,
1704            attempted_at: None,
1705            lane_seq,
1706            enqueue_shard: 0,
1707            created_at: run_at,
1708            unique_key: None,
1709            unique_states: None,
1710            payload: serde_json::json!({}),
1711        }
1712    }
1713
1714    #[test]
1715    fn ready_segments_split_on_run_at_boundaries() {
1716        let first = Utc
1717            .with_ymd_and_hms(2026, 6, 14, 12, 0, 0)
1718            .single()
1719            .expect("valid test timestamp");
1720        let second = first + Duration::seconds(1);
1721        let rows = vec![
1722            ready_row(1, first),
1723            ready_row(2, first),
1724            ready_row(3, second),
1725            ready_row(4, first),
1726        ];
1727
1728        let segments = QueueStorage::ready_segments_from_rows(&rows);
1729        let ranges: Vec<_> = segments
1730            .iter()
1731            .map(|segment| {
1732                (
1733                    segment.first_lane_seq,
1734                    segment.next_lane_seq,
1735                    segment.first_run_at,
1736                )
1737            })
1738            .collect();
1739
1740        assert_eq!(ranges, vec![(1, 3, first), (3, 4, second), (4, 5, first)]);
1741    }
1742}
1743
1744#[derive(Debug, Clone, sqlx::FromRow)]
1745struct DoneJobRow {
1746    ready_slot: i32,
1747    ready_generation: i64,
1748    job_id: i64,
1749    kind: String,
1750    queue: String,
1751    args: serde_json::Value,
1752    state: JobState,
1753    priority: i16,
1754    attempt: i16,
1755    run_lease: i64,
1756    max_attempts: i16,
1757    lane_seq: i64,
1758    /// Enqueue shard the row was claimed from. Part of the
1759    /// `done_entries` primary key so two shards' terminal rows at
1760    /// the same `(ready_slot, queue, priority, lane_seq)` do not
1761    /// collide.
1762    enqueue_shard: i16,
1763    run_at: DateTime<Utc>,
1764    attempted_at: Option<DateTime<Utc>>,
1765    finalized_at: DateTime<Utc>,
1766    created_at: DateTime<Utc>,
1767    unique_key: Option<Vec<u8>>,
1768    unique_states: Option<String>,
1769    payload: serde_json::Value,
1770}
1771
1772impl DoneJobRow {
1773    fn into_job_row(self) -> Result<JobRow, AwaError> {
1774        let payload = RuntimePayload::from_json(self.payload)?;
1775        Ok(JobRow {
1776            id: self.job_id,
1777            kind: self.kind,
1778            queue: self.queue,
1779            args: self.args,
1780            state: self.state,
1781            priority: self.priority,
1782            attempt: self.attempt,
1783            run_lease: self.run_lease,
1784            max_attempts: self.max_attempts,
1785            run_at: self.run_at,
1786            heartbeat_at: None,
1787            deadline_at: None,
1788            attempted_at: self.attempted_at,
1789            finalized_at: Some(self.finalized_at),
1790            created_at: self.created_at,
1791            errors: payload.errors_option(),
1792            metadata: payload.metadata,
1793            tags: payload.tags,
1794            unique_key: self.unique_key,
1795            unique_states: None,
1796            callback_id: None,
1797            callback_timeout_at: None,
1798            callback_filter: None,
1799            callback_on_complete: None,
1800            callback_on_fail: None,
1801            callback_transform: None,
1802            progress: payload.progress,
1803        })
1804    }
1805
1806    fn into_dlq_row(self, dlq_reason: String, dlq_at: DateTime<Utc>) -> DlqJobRow {
1807        DlqJobRow {
1808            job_id: self.job_id,
1809            kind: self.kind,
1810            queue: self.queue,
1811            args: self.args,
1812            state: self.state,
1813            priority: self.priority,
1814            attempt: self.attempt,
1815            run_lease: self.run_lease,
1816            max_attempts: self.max_attempts,
1817            run_at: self.run_at,
1818            attempted_at: self.attempted_at,
1819            finalized_at: self.finalized_at,
1820            created_at: self.created_at,
1821            unique_key: self.unique_key,
1822            unique_states: self.unique_states,
1823            payload: self.payload,
1824            dlq_reason,
1825            dlq_at,
1826            original_run_lease: self.run_lease,
1827        }
1828    }
1829}
1830
1831#[derive(Debug, Clone, sqlx::FromRow)]
1832struct DlqJobRow {
1833    job_id: i64,
1834    kind: String,
1835    queue: String,
1836    args: serde_json::Value,
1837    state: JobState,
1838    priority: i16,
1839    attempt: i16,
1840    run_lease: i64,
1841    max_attempts: i16,
1842    run_at: DateTime<Utc>,
1843    attempted_at: Option<DateTime<Utc>>,
1844    finalized_at: DateTime<Utc>,
1845    created_at: DateTime<Utc>,
1846    unique_key: Option<Vec<u8>>,
1847    unique_states: Option<String>,
1848    payload: serde_json::Value,
1849    dlq_reason: String,
1850    dlq_at: DateTime<Utc>,
1851    original_run_lease: i64,
1852}
1853
1854impl DlqJobRow {
1855    fn into_job_row(self) -> Result<JobRow, AwaError> {
1856        let payload = RuntimePayload::from_json(self.payload)?;
1857        Ok(JobRow {
1858            id: self.job_id,
1859            kind: self.kind,
1860            queue: self.queue,
1861            args: self.args,
1862            state: self.state,
1863            priority: self.priority,
1864            attempt: self.attempt,
1865            run_lease: self.run_lease,
1866            max_attempts: self.max_attempts,
1867            run_at: self.run_at,
1868            heartbeat_at: None,
1869            deadline_at: None,
1870            attempted_at: self.attempted_at,
1871            finalized_at: Some(self.finalized_at),
1872            created_at: self.created_at,
1873            errors: payload.errors_option(),
1874            metadata: payload.metadata,
1875            tags: payload.tags,
1876            unique_key: self.unique_key,
1877            unique_states: None,
1878            callback_id: None,
1879            callback_timeout_at: None,
1880            callback_filter: None,
1881            callback_on_complete: None,
1882            callback_on_fail: None,
1883            callback_transform: None,
1884            progress: payload.progress,
1885        })
1886    }
1887
1888    fn into_retry_ready_row(
1889        self,
1890        queue: String,
1891        priority: i16,
1892        run_at: DateTime<Utc>,
1893        payload: serde_json::Value,
1894    ) -> ExistingReadyRow {
1895        ExistingReadyRow {
1896            job_id: self.job_id,
1897            kind: self.kind,
1898            queue,
1899            args: self.args,
1900            priority,
1901            attempt: 0,
1902            run_lease: self.run_lease,
1903            max_attempts: self.max_attempts,
1904            run_at,
1905            attempted_at: None,
1906            created_at: self.created_at,
1907            unique_key: self.unique_key,
1908            unique_states: self.unique_states,
1909            payload,
1910        }
1911    }
1912
1913    fn into_retry_deferred_row(
1914        self,
1915        queue: String,
1916        priority: i16,
1917        run_at: DateTime<Utc>,
1918        payload: serde_json::Value,
1919    ) -> DeferredJobRow {
1920        DeferredJobRow {
1921            job_id: self.job_id,
1922            kind: self.kind,
1923            queue,
1924            args: self.args,
1925            state: JobState::Scheduled,
1926            priority,
1927            attempt: 0,
1928            run_lease: self.run_lease,
1929            max_attempts: self.max_attempts,
1930            run_at,
1931            attempted_at: None,
1932            finalized_at: None,
1933            created_at: self.created_at,
1934            unique_key: self.unique_key,
1935            unique_states: self.unique_states,
1936            payload,
1937        }
1938    }
1939}
1940
1941#[derive(Debug, Clone)]
1942struct ExistingReadyRow {
1943    job_id: i64,
1944    kind: String,
1945    queue: String,
1946    args: serde_json::Value,
1947    priority: i16,
1948    attempt: i16,
1949    run_lease: i64,
1950    max_attempts: i16,
1951    run_at: DateTime<Utc>,
1952    attempted_at: Option<DateTime<Utc>>,
1953    created_at: DateTime<Utc>,
1954    unique_key: Option<Vec<u8>>,
1955    unique_states: Option<String>,
1956    payload: serde_json::Value,
1957}
1958
1959#[derive(Debug, Clone, sqlx::FromRow)]
1960struct DeletedLeaseRow {
1961    ready_slot: i32,
1962    ready_generation: i64,
1963    job_id: i64,
1964    queue: String,
1965    state: JobState,
1966    priority: i16,
1967    attempt: i16,
1968    run_lease: i64,
1969    max_attempts: i16,
1970    lane_seq: i64,
1971    enqueue_shard: i16,
1972    attempted_at: Option<DateTime<Utc>>,
1973}
1974
1975#[derive(Debug, Clone, sqlx::FromRow)]
1976struct ReadySnapshotRow {
1977    ready_slot: i32,
1978    ready_generation: i64,
1979    job_id: i64,
1980    kind: String,
1981    queue: String,
1982    args: serde_json::Value,
1983    lane_seq: i64,
1984    enqueue_shard: i16,
1985    run_at: DateTime<Utc>,
1986    created_at: DateTime<Utc>,
1987    unique_key: Option<Vec<u8>>,
1988    unique_states: Option<String>,
1989    payload: serde_json::Value,
1990}
1991
1992#[derive(Debug, Clone, sqlx::FromRow)]
1993struct AttemptStateRow {
1994    job_id: i64,
1995    run_lease: i64,
1996    progress: Option<serde_json::Value>,
1997    callback_filter: Option<String>,
1998    callback_on_complete: Option<String>,
1999    callback_on_fail: Option<String>,
2000    callback_transform: Option<String>,
2001    callback_result: Option<serde_json::Value>,
2002}
2003
2004#[derive(Debug, Clone)]
2005struct LeaseTransitionRow {
2006    ready_slot: i32,
2007    ready_generation: i64,
2008    job_id: i64,
2009    kind: String,
2010    queue: String,
2011    args: serde_json::Value,
2012    state: JobState,
2013    priority: i16,
2014    attempt: i16,
2015    run_lease: i64,
2016    max_attempts: i16,
2017    lane_seq: i64,
2018    enqueue_shard: i16,
2019    run_at: DateTime<Utc>,
2020    attempted_at: Option<DateTime<Utc>>,
2021    created_at: DateTime<Utc>,
2022    unique_key: Option<Vec<u8>>,
2023    unique_states: Option<String>,
2024    payload: serde_json::Value,
2025    progress: Option<serde_json::Value>,
2026}
2027
2028impl LeaseTransitionRow {
2029    fn into_done_row(
2030        self,
2031        state: JobState,
2032        finalized_at: DateTime<Utc>,
2033        payload: serde_json::Value,
2034    ) -> DoneJobRow {
2035        DoneJobRow {
2036            ready_slot: self.ready_slot,
2037            ready_generation: self.ready_generation,
2038            job_id: self.job_id,
2039            kind: self.kind,
2040            queue: self.queue,
2041            args: self.args,
2042            state,
2043            priority: self.priority,
2044            attempt: self.attempt,
2045            run_lease: self.run_lease,
2046            max_attempts: self.max_attempts,
2047            lane_seq: self.lane_seq,
2048            enqueue_shard: self.enqueue_shard,
2049            run_at: self.run_at,
2050            attempted_at: self.attempted_at,
2051            finalized_at,
2052            created_at: self.created_at,
2053            unique_key: self.unique_key,
2054            unique_states: self.unique_states,
2055            payload,
2056        }
2057    }
2058
2059    fn into_deferred_row(
2060        self,
2061        state: JobState,
2062        run_at: DateTime<Utc>,
2063        finalized_at: Option<DateTime<Utc>>,
2064        payload: serde_json::Value,
2065    ) -> DeferredJobRow {
2066        DeferredJobRow {
2067            job_id: self.job_id,
2068            kind: self.kind,
2069            queue: self.queue,
2070            args: self.args,
2071            state,
2072            priority: self.priority,
2073            attempt: self.attempt,
2074            run_lease: self.run_lease,
2075            max_attempts: self.max_attempts,
2076            run_at,
2077            attempted_at: self.attempted_at,
2078            finalized_at,
2079            created_at: self.created_at,
2080            unique_key: self.unique_key,
2081            unique_states: self.unique_states,
2082            payload,
2083        }
2084    }
2085
2086    fn into_ready_row(self, run_at: DateTime<Utc>, payload: serde_json::Value) -> ExistingReadyRow {
2087        ExistingReadyRow {
2088            job_id: self.job_id,
2089            kind: self.kind,
2090            queue: self.queue,
2091            args: self.args,
2092            priority: self.priority,
2093            attempt: self.attempt,
2094            run_lease: self.run_lease,
2095            max_attempts: self.max_attempts,
2096            run_at,
2097            attempted_at: self.attempted_at,
2098            created_at: self.created_at,
2099            unique_key: self.unique_key,
2100            unique_states: self.unique_states,
2101            payload,
2102        }
2103    }
2104
2105    fn into_dlq_row(
2106        self,
2107        finalized_at: DateTime<Utc>,
2108        payload: serde_json::Value,
2109        dlq_reason: String,
2110        dlq_at: DateTime<Utc>,
2111    ) -> DlqJobRow {
2112        DlqJobRow {
2113            job_id: self.job_id,
2114            kind: self.kind,
2115            queue: self.queue,
2116            args: self.args,
2117            state: JobState::Failed,
2118            priority: self.priority,
2119            attempt: self.attempt,
2120            run_lease: self.run_lease,
2121            max_attempts: self.max_attempts,
2122            run_at: self.run_at,
2123            attempted_at: self.attempted_at,
2124            finalized_at,
2125            created_at: self.created_at,
2126            unique_key: self.unique_key,
2127            unique_states: self.unique_states,
2128            payload,
2129            dlq_reason,
2130            dlq_at,
2131            original_run_lease: self.run_lease,
2132        }
2133    }
2134}
2135
2136#[derive(Debug, Clone, sqlx::FromRow)]
2137struct LeaseJobRow {
2138    job_id: i64,
2139    kind: String,
2140    queue: String,
2141    args: serde_json::Value,
2142    state: JobState,
2143    priority: i16,
2144    attempt: i16,
2145    run_lease: i64,
2146    max_attempts: i16,
2147    run_at: DateTime<Utc>,
2148    heartbeat_at: Option<DateTime<Utc>>,
2149    deadline_at: Option<DateTime<Utc>>,
2150    attempted_at: Option<DateTime<Utc>>,
2151    finalized_at: Option<DateTime<Utc>>,
2152    created_at: DateTime<Utc>,
2153    unique_key: Option<Vec<u8>>,
2154    callback_id: Option<Uuid>,
2155    callback_timeout_at: Option<DateTime<Utc>>,
2156    callback_filter: Option<String>,
2157    callback_on_complete: Option<String>,
2158    callback_on_fail: Option<String>,
2159    callback_transform: Option<String>,
2160    payload: serde_json::Value,
2161    progress: Option<serde_json::Value>,
2162    callback_result: Option<serde_json::Value>,
2163}
2164
2165impl LeaseJobRow {
2166    fn into_job_row(self) -> Result<JobRow, AwaError> {
2167        let payload = QueueStorage::materialize_runtime_payload(
2168            self.payload,
2169            self.progress,
2170            self.callback_result,
2171        )?;
2172        Ok(JobRow {
2173            id: self.job_id,
2174            kind: self.kind,
2175            queue: self.queue,
2176            args: self.args,
2177            state: self.state,
2178            priority: self.priority,
2179            attempt: self.attempt,
2180            run_lease: self.run_lease,
2181            max_attempts: self.max_attempts,
2182            run_at: self.run_at,
2183            heartbeat_at: self.heartbeat_at,
2184            deadline_at: self.deadline_at,
2185            attempted_at: self.attempted_at,
2186            finalized_at: self.finalized_at,
2187            created_at: self.created_at,
2188            errors: payload.errors_option(),
2189            metadata: payload.metadata,
2190            tags: payload.tags,
2191            unique_key: self.unique_key,
2192            unique_states: None,
2193            callback_id: self.callback_id,
2194            callback_timeout_at: self.callback_timeout_at,
2195            callback_filter: self.callback_filter,
2196            callback_on_complete: self.callback_on_complete,
2197            callback_on_fail: self.callback_on_fail,
2198            callback_transform: self.callback_transform,
2199            progress: payload.progress,
2200        })
2201    }
2202}
2203
2204#[derive(Debug, Clone, sqlx::FromRow)]
2205struct DeferredJobRow {
2206    job_id: i64,
2207    kind: String,
2208    queue: String,
2209    args: serde_json::Value,
2210    state: JobState,
2211    priority: i16,
2212    attempt: i16,
2213    run_lease: i64,
2214    max_attempts: i16,
2215    run_at: DateTime<Utc>,
2216    attempted_at: Option<DateTime<Utc>>,
2217    finalized_at: Option<DateTime<Utc>>,
2218    created_at: DateTime<Utc>,
2219    unique_key: Option<Vec<u8>>,
2220    unique_states: Option<String>,
2221    payload: serde_json::Value,
2222}
2223
2224impl DeferredJobRow {
2225    fn into_job_row(self) -> Result<JobRow, AwaError> {
2226        let payload = RuntimePayload::from_json(self.payload)?;
2227        Ok(JobRow {
2228            id: self.job_id,
2229            kind: self.kind,
2230            queue: self.queue,
2231            args: self.args,
2232            state: self.state,
2233            priority: self.priority,
2234            attempt: self.attempt,
2235            run_lease: self.run_lease,
2236            max_attempts: self.max_attempts,
2237            run_at: self.run_at,
2238            heartbeat_at: None,
2239            deadline_at: None,
2240            attempted_at: self.attempted_at,
2241            finalized_at: self.finalized_at,
2242            created_at: self.created_at,
2243            errors: payload.errors_option(),
2244            metadata: payload.metadata,
2245            tags: payload.tags,
2246            unique_key: self.unique_key,
2247            unique_states: None,
2248            callback_id: None,
2249            callback_timeout_at: None,
2250            callback_filter: None,
2251            callback_on_complete: None,
2252            callback_on_fail: None,
2253            callback_transform: None,
2254            progress: payload.progress,
2255        })
2256    }
2257}
2258
2259/// Segmented queue storage backend.
2260///
2261/// Design goals:
2262/// - append-only queue segments in a rotated ring
2263/// - append-only completion segments keyed back to the queue segment
2264/// - a separate, faster rotating lease ring so delete churn is bounded by the
2265///   lease cycle rather than by queue retention
2266/// - hot mutable state restricted to queue cursors, narrow leases, and
2267///   optional per-attempt runtime state only when needed
2268///
2269#[derive(Debug)]
2270pub struct QueueStorage {
2271    config: QueueStorageConfig,
2272    next_stripe_probe: AtomicUsize,
2273    /// Per-store rotor that selects the enqueue shard for the next batch.
2274    /// Rotated once per `shard_for_enqueue` call so producers spread their
2275    /// writes across the per-`(queue, priority)` shard rows.
2276    shard_rotor: AtomicU16,
2277    /// Cache of `awa.queue_meta.enqueue_shards` per queue. Populated lazily
2278    /// on the first `shard_for_enqueue` call for a queue and invalidated by
2279    /// `reset()`. With the default `enqueue_shards = 1`, the cache holds 1
2280    /// and `pick_shard` returns 0 unconditionally.
2281    enqueue_shards_cache: Mutex<HashMap<String, i16>>,
2282    /// Lane-presence cache: `(physical_queue, priority, enqueue_shard)`
2283    /// triples whose three lane rows (queue_lanes, queue_enqueue_heads,
2284    /// queue_claim_heads) we have previously inserted. Skips the three
2285    /// `INSERT … ON CONFLICT DO NOTHING` round-trips on subsequent enqueue
2286    /// batches to a known lane/shard. Cleared on `reset()` because reset
2287    /// TRUNCATEs queue_lanes; `advance_enqueue_head` repairs a stale entry
2288    /// (head row gone after a rolled-back ensure_lane).
2289    ensured_lanes: Mutex<HashSet<(String, i16, i16)>>,
2290    prune_lock_timeout: Duration,
2291}
2292
2293impl QueueStorage {
2294    pub fn new(config: QueueStorageConfig) -> Result<Self, AwaError> {
2295        if config.queue_slot_count < 4 {
2296            return Err(AwaError::Validation(
2297                "queue storage requires at least 4 queue slots".into(),
2298            ));
2299        }
2300        if config.lease_slot_count < 2 {
2301            return Err(AwaError::Validation(
2302                "queue storage requires at least 2 lease slots".into(),
2303            ));
2304        }
2305        if config.claim_slot_count < 2 {
2306            return Err(AwaError::Validation(
2307                "queue storage requires at least 2 claim slots".into(),
2308            ));
2309        }
2310        if config.queue_stripe_count == 0 {
2311            return Err(AwaError::Validation(
2312                "queue storage requires at least 1 queue stripe".into(),
2313            ));
2314        }
2315        if config.schema == DEFAULT_SCHEMA
2316            && (config.queue_slot_count != DEFAULT_QUEUE_SLOT_COUNT
2317                || config.lease_slot_count != DEFAULT_LEASE_SLOT_COUNT
2318                || config.claim_slot_count != DEFAULT_CLAIM_SLOT_COUNT
2319                || !config.lease_claim_receipts)
2320        {
2321            return Err(AwaError::Validation(
2322                "the default `awa` queue-storage schema must use the default slot counts and \
2323                 lease_claim_receipts=true"
2324                    .into(),
2325            ));
2326        }
2327        validate_ident(&config.schema)?;
2328        Ok(Self {
2329            config,
2330            next_stripe_probe: AtomicUsize::new(0),
2331            shard_rotor: AtomicU16::new(0),
2332            enqueue_shards_cache: Mutex::new(HashMap::new()),
2333            ensured_lanes: Mutex::new(HashSet::new()),
2334            prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT,
2335        })
2336    }
2337
2338    #[doc(hidden)]
2339    pub fn with_prune_lock_timeout(mut self, timeout: Duration) -> Result<Self, AwaError> {
2340        if timeout.is_zero() {
2341            return Err(AwaError::Validation(
2342                "queue storage prune lock timeout must be greater than zero".into(),
2343            ));
2344        }
2345        self.prune_lock_timeout = timeout;
2346        Ok(self)
2347    }
2348
2349    pub fn from_existing_schema(schema: impl Into<String>) -> Result<Self, AwaError> {
2350        Self::new(QueueStorageConfig {
2351            schema: schema.into(),
2352            ..Default::default()
2353        })
2354    }
2355
2356    pub fn schema(&self) -> &str {
2357        &self.config.schema
2358    }
2359
2360    pub fn slot_count(&self) -> usize {
2361        self.queue_slot_count()
2362    }
2363
2364    pub fn queue_slot_count(&self) -> usize {
2365        self.config.queue_slot_count
2366    }
2367
2368    pub fn lease_slot_count(&self) -> usize {
2369        self.config.lease_slot_count
2370    }
2371
2372    pub fn claim_slot_count(&self) -> usize {
2373        self.config.claim_slot_count
2374    }
2375
2376    pub fn queue_stripe_count(&self) -> usize {
2377        self.config.queue_stripe_count
2378    }
2379
2380    pub fn lease_claim_receipts(&self) -> bool {
2381        self.config.lease_claim_receipts
2382    }
2383
2384    fn uses_queue_striping(&self) -> bool {
2385        self.queue_stripe_count() > 1
2386    }
2387
2388    fn is_physical_stripe_queue(&self, queue: &str) -> bool {
2389        self.uses_queue_striping()
2390            && queue
2391                .rsplit_once(QUEUE_STRIPE_DELIMITER)
2392                .is_some_and(|(_, suffix)| suffix.parse::<usize>().is_ok())
2393    }
2394
2395    fn physical_queue_for_stripe(&self, queue: &str, stripe: usize) -> String {
2396        format!("{queue}{QUEUE_STRIPE_DELIMITER}{stripe}")
2397    }
2398
2399    fn physical_queues_for_logical(&self, queue: &str) -> Vec<String> {
2400        if !self.uses_queue_striping() || self.is_physical_stripe_queue(queue) {
2401            return vec![queue.to_string()];
2402        }
2403        (0..self.queue_stripe_count())
2404            .map(|stripe| self.physical_queue_for_stripe(queue, stripe))
2405            .collect()
2406    }
2407
2408    fn stripe_probe_start(&self, stripe_count: usize) -> usize {
2409        if stripe_count <= 1 {
2410            return 0;
2411        }
2412        self.next_stripe_probe.fetch_add(1, Ordering::Relaxed) % stripe_count
2413    }
2414
2415    fn logical_queue_name<'a>(&self, queue: &'a str) -> &'a str {
2416        if !self.uses_queue_striping() {
2417            return queue;
2418        }
2419        queue
2420            .rsplit_once(QUEUE_STRIPE_DELIMITER)
2421            .and_then(|(prefix, suffix)| suffix.parse::<usize>().ok().map(|_| prefix))
2422            .unwrap_or(queue)
2423    }
2424
2425    fn queue_stripe_for_enqueue(
2426        &self,
2427        queue: &str,
2428        unique_key: &Option<Vec<u8>>,
2429        salt: i64,
2430    ) -> String {
2431        if !self.uses_queue_striping() || self.is_physical_stripe_queue(queue) {
2432            return queue.to_string();
2433        }
2434
2435        let stripe = if let Some(key) = unique_key {
2436            let mut hasher = DefaultHasher::new();
2437            key.hash(&mut hasher);
2438            (hasher.finish() as usize) % self.queue_stripe_count()
2439        } else {
2440            salt.rem_euclid(self.queue_stripe_count() as i64) as usize
2441        };
2442        self.physical_queue_for_stripe(queue, stripe)
2443    }
2444
2445    fn use_lease_claim_receipts_for_runtime(&self, _deadline_duration: Duration) -> bool {
2446        // Receipts mode now supports per-claim deadlines via
2447        // `lease_claims.deadline_at` (rescued by
2448        // `rescue_expired_receipt_deadlines_tx`), so receipts is the
2449        // live path whenever the engine is configured for receipts —
2450        // the queue's `deadline_duration` no longer disqualifies it.
2451        self.lease_claim_receipts()
2452    }
2453
2454    pub fn ready_child_relname(&self, slot: usize) -> String {
2455        format!("ready_entries_{slot}")
2456    }
2457
2458    pub fn done_child_relname(&self, slot: usize) -> String {
2459        format!("done_entries_{slot}")
2460    }
2461
2462    pub fn leases_relname(&self) -> &'static str {
2463        "leases"
2464    }
2465
2466    pub fn lease_claims_relname(&self) -> &'static str {
2467        "lease_claims"
2468    }
2469
2470    pub fn lease_claim_closures_relname(&self) -> &'static str {
2471        "lease_claim_closures"
2472    }
2473
2474    pub fn leases_child_relname(&self, slot: usize) -> String {
2475        format!("leases_{slot}")
2476    }
2477
2478    pub fn attempt_state_relname(&self) -> &'static str {
2479        "attempt_state"
2480    }
2481
2482    pub async fn active_schema(pool: &PgPool) -> Result<Option<String>, AwaError> {
2483        sqlx::query_scalar(
2484            "SELECT schema_name FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'",
2485        )
2486        .fetch_optional(pool)
2487        .await
2488        .map_err(map_sqlx_error)
2489    }
2490
2491    /// Transaction-aware variant of [`Self::active_schema`] — read the
2492    /// active queue-storage schema name inside the caller's transaction
2493    /// rather than acquiring a separate pool connection.
2494    pub async fn active_schema_in_tx(
2495        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2496    ) -> Result<Option<String>, AwaError> {
2497        sqlx::query_scalar(
2498            "SELECT schema_name FROM awa.runtime_storage_backends WHERE backend = 'queue_storage'",
2499        )
2500        .fetch_optional(tx.as_mut())
2501        .await
2502        .map_err(map_sqlx_error)
2503    }
2504
2505    fn materialize_runtime_payload(
2506        payload: serde_json::Value,
2507        progress: Option<serde_json::Value>,
2508        callback_result: Option<serde_json::Value>,
2509    ) -> Result<RuntimePayload, AwaError> {
2510        let mut payload = RuntimePayload::from_json(payload)?;
2511        if let Some(progress) = progress {
2512            payload.set_progress(Some(progress));
2513        }
2514        if let Some(callback_result) = callback_result {
2515            payload.insert_callback_result(Some(callback_result));
2516        }
2517        Ok(payload)
2518    }
2519
2520    fn payload_with_attempt_state(
2521        payload: serde_json::Value,
2522        progress: Option<serde_json::Value>,
2523    ) -> Result<serde_json::Value, AwaError> {
2524        let mut payload = RuntimePayload::from_json(payload)?;
2525        if let Some(progress) = progress {
2526            payload.set_progress(Some(progress));
2527        }
2528        Ok(payload.into_json())
2529    }
2530
2531    fn payload_from_parts(
2532        metadata: serde_json::Value,
2533        tags: Vec<String>,
2534        errors: Option<Vec<serde_json::Value>>,
2535        progress: Option<serde_json::Value>,
2536    ) -> Result<serde_json::Value, AwaError> {
2537        Ok(RuntimePayload {
2538            metadata,
2539            tags,
2540            errors: errors.unwrap_or_default(),
2541            progress,
2542        }
2543        .into_json())
2544    }
2545
2546    async fn sync_unique_claim<'a>(
2547        &self,
2548        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2549        job_id: i64,
2550        unique_key: &Option<Vec<u8>>,
2551        unique_states: Option<&str>,
2552        old_state: Option<JobState>,
2553        new_state: Option<JobState>,
2554    ) -> Result<(), AwaError> {
2555        let old_claim = old_state.is_some_and(|state| unique_state_claims(unique_states, state));
2556        let new_claim = new_state.is_some_and(|state| unique_state_claims(unique_states, state));
2557
2558        if old_claim && !new_claim {
2559            if let Some(key) = unique_key {
2560                sqlx::query(
2561                    "DELETE FROM awa.job_unique_claims WHERE unique_key = $1 AND job_id = $2",
2562                )
2563                .bind(key)
2564                .bind(job_id)
2565                .execute(tx.as_mut())
2566                .await
2567                .map_err(map_sqlx_error)?;
2568            }
2569        }
2570
2571        if new_claim && !old_claim {
2572            if let Some(key) = unique_key {
2573                let result = sqlx::query(
2574                    r#"
2575                    INSERT INTO awa.job_unique_claims (unique_key, job_id)
2576                    VALUES ($1, $2)
2577                    ON CONFLICT (unique_key)
2578                    DO UPDATE SET job_id = EXCLUDED.job_id
2579                    WHERE awa.job_unique_claims.job_id = EXCLUDED.job_id
2580                    "#,
2581                )
2582                .bind(key)
2583                .bind(job_id)
2584                .execute(tx.as_mut())
2585                .await
2586                .map_err(map_sqlx_error)?;
2587
2588                if result.rows_affected() == 0 {
2589                    return Err(AwaError::UniqueConflict {
2590                        constraint: Some("idx_awa_jobs_unique".to_string()),
2591                    });
2592                }
2593            }
2594        }
2595
2596        Ok(())
2597    }
2598
2599    async fn sync_enqueue_unique_claims<'a>(
2600        &self,
2601        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2602        claims: Vec<(Vec<u8>, i64)>,
2603    ) -> Result<(), AwaError> {
2604        if claims.is_empty() {
2605            return Ok(());
2606        }
2607
2608        let mut seen: HashSet<&[u8]> = HashSet::with_capacity(claims.len());
2609        for (key, _) in &claims {
2610            if !seen.insert(key.as_slice()) {
2611                return Err(AwaError::UniqueConflict {
2612                    constraint: Some("idx_awa_jobs_unique".to_string()),
2613                });
2614            }
2615        }
2616
2617        let (keys, job_ids): (Vec<Vec<u8>>, Vec<i64>) = claims.into_iter().unzip();
2618        let (requested, applied): (i64, i64) = sqlx::query_as(
2619            r#"
2620            WITH input(unique_key, job_id) AS (
2621                SELECT * FROM unnest($1::bytea[], $2::bigint[])
2622            ),
2623            inserted AS (
2624                INSERT INTO awa.job_unique_claims (unique_key, job_id)
2625                SELECT unique_key, job_id FROM input
2626                ON CONFLICT (unique_key)
2627                DO UPDATE SET job_id = EXCLUDED.job_id
2628                WHERE awa.job_unique_claims.job_id = EXCLUDED.job_id
2629                RETURNING unique_key
2630            )
2631            SELECT
2632                (SELECT count(*)::bigint FROM input) AS requested,
2633                (SELECT count(*)::bigint FROM inserted) AS applied
2634            "#,
2635        )
2636        .bind(keys)
2637        .bind(job_ids)
2638        .fetch_one(tx.as_mut())
2639        .await
2640        .map_err(map_sqlx_error)?;
2641
2642        if applied != requested {
2643            return Err(AwaError::UniqueConflict {
2644                constraint: Some("idx_awa_jobs_unique".to_string()),
2645            });
2646        }
2647
2648        Ok(())
2649    }
2650
2651    // Enqueue inserts have no prior storage state, so uniqueness only needs to
2652    // add claims for states included in the row's unique-state bitmask. State
2653    // transitions still use `sync_unique_claim`, which can release old claims.
2654    async fn sync_ready_enqueue_unique_claims<'a>(
2655        &self,
2656        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2657        rows: &[RuntimeReadyInsert],
2658    ) -> Result<(), AwaError> {
2659        let claims = rows
2660            .iter()
2661            .filter(|row| unique_state_claims(row.unique_states.as_deref(), JobState::Available))
2662            .filter_map(|row| row.unique_key.as_ref().map(|key| (key.clone(), row.job_id)))
2663            .collect();
2664        self.sync_enqueue_unique_claims(tx, claims).await
2665    }
2666
2667    async fn sync_deferred_enqueue_unique_claims<'a>(
2668        &self,
2669        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
2670        rows: &[DeferredJobRow],
2671    ) -> Result<(), AwaError> {
2672        let claims = rows
2673            .iter()
2674            .filter(|row| unique_state_claims(row.unique_states.as_deref(), row.state))
2675            .filter_map(|row| row.unique_key.as_ref().map(|key| (key.clone(), row.job_id)))
2676            .collect();
2677        self.sync_enqueue_unique_claims(tx, claims).await
2678    }
2679
2680    /// Idempotently install the queue-storage substrate (schema, tables,
2681    /// indexes, functions) for this configuration.
2682    ///
2683    /// Runs entirely inside one transaction on a single pooled connection,
2684    /// guarded by `pg_advisory_xact_lock` so concurrent worker startups
2685    /// serialize cleanly rather than fighting over pool slots. The lock
2686    /// auto-releases on COMMIT/ROLLBACK.
2687    ///
2688    /// **Note on index builds:** the `CREATE INDEX IF NOT EXISTS` calls
2689    /// below are not `CONCURRENTLY` (Postgres bans `CONCURRENTLY` inside a
2690    /// transaction). On a fresh install that's a no-op cost. On a redeploy
2691    /// against a database that already has rows in the queue tables, each
2692    /// fresh-index build takes `ShareUpdateExclusiveLock` on its table for
2693    /// the duration of the build, and concurrent writers to that table
2694    /// queue up. For typical operator scenarios this is bounded by the
2695    /// `IF NOT EXISTS` guard — only the *new* indexes added by a runtime
2696    /// upgrade get built; the rest are no-ops.
2697    #[tracing::instrument(skip(self, pool), name = "queue_storage.prepare_schema")]
2698    pub async fn prepare_schema(&self, pool: &PgPool) -> Result<(), AwaError> {
2699        let schema = self.schema();
2700        let install_lock_name = format!("awa.queue_storage.install:{schema}");
2701        let mut install_tx = pool.begin().await.map_err(map_sqlx_error)?;
2702
2703        sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
2704            .bind(&install_lock_name)
2705            .execute(install_tx.as_mut())
2706            .await
2707            .map_err(map_sqlx_error)?;
2708
2709        let install_result = async {
2710            // Helper-owned DDL. The helper takes the same per-schema advisory
2711            // xact lock we just acquired (re-entrant on the same session) and
2712            // performs the entire forward-only substrate install: sequences,
2713            // ring-state singletons, partitioned ready/done/lease tables,
2714            // lane indexes, claim_ready_runtime() function, and seed rows.
2715            // Validation inside the helper rejects non-default configuration
2716            // against the default `awa` schema (see migration v023 and #308).
2717            //
2718            // What stays here (constraint 7 of #308 PR 1):
2719            //   * `open_receipt_claims` non-empty rejection + drop (ADR-023).
2720            //   * Legacy rename of `lease_claims` / `lease_claim_closures`
2721            //     before the helper creates the partitioned parents, plus the
2722            //     post-helper copy + drop of the renamed-aside rows.
2723            //   * `queue_count_snapshots` legacy drop.
2724            //
2725            // `awa.runtime_storage_backends` is owned by migrations (v012);
2726            // activation paths only seed/update its row. The helper does NOT
2727            // touch it and does NOT change storage-transition state, so a call
2728            // to `prepare_schema` remains activation-neutral.
2729
2730            sqlx::query(&format!("CREATE SCHEMA IF NOT EXISTS {schema}"))
2731                .execute(install_tx.as_mut())
2732                .await
2733                .map_err(map_sqlx_error)?;
2734
2735            // The hot path reads "currently open" by anti-joining the
2736            // partitioned `lease_claims` / `lease_claim_closures` pair, so
2737            // `open_receipt_claims` is unused (see ADR-023). Drop it on every
2738            // install. Refuse to drop a non-empty table — non-empty here
2739            // means an operator rolled forward from an older build that
2740            // still wrote rows we don't want to silently delete.
2741            let open_receipt_claims_exists: bool = sqlx::query_scalar(
2742                r#"
2743                SELECT EXISTS (
2744                    SELECT 1 FROM pg_class c
2745                    JOIN pg_namespace n ON n.oid = c.relnamespace
2746                    WHERE n.nspname = $1 AND c.relname = 'open_receipt_claims'
2747                )
2748                "#,
2749            )
2750            .bind(schema)
2751            .fetch_one(install_tx.as_mut())
2752            .await
2753            .map_err(map_sqlx_error)?;
2754            if open_receipt_claims_exists {
2755                let row_count: i64 = sqlx::query_scalar(&format!(
2756                    "SELECT count(*)::bigint FROM {schema}.open_receipt_claims"
2757                ))
2758                .fetch_one(install_tx.as_mut())
2759                .await
2760                .map_err(map_sqlx_error)?;
2761                if row_count > 0 {
2762                    return Err(AwaError::Validation(format!(
2763                        "{schema}.open_receipt_claims has {row_count} rows but the runtime no \
2764                         longer reads or writes this table. Run the ADR-023 reverse migration \
2765                         (recreate from lease_claims minus durable closure evidence) to drain it, \
2766                         then re-run prepare_schema."
2767                    )));
2768                }
2769                sqlx::query(&format!(
2770                    "DROP TABLE IF EXISTS {schema}.open_receipt_claims CASCADE"
2771                ))
2772                .execute(install_tx.as_mut())
2773                .await
2774                .map_err(map_sqlx_error)?;
2775            }
2776
2777            // Detect the current shape of the legacy claim/closure parents.
2778            // The partitioned parents have relkind 'p'; regular tables have
2779            // 'r' and need to be renamed aside so the helper can create the
2780            // partitioned parents under the canonical name. Copying the data
2781            // back happens after the helper returns, all inside this single
2782            // transaction so a crash leaves the schema in one of exactly two
2783            // states (pre-migration or post-migration).
2784            let lease_claims_relkind: Option<String> = sqlx::query_scalar(
2785                r#"
2786                SELECT c.relkind::text
2787                FROM pg_class c
2788                JOIN pg_namespace n ON n.oid = c.relnamespace
2789                WHERE n.nspname = $1 AND c.relname = 'lease_claims'
2790                "#,
2791            )
2792            .bind(schema)
2793            .fetch_optional(install_tx.as_mut())
2794            .await
2795            .map_err(map_sqlx_error)?;
2796
2797            let closures_relkind: Option<String> = sqlx::query_scalar(
2798                r#"
2799                SELECT c.relkind::text
2800                FROM pg_class c
2801                JOIN pg_namespace n ON n.oid = c.relnamespace
2802                WHERE n.nspname = $1 AND c.relname = 'lease_claim_closures'
2803                "#,
2804            )
2805            .bind(schema)
2806            .fetch_optional(install_tx.as_mut())
2807            .await
2808            .map_err(map_sqlx_error)?;
2809
2810            if lease_claims_relkind.as_deref() == Some("r") {
2811                sqlx::query(&format!(
2812                    "ALTER TABLE {schema}.lease_claims RENAME TO lease_claims_legacy"
2813                ))
2814                .execute(install_tx.as_mut())
2815                .await
2816                .map_err(map_sqlx_error)?;
2817            }
2818            if closures_relkind.as_deref() == Some("r") {
2819                sqlx::query(&format!(
2820                    "ALTER TABLE {schema}.lease_claim_closures RENAME TO lease_claim_closures_legacy"
2821                ))
2822                .execute(install_tx.as_mut())
2823                .await
2824                .map_err(map_sqlx_error)?;
2825            }
2826
2827            // queue_count_snapshots was a staleness-cached counterpart of
2828            // queue_counts_exact. The dispatcher now derives the available
2829            // count directly from the head tables and nothing else needs the
2830            // snapshot. Drop it on every prepare_schema so an upgrade from an
2831            // older install reclaims the storage. Done before the helper
2832            // runs because the helper does not touch this legacy table.
2833            sqlx::query(&format!(
2834                "DROP TABLE IF EXISTS {schema}.queue_count_snapshots"
2835            ))
2836            .execute(install_tx.as_mut())
2837            .await
2838            .map_err(map_sqlx_error)?;
2839
2840            // The single forward-only DDL call. See the migration file
2841            // `awa-model/migrations/v023_install_queue_storage_substrate.sql`
2842            // for the function body. Default values match the constants in
2843            // this file (DEFAULT_QUEUE_SLOT_COUNT etc); the helper validates
2844            // that the default `awa` schema only ever gets default-shaped
2845            // installs.
2846            sqlx::query("SELECT awa.install_queue_storage_substrate($1, $2, $3, $4, $5)")
2847                .bind(schema)
2848                .bind(self.queue_slot_count() as i32)
2849                .bind(self.lease_slot_count() as i32)
2850                .bind(self.claim_slot_count() as i32)
2851                .bind(self.lease_claim_receipts())
2852                .execute(install_tx.as_mut())
2853                .await
2854                .map_err(map_sqlx_error)?;
2855
2856            // #169: apply receipt-plane fillfactor tunings. The install
2857            // helper above creates structure only; tunings live in
2858            // their own SQL function (see v024) and the orchestrator
2859            // here is the canonical place that composes them. This
2860            // means future perf knobs land additively in
2861            // `apply_receipt_plane_fillfactor` (or a sibling helper)
2862            // without touching the bigger install helper, and every
2863            // path that creates substrate — Rust prepare_schema, fresh
2864            // `awa migrate` (via v024's sweep), explicit operator call
2865            // against a custom schema — composes the same two pieces.
2866            sqlx::query("SELECT awa.apply_receipt_plane_fillfactor($1)")
2867                .bind(schema)
2868                .execute(install_tx.as_mut())
2869                .await
2870                .map_err(map_sqlx_error)?;
2871
2872            // Post-helper legacy fixups: copy any renamed-aside rows into the
2873            // newly partitioned parents created by the helper, then drop the
2874            // legacy table. ON CONFLICT DO NOTHING so a re-run after a
2875            // partial copy is idempotent. The outer transaction keeps the
2876            // copy + drop atomic so a crash between them leaves the schema
2877            // in one of exactly two states (pre or post migration); without
2878            // that the next `prepare_schema`'s ON CONFLICT would silently
2879            // mask the inconsistency.
2880            let lease_claims_legacy_exists: bool = sqlx::query_scalar(
2881                r#"
2882                SELECT EXISTS (
2883                    SELECT 1 FROM pg_class c
2884                    JOIN pg_namespace n ON n.oid = c.relnamespace
2885                    WHERE n.nspname = $1 AND c.relname = 'lease_claims_legacy'
2886                )
2887                "#,
2888            )
2889            .bind(schema)
2890            .fetch_one(install_tx.as_mut())
2891            .await
2892            .map_err(map_sqlx_error)?;
2893
2894            let closures_legacy_exists: bool = sqlx::query_scalar(
2895                r#"
2896                SELECT EXISTS (
2897                    SELECT 1 FROM pg_class c
2898                    JOIN pg_namespace n ON n.oid = c.relnamespace
2899                    WHERE n.nspname = $1 AND c.relname = 'lease_claim_closures_legacy'
2900                )
2901                "#,
2902            )
2903            .bind(schema)
2904            .fetch_one(install_tx.as_mut())
2905            .await
2906            .map_err(map_sqlx_error)?;
2907
2908            let legacy_claim_slot: Option<i32> =
2909                if lease_claims_legacy_exists || closures_legacy_exists {
2910                    Some(
2911                        sqlx::query_scalar(&format!(
2912                            "SELECT current_slot FROM {schema}.claim_ring_state WHERE singleton"
2913                        ))
2914                        .fetch_one(install_tx.as_mut())
2915                        .await
2916                        .map_err(map_sqlx_error)?,
2917                    )
2918                } else {
2919                    None
2920                };
2921
2922            if lease_claims_legacy_exists {
2923                sqlx::query(&format!(
2924                    "ALTER TABLE {schema}.lease_claims_legacy ADD COLUMN IF NOT EXISTS enqueue_shard SMALLINT NOT NULL DEFAULT 0"
2925                ))
2926                .execute(install_tx.as_mut())
2927                .await
2928                .map_err(map_sqlx_error)?;
2929                sqlx::query(&format!(
2930                    "ALTER TABLE {schema}.lease_claims_legacy ADD COLUMN IF NOT EXISTS deadline_at TIMESTAMPTZ"
2931                ))
2932                .execute(install_tx.as_mut())
2933                .await
2934                .map_err(map_sqlx_error)?;
2935
2936                sqlx::query(&format!(
2937                    r#"
2938                INSERT INTO {schema}.lease_claims (
2939                    claim_slot, job_id, run_lease, ready_slot, ready_generation,
2940                    queue, priority, attempt, max_attempts, lane_seq,
2941                    enqueue_shard, claimed_at, materialized_at, deadline_at
2942                )
2943                SELECT
2944                    $1,
2945                    job_id, run_lease, ready_slot, ready_generation,
2946                    queue, priority, attempt, max_attempts, lane_seq,
2947                    enqueue_shard, claimed_at, materialized_at, deadline_at
2948                FROM {schema}.lease_claims_legacy
2949                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
2950                "#
2951                ))
2952                .bind(legacy_claim_slot.expect("legacy claim slot should be present"))
2953                .execute(install_tx.as_mut())
2954                .await
2955                .map_err(map_sqlx_error)?;
2956
2957                sqlx::query(&format!(
2958                    "DROP TABLE {schema}.lease_claims_legacy"
2959                ))
2960                .execute(install_tx.as_mut())
2961                .await
2962                .map_err(map_sqlx_error)?;
2963            }
2964
2965            if closures_legacy_exists {
2966                sqlx::query(&format!(
2967                    r#"
2968                INSERT INTO {schema}.lease_claim_closures (
2969                    claim_slot, job_id, run_lease, outcome, closed_at
2970                )
2971                SELECT
2972                    $1,
2973                    job_id, run_lease, outcome, closed_at
2974                FROM {schema}.lease_claim_closures_legacy
2975                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
2976                "#
2977                ))
2978                .bind(legacy_claim_slot.expect("legacy claim slot should be present"))
2979                .execute(install_tx.as_mut())
2980                .await
2981                .map_err(map_sqlx_error)?;
2982
2983                sqlx::query(&format!(
2984                    "DROP TABLE {schema}.lease_claim_closures_legacy"
2985                ))
2986                .execute(install_tx.as_mut())
2987                .await
2988                .map_err(map_sqlx_error)?;
2989            }
2990
2991            Ok(())
2992        }
2993        .await;
2994
2995        match install_result {
2996            Ok(()) => install_tx.commit().await.map_err(map_sqlx_error),
2997            Err(err) => {
2998                let _ = install_tx.rollback().await;
2999                Err(err)
3000            }
3001        }
3002    }
3003
3004    #[tracing::instrument(skip(self, pool), name = "queue_storage.activate_backend")]
3005    pub async fn activate_backend(&self, pool: &PgPool) -> Result<(), AwaError> {
3006        let schema = self.schema();
3007        let details = serde_json::json!({ "schema": schema });
3008
3009        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
3010
3011        // Mark queue storage active only after the full schema, partitions,
3012        // indexes, and helper functions are in place. The explicit install
3013        // helper is used by tests and queue-storage-only setups, so it must
3014        // flip both the migration-owned routing registry row and the storage
3015        // transition state.
3016        sqlx::query(
3017            r#"
3018            INSERT INTO awa.runtime_storage_backends (backend, schema_name, updated_at)
3019            VALUES ('queue_storage', $1, now())
3020            ON CONFLICT (backend)
3021            DO UPDATE SET schema_name = EXCLUDED.schema_name, updated_at = EXCLUDED.updated_at
3022            "#,
3023        )
3024        .bind(schema)
3025        .execute(tx.as_mut())
3026        .await
3027        .map_err(map_sqlx_error)?;
3028
3029        let activation_result = sqlx::query(
3030            r#"
3031            UPDATE awa.storage_transition_state AS sts
3032            SET
3033                current_engine = 'queue_storage',
3034                prepared_engine = NULL,
3035                state = 'active',
3036                transition_epoch = CASE
3037                    WHEN sts.current_engine = 'queue_storage'
3038                     AND sts.prepared_engine IS NULL
3039                     AND sts.state = 'active'
3040                     AND sts.details = $1
3041                    THEN sts.transition_epoch
3042                    ELSE sts.transition_epoch + 1
3043                END,
3044                details = $1,
3045                entered_at = CASE
3046                    WHEN sts.current_engine = 'queue_storage'
3047                     AND sts.prepared_engine IS NULL
3048                     AND sts.state = 'active'
3049                     AND sts.details = $1
3050                    THEN sts.entered_at
3051                    ELSE now()
3052                END,
3053                updated_at = now(),
3054                finalized_at = CASE
3055                    WHEN sts.current_engine = 'queue_storage'
3056                     AND sts.prepared_engine IS NULL
3057                     AND sts.state = 'active'
3058                     AND sts.details = $1
3059                    THEN COALESCE(sts.finalized_at, now())
3060                    ELSE now()
3061                END
3062            WHERE sts.singleton
3063            "#,
3064        )
3065        .bind(details)
3066        .execute(tx.as_mut())
3067        .await
3068        .map_err(map_sqlx_error)?;
3069
3070        if activation_result.rows_affected() != 1 {
3071            return Err(AwaError::Validation(
3072                "queue storage activation requires the storage transition state row".into(),
3073            ));
3074        }
3075
3076        tx.commit().await.map_err(map_sqlx_error)?;
3077
3078        Ok(())
3079    }
3080
3081    #[tracing::instrument(skip(self, pool), name = "queue_storage.install")]
3082    pub async fn install(&self, pool: &PgPool) -> Result<(), AwaError> {
3083        self.prepare_schema(pool).await?;
3084        self.activate_backend(pool).await
3085    }
3086
3087    pub async fn reset(&self, pool: &PgPool) -> Result<(), AwaError> {
3088        let schema = self.schema();
3089        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
3090
3091        // Drop any partial-migration leftover tables before the main
3092        // TRUNCATE. If `prepare_schema` crashed mid-migration, the
3093        // schema may contain `lease_claims_legacy` /
3094        // `lease_claim_closures_legacy` alongside the partitioned
3095        // parents. `reset()` must clean these out, otherwise the next
3096        // `prepare_schema()` runs the legacy migration again on top of
3097        // the freshly-emptied parent and silently re-inserts old rows.
3098        sqlx::query(&format!(
3099            "DROP TABLE IF EXISTS {schema}.lease_claims_legacy"
3100        ))
3101        .execute(tx.as_mut())
3102        .await
3103        .map_err(map_sqlx_error)?;
3104
3105        sqlx::query(&format!(
3106            "DROP TABLE IF EXISTS {schema}.lease_claim_closures_legacy"
3107        ))
3108        .execute(tx.as_mut())
3109        .await
3110        .map_err(map_sqlx_error)?;
3111
3112        sqlx::query(&format!(
3113            r#"
3114            TRUNCATE
3115                {schema}.ready_entries,
3116                {schema}.ready_claim_attempt_batches,
3117                {schema}.ready_tombstones,
3118                {schema}.ready_segments,
3119                {schema}.done_entries,
3120                {schema}.receipt_completion_batches,
3121                {schema}.receipt_completion_tombstones,
3122                {schema}.dlq_entries,
3123                {schema}.leases,
3124                {schema}.lease_claims,
3125                {schema}.lease_claim_batches,
3126                {schema}.lease_claim_closures,
3127                {schema}.lease_claim_closure_batches,
3128                {schema}.attempt_state,
3129                {schema}.deferred_jobs,
3130                {schema}.queue_lanes,
3131                {schema}.queue_terminal_rollups,
3132                {schema}.queue_terminal_live_counts,
3133                {schema}.queue_terminal_count_deltas,
3134                {schema}.queue_claimer_leases,
3135                {schema}.queue_claimer_state,
3136                {schema}.queue_ring_slots,
3137                {schema}.lease_ring_slots,
3138                {schema}.claim_ring_slots
3139            "#
3140        ))
3141        .execute(tx.as_mut())
3142        .await
3143        .map_err(map_sqlx_error)?;
3144
3145        sqlx::query(&format!(
3146            "ALTER SEQUENCE {schema}.job_id_seq RESTART WITH 1"
3147        ))
3148        .execute(tx.as_mut())
3149        .await
3150        .map_err(map_sqlx_error)?;
3151
3152        sqlx::query(&format!(
3153            r#"
3154            UPDATE {schema}.queue_ring_state
3155            SET current_slot = 0,
3156                generation = 0,
3157                slot_count = $1
3158            WHERE singleton = TRUE
3159            "#
3160        ))
3161        .bind(self.queue_slot_count() as i32)
3162        .execute(tx.as_mut())
3163        .await
3164        .map_err(map_sqlx_error)?;
3165
3166        sqlx::query(&format!(
3167            r#"
3168            UPDATE {schema}.lease_ring_state
3169            SET current_slot = 0,
3170                generation = 0,
3171                slot_count = $1
3172            WHERE singleton = TRUE
3173            "#
3174        ))
3175        .bind(self.lease_slot_count() as i32)
3176        .execute(tx.as_mut())
3177        .await
3178        .map_err(map_sqlx_error)?;
3179
3180        sqlx::query(&format!(
3181            r#"
3182            UPDATE {schema}.claim_ring_state
3183            SET current_slot = 0,
3184                generation = 0,
3185                slot_count = $1
3186            WHERE singleton = TRUE
3187            "#
3188        ))
3189        .bind(self.claim_slot_count() as i32)
3190        .execute(tx.as_mut())
3191        .await
3192        .map_err(map_sqlx_error)?;
3193
3194        for slot in 0..self.queue_slot_count() {
3195            sqlx::query(&format!(
3196                r#"
3197                INSERT INTO {schema}.queue_ring_slots (slot, generation)
3198                VALUES ($1, $2)
3199                "#
3200            ))
3201            .bind(slot as i32)
3202            .bind(if slot == 0 { 0_i64 } else { -1_i64 })
3203            .execute(tx.as_mut())
3204            .await
3205            .map_err(map_sqlx_error)?;
3206        }
3207
3208        for slot in 0..self.lease_slot_count() {
3209            sqlx::query(&format!(
3210                r#"
3211                INSERT INTO {schema}.lease_ring_slots (slot, generation)
3212                VALUES ($1, $2)
3213                "#
3214            ))
3215            .bind(slot as i32)
3216            .bind(if slot == 0 { 0_i64 } else { -1_i64 })
3217            .execute(tx.as_mut())
3218            .await
3219            .map_err(map_sqlx_error)?;
3220        }
3221
3222        for slot in 0..self.claim_slot_count() {
3223            sqlx::query(&format!(
3224                r#"
3225                INSERT INTO {schema}.claim_ring_slots (slot, generation)
3226                VALUES ($1, $2)
3227                "#
3228            ))
3229            .bind(slot as i32)
3230            .bind(if slot == 0 { 0_i64 } else { -1_i64 })
3231            .execute(tx.as_mut())
3232            .await
3233            .map_err(map_sqlx_error)?;
3234        }
3235
3236        tx.commit().await.map_err(map_sqlx_error)?;
3237
3238        // queue_lanes was TRUNCATEd above and queue_meta may have a new
3239        // shard configuration for the next round. Clear both caches so
3240        // the next ensure_lane / shard_for_enqueue calls re-observe DB
3241        // state.
3242        self.clear_lane_cache();
3243        self.enqueue_shards_cache
3244            .lock()
3245            .expect("enqueue_shards_cache poisoned")
3246            .clear();
3247        Ok(())
3248    }
3249
3250    async fn ensure_lane<'a>(
3251        &self,
3252        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3253        queue: &str,
3254        priority: i16,
3255        enqueue_shard: i16,
3256    ) -> Result<(), AwaError> {
3257        // Fast path: this store has previously written the three lane
3258        // rows for this `(queue, priority, shard)` triple. The cache is
3259        // optimistic: another transaction may have marked the lane before
3260        // commit, or the marking transaction may later roll back. Verify the
3261        // head row is visible in this transaction before trusting the cache.
3262        if self.lane_is_cached(queue, priority, enqueue_shard) {
3263            let schema = self.schema();
3264            let visible: bool = sqlx::query_scalar(&format!(
3265                r#"
3266                SELECT EXISTS (
3267                    SELECT 1
3268                    FROM {schema}.queue_enqueue_heads
3269                    WHERE queue = $1
3270                      AND priority = $2
3271                      AND enqueue_shard = $3
3272                )
3273                "#
3274            ))
3275            .bind(queue)
3276            .bind(priority)
3277            .bind(enqueue_shard)
3278            .fetch_one(tx.as_mut())
3279            .await
3280            .map_err(map_sqlx_error)?;
3281
3282            if visible {
3283                return Ok(());
3284            }
3285
3286            self.invalidate_cached_lane(queue, priority, enqueue_shard);
3287        }
3288
3289        self.ensure_lane_inserts(tx, queue, priority, enqueue_shard)
3290            .await
3291    }
3292
3293    /// Run the three lane-row inserts unconditionally and mark the
3294    /// `(queue, priority)` pair as cached on success. Skips the cache
3295    /// fast path so callers in the rollback-recovery path can force a
3296    /// re-insert without racing another transaction that has marked
3297    /// the lane but not yet committed its inserts.
3298    async fn ensure_lane_inserts<'a>(
3299        &self,
3300        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3301        queue: &str,
3302        priority: i16,
3303        enqueue_shard: i16,
3304    ) -> Result<(), AwaError> {
3305        let schema = self.schema();
3306        let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}");
3307        sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
3308            .bind(lane_lock_key)
3309            .execute(tx.as_mut())
3310            .await
3311            .map_err(map_sqlx_error)?;
3312
3313        sqlx::query(&format!(
3314            r#"
3315            INSERT INTO {schema}.queue_lanes (queue, priority)
3316            VALUES ($1, $2)
3317            ON CONFLICT (queue, priority) DO NOTHING
3318            "#
3319        ))
3320        .bind(queue)
3321        .bind(priority)
3322        .execute(tx.as_mut())
3323        .await
3324        .map_err(map_sqlx_error)?;
3325
3326        sqlx::query(&format!(
3327            r#"
3328            INSERT INTO {schema}.queue_enqueue_heads (queue, priority, enqueue_shard)
3329            VALUES ($1, $2, $3)
3330            ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING
3331            "#
3332        ))
3333        .bind(queue)
3334        .bind(priority)
3335        .bind(enqueue_shard)
3336        .execute(tx.as_mut())
3337        .await
3338        .map_err(map_sqlx_error)?;
3339
3340        sqlx::query(&format!(
3341            r#"
3342            INSERT INTO {schema}.queue_claim_heads (queue, priority, enqueue_shard)
3343            VALUES ($1, $2, $3)
3344            ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING
3345            "#
3346        ))
3347        .bind(queue)
3348        .bind(priority)
3349        .bind(enqueue_shard)
3350        .execute(tx.as_mut())
3351        .await
3352        .map_err(map_sqlx_error)?;
3353
3354        sqlx::query(&format!(
3355            r#"
3356            SELECT {schema}.ensure_lane_sequences($1, $2, $3)
3357            "#
3358        ))
3359        .bind(queue)
3360        .bind(priority)
3361        .bind(enqueue_shard)
3362        .execute(tx.as_mut())
3363        .await
3364        .map_err(map_sqlx_error)?;
3365
3366        self.mark_lane_ensured(queue, priority, enqueue_shard);
3367        Ok(())
3368    }
3369
3370    /// Pick the enqueue shard for a row on `queue`.
3371    ///
3372    /// Reads `awa.queue_meta.enqueue_shards` (default 1, cached
3373    /// in-process). With `enqueue_shards = 1` the result is always 0.
3374    /// With `enqueue_shards > 1`:
3375    ///
3376    /// - If `ordering_key` is `Some(k)`, the shard is a stable hash
3377    ///   of `k` modulo the shard count. Two rows sharing an
3378    ///   ordering key always land on the same shard, which preserves
3379    ///   FIFO within the key.
3380    /// - If `ordering_key` is `None`, the shard comes from the
3381    ///   per-store rotor, which spreads consecutive picks across
3382    ///   shards.
3383    async fn shard_for_enqueue(
3384        &self,
3385        pool_executor: impl sqlx::PgExecutor<'_>,
3386        queue: &str,
3387        ordering_key: Option<&[u8]>,
3388    ) -> Result<i16, AwaError> {
3389        if let Some(cached) = self
3390            .enqueue_shards_cache
3391            .lock()
3392            .expect("enqueue_shards_cache poisoned")
3393            .get(queue)
3394            .copied()
3395        {
3396            return Ok(self.pick_shard(cached, ordering_key));
3397        }
3398
3399        let shards: i16 = sqlx::query_scalar(
3400            r#"
3401            SELECT COALESCE(MAX(enqueue_shards), 1)::smallint
3402            FROM awa.queue_meta
3403            WHERE queue = $1
3404            "#,
3405        )
3406        .bind(queue)
3407        .fetch_one(pool_executor)
3408        .await
3409        .map_err(map_sqlx_error)?;
3410
3411        let shards = shards.max(1);
3412        self.enqueue_shards_cache
3413            .lock()
3414            .expect("enqueue_shards_cache poisoned")
3415            .insert(queue.to_string(), shards);
3416
3417        Ok(self.pick_shard(shards, ordering_key))
3418    }
3419
3420    /// Map a row (or a no-key sub-batch) to its shard.
3421    ///
3422    /// At `shards <= 1` every row goes to shard 0. Otherwise the
3423    /// caller-supplied `ordering_key` hashes deterministically into
3424    /// `[0, shards)` via [`shard_for_ordering_key`]; an absent key
3425    /// advances the per-store rotor once and returns the next shard.
3426    /// Callers route a whole no-key sub-batch through a single
3427    /// `pick_shard(_, None)` so the rotor amortises across the batch
3428    /// rather than firing per row.
3429    fn pick_shard(&self, shards: i16, ordering_key: Option<&[u8]>) -> i16 {
3430        if shards <= 1 {
3431            return 0;
3432        }
3433        match ordering_key {
3434            Some(key) => shard_for_ordering_key(key, shards),
3435            None => {
3436                let raw = self.shard_rotor.fetch_add(1, Ordering::Relaxed) as i32;
3437                raw.rem_euclid(shards as i32) as i16
3438            }
3439        }
3440    }
3441
3442    fn lane_is_cached(&self, queue: &str, priority: i16, enqueue_shard: i16) -> bool {
3443        let cache = self.ensured_lanes.lock().expect("ensured_lanes mutex");
3444        cache.contains(&(queue.to_string(), priority, enqueue_shard))
3445    }
3446
3447    fn mark_lane_ensured(&self, queue: &str, priority: i16, enqueue_shard: i16) {
3448        self.ensured_lanes
3449            .lock()
3450            .expect("ensured_lanes mutex")
3451            .insert((queue.to_string(), priority, enqueue_shard));
3452    }
3453
3454    fn invalidate_cached_lane(&self, queue: &str, priority: i16, enqueue_shard: i16) {
3455        self.ensured_lanes
3456            .lock()
3457            .expect("ensured_lanes mutex")
3458            .remove(&(queue.to_string(), priority, enqueue_shard));
3459    }
3460
3461    fn clear_lane_cache(&self) {
3462        self.ensured_lanes
3463            .lock()
3464            .expect("ensured_lanes mutex")
3465            .clear();
3466    }
3467
3468    // Reserve lane sequence numbers for a specific
3469    // `(queue, priority, shard)` triple and return the lane sequence at
3470    // which the caller's range starts. If the head row is missing —
3471    // typically because a previous ensure_lane ran inside a transaction
3472    // that ultimately rolled back, leaving a stale cache entry behind —
3473    // the cache entry is invalidated, ensure_lane is re-run, and the reserve
3474    // is retried exactly once. A second miss surfaces as
3475    // RowNotFound.
3476    async fn advance_enqueue_head<'a>(
3477        &self,
3478        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3479        queue: &str,
3480        priority: i16,
3481        enqueue_shard: i16,
3482        count: i64,
3483    ) -> Result<i64, AwaError> {
3484        let schema = self.schema();
3485        let sql = format!(
3486            r#"
3487            SELECT {schema}.reserve_enqueue_seq($1, $2, $3, $4)
3488            FROM {schema}.queue_enqueue_heads
3489            WHERE queue = $1 AND priority = $2 AND enqueue_shard = $3
3490            "#
3491        );
3492
3493        let maybe_start: Option<i64> = sqlx::query_scalar(&sql)
3494            .bind(queue)
3495            .bind(priority)
3496            .bind(enqueue_shard)
3497            .bind(count)
3498            .fetch_optional(tx.as_mut())
3499            .await
3500            .map_err(map_sqlx_error)?;
3501
3502        if let Some(start) = maybe_start {
3503            return Ok(start);
3504        }
3505
3506        // Recovery path: a prior ensure_lane marked the cache and
3507        // then rolled back, leaving the head row missing. Bypass the
3508        // cache fast path here and run the inserts unconditionally —
3509        // calling `ensure_lane` would re-take the fast path if
3510        // another concurrent transaction has re-marked the cache but
3511        // not yet committed its inserts, leaving us in the same
3512        // failure state.
3513        self.invalidate_cached_lane(queue, priority, enqueue_shard);
3514        self.ensure_lane_inserts(tx, queue, priority, enqueue_shard)
3515            .await?;
3516        let start: i64 = sqlx::query_scalar(&sql)
3517            .bind(queue)
3518            .bind(priority)
3519            .bind(enqueue_shard)
3520            .bind(count)
3521            .fetch_one(tx.as_mut())
3522            .await
3523            .map_err(map_sqlx_error)?;
3524        Ok(start)
3525    }
3526
3527    async fn current_queue_ring<'a>(
3528        &self,
3529        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3530    ) -> Result<(i32, i64), AwaError> {
3531        let schema = self.schema();
3532        sqlx::query_as(&format!(
3533            r#"
3534            SELECT current_slot, generation
3535            FROM {schema}.queue_ring_state
3536            WHERE singleton = TRUE
3537            "#
3538        ))
3539        .fetch_one(tx.as_mut())
3540        .await
3541        .map_err(map_sqlx_error)
3542    }
3543
3544    async fn next_job_ids<'a>(
3545        &self,
3546        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3547        count: usize,
3548    ) -> Result<Vec<i64>, AwaError> {
3549        if count == 0 {
3550            return Ok(Vec::new());
3551        }
3552
3553        let query = format!(
3554            "SELECT nextval('{}')::bigint FROM generate_series(1, $1::int)",
3555            self.job_id_sequence()
3556        );
3557
3558        sqlx::query_scalar(&query)
3559            .bind(count as i32)
3560            .fetch_all(tx.as_mut())
3561            .await
3562            .map_err(map_sqlx_error)
3563    }
3564
3565    async fn current_timestamp_tx<'a>(
3566        &self,
3567        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3568    ) -> Result<DateTime<Utc>, AwaError> {
3569        sqlx::query_scalar("SELECT clock_timestamp()")
3570            .fetch_one(tx.as_mut())
3571            .await
3572            .map_err(map_sqlx_error)
3573    }
3574
3575    async fn claim_ready_rows_tx<'a>(
3576        &self,
3577        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3578        queue: &str,
3579        max_batch: i64,
3580        deadline_duration: Duration,
3581        aging_interval: Duration,
3582    ) -> Result<Vec<ReadyJobLeaseRow>, AwaError> {
3583        let schema = self.schema();
3584        sqlx::query_as(&format!(
3585            r#"
3586            SELECT
3587                ready_slot,
3588                ready_generation,
3589                lane_seq,
3590                enqueue_shard,
3591                lease_slot,
3592                lease_generation,
3593                claim_slot,
3594                receipt_id,
3595                claim_batch_id,
3596                claim_batch_index,
3597                job_id,
3598                kind,
3599                queue,
3600                args,
3601                lane_priority,
3602                priority,
3603                attempt,
3604                run_lease,
3605                max_attempts,
3606                run_at,
3607                heartbeat_at,
3608                deadline_at,
3609                attempted_at,
3610                created_at,
3611                unique_key,
3612                unique_states,
3613                COALESCE(payload, '{{}}'::jsonb) AS payload
3614            FROM {schema}.claim_ready_runtime($1, $2, $3, $4)
3615            "#
3616        ))
3617        .bind(queue)
3618        .bind(max_batch)
3619        .bind(deadline_duration.as_secs_f64())
3620        .bind(aging_interval.as_secs_f64())
3621        .fetch_all(tx.as_mut())
3622        .await
3623        .map_err(map_sqlx_error)
3624    }
3625
3626    fn claim_cursor_advances(rows: &[ReadyJobLeaseRow]) -> Vec<ClaimCursorAdvance> {
3627        let mut next_by_lane: BTreeMap<ClaimCursorLaneKey, i64> = BTreeMap::new();
3628        for row in rows {
3629            let key = (row.queue.clone(), row.lane_priority, row.enqueue_shard);
3630            let next = row.lane_seq + 1;
3631            next_by_lane
3632                .entry(key)
3633                .and_modify(|current| *current = (*current).max(next))
3634                .or_insert(next);
3635        }
3636
3637        next_by_lane
3638            .into_iter()
3639            .map(
3640                |((queue, priority, enqueue_shard), next_seq)| ClaimCursorAdvance {
3641                    queue,
3642                    priority,
3643                    enqueue_shard,
3644                    next_seq,
3645                    only_if_current: None,
3646                },
3647            )
3648            .collect()
3649    }
3650
3651    fn normalize_claim_cursor_advances(advances: &[ClaimCursorAdvance]) -> Vec<ClaimCursorAdvance> {
3652        let mut grouped: GroupedClaimCursorAdvances = BTreeMap::new();
3653
3654        for advance in advances {
3655            let key = (
3656                advance.queue.clone(),
3657                advance.priority,
3658                advance.enqueue_shard,
3659            );
3660            let (unconditional, conditional) = grouped.entry(key).or_default();
3661            if let Some(only_if_current) = advance.only_if_current {
3662                conditional
3663                    .entry(only_if_current)
3664                    .and_modify(|next| *next = (*next).max(advance.next_seq))
3665                    .or_insert(advance.next_seq);
3666            } else {
3667                *unconditional = Some(
3668                    unconditional
3669                        .map(|next| next.max(advance.next_seq))
3670                        .unwrap_or(advance.next_seq),
3671                );
3672            }
3673        }
3674
3675        let mut normalized = Vec::with_capacity(advances.len());
3676        for ((queue, priority, enqueue_shard), (unconditional, conditional)) in grouped {
3677            if let Some(next_seq) = unconditional {
3678                normalized.push(ClaimCursorAdvance {
3679                    queue,
3680                    priority,
3681                    enqueue_shard,
3682                    next_seq,
3683                    only_if_current: None,
3684                });
3685                continue;
3686            }
3687
3688            for (only_if_current, next_seq) in conditional {
3689                normalized.push(ClaimCursorAdvance {
3690                    queue: queue.clone(),
3691                    priority,
3692                    enqueue_shard,
3693                    next_seq,
3694                    only_if_current: Some(only_if_current),
3695                });
3696            }
3697        }
3698
3699        normalized
3700    }
3701
3702    async fn advance_claim_cursors(&self, pool: &PgPool, advances: &[ClaimCursorAdvance]) {
3703        let advances = Self::normalize_claim_cursor_advances(advances);
3704        // PostgreSQL sequence state is not rolled back with the surrounding
3705        // transaction. Keep claim cursors lagging rather than risking a cursor
3706        // that gets ahead of ready rows whose claim/cancel transaction aborts.
3707        for attempt in 1..=3 {
3708            match self.advance_claim_cursors_strict(pool, &advances).await {
3709                Ok(()) => return,
3710                Err(err) if attempt < 3 => {
3711                    tracing::warn!(
3712                        error = ?err,
3713                        lanes = advances.len(),
3714                        attempt,
3715                        "failed to advance queue-storage claim cursors after committed state change; retrying"
3716                    );
3717                    tokio::time::sleep(Duration::from_millis(25 * attempt as u64)).await;
3718                }
3719                Err(err) => {
3720                    tracing::warn!(
3721                        error = ?err,
3722                        lanes = advances.len(),
3723                        attempts = attempt,
3724                        "failed to advance queue-storage claim cursors after committed state change"
3725                    );
3726                    return;
3727                }
3728            }
3729        }
3730    }
3731
3732    async fn advance_claim_cursors_strict(
3733        &self,
3734        pool: &PgPool,
3735        advances: &[ClaimCursorAdvance],
3736    ) -> Result<(), AwaError> {
3737        if advances.is_empty() {
3738            return Ok(());
3739        }
3740
3741        let schema = self.schema();
3742        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
3743        for advance in advances {
3744            sqlx::query(&format!(
3745                r#"
3746                WITH head AS MATERIALIZED (
3747                    SELECT seq_name
3748                    FROM {schema}.queue_claim_heads
3749                    WHERE queue = $1
3750                      AND priority = $2
3751                      AND enqueue_shard = $3
3752                    FOR UPDATE
3753                )
3754                SELECT {schema}.set_sequence_next(seq_name, $4)
3755                FROM head
3756                WHERE $5::bigint IS NULL
3757                   OR {schema}.sequence_next_value(seq_name) = $5
3758                "#
3759            ))
3760            .bind(&advance.queue)
3761            .bind(advance.priority)
3762            .bind(advance.enqueue_shard)
3763            .bind(advance.next_seq)
3764            .bind(advance.only_if_current)
3765            .execute(tx.as_mut())
3766            .await
3767            .map_err(map_sqlx_error)?;
3768        }
3769
3770        tx.commit().await.map_err(map_sqlx_error)?;
3771        Ok(())
3772    }
3773
3774    async fn execute_ready_inserts_tx<'a>(
3775        &self,
3776        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3777        ring: (i32, i64),
3778        rows: &[RuntimeReadyInsert],
3779    ) -> Result<usize, AwaError> {
3780        if rows.is_empty() {
3781            return Ok(0);
3782        }
3783
3784        let schema = self.schema();
3785        let mut builder = QueryBuilder::<Postgres>::new(format!(
3786            "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) "
3787        ));
3788        builder.push_values(rows.iter(), |mut b, row| {
3789            b.push_bind(ring.0)
3790                .push_bind(ring.1)
3791                .push_bind(row.job_id)
3792                .push_bind(&row.kind)
3793                .push_bind(&row.queue)
3794                .push_bind(&row.args)
3795                .push_bind(row.priority)
3796                .push_bind(row.attempt)
3797                .push_bind(row.run_lease)
3798                .push_bind(row.max_attempts)
3799                .push_bind(row.lane_seq)
3800                .push_bind(row.enqueue_shard)
3801                .push_bind(row.run_at)
3802                .push_bind(row.attempted_at)
3803                .push_bind(row.created_at)
3804                .push_bind(&row.unique_key)
3805                .push_bind(&row.unique_states)
3806                .push_bind(storage_payload(&row.payload));
3807        });
3808        builder
3809            .build()
3810            .execute(tx.as_mut())
3811            .await
3812            .map_err(map_sqlx_error)?;
3813
3814        Ok(rows.len())
3815    }
3816
3817    fn ready_segments_from_rows(rows: &[RuntimeReadyInsert]) -> Vec<ReadySegmentInsert> {
3818        if rows.is_empty() {
3819            return Vec::new();
3820        }
3821
3822        let mut segments = Vec::new();
3823        let mut start_index = 0usize;
3824
3825        for index in 1..=rows.len() {
3826            let split = if index == rows.len() {
3827                true
3828            } else {
3829                let previous = &rows[index - 1];
3830                let current = &rows[index];
3831                previous.queue != current.queue
3832                    || previous.priority != current.priority
3833                    || previous.enqueue_shard != current.enqueue_shard
3834                    || previous.lane_seq + 1 != current.lane_seq
3835                    || previous.run_at != current.run_at
3836            };
3837
3838            if split {
3839                let first = &rows[start_index];
3840                let last = &rows[index - 1];
3841                segments.push(ReadySegmentInsert {
3842                    queue: first.queue.clone(),
3843                    priority: first.priority,
3844                    enqueue_shard: first.enqueue_shard,
3845                    first_lane_seq: first.lane_seq,
3846                    next_lane_seq: last.lane_seq + 1,
3847                    first_run_at: first.run_at,
3848                });
3849                start_index = index;
3850            }
3851        }
3852
3853        segments
3854    }
3855
3856    async fn insert_ready_segments_tx<'a>(
3857        &self,
3858        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3859        ring: (i32, i64),
3860        rows: &[RuntimeReadyInsert],
3861    ) -> Result<usize, AwaError> {
3862        let segments = Self::ready_segments_from_rows(rows);
3863        if segments.is_empty() {
3864            return Ok(0);
3865        }
3866
3867        let schema = self.schema();
3868        let mut builder = QueryBuilder::<Postgres>::new(format!(
3869            "INSERT INTO {schema}.ready_segments (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq, next_lane_seq, first_run_at) "
3870        ));
3871        builder.push_values(segments.iter(), |mut b, segment| {
3872            b.push_bind(ring.0)
3873                .push_bind(ring.1)
3874                .push_bind(&segment.queue)
3875                .push_bind(segment.priority)
3876                .push_bind(segment.enqueue_shard)
3877                .push_bind(segment.first_lane_seq)
3878                .push_bind(segment.next_lane_seq)
3879                .push_bind(segment.first_run_at);
3880        });
3881        builder.push(
3882            " ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) DO NOTHING",
3883        );
3884        builder
3885            .build()
3886            .execute(tx.as_mut())
3887            .await
3888            .map_err(map_sqlx_error)?;
3889
3890        Ok(segments.len())
3891    }
3892
3893    async fn execute_ready_copy_tx<'a>(
3894        &self,
3895        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3896        ring: (i32, i64),
3897        rows: &[RuntimeReadyInsert],
3898    ) -> Result<usize, AwaError> {
3899        if rows.is_empty() {
3900            return Ok(0);
3901        }
3902
3903        let schema = self.schema();
3904        let copy_sql = format!(
3905            "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}')"
3906        );
3907        let mut copy_in = tx
3908            .as_mut()
3909            .copy_in_raw(&copy_sql)
3910            .await
3911            .map_err(map_sqlx_error)?;
3912        // 320 bytes/row is only a rough starting point; large JSON payloads
3913        // are bounded by chunked COPY sends below rather than by this reserve.
3914        let mut csv_buf = Vec::with_capacity(rows.len().min(1024) * 320);
3915        for row in rows {
3916            write_ready_copy_row(&mut csv_buf, ring.0, ring.1, row);
3917            if csv_buf.len() >= COPY_CHUNK_TARGET_BYTES {
3918                let chunk =
3919                    std::mem::replace(&mut csv_buf, Vec::with_capacity(COPY_CHUNK_TARGET_BYTES));
3920                copy_in.send(chunk).await.map_err(map_sqlx_error)?;
3921            }
3922        }
3923        if !csv_buf.is_empty() {
3924            copy_in.send(csv_buf).await.map_err(map_sqlx_error)?;
3925        }
3926        copy_in.finish().await.map_err(map_sqlx_error)?;
3927
3928        Ok(rows.len())
3929    }
3930
3931    async fn insert_ready_rows_tx<'a>(
3932        &self,
3933        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
3934        rows: Vec<RuntimeReadyRow>,
3935    ) -> Result<usize, AwaError> {
3936        if rows.is_empty() {
3937            return Ok(0);
3938        }
3939
3940        let grouped = self.group_ready_rows_by_shard(tx, rows).await?;
3941        let total_rows: usize = grouped.values().map(Vec::len).sum();
3942        let job_ids = self.next_job_ids(tx, total_rows).await?;
3943        let mut job_id_iter = job_ids.into_iter();
3944
3945        let mut ready_rows = Vec::with_capacity(total_rows);
3946        let mut lane_ranges = Vec::with_capacity(grouped.len());
3947
3948        for ((queue, priority, enqueue_shard), lane_rows) in grouped {
3949            let range_start = ready_rows.len();
3950
3951            for row in lane_rows {
3952                let job_id = job_id_iter.next().ok_or_else(|| {
3953                    AwaError::Validation("queue storage job id allocation underflow".to_string())
3954                })?;
3955                ready_rows.push(RuntimeReadyInsert {
3956                    job_id,
3957                    kind: row.kind,
3958                    queue: row.queue,
3959                    args: row.args,
3960                    priority: row.priority,
3961                    attempt: row.attempt,
3962                    run_lease: row.run_lease,
3963                    max_attempts: row.max_attempts,
3964                    run_at: row.run_at,
3965                    attempted_at: row.attempted_at,
3966                    lane_seq: 0,
3967                    enqueue_shard,
3968                    created_at: row.created_at,
3969                    unique_key: row.unique_key,
3970                    unique_states: row.unique_states,
3971                    payload: row.payload,
3972                });
3973            }
3974            lane_ranges.push((
3975                queue,
3976                priority,
3977                enqueue_shard,
3978                range_start,
3979                ready_rows.len(),
3980            ));
3981        }
3982
3983        self.sync_ready_enqueue_unique_claims(tx, &ready_rows)
3984            .await?;
3985        for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges {
3986            self.ensure_lane(tx, &queue, priority, enqueue_shard)
3987                .await?;
3988
3989            let count = (range_end - range_start) as i64;
3990            let start_seq = self
3991                .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count)
3992                .await?;
3993
3994            for (offset, row) in ready_rows[range_start..range_end].iter_mut().enumerate() {
3995                row.lane_seq = start_seq + offset as i64;
3996            }
3997        }
3998        let ring = self.current_queue_ring(tx).await?;
3999        self.execute_ready_inserts_tx(tx, ring, &ready_rows).await?;
4000        self.insert_ready_segments_tx(tx, ring, &ready_rows).await?;
4001        Ok(total_rows)
4002    }
4003
4004    /// Re-group rows by `(queue, priority, enqueue_shard)` so each
4005    /// resulting bucket targets exactly one head row.
4006    ///
4007    /// The two routing modes are amortised differently:
4008    ///
4009    /// - **Rows with `ordering_key`** are hashed per row via
4010    ///   `shard_for_ordering_key`, so jobs that share a key always
4011    ///   land on the same shard. Each distinct key may produce its
4012    ///   own sub-bucket inside one batch.
4013    /// - **Rows without `ordering_key`** share **one rotor pick per
4014    ///   `(queue, priority)` per call**. A batch of 500 rotor-routed
4015    ///   rows produces one `advance_enqueue_head` UPDATE and one
4016    ///   INSERT, not 500. The rotor still advances per batch, so
4017    ///   successive batches spread across shards; the per-batch
4018    ///   amortisation is what makes `enqueue_shards > 1` net-faster
4019    ///   than `S = 1` at moderate concurrency.
4020    ///
4021    /// Mixing keyed and rotor rows in one call is supported — the
4022    /// keyed rows fan out by key while the rotor rows collapse into
4023    /// a single shard-sub-batch.
4024    async fn group_ready_rows_by_shard<'a>(
4025        &self,
4026        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4027        rows: Vec<RuntimeReadyRow>,
4028    ) -> Result<BTreeMap<(String, i16, i16), Vec<RuntimeReadyRow>>, AwaError> {
4029        // Partition by (queue, priority) first so the rotor pick for
4030        // the no-key sub-batch can amortise over every row that shares
4031        // its destination lane.
4032        let mut by_queue_priority: BTreeMap<(String, i16), Vec<RuntimeReadyRow>> = BTreeMap::new();
4033        for row in rows {
4034            by_queue_priority
4035                .entry((row.queue.clone(), row.priority))
4036                .or_default()
4037                .push(row);
4038        }
4039
4040        let mut grouped: BTreeMap<(String, i16, i16), Vec<RuntimeReadyRow>> = BTreeMap::new();
4041        for ((queue, priority), bucket) in by_queue_priority {
4042            let mut rotor_rows: Vec<RuntimeReadyRow> = Vec::with_capacity(bucket.len());
4043            for row in bucket {
4044                if row.ordering_key.is_some() {
4045                    let shard = self
4046                        .shard_for_enqueue(tx.as_mut(), &queue, row.ordering_key.as_deref())
4047                        .await?;
4048                    grouped
4049                        .entry((queue.clone(), priority, shard))
4050                        .or_default()
4051                        .push(row);
4052                } else {
4053                    rotor_rows.push(row);
4054                }
4055            }
4056            if !rotor_rows.is_empty() {
4057                let shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?;
4058                grouped
4059                    .entry((queue.clone(), priority, shard))
4060                    .or_default()
4061                    .extend(rotor_rows);
4062            }
4063        }
4064        Ok(grouped)
4065    }
4066
4067    async fn insert_ready_rows_copy_tx<'a>(
4068        &self,
4069        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4070        rows: Vec<RuntimeReadyRow>,
4071        job_ids: Vec<i64>,
4072    ) -> Result<usize, AwaError> {
4073        if rows.is_empty() {
4074            return Ok(0);
4075        }
4076
4077        let grouped = self.group_ready_rows_by_shard(tx, rows).await?;
4078
4079        let total_rows: usize = grouped.values().map(Vec::len).sum();
4080        if job_ids.len() != total_rows {
4081            return Err(AwaError::Validation(
4082                "queue storage job id allocation count mismatch".to_string(),
4083            ));
4084        }
4085        let mut job_id_iter = job_ids.into_iter();
4086
4087        let mut ready_rows = Vec::with_capacity(total_rows);
4088        let mut lane_ranges = Vec::with_capacity(grouped.len());
4089
4090        for ((queue, priority, enqueue_shard), lane_rows) in grouped {
4091            let range_start = ready_rows.len();
4092
4093            for row in lane_rows {
4094                let job_id = job_id_iter.next().ok_or_else(|| {
4095                    AwaError::Validation("queue storage job id allocation underflow".to_string())
4096                })?;
4097                ready_rows.push(RuntimeReadyInsert {
4098                    job_id,
4099                    kind: row.kind,
4100                    queue: row.queue,
4101                    args: row.args,
4102                    priority: row.priority,
4103                    attempt: row.attempt,
4104                    run_lease: row.run_lease,
4105                    max_attempts: row.max_attempts,
4106                    run_at: row.run_at,
4107                    attempted_at: row.attempted_at,
4108                    lane_seq: 0,
4109                    enqueue_shard,
4110                    created_at: row.created_at,
4111                    unique_key: row.unique_key,
4112                    unique_states: row.unique_states,
4113                    payload: row.payload,
4114                });
4115            }
4116            lane_ranges.push((
4117                queue,
4118                priority,
4119                enqueue_shard,
4120                range_start,
4121                ready_rows.len(),
4122            ));
4123        }
4124
4125        self.sync_ready_enqueue_unique_claims(tx, &ready_rows)
4126            .await?;
4127        for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges {
4128            self.ensure_lane(tx, &queue, priority, enqueue_shard)
4129                .await?;
4130
4131            let count = (range_end - range_start) as i64;
4132            let start_seq = self
4133                .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count)
4134                .await?;
4135
4136            for (offset, row) in ready_rows[range_start..range_end].iter_mut().enumerate() {
4137                row.lane_seq = start_seq + offset as i64;
4138            }
4139        }
4140        let ring = self.current_queue_ring(tx).await?;
4141        self.execute_ready_copy_tx(tx, ring, &ready_rows).await?;
4142        self.insert_ready_segments_tx(tx, ring, &ready_rows).await?;
4143        Ok(total_rows)
4144    }
4145
4146    async fn insert_existing_ready_rows_tx<'a>(
4147        &self,
4148        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4149        rows: Vec<ExistingReadyRow>,
4150        old_state: Option<JobState>,
4151    ) -> Result<usize, AwaError> {
4152        if rows.is_empty() {
4153            return Ok(0);
4154        }
4155
4156        let mut grouped: BTreeMap<(String, i16), Vec<ExistingReadyRow>> = BTreeMap::new();
4157        for row in rows {
4158            grouped
4159                .entry((row.queue.clone(), row.priority))
4160                .or_default()
4161                .push(row);
4162        }
4163
4164        let total_rows: usize = grouped.values().map(Vec::len).sum();
4165        let mut ready_rows = Vec::with_capacity(total_rows);
4166
4167        for ((queue, priority), lane_rows) in grouped {
4168            // Re-enqueue paths (retry-after, age-waiting, DLQ retry,
4169            // callback resume) do not carry the producer's original
4170            // `ordering_key`: the row came back from terminal storage
4171            // or from a deferred row, where the key was not retained.
4172            // Fall back to the rotor so the retry batch picks a shard
4173            // uniformly. Workloads that need ordering preserved across
4174            // retries must re-enqueue with the original key from the
4175            // application layer.
4176            let enqueue_shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?;
4177            self.ensure_lane(tx, &queue, priority, enqueue_shard)
4178                .await?;
4179
4180            let count = lane_rows.len() as i64;
4181            let start_seq = self
4182                .advance_enqueue_head(tx, &queue, priority, enqueue_shard, count)
4183                .await?;
4184
4185            for (offset, row) in lane_rows.into_iter().enumerate() {
4186                self.sync_unique_claim(
4187                    tx,
4188                    row.job_id,
4189                    &row.unique_key,
4190                    row.unique_states.as_deref(),
4191                    old_state,
4192                    Some(JobState::Available),
4193                )
4194                .await?;
4195                ready_rows.push(RuntimeReadyInsert {
4196                    job_id: row.job_id,
4197                    kind: row.kind,
4198                    queue: row.queue,
4199                    args: row.args,
4200                    priority: row.priority,
4201                    attempt: row.attempt,
4202                    run_lease: row.run_lease,
4203                    max_attempts: row.max_attempts,
4204                    run_at: row.run_at,
4205                    attempted_at: row.attempted_at,
4206                    lane_seq: start_seq + offset as i64,
4207                    enqueue_shard,
4208                    created_at: row.created_at,
4209                    unique_key: row.unique_key,
4210                    unique_states: row.unique_states,
4211                    payload: row.payload,
4212                });
4213            }
4214        }
4215
4216        let ring = self.current_queue_ring(tx).await?;
4217        self.execute_ready_inserts_tx(tx, ring, &ready_rows).await?;
4218        self.insert_ready_segments_tx(tx, ring, &ready_rows).await?;
4219        Ok(total_rows)
4220    }
4221
4222    async fn insert_deferred_rows_tx<'a>(
4223        &self,
4224        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4225        rows: Vec<DeferredJobRow>,
4226        old_state: Option<JobState>,
4227    ) -> Result<usize, AwaError> {
4228        if rows.is_empty() {
4229            return Ok(0);
4230        }
4231
4232        if old_state.is_none() {
4233            self.sync_deferred_enqueue_unique_claims(tx, &rows).await?;
4234        } else {
4235            for row in &rows {
4236                self.sync_unique_claim(
4237                    tx,
4238                    row.job_id,
4239                    &row.unique_key,
4240                    row.unique_states.as_deref(),
4241                    old_state,
4242                    Some(row.state),
4243                )
4244                .await?;
4245            }
4246        }
4247
4248        let schema = self.schema();
4249        let mut builder = QueryBuilder::<Postgres>::new(format!(
4250            "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) "
4251        ));
4252        builder.push_values(rows.iter(), |mut b, row| {
4253            b.push_bind(row.job_id)
4254                .push_bind(&row.kind)
4255                .push_bind(&row.queue)
4256                .push_bind(&row.args)
4257                .push_bind(row.state)
4258                .push_bind(row.priority)
4259                .push_bind(row.attempt)
4260                .push_bind(row.run_lease)
4261                .push_bind(row.max_attempts)
4262                .push_bind(row.run_at)
4263                .push_bind(row.attempted_at)
4264                .push_bind(row.finalized_at)
4265                .push_bind(row.created_at)
4266                .push_bind(&row.unique_key)
4267                .push_bind(&row.unique_states)
4268                .push_bind(storage_payload(&row.payload));
4269        });
4270        builder
4271            .build()
4272            .execute(tx.as_mut())
4273            .await
4274            .map_err(map_sqlx_error)?;
4275
4276        Ok(rows.len())
4277    }
4278
4279    async fn insert_deferred_rows_copy_tx<'a>(
4280        &self,
4281        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4282        rows: Vec<DeferredJobRow>,
4283    ) -> Result<usize, AwaError> {
4284        if rows.is_empty() {
4285            return Ok(0);
4286        }
4287
4288        self.sync_deferred_enqueue_unique_claims(tx, &rows).await?;
4289
4290        let schema = self.schema();
4291        let copy_sql = format!(
4292            "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}')"
4293        );
4294        let mut copy_in = tx
4295            .as_mut()
4296            .copy_in_raw(&copy_sql)
4297            .await
4298            .map_err(map_sqlx_error)?;
4299        // 320 bytes/row is only a rough starting point; large JSON payloads
4300        // are bounded by chunked COPY sends below rather than by this reserve.
4301        let mut csv_buf = Vec::with_capacity(rows.len().min(1024) * 320);
4302        for row in &rows {
4303            write_deferred_copy_row(&mut csv_buf, row);
4304            if csv_buf.len() >= COPY_CHUNK_TARGET_BYTES {
4305                let chunk =
4306                    std::mem::replace(&mut csv_buf, Vec::with_capacity(COPY_CHUNK_TARGET_BYTES));
4307                copy_in.send(chunk).await.map_err(map_sqlx_error)?;
4308            }
4309        }
4310        if !csv_buf.is_empty() {
4311            copy_in.send(csv_buf).await.map_err(map_sqlx_error)?;
4312        }
4313        copy_in.finish().await.map_err(map_sqlx_error)?;
4314
4315        Ok(rows.len())
4316    }
4317
4318    async fn insert_done_rows_tx<'a>(
4319        &self,
4320        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4321        rows: &[DoneJobRow],
4322        old_state: Option<JobState>,
4323    ) -> Result<usize, AwaError> {
4324        if rows.is_empty() {
4325            return Ok(0);
4326        }
4327
4328        for row in rows {
4329            self.sync_unique_claim(
4330                tx,
4331                row.job_id,
4332                &row.unique_key,
4333                row.unique_states.as_deref(),
4334                old_state,
4335                Some(row.state),
4336            )
4337            .await?;
4338        }
4339
4340        let schema = self.schema();
4341        let keep_ready_backing = matches!(
4342            old_state,
4343            Some(JobState::Running | JobState::WaitingExternal)
4344        );
4345        let ready_payloads = if rows
4346            .iter()
4347            .any(|row| !is_storage_payload_empty(&row.payload))
4348        {
4349            self.ready_payloads_for_done_rows_tx(tx, rows).await?
4350        } else {
4351            HashMap::new()
4352        };
4353        let mut ordered_rows: Vec<&DoneJobRow> = rows.iter().collect();
4354        ordered_rows.sort_unstable_by_key(|row| {
4355            (
4356                row.ready_slot,
4357                row.ready_generation,
4358                row.queue.as_str(),
4359                row.priority,
4360                row.enqueue_shard,
4361                row.lane_seq,
4362                row.job_id,
4363            )
4364        });
4365        let (ready_backed, synthetic): (Vec<_>, Vec<_>) = ordered_rows
4366            .into_iter()
4367            .partition(|row| keep_ready_backing && row.lane_seq >= 0);
4368
4369        if !ready_backed.is_empty() {
4370            let mut builder = QueryBuilder::<Postgres>::new(format!(
4371                "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) "
4372            ));
4373            builder.push_values(ready_backed, |mut b, row| {
4374                let ready_key = (
4375                    row.ready_slot,
4376                    row.ready_generation,
4377                    row.queue.as_str(),
4378                    row.priority,
4379                    row.enqueue_shard,
4380                    row.lane_seq,
4381                );
4382                let ready_payload = ready_payloads.get(&ready_key);
4383                b.push_bind(row.ready_slot)
4384                    .push_bind(row.ready_generation)
4385                    .push_bind(row.job_id)
4386                    .push_bind(&row.kind)
4387                    .push_bind(&row.queue)
4388                    .push_bind(row.state)
4389                    .push_bind(row.priority)
4390                    .push_bind(row.attempt)
4391                    .push_bind(row.run_lease)
4392                    .push_bind(row.lane_seq)
4393                    .push_bind(row.enqueue_shard)
4394                    .push_bind(row.attempted_at)
4395                    .push_bind(row.finalized_at)
4396                    .push_bind(terminal_storage_payload(&row.payload, ready_payload));
4397            });
4398            builder
4399                .build()
4400                .execute(tx.as_mut())
4401                .await
4402                .map_err(map_sqlx_error)?;
4403        }
4404
4405        if !synthetic.is_empty() {
4406            let mut builder = QueryBuilder::<Postgres>::new(format!(
4407                "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) "
4408            ));
4409            builder.push_values(synthetic, |mut b, row| {
4410                let ready_key = (
4411                    row.ready_slot,
4412                    row.ready_generation,
4413                    row.queue.as_str(),
4414                    row.priority,
4415                    row.enqueue_shard,
4416                    row.lane_seq,
4417                );
4418                let ready_payload = ready_payloads.get(&ready_key);
4419                b.push_bind(row.ready_slot)
4420                    .push_bind(row.ready_generation)
4421                    .push_bind(row.job_id)
4422                    .push_bind(&row.kind)
4423                    .push_bind(&row.queue)
4424                    .push_bind(&row.args)
4425                    .push_bind(row.state)
4426                    .push_bind(row.priority)
4427                    .push_bind(row.attempt)
4428                    .push_bind(row.run_lease)
4429                    .push_bind(row.max_attempts)
4430                    .push_bind(row.lane_seq)
4431                    .push_bind(row.enqueue_shard)
4432                    .push_bind(row.run_at)
4433                    .push_bind(row.attempted_at)
4434                    .push_bind(row.finalized_at)
4435                    .push_bind(row.created_at)
4436                    .push_bind(&row.unique_key)
4437                    .push_bind(&row.unique_states)
4438                    .push_bind(terminal_storage_payload(&row.payload, ready_payload));
4439            });
4440            builder
4441                .build()
4442                .execute(tx.as_mut())
4443                .await
4444                .map_err(map_sqlx_error)?;
4445        }
4446
4447        // Terminal counts stay exact without mutating the hot live-counter
4448        // rows on every completion. The hot path appends signed deltas; the
4449        // maintenance rollup folds them into queue_terminal_live_counts later.
4450        self.increment_live_terminal_counters_tx(tx, rows).await?;
4451
4452        Ok(rows.len())
4453    }
4454
4455    /// Append positive terminal-count deltas for newly inserted `done_entries`
4456    /// terminal rows.
4457    ///
4458    /// For done-entry terminal rows, the exact invariant is:
4459    ///
4460    /// `SUM(queue_terminal_live_counts) + SUM(queue_terminal_count_deltas)
4461    /// == count(*) FROM done_entries`.
4462    ///
4463    /// Compact receipt completions are already append-only batches, so
4464    /// `queue_counts_exact` counts retained `receipt_completion_batches`
4465    /// directly instead of adding another hot-path count-delta write.
4466    ///
4467    /// Maintenance asynchronously folds delta rows into the live-counter table
4468    /// and truncates the delta segment. The hot path therefore performs only
4469    /// append-only writes.
4470    async fn increment_live_terminal_counters_tx<'a>(
4471        &self,
4472        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4473        rows: &[DoneJobRow],
4474    ) -> Result<(), AwaError> {
4475        if rows.is_empty() {
4476            return Ok(());
4477        }
4478        let mut by_group: BTreeMap<TerminalCounterKey, i64> = BTreeMap::new();
4479        for row in rows {
4480            let key = (
4481                row.ready_slot,
4482                row.ready_generation,
4483                row.queue.clone(),
4484                row.priority,
4485                row.enqueue_shard,
4486                terminal_counter_bucket(row.job_id),
4487            );
4488            *by_group.entry(key).or_insert(0) += 1;
4489        }
4490        self.append_terminal_count_deltas_tx(tx, by_group).await
4491    }
4492
4493    /// Append negative terminal-count deltas for deleted `done_entries`
4494    /// terminal rows.
4495    ///
4496    /// This is used by retry-from-terminal, DLQ moves, discard paths, and the
4497    /// SQL compatibility delete function. The negative row may cancel a
4498    /// not-yet-rolled positive delta or reduce an already-folded live counter;
4499    /// exact reads sum both sources, so either order is correct.
4500    async fn decrement_live_terminal_counters_tx<'a>(
4501        &self,
4502        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4503        rows: &[TerminalCounterKey],
4504    ) -> Result<(), AwaError> {
4505        if rows.is_empty() {
4506            return Ok(());
4507        }
4508        let mut by_group: BTreeMap<TerminalCounterKey, i64> = BTreeMap::new();
4509        for key in rows {
4510            *by_group.entry(key.clone()).or_insert(0) -= 1;
4511        }
4512        self.append_terminal_count_deltas_tx(tx, by_group).await
4513    }
4514
4515    async fn append_terminal_count_deltas_tx<'a>(
4516        &self,
4517        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4518        by_group: BTreeMap<TerminalCounterKey, i64>,
4519    ) -> Result<(), AwaError> {
4520        if by_group.is_empty() {
4521            return Ok(());
4522        }
4523
4524        let mut ready_slots: Vec<i32> = Vec::with_capacity(by_group.len());
4525        let mut ready_generations: Vec<i64> = Vec::with_capacity(by_group.len());
4526        let mut queues: Vec<String> = Vec::with_capacity(by_group.len());
4527        let mut priorities: Vec<i16> = Vec::with_capacity(by_group.len());
4528        let mut enqueue_shards: Vec<i16> = Vec::with_capacity(by_group.len());
4529        let mut counter_buckets: Vec<i16> = Vec::with_capacity(by_group.len());
4530        let mut deltas: Vec<i64> = Vec::with_capacity(by_group.len());
4531        for ((slot, generation, queue, prio, shard, bucket), delta) in by_group {
4532            if delta == 0 {
4533                continue;
4534            }
4535            ready_slots.push(slot);
4536            ready_generations.push(generation);
4537            queues.push(queue);
4538            priorities.push(prio);
4539            enqueue_shards.push(shard);
4540            counter_buckets.push(bucket);
4541            deltas.push(delta);
4542        }
4543
4544        if deltas.is_empty() {
4545            return Ok(());
4546        }
4547
4548        let schema = self.schema();
4549        sqlx::query(&format!(
4550            r#"
4551            INSERT INTO {schema}.queue_terminal_count_deltas (
4552                ready_slot,
4553                ready_generation,
4554                queue,
4555                priority,
4556                enqueue_shard,
4557                counter_bucket,
4558                terminal_delta
4559            )
4560            SELECT
4561                ready_slot,
4562                ready_generation,
4563                queue,
4564                priority,
4565                enqueue_shard,
4566                counter_bucket,
4567                terminal_delta
4568            FROM unnest(
4569                $1::int[],
4570                $2::bigint[],
4571                $3::text[],
4572                $4::smallint[],
4573                $5::smallint[],
4574                $6::smallint[],
4575                $7::bigint[]
4576            ) AS d(
4577                ready_slot,
4578                ready_generation,
4579                queue,
4580                priority,
4581                enqueue_shard,
4582                counter_bucket,
4583                terminal_delta
4584            )
4585            ORDER BY
4586                ready_slot,
4587                ready_generation,
4588                queue,
4589                priority,
4590                enqueue_shard,
4591                counter_bucket
4592            "#
4593        ))
4594        .bind(&ready_slots)
4595        .bind(&ready_generations)
4596        .bind(&queues)
4597        .bind(&priorities)
4598        .bind(&enqueue_shards)
4599        .bind(&counter_buckets)
4600        .bind(&deltas)
4601        .execute(tx.as_mut())
4602        .await
4603        .map_err(map_sqlx_error)?;
4604        Ok(())
4605    }
4606
4607    /// Build the row-key vector for `decrement_live_terminal_counters_tx`
4608    /// from a slice of `DoneJobRow` (used by retry-from-terminal,
4609    /// DLQ move, and discard, all of which DELETE FROM done_entries
4610    /// with `RETURNING *` materialised into `DoneJobRow`).
4611    fn done_rows_to_counter_keys(rows: &[DoneJobRow]) -> Vec<TerminalCounterKey> {
4612        rows.iter()
4613            .map(|row| {
4614                (
4615                    row.ready_slot,
4616                    row.ready_generation,
4617                    row.queue.clone(),
4618                    row.priority,
4619                    row.enqueue_shard,
4620                    terminal_counter_bucket(row.job_id),
4621                )
4622            })
4623            .collect()
4624    }
4625
4626    async fn ensure_terminal_removed_receipt_closures_tx<'a>(
4627        &self,
4628        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4629        rows: &[DoneJobRow],
4630    ) -> Result<(), AwaError> {
4631        if rows.is_empty() || !self.lease_claim_receipts() {
4632            return Ok(());
4633        }
4634
4635        let receipt_pairs: Vec<(i64, i64)> =
4636            rows.iter().map(|row| (row.job_id, row.run_lease)).collect();
4637        self.lock_receipt_attempts_tx(tx, &receipt_pairs).await?;
4638
4639        let schema = self.schema();
4640        let job_ids: Vec<i64> = rows.iter().map(|row| row.job_id).collect();
4641        let run_leases: Vec<i64> = rows.iter().map(|row| row.run_lease).collect();
4642        sqlx::query(&format!(
4643            r#"
4644            WITH refs(job_id, run_lease) AS (
4645                SELECT * FROM unnest($1::bigint[], $2::bigint[])
4646            ),
4647            inserted AS (
4648                INSERT INTO {schema}.lease_claim_closures (
4649                    claim_slot,
4650                    job_id,
4651                    run_lease,
4652                    outcome,
4653                    closed_at
4654                )
4655                SELECT
4656                    claims.claim_slot,
4657                    claims.job_id,
4658                    claims.run_lease,
4659                    'terminal_removed',
4660                    clock_timestamp()
4661                FROM {schema}.lease_claims AS claims
4662                JOIN refs
4663                  ON refs.job_id = claims.job_id
4664                 AND refs.run_lease = claims.run_lease
4665                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
4666                RETURNING claim_slot, job_id, run_lease, closed_at
4667            ),
4668            marked AS (
4669                UPDATE {schema}.lease_claims AS claims
4670                SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
4671                FROM inserted
4672                WHERE claims.claim_slot = inserted.claim_slot
4673                  AND claims.job_id = inserted.job_id
4674                  AND claims.run_lease = inserted.run_lease
4675                RETURNING claims.job_id
4676            )
4677            SELECT count(*) FROM marked
4678            "#
4679        ))
4680        .bind(&job_ids)
4681        .bind(&run_leases)
4682        .execute(tx.as_mut())
4683        .await
4684        .map_err(map_sqlx_error)?;
4685        Ok(())
4686    }
4687
4688    async fn ready_payloads_for_done_rows_tx<'a, 'r>(
4689        &self,
4690        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4691        rows: &'r [DoneJobRow],
4692    ) -> Result<HashMap<(i32, i64, &'r str, i16, i16, i64), serde_json::Value>, AwaError> {
4693        if rows.is_empty() {
4694            return Ok(HashMap::new());
4695        }
4696
4697        let schema = self.schema();
4698        let ready_slots: Vec<i32> = rows.iter().map(|row| row.ready_slot).collect();
4699        let ready_generations: Vec<i64> = rows.iter().map(|row| row.ready_generation).collect();
4700        let queues: Vec<&str> = rows.iter().map(|row| row.queue.as_str()).collect();
4701        let priorities: Vec<i16> = rows.iter().map(|row| row.priority).collect();
4702        let enqueue_shards: Vec<i16> = rows.iter().map(|row| row.enqueue_shard).collect();
4703        let lane_seqs: Vec<i64> = rows.iter().map(|row| row.lane_seq).collect();
4704
4705        let payload_rows: Vec<(i32, i64, String, i16, i16, i64, serde_json::Value)> =
4706            sqlx::query_as(&format!(
4707                r#"
4708                WITH refs(ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq) AS (
4709                    SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::smallint[], $6::bigint[])
4710                )
4711                SELECT
4712                    ready.ready_slot,
4713                    ready.ready_generation,
4714                    ready.queue,
4715                    ready.priority,
4716                    ready.enqueue_shard,
4717                    ready.lane_seq,
4718                    COALESCE(ready.payload, '{{}}'::jsonb) AS payload
4719                FROM refs
4720                JOIN {schema}.ready_entries AS ready
4721                  ON ready.ready_slot = refs.ready_slot
4722                 AND ready.ready_generation = refs.ready_generation
4723                 AND ready.queue = refs.queue
4724                 AND ready.priority = refs.priority
4725                 AND ready.enqueue_shard = refs.enqueue_shard
4726                 AND ready.lane_seq = refs.lane_seq
4727                "#
4728            ))
4729            .bind(&ready_slots)
4730            .bind(&ready_generations)
4731            .bind(&queues)
4732            .bind(&priorities)
4733            .bind(&enqueue_shards)
4734            .bind(&lane_seqs)
4735            .fetch_all(tx.as_mut())
4736            .await
4737            .map_err(map_sqlx_error)?;
4738
4739        let mut payload_by_owned_key = HashMap::with_capacity(payload_rows.len());
4740        for (ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, payload) in
4741            payload_rows
4742        {
4743            payload_by_owned_key.insert(
4744                (
4745                    ready_slot,
4746                    ready_generation,
4747                    queue,
4748                    priority,
4749                    enqueue_shard,
4750                    lane_seq,
4751                ),
4752                payload,
4753            );
4754        }
4755
4756        let mut payloads = HashMap::with_capacity(payload_by_owned_key.len());
4757        for row in rows {
4758            if let Some(payload) = payload_by_owned_key.remove(&(
4759                row.ready_slot,
4760                row.ready_generation,
4761                row.queue.clone(),
4762                row.priority,
4763                row.enqueue_shard,
4764                row.lane_seq,
4765            )) {
4766                payloads.insert(
4767                    (
4768                        row.ready_slot,
4769                        row.ready_generation,
4770                        row.queue.as_str(),
4771                        row.priority,
4772                        row.enqueue_shard,
4773                        row.lane_seq,
4774                    ),
4775                    payload,
4776                );
4777            }
4778        }
4779        Ok(payloads)
4780    }
4781
4782    async fn insert_dlq_rows_tx<'a>(
4783        &self,
4784        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4785        rows: &[DlqJobRow],
4786        old_state: Option<JobState>,
4787    ) -> Result<usize, AwaError> {
4788        if rows.is_empty() {
4789            return Ok(0);
4790        }
4791
4792        for row in rows {
4793            self.sync_unique_claim(
4794                tx,
4795                row.job_id,
4796                &row.unique_key,
4797                row.unique_states.as_deref(),
4798                old_state,
4799                Some(JobState::Failed),
4800            )
4801            .await?;
4802        }
4803
4804        let schema = self.schema();
4805        let mut builder = QueryBuilder::<Postgres>::new(format!(
4806            "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) "
4807        ));
4808        builder.push_values(rows.iter(), |mut b, row| {
4809            b.push_bind(row.job_id)
4810                .push_bind(&row.kind)
4811                .push_bind(&row.queue)
4812                .push_bind(&row.args)
4813                .push_bind(row.state)
4814                .push_bind(row.priority)
4815                .push_bind(row.attempt)
4816                .push_bind(row.run_lease)
4817                .push_bind(row.max_attempts)
4818                .push_bind(row.run_at)
4819                .push_bind(row.attempted_at)
4820                .push_bind(row.finalized_at)
4821                .push_bind(row.created_at)
4822                .push_bind(&row.unique_key)
4823                .push_bind(&row.unique_states)
4824                .push_bind(storage_payload(&row.payload))
4825                .push_bind(&row.dlq_reason)
4826                .push_bind(row.dlq_at)
4827                .push_bind(row.original_run_lease);
4828        });
4829        builder
4830            .build()
4831            .execute(tx.as_mut())
4832            .await
4833            .map_err(map_sqlx_error)?;
4834
4835        Ok(rows.len())
4836    }
4837
4838    async fn adjust_terminal_rollups_batch<'a, I>(
4839        &self,
4840        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
4841        deltas: I,
4842    ) -> Result<(), AwaError>
4843    where
4844        I: IntoIterator<Item = (String, i16, i64, i64)>,
4845    {
4846        let mut grouped: BTreeMap<(String, i16), (i64, i64)> = BTreeMap::new();
4847        for (queue, priority, pruned_completed_delta, pruned_failed_delta) in deltas {
4848            if pruned_completed_delta == 0 && pruned_failed_delta == 0 {
4849                continue;
4850            }
4851            let entry = grouped.entry((queue, priority)).or_insert((0_i64, 0_i64));
4852            entry.0 += pruned_completed_delta;
4853            entry.1 += pruned_failed_delta;
4854        }
4855
4856        if grouped.is_empty() {
4857            return Ok(());
4858        }
4859
4860        let schema = self.schema();
4861        let mut queues = Vec::with_capacity(grouped.len());
4862        let mut priorities = Vec::with_capacity(grouped.len());
4863        let mut pruned_completed_deltas = Vec::with_capacity(grouped.len());
4864        let mut pruned_failed_deltas = Vec::with_capacity(grouped.len());
4865
4866        for ((queue, priority), (pruned_completed_delta, pruned_failed_delta)) in grouped {
4867            queues.push(queue);
4868            priorities.push(priority);
4869            pruned_completed_deltas.push(pruned_completed_delta);
4870            pruned_failed_deltas.push(pruned_failed_delta);
4871        }
4872
4873        sqlx::query(&format!(
4874            r#"
4875            WITH deltas(queue, priority, pruned_completed_delta, pruned_failed_delta) AS (
4876                SELECT *
4877                FROM unnest(
4878                    $1::text[],
4879                    $2::smallint[],
4880                    $3::bigint[],
4881                    $4::bigint[]
4882                )
4883            )
4884            INSERT INTO {schema}.queue_terminal_rollups AS rollups (
4885                queue,
4886                priority,
4887                pruned_completed_count,
4888                pruned_failed_count
4889            )
4890            SELECT
4891                deltas.queue,
4892                deltas.priority,
4893                deltas.pruned_completed_delta,
4894                deltas.pruned_failed_delta
4895            FROM deltas
4896            ON CONFLICT (queue, priority) DO UPDATE
4897            SET pruned_completed_count = GREATEST(
4898                    0,
4899                    rollups.pruned_completed_count + EXCLUDED.pruned_completed_count
4900                ),
4901                pruned_failed_count = GREATEST(
4902                    0,
4903                    rollups.pruned_failed_count + EXCLUDED.pruned_failed_count
4904                )
4905            "#
4906        ))
4907        .bind(&queues)
4908        .bind(&priorities)
4909        .bind(&pruned_completed_deltas)
4910        .bind(&pruned_failed_deltas)
4911        .execute(tx.as_mut())
4912        .await
4913        .map_err(map_sqlx_error)?;
4914        Ok(())
4915    }
4916
4917    async fn enqueue_runtime_rows(
4918        &self,
4919        pool: &PgPool,
4920        rows: Vec<RuntimeReadyRow>,
4921    ) -> Result<usize, AwaError> {
4922        if rows.is_empty() {
4923            return Ok(0);
4924        }
4925
4926        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
4927        let total_rows = self.insert_ready_rows_tx(&mut tx, rows.clone()).await?;
4928
4929        let queues_to_notify: Vec<String> = rows.iter().map(|row| row.queue.clone()).collect();
4930        self.notify_queues_tx(&mut tx, queues_to_notify).await?;
4931
4932        tx.commit().await.map_err(map_sqlx_error)?;
4933        Ok(total_rows)
4934    }
4935
4936    pub async fn enqueue_batch(
4937        &self,
4938        pool: &PgPool,
4939        queue: &str,
4940        priority: i16,
4941        count: i64,
4942    ) -> Result<i64, AwaError> {
4943        if count <= 0 {
4944            return Ok(0);
4945        }
4946
4947        let rows: Vec<_> = (0..count)
4948            .map(|seq| RuntimeReadyRow {
4949                kind: "bench_job".to_string(),
4950                queue: if self.uses_queue_striping() && !self.is_physical_stripe_queue(queue) {
4951                    self.physical_queue_for_stripe(
4952                        queue,
4953                        seq.rem_euclid(self.queue_stripe_count() as i64) as usize,
4954                    )
4955                } else {
4956                    queue.to_string()
4957                },
4958                args: serde_json::json!({ "seq": seq }),
4959                priority,
4960                attempt: 0,
4961                run_lease: 0,
4962                max_attempts: 25,
4963                run_at: Utc::now(),
4964                attempted_at: None,
4965                created_at: Utc::now(),
4966                unique_key: None,
4967                unique_states: None,
4968                payload: RuntimePayload::default().into_json(),
4969                ordering_key: None,
4970            })
4971            .collect();
4972        self.enqueue_runtime_rows(pool, rows)
4973            .await
4974            .map(|count| count as i64)
4975    }
4976
4977    pub async fn enqueue_params_batch(
4978        &self,
4979        pool: &PgPool,
4980        jobs: &[InsertParams],
4981    ) -> Result<usize, AwaError> {
4982        if jobs.is_empty() {
4983            return Ok(0);
4984        }
4985
4986        let now = Utc::now();
4987        let mut ready_rows = Vec::new();
4988        let mut deferred_rows = Vec::new();
4989
4990        for (idx, job) in jobs.iter().enumerate() {
4991            let prepared = prepare_row_raw(job.kind.clone(), job.args.clone(), job.opts.clone())?;
4992            let payload = Self::payload_from_parts(prepared.metadata, prepared.tags, None, None)?;
4993            let queue =
4994                self.queue_stripe_for_enqueue(&prepared.queue, &prepared.unique_key, idx as i64);
4995
4996            let ready_row = RuntimeReadyRow {
4997                kind: prepared.kind,
4998                queue: queue.clone(),
4999                args: prepared.args,
5000                priority: prepared.priority,
5001                attempt: 0,
5002                run_lease: 0,
5003                max_attempts: prepared.max_attempts,
5004                run_at: prepared.run_at.unwrap_or(now),
5005                attempted_at: None,
5006                created_at: now,
5007                unique_key: prepared.unique_key,
5008                unique_states: prepared.unique_states,
5009                payload: payload.clone(),
5010                ordering_key: prepared.ordering_key,
5011            };
5012
5013            match prepared.state {
5014                JobState::Available => ready_rows.push(ready_row),
5015                JobState::Scheduled => deferred_rows.push(DeferredJobRow {
5016                    job_id: 0,
5017                    kind: ready_row.kind,
5018                    queue,
5019                    args: ready_row.args,
5020                    state: JobState::Scheduled,
5021                    priority: ready_row.priority,
5022                    attempt: ready_row.attempt,
5023                    run_lease: ready_row.run_lease,
5024                    max_attempts: ready_row.max_attempts,
5025                    run_at: ready_row.run_at,
5026                    attempted_at: ready_row.attempted_at,
5027                    finalized_at: None,
5028                    created_at: ready_row.created_at,
5029                    unique_key: ready_row.unique_key,
5030                    unique_states: ready_row.unique_states,
5031                    payload: payload.clone(),
5032                }),
5033                other => {
5034                    return Err(AwaError::Validation(format!(
5035                        "queue storage does not support initial state {other}"
5036                    )));
5037                }
5038            }
5039        }
5040
5041        let queues_to_notify: Vec<String> =
5042            ready_rows.iter().map(|row| row.queue.clone()).collect();
5043
5044        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5045        let mut total = 0usize;
5046        if !ready_rows.is_empty() {
5047            total += self
5048                .insert_ready_rows_tx(&mut tx, ready_rows.clone())
5049                .await?;
5050        }
5051        if !deferred_rows.is_empty() {
5052            let ids = self.next_job_ids(&mut tx, deferred_rows.len()).await?;
5053            let deferred_rows: Vec<_> = deferred_rows
5054                .into_iter()
5055                .zip(ids)
5056                .map(|(row, id)| DeferredJobRow { job_id: id, ..row })
5057                .collect();
5058            total += self
5059                .insert_deferred_rows_tx(&mut tx, deferred_rows, None)
5060                .await?;
5061        }
5062
5063        self.notify_queues_tx(&mut tx, queues_to_notify).await?;
5064
5065        tx.commit().await.map_err(map_sqlx_error)?;
5066        Ok(total)
5067    }
5068
5069    /// Enqueue prepared jobs into queue storage using PostgreSQL COPY.
5070    ///
5071    /// This follows the same preparation, queue striping, lane sequencing,
5072    /// uniqueness, and notification semantics as [`Self::enqueue_params_batch`],
5073    /// but streams materialized rows directly into `ready_entries` and
5074    /// `deferred_jobs` instead of building multi-row `INSERT` statements.
5075    #[tracing::instrument(skip(self, pool, jobs), fields(job.count = jobs.len()), name = "queue_storage.enqueue_params_copy")]
5076    pub async fn enqueue_params_copy(
5077        &self,
5078        pool: &PgPool,
5079        jobs: &[InsertParams],
5080    ) -> Result<usize, AwaError> {
5081        if jobs.is_empty() {
5082            return Ok(0);
5083        }
5084
5085        let now = Utc::now();
5086        let mut ready_rows = Vec::new();
5087        let mut deferred_rows = Vec::new();
5088
5089        for (idx, job) in jobs.iter().enumerate() {
5090            let prepared = prepare_row_raw(job.kind.clone(), job.args.clone(), job.opts.clone())?;
5091            let payload = Self::payload_from_parts(prepared.metadata, prepared.tags, None, None)?;
5092            let queue =
5093                self.queue_stripe_for_enqueue(&prepared.queue, &prepared.unique_key, idx as i64);
5094
5095            let ready_row = RuntimeReadyRow {
5096                kind: prepared.kind,
5097                queue: queue.clone(),
5098                args: prepared.args,
5099                priority: prepared.priority,
5100                attempt: 0,
5101                run_lease: 0,
5102                max_attempts: prepared.max_attempts,
5103                run_at: prepared.run_at.unwrap_or(now),
5104                attempted_at: None,
5105                created_at: now,
5106                unique_key: prepared.unique_key,
5107                unique_states: prepared.unique_states,
5108                payload: payload.clone(),
5109                ordering_key: prepared.ordering_key,
5110            };
5111
5112            match prepared.state {
5113                JobState::Available => ready_rows.push(ready_row),
5114                JobState::Scheduled => deferred_rows.push(DeferredJobRow {
5115                    job_id: 0,
5116                    kind: ready_row.kind,
5117                    queue,
5118                    args: ready_row.args,
5119                    state: JobState::Scheduled,
5120                    priority: ready_row.priority,
5121                    attempt: ready_row.attempt,
5122                    run_lease: ready_row.run_lease,
5123                    max_attempts: ready_row.max_attempts,
5124                    run_at: ready_row.run_at,
5125                    attempted_at: ready_row.attempted_at,
5126                    finalized_at: None,
5127                    created_at: ready_row.created_at,
5128                    unique_key: ready_row.unique_key,
5129                    unique_states: ready_row.unique_states,
5130                    payload: payload.clone(),
5131                }),
5132                other => {
5133                    return Err(AwaError::Validation(format!(
5134                        "queue storage does not support initial state {other}"
5135                    )));
5136                }
5137            }
5138        }
5139
5140        let queues_to_notify: Vec<String> =
5141            ready_rows.iter().map(|row| row.queue.clone()).collect();
5142
5143        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5144        let mut total = 0usize;
5145        let job_ids = self
5146            .next_job_ids(&mut tx, ready_rows.len() + deferred_rows.len())
5147            .await?;
5148        let (ready_job_ids, deferred_job_ids) = job_ids.split_at(ready_rows.len());
5149        if !ready_rows.is_empty() {
5150            total += self
5151                .insert_ready_rows_copy_tx(&mut tx, ready_rows, ready_job_ids.to_vec())
5152                .await?;
5153        }
5154        if !deferred_rows.is_empty() {
5155            let deferred_rows: Vec<_> = deferred_rows
5156                .into_iter()
5157                .zip(deferred_job_ids.iter().copied())
5158                .map(|(row, id)| DeferredJobRow { job_id: id, ..row })
5159                .collect();
5160            total += self
5161                .insert_deferred_rows_copy_tx(&mut tx, deferred_rows)
5162                .await?;
5163        }
5164
5165        self.notify_queues_tx(&mut tx, queues_to_notify).await?;
5166
5167        tx.commit().await.map_err(map_sqlx_error)?;
5168        Ok(total)
5169    }
5170
5171    #[tracing::instrument(skip(self, pool), name = "queue_storage.claim_batch")]
5172    pub async fn claim_batch(
5173        &self,
5174        pool: &PgPool,
5175        queue: &str,
5176        max_batch: i64,
5177    ) -> Result<Vec<ClaimedEntry>, AwaError> {
5178        if max_batch <= 0 {
5179            return Ok(Vec::new());
5180        }
5181
5182        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5183        let mut claimed_rows = Vec::new();
5184        let stripe_queues = self.physical_queues_for_logical(queue);
5185        let start = self.stripe_probe_start(stripe_queues.len());
5186        for offset in 0..stripe_queues.len() {
5187            if claimed_rows.len() >= max_batch as usize {
5188                break;
5189            }
5190            let stripe_queue = &stripe_queues[(start + offset) % stripe_queues.len()];
5191            let remaining = max_batch - claimed_rows.len() as i64;
5192            claimed_rows.extend(
5193                self.claim_ready_rows_tx(
5194                    &mut tx,
5195                    stripe_queue,
5196                    remaining,
5197                    Duration::ZERO,
5198                    Duration::ZERO,
5199                )
5200                .await?,
5201            );
5202        }
5203        let claim_cursor_advances = Self::claim_cursor_advances(&claimed_rows);
5204        let claimed = claimed_rows
5205            .into_iter()
5206            .map(|row| row.claim_ref(self.lease_claim_receipts()))
5207            .collect();
5208
5209        tx.commit().await.map_err(map_sqlx_error)?;
5210        self.advance_claim_cursors(pool, &claim_cursor_advances)
5211            .await;
5212        Ok(claimed)
5213    }
5214
5215    #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.claim_runtime_batch")]
5216    pub async fn claim_runtime_batch(
5217        &self,
5218        pool: &PgPool,
5219        queue: &str,
5220        max_batch: i64,
5221        deadline_duration: Duration,
5222    ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5223        self.claim_runtime_batch_with_aging(
5224            pool,
5225            queue,
5226            max_batch,
5227            deadline_duration,
5228            Duration::ZERO,
5229        )
5230        .await
5231    }
5232
5233    #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.claim_runtime_batch_with_aging")]
5234    pub async fn claim_runtime_batch_with_aging(
5235        &self,
5236        pool: &PgPool,
5237        queue: &str,
5238        max_batch: i64,
5239        deadline_duration: Duration,
5240        aging_interval: Duration,
5241    ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5242        if max_batch <= 0 {
5243            return Ok(Vec::new());
5244        }
5245
5246        let stripe_queues = self.physical_queues_for_logical(queue);
5247        if stripe_queues.len() > 1 {
5248            let mut claimed = Vec::new();
5249            let start = self.stripe_probe_start(stripe_queues.len());
5250            for offset in 0..stripe_queues.len() {
5251                if claimed.len() >= max_batch as usize {
5252                    break;
5253                }
5254                let stripe_queue = &stripe_queues[(start + offset) % stripe_queues.len()];
5255                let remaining = max_batch - claimed.len() as i64;
5256                match self
5257                    .claim_runtime_batch_with_aging_physical(
5258                        pool,
5259                        stripe_queue,
5260                        remaining,
5261                        deadline_duration,
5262                        aging_interval,
5263                    )
5264                    .await
5265                {
5266                    Ok(stripe_claims) => claimed.extend(stripe_claims),
5267                    Err(err) if claimed.is_empty() => return Err(err),
5268                    Err(err) => {
5269                        tracing::warn!(
5270                            queue = %queue,
5271                            stripe_queue = %stripe_queue,
5272                            claimed = claimed.len(),
5273                            error = ?err,
5274                            "returning already-claimed runtime jobs after striped claim error"
5275                        );
5276                        break;
5277                    }
5278                }
5279            }
5280            return Ok(claimed);
5281        }
5282
5283        self.claim_runtime_batch_with_aging_physical(
5284            pool,
5285            &stripe_queues[0],
5286            max_batch,
5287            deadline_duration,
5288            aging_interval,
5289        )
5290        .await
5291    }
5292
5293    async fn claim_runtime_batch_with_aging_physical(
5294        &self,
5295        pool: &PgPool,
5296        queue: &str,
5297        max_batch: i64,
5298        deadline_duration: Duration,
5299        aging_interval: Duration,
5300    ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5301        if max_batch <= 0 {
5302            return Ok(Vec::new());
5303        }
5304
5305        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5306        let mut claimed = Vec::new();
5307        claimed.extend(
5308            self.claim_ready_rows_tx(&mut tx, queue, max_batch, deadline_duration, aging_interval)
5309                .await?,
5310        );
5311
5312        for row in &claimed {
5313            self.sync_unique_claim(
5314                &mut tx,
5315                row.job_id,
5316                &row.unique_key,
5317                row.unique_states.as_deref(),
5318                Some(JobState::Available),
5319                Some(JobState::Running),
5320            )
5321            .await?;
5322        }
5323
5324        let use_lease_claim_receipts = self.use_lease_claim_receipts_for_runtime(deadline_duration);
5325        if !use_lease_claim_receipts && deadline_duration.is_zero() {
5326            // Legacy zero-deadline claims have no heartbeat/deadline rescue
5327            // timestamp, so a post-commit conversion error would strand the
5328            // materialized lease indefinitely. Keep the old rollback semantics
5329            // for that path.
5330            let converted = claimed
5331                .iter()
5332                .cloned()
5333                .map(|row| row.into_claimed_runtime_job(use_lease_claim_receipts))
5334                .collect::<Result<Vec<_>, _>>()?;
5335            let claim_cursor_advances = Self::claim_cursor_advances(&claimed);
5336            tx.commit().await.map_err(map_sqlx_error)?;
5337            self.advance_claim_cursors(pool, &claim_cursor_advances)
5338                .await;
5339            return Ok(converted);
5340        }
5341
5342        let claim_cursor_advances = Self::claim_cursor_advances(&claimed);
5343
5344        // Release claim locks before doing Rust-side payload conversion; this
5345        // keeps the hot claim transaction focused on database state changes.
5346        tx.commit().await.map_err(map_sqlx_error)?;
5347        self.advance_claim_cursors(pool, &claim_cursor_advances)
5348            .await;
5349
5350        claimed
5351            .into_iter()
5352            .map(|row| row.into_claimed_runtime_job(use_lease_claim_receipts))
5353            .collect()
5354    }
5355
5356    #[tracing::instrument(skip(self, pool), fields(queue = %queue, instance_id = %instance_id), name = "queue_storage.acquire_queue_claimer")]
5357    pub async fn acquire_queue_claimer(
5358        &self,
5359        pool: &PgPool,
5360        queue: &str,
5361        instance_id: Uuid,
5362        max_claimers: i16,
5363        lease_ttl: Duration,
5364        idle_threshold: Duration,
5365    ) -> Result<Option<QueueClaimerLease>, AwaError> {
5366        Ok(self
5367            .acquire_queue_claimer_row(
5368                pool,
5369                queue,
5370                instance_id,
5371                max_claimers,
5372                lease_ttl,
5373                idle_threshold,
5374            )
5375            .await?
5376            .map(QueueClaimerLeaseRow::lease))
5377    }
5378
5379    async fn acquire_queue_claimer_row(
5380        &self,
5381        pool: &PgPool,
5382        queue: &str,
5383        instance_id: Uuid,
5384        max_claimers: i16,
5385        lease_ttl: Duration,
5386        idle_threshold: Duration,
5387    ) -> Result<Option<QueueClaimerLeaseRow>, AwaError> {
5388        if max_claimers <= 0 {
5389            return Ok(None);
5390        }
5391
5392        let schema = self.schema();
5393        let now = Utc::now();
5394        let expires_at = now
5395            + TimeDelta::from_std(lease_ttl)
5396                .map_err(|err| AwaError::Validation(format!("invalid claimer lease ttl: {err}")))?;
5397        let idle_cutoff = now
5398            - TimeDelta::from_std(idle_threshold).map_err(|err| {
5399                AwaError::Validation(format!("invalid claimer idle threshold: {err}"))
5400            })?;
5401        let probe_start = if max_claimers > 1 {
5402            ((instance_id.as_u128() ^ (now.timestamp_millis() as u128)) % (max_claimers as u128))
5403                as i16
5404        } else {
5405            0
5406        };
5407
5408        if let Some(owned) = sqlx::query_as::<_, QueueClaimerLeaseRow>(&format!(
5409            r#"
5410            SELECT claimer_slot, lease_epoch, last_claimed_at, expires_at
5411            FROM {schema}.queue_claimer_leases
5412            WHERE queue = $1
5413              AND owner_instance_id = $2
5414              AND expires_at > $3
5415            ORDER BY claimer_slot
5416            LIMIT 1
5417            "#
5418        ))
5419        .bind(queue)
5420        .bind(instance_id)
5421        .bind(now)
5422        .fetch_optional(pool)
5423        .await
5424        .map_err(map_sqlx_error)?
5425        {
5426            return Ok(Some(owned));
5427        }
5428
5429        for offset in 0..max_claimers {
5430            let slot = (probe_start + offset) % max_claimers;
5431            if let Some(updated) = sqlx::query_as::<_, QueueClaimerLeaseRow>(&format!(
5432                r#"
5433                UPDATE {schema}.queue_claimer_leases
5434                SET owner_instance_id = $3,
5435                    lease_epoch = CASE
5436                        WHEN owner_instance_id = $3 THEN lease_epoch
5437                        ELSE lease_epoch + 1
5438                    END,
5439                    leased_at = $4,
5440                    last_claimed_at = $4,
5441                    expires_at = $5
5442                WHERE queue = $1
5443                  AND claimer_slot = $2
5444                  AND (
5445                        owner_instance_id = $3
5446                     OR expires_at <= $4
5447                     OR last_claimed_at <= $6
5448                  )
5449                RETURNING claimer_slot, lease_epoch, last_claimed_at, expires_at
5450                "#
5451            ))
5452            .bind(queue)
5453            .bind(slot)
5454            .bind(instance_id)
5455            .bind(now)
5456            .bind(expires_at)
5457            .bind(idle_cutoff)
5458            .fetch_optional(pool)
5459            .await
5460            .map_err(map_sqlx_error)?
5461            {
5462                return Ok(Some(updated));
5463            }
5464
5465            if let Some(inserted) = sqlx::query_as::<_, QueueClaimerLeaseRow>(&format!(
5466                r#"
5467                INSERT INTO {schema}.queue_claimer_leases (
5468                    queue,
5469                    claimer_slot,
5470                    owner_instance_id,
5471                    lease_epoch,
5472                    leased_at,
5473                    last_claimed_at,
5474                    expires_at
5475                )
5476                VALUES ($1, $2, $3, 0, $4, $4, $5)
5477                ON CONFLICT (queue, claimer_slot) DO NOTHING
5478                RETURNING claimer_slot, lease_epoch, last_claimed_at, expires_at
5479                "#
5480            ))
5481            .bind(queue)
5482            .bind(slot)
5483            .bind(instance_id)
5484            .bind(now)
5485            .bind(expires_at)
5486            .fetch_optional(pool)
5487            .await
5488            .map_err(map_sqlx_error)?
5489            {
5490                return Ok(Some(inserted));
5491            }
5492        }
5493
5494        Ok(None)
5495    }
5496
5497    #[tracing::instrument(skip(self, pool), fields(queue = %queue, instance_id = %instance_id, claimer_slot = lease.claimer_slot), name = "queue_storage.mark_queue_claimer_active")]
5498    pub async fn mark_queue_claimer_active(
5499        &self,
5500        pool: &PgPool,
5501        queue: &str,
5502        instance_id: Uuid,
5503        lease: QueueClaimerLease,
5504        lease_ttl: Duration,
5505    ) -> Result<bool, AwaError> {
5506        let schema = self.schema();
5507        let now = Utc::now();
5508        let expires_at = now
5509            + TimeDelta::from_std(lease_ttl)
5510                .map_err(|err| AwaError::Validation(format!("invalid claimer lease ttl: {err}")))?;
5511
5512        let result = sqlx::query(&format!(
5513            r#"
5514            UPDATE {schema}.queue_claimer_leases
5515            SET last_claimed_at = $5,
5516                expires_at = $6
5517            WHERE queue = $1
5518              AND claimer_slot = $2
5519              AND owner_instance_id = $3
5520              AND lease_epoch = $4
5521            "#
5522        ))
5523        .bind(queue)
5524        .bind(lease.claimer_slot)
5525        .bind(instance_id)
5526        .bind(lease.lease_epoch)
5527        .bind(now)
5528        .bind(expires_at)
5529        .execute(pool)
5530        .await
5531        .map_err(map_sqlx_error)?;
5532
5533        Ok(result.rows_affected() == 1)
5534    }
5535
5536    fn desired_queue_claimer_target(
5537        &self,
5538        current_target: Option<i16>,
5539        signal: &AvailableSignal,
5540        max_claimers: i16,
5541    ) -> i16 {
5542        // The signal source counts sequence positions reserved for enqueue
5543        // but not yet claimed. It can over-count deleted or uncommitted
5544        // positions, but already-claimed rows are excluded by the claim cursor
5545        // advance.
5546        // `backlog` is retained as a separate name in the threshold
5547        // table to keep room for shape tweaks that diverge from
5548        // `available` later.
5549        let available = signal.available.max(0) as u64;
5550        let backlog = available;
5551        let current = current_target.unwrap_or(1).clamp(1, max_claimers.max(1));
5552        let max_four = 4.min(max_claimers.max(1));
5553        let max_two = 2.min(max_claimers.max(1));
5554
5555        match current {
5556            4.. => {
5557                if available >= 32 || backlog >= 16 {
5558                    max_four
5559                } else if available >= 8 || backlog >= 4 {
5560                    max_two
5561                } else {
5562                    1
5563                }
5564            }
5565            2..=3 => {
5566                if available >= 128 || backlog >= 64 {
5567                    max_four
5568                } else if available >= 4 || backlog >= 2 {
5569                    max_two
5570                } else {
5571                    1
5572                }
5573            }
5574            _ => {
5575                if available >= 64 || backlog >= 32 {
5576                    max_four
5577                } else if available >= 8 || backlog >= 4 {
5578                    max_two
5579                } else {
5580                    1
5581                }
5582            }
5583        }
5584    }
5585
5586    async fn queue_claimer_target(
5587        &self,
5588        pool: &PgPool,
5589        queue: &str,
5590        max_claimers: i16,
5591        control_interval: Duration,
5592    ) -> Result<i16, AwaError> {
5593        let schema = self.schema();
5594        let now = Utc::now();
5595        let stale_cutoff = now
5596            - TimeDelta::from_std(control_interval).map_err(|err| {
5597                AwaError::Validation(format!("invalid claimer control interval: {err}"))
5598            })?;
5599
5600        if let Some(target) = sqlx::query_scalar::<_, i16>(&format!(
5601            r#"
5602            SELECT target_claimers
5603            FROM {schema}.queue_claimer_state
5604            WHERE queue = $1
5605              AND updated_at > $2
5606            "#
5607        ))
5608        .bind(queue)
5609        .bind(stale_cutoff)
5610        .fetch_optional(pool)
5611        .await
5612        .map_err(map_sqlx_error)?
5613        {
5614            return Ok(target.clamp(1, max_claimers.max(1)));
5615        }
5616
5617        let current_target = sqlx::query_scalar::<_, i16>(&format!(
5618            r#"
5619            SELECT target_claimers
5620            FROM {schema}.queue_claimer_state
5621            WHERE queue = $1
5622            "#
5623        ))
5624        .bind(queue)
5625        .fetch_optional(pool)
5626        .await
5627        .map_err(map_sqlx_error)?;
5628
5629        let signal = self.queue_claimer_signal(pool, queue).await?;
5630        let desired = self.desired_queue_claimer_target(current_target, &signal, max_claimers);
5631
5632        if let Some(updated) = sqlx::query_scalar::<_, i16>(&format!(
5633            r#"
5634            INSERT INTO {schema}.queue_claimer_state (queue, target_claimers, updated_at)
5635            VALUES ($1, $2, $3)
5636            ON CONFLICT (queue) DO UPDATE
5637            SET target_claimers = EXCLUDED.target_claimers,
5638                updated_at = EXCLUDED.updated_at
5639            WHERE {schema}.queue_claimer_state.updated_at <= $4
5640            RETURNING target_claimers
5641            "#
5642        ))
5643        .bind(queue)
5644        .bind(desired)
5645        .bind(now)
5646        .bind(stale_cutoff)
5647        .fetch_optional(pool)
5648        .await
5649        .map_err(map_sqlx_error)?
5650        {
5651            return Ok(updated.clamp(1, max_claimers.max(1)));
5652        }
5653
5654        Ok(current_target
5655            .unwrap_or(desired)
5656            .clamp(1, max_claimers.max(1)))
5657    }
5658
5659    /// Cheap, dispatcher-grade available-count signal.
5660    ///
5661    /// Sums `enqueue_cursor - claim_cursor` across the queue's
5662    /// (queue, priority) lanes — one PK read into each of the two head tables per lane
5663    /// (typically ≤ 4 lanes per logical queue). The difference is an
5664    /// upper bound on the count of unclaimed ready rows: admin DELETEs
5665    /// of unclaimed jobs and in-flight enqueue reservations can leave a
5666    /// gap between `claim_cursor` and `enqueue_cursor`. The dispatcher
5667    /// tolerates this drift because the worst case is a wasted claim
5668    /// attempt that finds no committed rows.
5669    ///
5670    /// For an exact count, use [`Self::queue_counts_exact`], which
5671    /// scans `ready_entries`.
5672    async fn queue_claimer_signal(
5673        &self,
5674        pool: &PgPool,
5675        queue: &str,
5676    ) -> Result<AvailableSignal, AwaError> {
5677        let schema = self.schema();
5678        let queues = self.physical_queues_for_logical(queue);
5679        let available: i64 = sqlx::query_scalar(&format!(
5680            r#"
5681            SELECT COALESCE(
5682                sum(GREATEST(
5683                    {schema}.sequence_next_value(qe.seq_name)
5684                        - {schema}.sequence_next_value(qc.seq_name),
5685                    0
5686                )),
5687                0
5688            )::bigint
5689            FROM {schema}.queue_enqueue_heads AS qe
5690            JOIN {schema}.queue_claim_heads AS qc
5691              ON qc.queue = qe.queue
5692             AND qc.priority = qe.priority
5693             AND qc.enqueue_shard = qe.enqueue_shard
5694            WHERE qe.queue = ANY($1)
5695            "#
5696        ))
5697        .bind(&queues)
5698        .fetch_one(pool)
5699        .await
5700        .map_err(map_sqlx_error)?;
5701
5702        Ok(AvailableSignal { available })
5703    }
5704
5705    #[allow(clippy::too_many_arguments)]
5706    #[tracing::instrument(skip(self, pool), fields(queue = %queue, instance_id = %instance_id), name = "queue_storage.claim_runtime_batch_with_aging_for_instance")]
5707    pub async fn claim_runtime_batch_with_aging_for_instance(
5708        &self,
5709        pool: &PgPool,
5710        queue: &str,
5711        max_batch: i64,
5712        deadline_duration: Duration,
5713        aging_interval: Duration,
5714        instance_id: Uuid,
5715        max_claimers: i16,
5716        lease_ttl: Duration,
5717        idle_threshold: Duration,
5718    ) -> Result<Vec<ClaimedRuntimeJob>, AwaError> {
5719        let target_claimers = self
5720            .queue_claimer_target(pool, queue, max_claimers, Duration::from_millis(500))
5721            .await?;
5722
5723        let Some(lease) = self
5724            .acquire_queue_claimer_row(
5725                pool,
5726                queue,
5727                instance_id,
5728                target_claimers,
5729                lease_ttl,
5730                idle_threshold,
5731            )
5732            .await?
5733        else {
5734            return Ok(Vec::new());
5735        };
5736
5737        let claimed = self
5738            .claim_runtime_batch_with_aging(
5739                pool,
5740                queue,
5741                max_batch,
5742                deadline_duration,
5743                aging_interval,
5744            )
5745            .await?;
5746
5747        if !claimed.is_empty() && lease.needs_refresh(Utc::now(), lease_ttl, idle_threshold) {
5748            let _ = self
5749                .mark_queue_claimer_active(pool, queue, instance_id, lease.lease(), lease_ttl)
5750                .await?;
5751        }
5752
5753        Ok(claimed)
5754    }
5755
5756    #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.claim_job_batch")]
5757    pub async fn claim_job_batch(
5758        &self,
5759        pool: &PgPool,
5760        queue: &str,
5761        max_batch: i64,
5762        deadline_duration: Duration,
5763    ) -> Result<Vec<JobRow>, AwaError> {
5764        self.claim_runtime_batch(pool, queue, max_batch, deadline_duration)
5765            .await
5766            .map(|claimed| claimed.into_iter().map(|row| row.job).collect())
5767    }
5768
5769    #[tracing::instrument(skip(self, pool, claimed), name = "queue_storage.complete_batch")]
5770    pub async fn complete_batch(
5771        &self,
5772        pool: &PgPool,
5773        claimed: &[ClaimedEntry],
5774    ) -> Result<usize, AwaError> {
5775        self.complete_claimed_batch(pool, claimed)
5776            .await
5777            .map(|updated| updated.len())
5778    }
5779
5780    #[tracing::instrument(
5781        skip(self, pool, claimed),
5782        name = "queue_storage.complete_claimed_batch"
5783    )]
5784    pub async fn complete_claimed_batch(
5785        &self,
5786        pool: &PgPool,
5787        claimed: &[ClaimedEntry],
5788    ) -> Result<Vec<(i64, i64)>, AwaError> {
5789        if claimed.is_empty() {
5790            return Ok(Vec::new());
5791        }
5792
5793        let schema = self.schema();
5794        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5795
5796        let lease_slots: Vec<i32> = claimed.iter().map(|entry| entry.lease_slot).collect();
5797        let queues: Vec<String> = claimed.iter().map(|entry| entry.queue.clone()).collect();
5798        let priorities: Vec<i16> = claimed.iter().map(|entry| entry.priority).collect();
5799        let enqueue_shards: Vec<i16> = claimed.iter().map(|entry| entry.enqueue_shard).collect();
5800        let lane_seqs: Vec<i64> = claimed.iter().map(|entry| entry.lane_seq).collect();
5801
5802        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
5803            r#"
5804            WITH completed(lease_slot, queue, priority, enqueue_shard, lane_seq) AS (
5805                SELECT * FROM unnest($1::int[], $2::text[], $3::smallint[], $4::smallint[], $5::bigint[])
5806            )
5807            DELETE FROM {schema}.leases AS leases
5808            USING completed
5809            WHERE leases.lease_slot = completed.lease_slot
5810              AND leases.queue = completed.queue
5811              AND leases.priority = completed.priority
5812              AND leases.enqueue_shard = completed.enqueue_shard
5813              AND leases.lane_seq = completed.lane_seq
5814            RETURNING
5815                leases.ready_slot,
5816                leases.ready_generation,
5817                leases.job_id,
5818                leases.queue,
5819                leases.state,
5820                leases.priority,
5821                leases.attempt,
5822                leases.run_lease,
5823                leases.max_attempts,
5824                leases.lane_seq,
5825                leases.enqueue_shard,
5826                leases.heartbeat_at,
5827                leases.deadline_at,
5828                leases.attempted_at,
5829                leases.callback_id,
5830                leases.callback_timeout_at
5831            "#
5832        ))
5833        .bind(&lease_slots)
5834        .bind(&queues)
5835        .bind(&priorities)
5836        .bind(&enqueue_shards)
5837        .bind(&lane_seqs)
5838        .fetch_all(tx.as_mut())
5839        .await
5840        .map_err(map_sqlx_error)?;
5841
5842        if deleted.is_empty() {
5843            tx.commit().await.map_err(map_sqlx_error)?;
5844            return Ok(Vec::new());
5845        }
5846
5847        let completed_pairs: Vec<(i64, i64)> = deleted
5848            .iter()
5849            .map(|row| (row.job_id, row.run_lease))
5850            .collect();
5851        self.close_receipt_pairs_tx(&mut tx, &completed_pairs, "completed")
5852            .await?;
5853
5854        let moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
5855
5856        let finalized_at = Utc::now();
5857        let mut done_rows = Vec::with_capacity(moved.len());
5858        for entry in moved.iter().cloned() {
5859            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
5860                entry.payload.clone(),
5861                entry.progress.clone(),
5862            )?)?;
5863            payload.set_progress(None);
5864            done_rows.push(entry.into_done_row(
5865                JobState::Completed,
5866                finalized_at,
5867                payload.into_json(),
5868            ));
5869        }
5870
5871        self.insert_done_rows_tx(&mut tx, &done_rows, Some(JobState::Running))
5872            .await?;
5873        tx.commit().await.map_err(map_sqlx_error)?;
5874        Ok(moved
5875            .into_iter()
5876            .map(|entry| (entry.job_id, entry.run_lease))
5877            .collect())
5878    }
5879
5880    fn receipt_fast_complete_candidate(entry: &ClaimedRuntimeJob) -> bool {
5881        entry.claim.lease_claim_receipt
5882            && entry.claim.receipt_id.is_some()
5883            && entry.job.unique_key.is_none()
5884            && is_compact_receipt_completion_metadata(&entry.job.metadata)
5885            && entry.job.tags.is_empty()
5886            && entry.job.errors.as_ref().is_none_or(Vec::is_empty)
5887    }
5888
5889    async fn complete_receipt_runtime_batch_fast(
5890        &self,
5891        pool: &PgPool,
5892        claimed: &[ClaimedRuntimeJob],
5893    ) -> Result<Vec<(i64, i64)>, AwaError> {
5894        if claimed.is_empty() {
5895            return Ok(Vec::new());
5896        }
5897
5898        let schema = self.schema();
5899        let finalized_at = Utc::now();
5900        let job_ids: Vec<i64> = claimed.iter().map(|entry| entry.job.id).collect();
5901        let run_leases: Vec<i64> = claimed.iter().map(|entry| entry.job.run_lease).collect();
5902        let mut by_partition: BTreeMap<usize, Vec<&ClaimedRuntimeJob>> = BTreeMap::new();
5903        for entry in claimed {
5904            let claim_slot_index =
5905                ring_slot_index(entry.claim.claim_slot, self.claim_slot_count(), "claim")?;
5906            by_partition
5907                .entry(claim_slot_index)
5908                .or_default()
5909                .push(entry);
5910        }
5911
5912        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
5913
5914        // Without the per-job closure-row insert,
5915        // lease_claim_closure_batches become the idempotence evidence.
5916        // Serialize duplicate completion attempts before the main statement
5917        // so a waiter takes its statement snapshot after the first
5918        // transaction commits and sees that evidence.
5919        let receipt_pairs: Vec<(i64, i64)> = job_ids
5920            .iter()
5921            .zip(run_leases.iter())
5922            .map(|(job_id, run_lease)| (*job_id, *run_lease))
5923            .collect();
5924        if let Err(err) = self.lock_receipt_attempts_tx(&mut tx, &receipt_pairs).await {
5925            let _ = tx.rollback().await;
5926            return Err(err);
5927        }
5928
5929        let mut rows = Vec::with_capacity(claimed.len());
5930        for (claim_slot_index, group) in by_partition {
5931            let claim_rel = claim_child_name(schema, claim_slot_index);
5932            let claim_batch_rel = claim_batch_child_name(schema, claim_slot_index);
5933            let closure_rel = closure_child_name(schema, claim_slot_index);
5934            let closure_batch_rel = claim_closure_batch_child_name(schema, claim_slot_index);
5935            let receipt_batch_rel = format!("{schema}.receipt_completion_batches");
5936            let closed_evidence =
5937                receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
5938
5939            let claim_slots: Vec<i32> = group.iter().map(|entry| entry.claim.claim_slot).collect();
5940            let ready_slots: Vec<i32> = group.iter().map(|entry| entry.claim.ready_slot).collect();
5941            let ready_generations: Vec<i64> = group
5942                .iter()
5943                .map(|entry| entry.claim.ready_generation)
5944                .collect();
5945            let job_ids: Vec<i64> = group.iter().map(|entry| entry.job.id).collect();
5946            let queues: Vec<String> = group.iter().map(|entry| entry.job.queue.clone()).collect();
5947            let priorities: Vec<i16> = group.iter().map(|entry| entry.claim.priority).collect();
5948            let attempts: Vec<i16> = group.iter().map(|entry| entry.job.attempt).collect();
5949            let run_leases: Vec<i64> = group.iter().map(|entry| entry.job.run_lease).collect();
5950            let receipt_ids: Vec<i64> = group
5951                .iter()
5952                .map(|entry| {
5953                    entry
5954                        .claim
5955                        .receipt_id
5956                        .expect("receipt fast completion requires receipt_id")
5957                })
5958                .collect();
5959            let claim_batch_ids: Vec<Option<i64>> = group
5960                .iter()
5961                .map(|entry| entry.claim.claim_batch_id)
5962                .collect();
5963            let claim_batch_indices: Vec<Option<i32>> = group
5964                .iter()
5965                .map(|entry| entry.claim.claim_batch_index)
5966                .collect();
5967            let lane_seqs: Vec<i64> = group.iter().map(|entry| entry.claim.lane_seq).collect();
5968            let enqueue_shards: Vec<i16> = group
5969                .iter()
5970                .map(|entry| entry.claim.enqueue_shard)
5971                .collect();
5972            let attempted_ats: Vec<Option<DateTime<Utc>>> =
5973                group.iter().map(|entry| entry.job.attempted_at).collect();
5974            let finalized_ats: Vec<DateTime<Utc>> = vec![finalized_at; group.len()];
5975
5976            let completed: Vec<(i64, i64)> = match sqlx::query_as(&format!(
5977                r#"
5978                WITH completed(
5979                    claim_slot,
5980                    ready_slot,
5981                    ready_generation,
5982                    job_id,
5983                    queue,
5984                    priority,
5985                    attempt,
5986                    run_lease,
5987                    receipt_id,
5988                    claim_batch_id,
5989                    claim_batch_index,
5990                    lane_seq,
5991                    enqueue_shard,
5992                    attempted_at,
5993                    finalized_at
5994                ) AS (
5995                    SELECT *
5996                    FROM unnest(
5997                        $1::int[],
5998                        $2::int[],
5999                        $3::bigint[],
6000                        $4::bigint[],
6001                        $5::text[],
6002                        $6::smallint[],
6003                        $7::smallint[],
6004                        $8::bigint[],
6005                        $9::bigint[],
6006                        $10::bigint[],
6007                        $11::int[],
6008                        $12::bigint[],
6009                        $13::smallint[],
6010                        $14::timestamptz[],
6011                        $15::timestamptz[]
6012                    )
6013                ),
6014                row_claim_refs AS (
6015                    SELECT claims.claim_slot, claims.job_id, claims.run_lease, claims.receipt_id
6016                    FROM {claim_rel} AS claims
6017                    JOIN completed
6018                      ON completed.claim_slot = claims.claim_slot
6019                     AND completed.job_id = claims.job_id
6020                     AND completed.run_lease = claims.run_lease
6021                     AND completed.receipt_id = claims.receipt_id
6022                     AND completed.claim_batch_id IS NULL
6023                    WHERE NOT {closed_evidence}
6024                      AND NOT EXISTS (
6025                          SELECT 1
6026                          FROM {schema}.leases AS lease
6027                        WHERE lease.job_id = claims.job_id
6028                            AND lease.run_lease = claims.run_lease
6029                      )
6030                ),
6031                batch_claim_refs AS (
6032                    SELECT
6033                        claim_batches.claim_slot,
6034                        items.job_id,
6035                        items.run_lease,
6036                        items.receipt_id
6037                    FROM completed
6038                    JOIN {claim_batch_rel} AS claim_batches
6039                      ON claim_batches.claim_slot = completed.claim_slot
6040                     AND claim_batches.batch_id = completed.claim_batch_id
6041                    CROSS JOIN LATERAL (
6042                        SELECT
6043                            claim_batches.job_ids[completed.claim_batch_index] AS job_id,
6044                            claim_batches.run_leases[completed.claim_batch_index] AS run_lease,
6045                            claim_batches.receipt_ids[completed.claim_batch_index] AS receipt_id
6046                    ) AS items
6047                    WHERE completed.claim_batch_id IS NOT NULL
6048                      AND completed.claim_batch_index IS NOT NULL
6049                      AND completed.claim_batch_index BETWEEN 1 AND claim_batches.claimed_count
6050                      AND completed.job_id = items.job_id
6051                      AND completed.run_lease = items.run_lease
6052                      AND completed.receipt_id = items.receipt_id
6053                      AND NOT EXISTS (
6054                          SELECT 1
6055                          FROM {closure_rel} AS closures
6056                          WHERE closures.claim_slot = claim_batches.claim_slot
6057                            AND closures.job_id = items.job_id
6058                            AND closures.run_lease = items.run_lease
6059                      )
6060                      -- Compact successful completion treats receipt closure
6061                      -- evidence as the same-attempt disposition fence. The
6062                      -- non-success paths write explicit closure rows before
6063                      -- moving the job to done/deferred/DLQ, so probing those
6064                      -- terminal ledgers here only adds hot-path work.
6065                      AND NOT EXISTS (
6066                          SELECT 1
6067                          FROM {closure_batch_rel} AS closure_batches
6068                          WHERE closure_batches.receipt_ranges @> items.receipt_id
6069                      )
6070                      AND NOT EXISTS (
6071                          SELECT 1
6072                          FROM {schema}.leases AS lease
6073                          WHERE lease.job_id = items.job_id
6074                            AND lease.run_lease = items.run_lease
6075                      )
6076                ),
6077                claim_refs AS (
6078                    SELECT claim_slot, job_id, run_lease, receipt_id FROM row_claim_refs
6079                    UNION ALL
6080                    SELECT claim_slot, job_id, run_lease, receipt_id FROM batch_claim_refs
6081                ),
6082                deleted_attempts AS (
6083                    DELETE FROM {schema}.attempt_state AS attempt
6084                    USING claim_refs
6085                    WHERE attempt.job_id = claim_refs.job_id
6086                      AND attempt.run_lease = claim_refs.run_lease
6087                    RETURNING attempt.job_id
6088                ),
6089                completed_rows AS (
6090                    SELECT completed.*
6091                    FROM completed
6092                    JOIN claim_refs
6093                      ON claim_refs.claim_slot = completed.claim_slot
6094                     AND claim_refs.job_id = completed.job_id
6095                     AND claim_refs.run_lease = completed.run_lease
6096                     AND claim_refs.receipt_id = completed.receipt_id
6097                ),
6098                claim_closure_batches AS (
6099                    INSERT INTO {closure_batch_rel} (
6100                        claim_slot,
6101                        ready_slot,
6102                        ready_generation,
6103                        outcome,
6104                        closed_count,
6105                        receipt_ids,
6106                        receipt_ranges,
6107                        closed_at
6108                    )
6109                    SELECT
6110                        completed.claim_slot,
6111                        completed.ready_slot,
6112                        completed.ready_generation,
6113                        'completed',
6114                        count(*)::int AS closed_count,
6115                        array_agg(completed.receipt_id ORDER BY completed.lane_seq, completed.job_id),
6116                        range_agg(int8range(completed.receipt_id, completed.receipt_id + 1, '[)') ORDER BY completed.receipt_id),
6117                        max(completed.finalized_at)
6118                    FROM completed_rows AS completed
6119                    GROUP BY
6120                        completed.claim_slot,
6121                        completed.ready_slot,
6122                        completed.ready_generation
6123                    RETURNING claim_slot, ready_slot, ready_generation, receipt_ids
6124                ),
6125                terminal AS (
6126                    INSERT INTO {receipt_batch_rel} (
6127                        ready_slot,
6128                        ready_generation,
6129                        claim_slot,
6130                        queue,
6131                        priority,
6132                        enqueue_shard,
6133                        completed_count,
6134                        job_ids,
6135                        run_leases,
6136                        lane_seqs,
6137                        attempts,
6138                        attempted_ats,
6139                        finalized_at
6140                    )
6141                    SELECT
6142                        completed.ready_slot,
6143                        completed.ready_generation,
6144                        completed.claim_slot,
6145                        completed.queue,
6146                        completed.priority,
6147                        completed.enqueue_shard,
6148                        count(*)::int AS completed_count,
6149                        array_agg(completed.job_id ORDER BY completed.lane_seq, completed.job_id),
6150                        array_agg(completed.run_lease ORDER BY completed.lane_seq, completed.job_id),
6151                        array_agg(completed.lane_seq ORDER BY completed.lane_seq, completed.job_id),
6152                        array_agg(completed.attempt ORDER BY completed.lane_seq, completed.job_id),
6153                        array_agg(completed.attempted_at ORDER BY completed.lane_seq, completed.job_id),
6154                        max(completed.finalized_at)
6155                    FROM completed_rows AS completed
6156                    GROUP BY
6157                        completed.ready_slot,
6158                        completed.ready_generation,
6159                        completed.claim_slot,
6160                        completed.queue,
6161                        completed.priority,
6162                        completed.enqueue_shard
6163                    RETURNING
6164                        ready_slot,
6165                        ready_generation,
6166                        claim_slot,
6167                        queue,
6168                        priority,
6169                        enqueue_shard,
6170                        job_ids,
6171                        run_leases
6172                )
6173                SELECT completed_rows.job_id, completed_rows.run_lease
6174                FROM completed_rows
6175                CROSS JOIN (SELECT count(*) FROM claim_closure_batches) AS closure_batch_write
6176                CROSS JOIN (SELECT count(*) FROM terminal) AS terminal_write
6177                "#
6178            ))
6179            .bind(&claim_slots)
6180            .bind(&ready_slots)
6181            .bind(&ready_generations)
6182            .bind(&job_ids)
6183            .bind(&queues)
6184            .bind(&priorities)
6185            .bind(&attempts)
6186            .bind(&run_leases)
6187            .bind(&receipt_ids)
6188            .bind(&claim_batch_ids)
6189            .bind(&claim_batch_indices)
6190            .bind(&lane_seqs)
6191            .bind(&enqueue_shards)
6192            .bind(&attempted_ats)
6193            .bind(&finalized_ats)
6194            .fetch_all(tx.as_mut())
6195            .await
6196            {
6197                Ok(rows) => rows,
6198                Err(err) => {
6199                    let _ = tx.rollback().await;
6200                    return Err(map_sqlx_error(err));
6201                }
6202            };
6203
6204            rows.extend(completed);
6205        }
6206
6207        tx.commit().await.map_err(map_sqlx_error)?;
6208        Ok(rows)
6209    }
6210
6211    #[tracing::instrument(
6212        skip(self, pool, claimed),
6213        name = "queue_storage.complete_runtime_batch"
6214    )]
6215    pub async fn complete_runtime_batch(
6216        &self,
6217        pool: &PgPool,
6218        claimed: &[ClaimedRuntimeJob],
6219    ) -> Result<Vec<(i64, i64)>, AwaError> {
6220        if claimed.is_empty() {
6221            return Ok(Vec::new());
6222        }
6223
6224        if self.lease_claim_receipts() && claimed.iter().all(Self::receipt_fast_complete_candidate)
6225        {
6226            let mut updated = self
6227                .complete_receipt_runtime_batch_fast(pool, claimed)
6228                .await?;
6229            if updated.len() == claimed.len() {
6230                return Ok(updated);
6231            }
6232
6233            let updated_pairs: BTreeSet<(i64, i64)> = updated.iter().copied().collect();
6234            let missed: Vec<_> = claimed
6235                .iter()
6236                .filter(|entry| !updated_pairs.contains(&(entry.job.id, entry.job.run_lease)))
6237                .cloned()
6238                .collect();
6239            if !missed.is_empty() {
6240                updated.extend(self.complete_runtime_batch_slow(pool, &missed).await?);
6241            }
6242            return Ok(updated);
6243        }
6244
6245        self.complete_runtime_batch_slow(pool, claimed).await
6246    }
6247
6248    async fn complete_runtime_batch_slow(
6249        &self,
6250        pool: &PgPool,
6251        claimed: &[ClaimedRuntimeJob],
6252    ) -> Result<Vec<(i64, i64)>, AwaError> {
6253        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
6254        let result = self
6255            .complete_runtime_batch_slow_in_tx(&mut tx, claimed)
6256            .await?;
6257        tx.commit().await.map_err(map_sqlx_error)?;
6258        Ok(result)
6259    }
6260
6261    /// Same as [`Self::complete_runtime_batch_slow`] but runs on the caller's
6262    /// transaction so additional writes — for example ADR-029 follow-up job
6263    /// inserts — can join the same commit. The caller is responsible for
6264    /// `commit()` / `rollback()`. Handles both receipt-claimed and
6265    /// materialised leases.
6266    pub async fn complete_runtime_batch_slow_in_tx(
6267        &self,
6268        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
6269        claimed: &[ClaimedRuntimeJob],
6270    ) -> Result<Vec<(i64, i64)>, AwaError> {
6271        if claimed.is_empty() {
6272            return Ok(Vec::new());
6273        }
6274
6275        let schema = self.schema();
6276
6277        let claimed_map: BTreeMap<(i64, i64), ClaimedRuntimeJob> = claimed
6278            .iter()
6279            .cloned()
6280            .map(|entry| ((entry.job.id, entry.job.run_lease), entry))
6281            .collect();
6282
6283        if self.lease_claim_receipts() {
6284            let (mut receipt_claimed, mut materialized_claimed): (Vec<_>, Vec<_>) = claimed
6285                .iter()
6286                .cloned()
6287                .partition(|entry| entry.claim.lease_claim_receipt);
6288            let mut updated_all = Vec::new();
6289
6290            if !receipt_claimed.is_empty() {
6291                let receipt_pairs: Vec<(i64, i64)> = receipt_claimed
6292                    .iter()
6293                    .map(|entry| (entry.job.id, entry.job.run_lease))
6294                    .collect();
6295                self.lock_receipt_attempts_tx(tx, &receipt_pairs).await?;
6296
6297                // claim_slot rides along on `ClaimedEntry`, so receipt
6298                // completion can validate exact claim evidence and route the
6299                // explicit completed closure without an extra lookup.
6300                let receipt_claim_slots: Vec<i32> = receipt_claimed
6301                    .iter()
6302                    .map(|entry| entry.claim.claim_slot)
6303                    .collect();
6304                let receipt_job_ids: Vec<i64> =
6305                    receipt_claimed.iter().map(|entry| entry.job.id).collect();
6306                let receipt_run_leases: Vec<i64> = receipt_claimed
6307                    .iter()
6308                    .map(|entry| entry.job.run_lease)
6309                    .collect();
6310                let receipt_receipt_ids: Vec<i64> = receipt_claimed
6311                    .iter()
6312                    .map(|entry| {
6313                        entry
6314                            .claim
6315                            .receipt_id
6316                            .expect("receipt-backed slow completion requires receipt_id")
6317                    })
6318                    .collect();
6319                let closure_rel = format!("{schema}.lease_claim_closures");
6320                let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
6321                let closed_evidence =
6322                    receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
6323                let updated: Vec<(i64, i64)> = sqlx::query_as(&format!(
6324                    r#"
6325                    WITH completed(claim_slot, job_id, run_lease, receipt_id) AS (
6326                        SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[])
6327                    ),
6328                    locked_row_claims AS (
6329                        SELECT claims.claim_slot, claims.job_id, claims.run_lease
6330                        FROM {schema}.lease_claims AS claims
6331                        JOIN completed
6332                          ON completed.claim_slot = claims.claim_slot
6333                         AND completed.job_id = claims.job_id
6334                         AND completed.run_lease = claims.run_lease
6335                         AND completed.receipt_id = claims.receipt_id
6336                        WHERE NOT {closed_evidence}
6337                          AND NOT EXISTS (
6338                              SELECT 1
6339                              FROM {schema}.leases AS lease
6340                              WHERE lease.job_id = claims.job_id
6341                                AND lease.run_lease = claims.run_lease
6342                              )
6343                        FOR UPDATE OF claims
6344                    ),
6345                    locked_batch_claims AS (
6346                        SELECT
6347                            claim_batches.claim_slot,
6348                            claim_batches.ready_slot,
6349                            claim_batches.ready_generation,
6350                            items.job_id,
6351                            items.run_lease,
6352                            items.receipt_id
6353                        FROM {schema}.lease_claim_batches AS claim_batches
6354                        CROSS JOIN LATERAL unnest(
6355                            claim_batches.job_ids,
6356                            claim_batches.run_leases,
6357                            claim_batches.receipt_ids
6358                        ) AS items(job_id, run_lease, receipt_id)
6359                        JOIN completed
6360                          ON completed.claim_slot = claim_batches.claim_slot
6361                         AND completed.job_id = items.job_id
6362                         AND completed.run_lease = items.run_lease
6363                         AND completed.receipt_id = items.receipt_id
6364                        WHERE NOT EXISTS (
6365                              SELECT 1
6366                              FROM {schema}.lease_claim_closures AS closures
6367                              WHERE closures.claim_slot = claim_batches.claim_slot
6368                                AND closures.job_id = items.job_id
6369                                AND closures.run_lease = items.run_lease
6370                          )
6371                          AND NOT EXISTS (
6372                              SELECT 1
6373                              FROM {schema}.lease_claim_closure_batches AS closure_batches
6374                              WHERE closure_batches.claim_slot = claim_batches.claim_slot
6375                                AND closure_batches.receipt_ranges @> items.receipt_id
6376                          )
6377                          AND NOT EXISTS (
6378                              SELECT 1
6379                              FROM {schema}.leases AS lease
6380                              WHERE lease.job_id = items.job_id
6381                                AND lease.run_lease = items.run_lease
6382                          )
6383                          AND NOT EXISTS (
6384                              SELECT 1
6385                              FROM {schema}.done_entries AS done
6386                              WHERE done.job_id = items.job_id
6387                                AND done.run_lease = items.run_lease
6388                          )
6389                          AND NOT EXISTS (
6390                              SELECT 1
6391                              FROM {schema}.deferred_jobs AS deferred
6392                              WHERE deferred.job_id = items.job_id
6393                                AND deferred.run_lease = items.run_lease
6394                          )
6395                          AND NOT EXISTS (
6396                              SELECT 1
6397                              FROM {schema}.dlq_entries AS dlq
6398                              WHERE dlq.job_id = items.job_id
6399                                AND dlq.run_lease = items.run_lease
6400                          )
6401                        FOR UPDATE OF claim_batches
6402                    ),
6403                    locked_claims AS (
6404                        SELECT claim_slot, job_id, run_lease FROM locked_row_claims
6405                        UNION ALL
6406                        SELECT claim_slot, job_id, run_lease FROM locked_batch_claims
6407                    ),
6408                    deleted_attempts AS (
6409                        DELETE FROM {schema}.attempt_state AS attempt
6410                        USING locked_claims
6411                        WHERE attempt.job_id = locked_claims.job_id
6412                          AND attempt.run_lease = locked_claims.run_lease
6413                        RETURNING attempt.job_id
6414                    ),
6415                    -- Row-sourced (deadline lease) claims close into the
6416                    -- explicit ledger so the prune gates balance them
6417                    -- against the lease_claims row via the closure JOIN.
6418                    closed_claims AS (
6419                        INSERT INTO {schema}.lease_claim_closures (
6420                            claim_slot,
6421                            job_id,
6422                            run_lease,
6423                            outcome,
6424                            closed_at
6425                        )
6426                        SELECT
6427                            locked_row_claims.claim_slot,
6428                            locked_row_claims.job_id,
6429                            locked_row_claims.run_lease,
6430                            'completed',
6431                            clock_timestamp()
6432                        FROM locked_row_claims
6433                        ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
6434                        RETURNING claim_slot, job_id, run_lease, closed_at
6435                    ),
6436                    -- Compact batch-sourced claims have no lease_claims row
6437                    -- to JOIN, so close them into the batch ledger that the
6438                    -- queue prune gate counts via compact_count.
6439                    closed_batches AS (
6440                        INSERT INTO {schema}.lease_claim_closure_batches (
6441                            claim_slot,
6442                            ready_slot,
6443                            ready_generation,
6444                            outcome,
6445                            closed_count,
6446                            receipt_ids,
6447                            receipt_ranges,
6448                            closed_at
6449                        )
6450                        SELECT
6451                            locked_batch_claims.claim_slot,
6452                            locked_batch_claims.ready_slot,
6453                            locked_batch_claims.ready_generation,
6454                            'completed',
6455                            count(*)::int,
6456                            array_agg(locked_batch_claims.receipt_id ORDER BY locked_batch_claims.receipt_id),
6457                            range_agg(int8range(locked_batch_claims.receipt_id, locked_batch_claims.receipt_id + 1, '[)') ORDER BY locked_batch_claims.receipt_id),
6458                            clock_timestamp()
6459                        FROM locked_batch_claims
6460                        GROUP BY
6461                            locked_batch_claims.claim_slot,
6462                            locked_batch_claims.ready_slot,
6463                            locked_batch_claims.ready_generation
6464                        RETURNING claim_slot
6465                    ),
6466                    marked_claims AS (
6467                        UPDATE {schema}.lease_claims AS claims
6468                        SET closed_at = COALESCE(claims.closed_at, closed_claims.closed_at)
6469                        FROM closed_claims
6470                        WHERE claims.claim_slot = closed_claims.claim_slot
6471                          AND claims.job_id = closed_claims.job_id
6472                          AND claims.run_lease = closed_claims.run_lease
6473                        RETURNING claims.job_id
6474                    ),
6475                    -- Force the batch insert to run and report the compact
6476                    -- claims it closed so the caller finalizes them too.
6477                    closed_batch_pairs AS (
6478                        SELECT locked_batch_claims.job_id, locked_batch_claims.run_lease
6479                        FROM locked_batch_claims
6480                        WHERE EXISTS (SELECT 1 FROM closed_batches)
6481                    )
6482                    SELECT job_id, run_lease
6483                    FROM closed_claims
6484                    UNION ALL
6485                    SELECT job_id, run_lease
6486                    FROM closed_batch_pairs
6487                    "#
6488                ))
6489                .bind(&receipt_claim_slots)
6490                .bind(&receipt_job_ids)
6491                .bind(&receipt_run_leases)
6492                .bind(&receipt_receipt_ids)
6493                .fetch_all(tx.as_mut())
6494                .await
6495                .map_err(map_sqlx_error)?;
6496
6497                if !updated.is_empty() {
6498                    let finalized_at = Utc::now();
6499                    let mut done_rows = Vec::with_capacity(updated.len());
6500                    for (job_id, run_lease) in &updated {
6501                        if let Some(runtime_job) = claimed_map.get(&(*job_id, *run_lease)).cloned()
6502                        {
6503                            done_rows.push(runtime_job.into_done_row(finalized_at)?);
6504                        }
6505                    }
6506
6507                    self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6508                        .await?;
6509                    updated_all.extend(updated);
6510                }
6511
6512                let updated_pairs: BTreeSet<(i64, i64)> = updated_all.iter().copied().collect();
6513                let mut escalated_receipts = Vec::new();
6514                for entry in receipt_claimed.drain(..) {
6515                    if !updated_pairs.contains(&(entry.job.id, entry.job.run_lease)) {
6516                        escalated_receipts.push(entry);
6517                    }
6518                }
6519                materialized_claimed.extend(escalated_receipts);
6520            }
6521
6522            if !materialized_claimed.is_empty() {
6523                let ready_slots: Vec<i32> = materialized_claimed
6524                    .iter()
6525                    .map(|entry| entry.claim.ready_slot)
6526                    .collect();
6527                let ready_generations: Vec<i64> = materialized_claimed
6528                    .iter()
6529                    .map(|entry| entry.claim.ready_generation)
6530                    .collect();
6531                let job_ids: Vec<i64> = materialized_claimed
6532                    .iter()
6533                    .map(|entry| entry.job.id)
6534                    .collect();
6535                let queues: Vec<String> = materialized_claimed
6536                    .iter()
6537                    .map(|entry| entry.claim.queue.clone())
6538                    .collect();
6539                let priorities: Vec<i16> = materialized_claimed
6540                    .iter()
6541                    .map(|entry| entry.claim.priority)
6542                    .collect();
6543                let enqueue_shards: Vec<i16> = materialized_claimed
6544                    .iter()
6545                    .map(|entry| entry.claim.enqueue_shard)
6546                    .collect();
6547                let lane_seqs: Vec<i64> = materialized_claimed
6548                    .iter()
6549                    .map(|entry| entry.claim.lane_seq)
6550                    .collect();
6551                let run_leases: Vec<i64> = materialized_claimed
6552                    .iter()
6553                    .map(|entry| entry.job.run_lease)
6554                    .collect();
6555
6556                // CTE-as-DML: delete the leases and the matching attempt_state
6557                // rows in one round-trip. Receipt claims may materialize a
6558                // lease after the original claim, so the materialized lease can
6559                // live in a newer lease slot than the claim carried. Match on
6560                // the stable ready-lane and attempt identity instead.
6561                let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
6562                    r#"
6563                    WITH completed(ready_slot, ready_generation, job_id, queue, priority, enqueue_shard, lane_seq, run_lease) AS (
6564                        SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::text[], $5::smallint[], $6::smallint[], $7::bigint[], $8::bigint[])
6565                    ),
6566                    deleted AS (
6567                        DELETE FROM {schema}.leases AS leases
6568                        USING completed
6569                        WHERE leases.ready_slot = completed.ready_slot
6570                          AND leases.ready_generation = completed.ready_generation
6571                          AND leases.job_id = completed.job_id
6572                          AND leases.queue = completed.queue
6573                          AND leases.priority = completed.priority
6574                          AND leases.enqueue_shard = completed.enqueue_shard
6575                          AND leases.lane_seq = completed.lane_seq
6576                          AND leases.run_lease = completed.run_lease
6577                        RETURNING
6578                            leases.ready_slot,
6579                            leases.ready_generation,
6580                            leases.job_id,
6581                            leases.queue,
6582                            leases.state,
6583                            leases.priority,
6584                            leases.attempt,
6585                            leases.run_lease,
6586                            leases.max_attempts,
6587                            leases.lane_seq,
6588                            leases.enqueue_shard,
6589                            leases.heartbeat_at,
6590                            leases.deadline_at,
6591                            leases.attempted_at,
6592                            leases.callback_id,
6593                            leases.callback_timeout_at
6594                    ),
6595                    del_attempts AS (
6596                        DELETE FROM {schema}.attempt_state AS attempt
6597                        USING deleted
6598                        WHERE attempt.job_id = deleted.job_id
6599                          AND attempt.run_lease = deleted.run_lease
6600                        RETURNING attempt.job_id
6601                    )
6602                    SELECT
6603                        ready_slot,
6604                        ready_generation,
6605                        job_id,
6606                        queue,
6607                        state,
6608                        priority,
6609                        attempt,
6610                        run_lease,
6611                        max_attempts,
6612                        lane_seq,
6613                        enqueue_shard,
6614                        heartbeat_at,
6615                        deadline_at,
6616                        attempted_at,
6617                        callback_id,
6618                        callback_timeout_at
6619                    FROM deleted
6620                    "#
6621                ))
6622                .bind(&ready_slots)
6623                .bind(&ready_generations)
6624                .bind(&job_ids)
6625                .bind(&queues)
6626                .bind(&priorities)
6627                .bind(&enqueue_shards)
6628                .bind(&lane_seqs)
6629                .bind(&run_leases)
6630                .fetch_all(tx.as_mut())
6631                .await
6632                .map_err(map_sqlx_error)?;
6633
6634                if !deleted.is_empty() {
6635                    let completed_pairs: Vec<(i64, i64)> = deleted
6636                        .iter()
6637                        .map(|row| (row.job_id, row.run_lease))
6638                        .collect();
6639                    self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
6640                        .await?;
6641
6642                    let finalized_at = Utc::now();
6643                    let mut done_rows = Vec::with_capacity(deleted.len());
6644                    for deleted_row in deleted {
6645                        if let Some(runtime_job) = claimed_map
6646                            .get(&(deleted_row.job_id, deleted_row.run_lease))
6647                            .cloned()
6648                        {
6649                            done_rows.push(runtime_job.into_done_row(finalized_at)?);
6650                            updated_all.push((deleted_row.job_id, deleted_row.run_lease));
6651                        }
6652                    }
6653
6654                    self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6655                        .await?;
6656                }
6657            }
6658
6659            return Ok(updated_all);
6660        }
6661
6662        let lease_slots: Vec<i32> = claimed.iter().map(|entry| entry.claim.lease_slot).collect();
6663        let queues: Vec<String> = claimed
6664            .iter()
6665            .map(|entry| entry.claim.queue.clone())
6666            .collect();
6667        let priorities: Vec<i16> = claimed.iter().map(|entry| entry.claim.priority).collect();
6668        let enqueue_shards: Vec<i16> = claimed
6669            .iter()
6670            .map(|entry| entry.claim.enqueue_shard)
6671            .collect();
6672        let lane_seqs: Vec<i64> = claimed.iter().map(|entry| entry.claim.lane_seq).collect();
6673        let run_leases: Vec<i64> = claimed.iter().map(|entry| entry.job.run_lease).collect();
6674
6675        // Single CTE-as-DML statement: delete the leases and the matching
6676        // attempt_state rows in one round-trip. The `deleted` CTE materialises
6677        // the lease deletion (so its RETURNING is observable to the
6678        // attempt-state delete and to the final SELECT), and `del_attempts`
6679        // hangs off it. Saves one round-trip per completion batch versus
6680        // issuing the attempt-state delete as a separate statement.
6681        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
6682            r#"
6683            WITH completed(lease_slot, queue, priority, enqueue_shard, lane_seq, run_lease) AS (
6684                SELECT * FROM unnest($1::int[], $2::text[], $3::smallint[], $4::smallint[], $5::bigint[], $6::bigint[])
6685            ),
6686            deleted AS (
6687                DELETE FROM {schema}.leases AS leases
6688                USING completed
6689                WHERE leases.lease_slot = completed.lease_slot
6690                  AND leases.queue = completed.queue
6691                  AND leases.priority = completed.priority
6692                  AND leases.enqueue_shard = completed.enqueue_shard
6693                  AND leases.lane_seq = completed.lane_seq
6694                  AND leases.run_lease = completed.run_lease
6695                RETURNING
6696                    leases.ready_slot,
6697                    leases.ready_generation,
6698                    leases.job_id,
6699                    leases.queue,
6700                    leases.state,
6701                    leases.priority,
6702                    leases.attempt,
6703                    leases.run_lease,
6704                    leases.max_attempts,
6705                    leases.lane_seq,
6706                    leases.enqueue_shard,
6707                    leases.heartbeat_at,
6708                    leases.deadline_at,
6709                    leases.attempted_at,
6710                    leases.callback_id,
6711                    leases.callback_timeout_at
6712            ),
6713            del_attempts AS (
6714                DELETE FROM {schema}.attempt_state AS attempt
6715                USING deleted
6716                WHERE attempt.job_id = deleted.job_id
6717                  AND attempt.run_lease = deleted.run_lease
6718                RETURNING attempt.job_id
6719            )
6720            SELECT
6721                ready_slot,
6722                ready_generation,
6723                job_id,
6724                queue,
6725                state,
6726                priority,
6727                attempt,
6728                run_lease,
6729                max_attempts,
6730                lane_seq,
6731                enqueue_shard,
6732                heartbeat_at,
6733                deadline_at,
6734                attempted_at,
6735                callback_id,
6736                callback_timeout_at
6737            FROM deleted
6738            "#
6739        ))
6740        .bind(&lease_slots)
6741        .bind(&queues)
6742        .bind(&priorities)
6743        .bind(&enqueue_shards)
6744        .bind(&lane_seqs)
6745        .bind(&run_leases)
6746        .fetch_all(tx.as_mut())
6747        .await
6748        .map_err(map_sqlx_error)?;
6749
6750        if deleted.is_empty() {
6751            return Ok(Vec::new());
6752        }
6753
6754        let completed_pairs: Vec<(i64, i64)> = deleted
6755            .iter()
6756            .map(|row| (row.job_id, row.run_lease))
6757            .collect();
6758        self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
6759            .await?;
6760
6761        let finalized_at = Utc::now();
6762        let mut done_rows = Vec::with_capacity(deleted.len());
6763        let mut updated = Vec::with_capacity(deleted.len());
6764        for deleted_row in deleted {
6765            if let Some(runtime_job) = claimed_map
6766                .get(&(deleted_row.job_id, deleted_row.run_lease))
6767                .cloned()
6768            {
6769                done_rows.push(runtime_job.into_done_row(finalized_at)?);
6770                updated.push((deleted_row.job_id, deleted_row.run_lease));
6771            }
6772        }
6773
6774        self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6775            .await?;
6776        Ok(updated)
6777    }
6778
6779    #[tracing::instrument(
6780        skip(self, pool, completions),
6781        name = "queue_storage.complete_job_batch_by_id"
6782    )]
6783    pub async fn complete_job_batch_by_id(
6784        &self,
6785        pool: &PgPool,
6786        completions: &[(i64, i64)],
6787    ) -> Result<Vec<(i64, i64)>, AwaError> {
6788        if completions.is_empty() {
6789            return Ok(Vec::new());
6790        }
6791        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
6792        let result = self
6793            .complete_job_batch_by_id_in_tx(&mut tx, completions)
6794            .await?;
6795        tx.commit().await.map_err(map_sqlx_error)?;
6796        Ok(result)
6797    }
6798
6799    /// Same as [`Self::complete_job_batch_by_id`] but runs on the caller's
6800    /// transaction so additional writes — for example ADR-029 follow-up job
6801    /// inserts — can join the same commit. The caller is responsible for
6802    /// `commit()` / `rollback()`.
6803    pub async fn complete_job_batch_by_id_in_tx(
6804        &self,
6805        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
6806        completions: &[(i64, i64)],
6807    ) -> Result<Vec<(i64, i64)>, AwaError> {
6808        if completions.is_empty() {
6809            return Ok(Vec::new());
6810        }
6811
6812        let schema = self.schema();
6813
6814        let job_ids: Vec<i64> = completions.iter().map(|(job_id, _)| *job_id).collect();
6815        let run_leases: Vec<i64> = completions
6816            .iter()
6817            .map(|(_, run_lease)| *run_lease)
6818            .collect();
6819
6820        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
6821            r#"
6822            WITH completed(job_id, run_lease) AS (
6823                SELECT * FROM unnest($1::bigint[], $2::bigint[])
6824            )
6825            DELETE FROM {schema}.leases AS leases
6826            USING completed
6827            WHERE leases.job_id = completed.job_id
6828              AND leases.run_lease = completed.run_lease
6829            RETURNING
6830                leases.ready_slot,
6831                leases.ready_generation,
6832                leases.job_id,
6833                leases.queue,
6834                leases.state,
6835                leases.priority,
6836                leases.attempt,
6837                leases.run_lease,
6838                leases.max_attempts,
6839                leases.lane_seq,
6840                leases.enqueue_shard,
6841                leases.heartbeat_at,
6842                leases.deadline_at,
6843                leases.attempted_at,
6844                leases.callback_id,
6845                leases.callback_timeout_at
6846            "#
6847        ))
6848        .bind(&job_ids)
6849        .bind(&run_leases)
6850        .fetch_all(tx.as_mut())
6851        .await
6852        .map_err(map_sqlx_error)?;
6853
6854        if deleted.is_empty() {
6855            return Ok(Vec::new());
6856        }
6857
6858        let completed_pairs: Vec<(i64, i64)> = deleted
6859            .iter()
6860            .map(|row| (row.job_id, row.run_lease))
6861            .collect();
6862        self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
6863            .await?;
6864
6865        let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
6866
6867        let finalized_at = Utc::now();
6868        let mut done_rows = Vec::with_capacity(moved.len());
6869        for entry in moved.iter().cloned() {
6870            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
6871                entry.payload.clone(),
6872                entry.progress.clone(),
6873            )?)?;
6874            payload.set_progress(None);
6875            done_rows.push(entry.into_done_row(
6876                JobState::Completed,
6877                finalized_at,
6878                payload.into_json(),
6879            ));
6880        }
6881
6882        self.insert_done_rows_tx(tx, &done_rows, Some(JobState::Running))
6883            .await?;
6884        Ok(moved
6885            .into_iter()
6886            .map(|entry| (entry.job_id, entry.run_lease))
6887            .collect())
6888    }
6889
6890    /// Exact admin/UI-grade queue counts. Scans `ready_entries` rather
6891    /// than reading the head tables — slower, but unaffected by the
6892    /// transient gap between sequence reservations and committed,
6893    /// still-claimable ready rows. Use [`Self::queue_claimer_signal`] for the
6894    /// dispatcher hot path.
6895    ///
6896    /// The live-terminal portion reads retained compact receipt batches
6897    /// directly. Done-entry terminal rows read from folded
6898    /// `queue_terminal_live_counts` plus unrolled
6899    /// `queue_terminal_count_deltas` when [`Self::terminal_counter_trusted`]
6900    /// returns true, and fall back to a `count(*) FROM terminal_jobs` scan
6901    /// when not. The fallback exists for the rolling-upgrade window: older
6902    /// binaries may have written terminal rows without maintaining the
6903    /// counter/delta contract, so reads stay honest until the operator runs
6904    /// `awa storage rebuild-terminal-counters`.
6905    async fn queue_counts_exact(
6906        &self,
6907        pool: &PgPool,
6908        queue: &str,
6909    ) -> Result<QueueCounts, AwaError> {
6910        let schema = self.schema();
6911        let queues = self.physical_queues_for_logical(queue);
6912        let closure_rel = format!("{schema}.lease_claim_closures");
6913        let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
6914        let closed_evidence =
6915            receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
6916        let counter_trusted = self.terminal_counter_trusted(pool).await?;
6917        // The live-terminal CTE swaps between counter-fed and
6918        // scan-fed depending on trust. Build it as a string so the
6919        // outer query plan is otherwise identical between the two
6920        // paths.
6921        let live_terminal_cte = if counter_trusted {
6922            format!(
6923                "live_terminal AS (
6924                    SELECT GREATEST(
6925                        0,
6926                        COALESCE((
6927                            SELECT SUM(live_terminal_count)
6928                            FROM {schema}.queue_terminal_live_counts
6929                            WHERE queue = ANY($1)
6930                        ), 0)
6931                        +
6932                        COALESCE((
6933                            SELECT SUM(terminal_delta)
6934                            FROM {schema}.queue_terminal_count_deltas
6935                            WHERE queue = ANY($1)
6936                        ), 0)
6937                        +
6938                        COALESCE((
6939                            SELECT SUM(completed_count)
6940                            FROM {schema}.receipt_completion_batches
6941                            WHERE queue = ANY($1)
6942                        ), 0)
6943                        -
6944                        COALESCE((
6945                            SELECT count(*)::bigint
6946                            FROM {schema}.receipt_completion_tombstones
6947                            WHERE queue = ANY($1)
6948                        ), 0)
6949                    )::bigint AS terminal
6950                )"
6951            )
6952        } else {
6953            format!(
6954                "live_terminal AS (
6955                    SELECT count(*)::bigint AS terminal
6956                    FROM {schema}.terminal_jobs
6957                    WHERE queue = ANY($1)
6958                )"
6959            )
6960        };
6961        let row: (i64, i64, i64, i64) = sqlx::query_as(&format!(
6962            r#"
6963            WITH lane_counts AS (
6964                -- Exact count: a ready row is available iff its
6965                -- lane_seq has not yet been passed by the lane's
6966                -- claim sequence cursor. Each shard within a (queue,
6967                -- priority) lane carries its own sequence, so the
6968                -- join matches on shard too — otherwise a ready row
6969                -- in shard A could be incorrectly compared against
6970                -- shard B's claim cursor.
6971                SELECT COALESCE(count(*)::bigint, 0) AS available
6972                FROM {schema}.ready_entries AS ready
6973                JOIN {schema}.queue_claim_heads AS claims
6974                  ON claims.queue = ready.queue
6975                 AND claims.priority = ready.priority
6976                 AND claims.enqueue_shard = ready.enqueue_shard
6977                WHERE ready.queue = ANY($1)
6978                  AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
6979                  AND NOT EXISTS (
6980                      SELECT 1 FROM {schema}.ready_tombstones AS tomb
6981                      WHERE tomb.queue = ready.queue
6982                        AND tomb.priority = ready.priority
6983                        AND tomb.enqueue_shard = ready.enqueue_shard
6984                        AND tomb.lane_seq = ready.lane_seq
6985                        AND tomb.ready_slot = ready.ready_slot
6986                        AND tomb.ready_generation = ready.ready_generation
6987                  )
6988            ),
6989            pruned_terminal AS (
6990                -- The GREATEST legacy dedupe applies to the completed
6991                -- column only: queue_lanes never carried a failed
6992                -- column, so failed counts come from the rollups alone.
6993                SELECT
6994                    COALESCE(
6995                        sum(
6996                            GREATEST(
6997                                COALESCE(lanes.pruned_completed_count, 0),
6998                                COALESCE(rollups.pruned_completed_count, 0)
6999                            )
7000                        ),
7001                        0
7002                    )::bigint AS completed,
7003                    COALESCE(
7004                        sum(COALESCE(rollups.pruned_failed_count, 0)),
7005                        0
7006                    )::bigint AS failed
7007                FROM (
7008                    SELECT queue, priority, pruned_completed_count
7009                    FROM {schema}.queue_lanes
7010                    WHERE queue = ANY($1)
7011                ) AS lanes
7012                FULL OUTER JOIN (
7013                    SELECT queue, priority, pruned_completed_count, pruned_failed_count
7014                    FROM {schema}.queue_terminal_rollups
7015                    WHERE queue = ANY($1)
7016                ) AS rollups
7017                USING (queue, priority)
7018            ),
7019            live_running AS (
7020                SELECT (
7021                    COALESCE((
7022                        SELECT count(*)::bigint
7023                        FROM {schema}.leases
7024                        WHERE queue = ANY($1)
7025                          AND state = 'running'
7026                    ), 0)
7027                    +
7028                    -- Derive the receipt-backed running count from
7029                    -- lease_claims anti-joined with every durable
7030                    -- closure evidence shape.
7031                    COALESCE((
7032                        SELECT count(*)::bigint
7033                        FROM {schema}.lease_claims AS claims
7034                        WHERE claims.queue = ANY($1)
7035                          AND NOT {closed_evidence}
7036                          AND NOT EXISTS (
7037                              SELECT 1
7038                              FROM {schema}.leases AS lease
7039	                              WHERE lease.job_id = claims.job_id
7040	                                AND lease.run_lease = claims.run_lease
7041	                          )
7042	                    ), 0)
7043	                    +
7044	                    -- Zero-deadline compact receipt claims live in
7045	                    -- lease_claim_batches until they complete or a cold
7046	                    -- path materializes them. Expand only for exact
7047	                    -- admin-grade counts.
7048	                    COALESCE((
7049	                        SELECT count(*)::bigint
7050	                        FROM {schema}.lease_claim_batches AS batches
7051	                        CROSS JOIN LATERAL unnest(
7052	                            batches.job_ids,
7053	                            batches.run_leases,
7054	                            batches.receipt_ids
7055	                        ) AS items(job_id, run_lease, receipt_id)
7056	                        WHERE batches.queue = ANY($1)
7057	                          AND NOT EXISTS (
7058	                              SELECT 1
7059	                              FROM {schema}.lease_claim_closures AS closures
7060	                              WHERE closures.claim_slot = batches.claim_slot
7061	                                AND closures.job_id = items.job_id
7062	                                AND closures.run_lease = items.run_lease
7063	                          )
7064	                          AND NOT EXISTS (
7065	                              SELECT 1
7066	                              FROM {schema}.lease_claim_closure_batches AS closure_batches
7067	                              WHERE closure_batches.claim_slot = batches.claim_slot
7068	                                AND closure_batches.receipt_ranges @> items.receipt_id
7069	                          )
7070	                          AND NOT EXISTS (
7071	                              SELECT 1
7072	                              FROM {schema}.leases AS lease
7073	                              WHERE lease.job_id = items.job_id
7074	                                AND lease.run_lease = items.run_lease
7075	                          )
7076	                          AND NOT EXISTS (
7077	                              SELECT 1 FROM {schema}.done_entries AS done
7078	                              WHERE done.job_id = items.job_id
7079	                                AND done.run_lease = items.run_lease
7080	                          )
7081	                          AND NOT EXISTS (
7082	                              SELECT 1 FROM {schema}.deferred_jobs AS deferred
7083	                              WHERE deferred.job_id = items.job_id
7084	                                AND deferred.run_lease = items.run_lease
7085	                          )
7086	                          AND NOT EXISTS (
7087	                              SELECT 1 FROM {schema}.dlq_entries AS dlq
7088	                              WHERE dlq.job_id = items.job_id
7089	                                AND dlq.run_lease = items.run_lease
7090	                          )
7091	                    ), 0)
7092	                )::bigint AS running
7093	            ),
7094            {live_terminal_cte}
7095            SELECT
7096                lane_counts.available,
7097                live_running.running,
7098                pruned_terminal.completed + pruned_terminal.failed
7099                    + live_terminal.terminal AS terminal,
7100                pruned_terminal.failed AS pruned_failed
7101            FROM lane_counts
7102            CROSS JOIN pruned_terminal
7103            CROSS JOIN live_running
7104            CROSS JOIN live_terminal
7105            "#
7106        ))
7107        .bind(&queues)
7108        .fetch_one(pool)
7109        .await
7110        .map_err(map_sqlx_error)?;
7111
7112        let (available, running, terminal, pruned_failed) = row;
7113        Ok(QueueCounts {
7114            available,
7115            running,
7116            terminal,
7117            pruned_failed,
7118        })
7119    }
7120
7121    #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.queue_counts")]
7122    pub async fn queue_counts(&self, pool: &PgPool, queue: &str) -> Result<QueueCounts, AwaError> {
7123        self.queue_counts_exact(pool, queue).await
7124    }
7125
7126    /// Index-only queue depth probe — for observability / depth-target
7127    /// throttling. Returns the same shape as [`Self::queue_counts`] but
7128    /// skips the table scans that [`Self::queue_counts_exact`] needs for
7129    /// exact terminal accounting:
7130    ///
7131    /// - **available** is the same as the dispatcher's claim signal:
7132    ///   `sum(GREATEST(enqueue_seq - claim_seq, 0))` over the shard
7133    ///   head tables. No scan of `ready_entries`. This is an upper
7134    ///   bound: admin DELETEs, committed gaps, and uncommitted enqueue
7135    ///   reservations can leave the enqueue sequence ahead of the actual
7136    ///   ready row count. Acceptable for depth-target throttling and
7137    ///   dashboards; not suitable for exact billing-style counts.
7138    /// - **running** matches [`Self::queue_counts`]'s strict definition:
7139    ///   `leases.state = 'running'` only. Receipt-plane claims that
7140    ///   have not yet materialised a lease row are omitted (the
7141    ///   exact path catches them via the `lease_claims` anti-join, but
7142    ///   that anti-join is what this fast variant exists to avoid).
7143    ///   `waiting_external` is *not* included — it's reported as part
7144    ///   of admin's parked-callback view, not running.
7145    /// - **terminal** is read from the persisted
7146    ///   `queue_terminal_rollups` denormaliser
7147    ///   (`pruned_completed_count + pruned_failed_count`).
7148    ///   Rows currently in `done_entries` or
7149    ///   `receipt_completion_batches` that have not yet rolled up are
7150    ///   not included. Strictly a lower bound; converges to the exact
7151    ///   count when rotation prunes the live queue segment. (The name
7152    ///   `terminal` is honest — this number counts `completed`,
7153    ///   `failed`, and `cancelled` terminal facts with the same
7154    ///   semantics as [`Self::queue_counts_exact`]; renamed from
7155    ///   `completed` in #290.)
7156    ///
7157    /// All three counters are O(num shards) lookups against small head
7158    /// tables and `leases` index probes. Use this for high-cadence
7159    /// pollers (admin dashboards, depth-target producers, soak
7160    /// observability); use [`Self::queue_counts`] for admin tooling
7161    /// that needs the exact terminal count.
7162    #[tracing::instrument(skip(self, pool), fields(queue = %queue), name = "queue_storage.queue_counts_fast")]
7163    pub async fn queue_counts_fast(
7164        &self,
7165        pool: &PgPool,
7166        queue: &str,
7167    ) -> Result<QueueCounts, AwaError> {
7168        let schema = self.schema();
7169        let queues = self.physical_queues_for_logical(queue);
7170        // available: dispatcher signal — already an index-only sum over
7171        // the (queue, priority, enqueue_shard) head tables.
7172        let available = self.queue_claimer_signal(pool, queue).await?.available;
7173        // running: leases.state = 'running' only, matching
7174        // queue_counts_exact's strict definition. Receipt-plane claims
7175        // that haven't materialised a lease row yet are documented as
7176        // omitted in the method-level doc.
7177        let running: i64 = sqlx::query_scalar(&format!(
7178            r#"
7179            SELECT COALESCE(count(*)::bigint, 0)
7180            FROM {schema}.leases
7181            WHERE queue = ANY($1)
7182              AND state = 'running'
7183            "#
7184        ))
7185        .bind(&queues)
7186        .fetch_one(pool)
7187        .await
7188        .map_err(map_sqlx_error)?;
7189        // terminal: denormalised rollup only. Live (un-rolled-up)
7190        // done_entries rows are excluded — see method-level docs. The
7191        // GREATEST legacy dedupe applies to the completed column only:
7192        // queue_lanes never carried a failed column.
7193        let (pruned_completed, pruned_failed): (i64, i64) = sqlx::query_as(&format!(
7194            r#"
7195            SELECT
7196                COALESCE(sum(GREATEST(
7197                    COALESCE(lanes.pruned_completed_count, 0),
7198                    COALESCE(rollups.pruned_completed_count, 0)
7199                )), 0)::bigint,
7200                COALESCE(sum(COALESCE(rollups.pruned_failed_count, 0)), 0)::bigint
7201            FROM (
7202                SELECT queue, priority, pruned_completed_count
7203                FROM {schema}.queue_lanes
7204                WHERE queue = ANY($1)
7205            ) AS lanes
7206            FULL OUTER JOIN (
7207                SELECT queue, priority, pruned_completed_count, pruned_failed_count
7208                FROM {schema}.queue_terminal_rollups
7209                WHERE queue = ANY($1)
7210            ) AS rollups
7211            USING (queue, priority)
7212            "#
7213        ))
7214        .bind(&queues)
7215        .fetch_one(pool)
7216        .await
7217        .map_err(map_sqlx_error)?;
7218        Ok(QueueCounts {
7219            available,
7220            running,
7221            terminal: pruned_completed + pruned_failed,
7222            pruned_failed,
7223        })
7224    }
7225
7226    /// Cumulative count of `failed` terminal rows pruned past the
7227    /// retention floor for a queue, summed over the queue's
7228    /// `queue_terminal_rollups` priorities (and stripes, for striped
7229    /// queues). Monotonic — rollups never decrease. These rows no
7230    /// longer exist in `done_entries` and cannot be retried.
7231    pub async fn pruned_failed_count_for_queue(
7232        &self,
7233        pool: &PgPool,
7234        queue: &str,
7235    ) -> Result<u64, AwaError> {
7236        let schema = self.schema();
7237        let queues = self.physical_queues_for_logical(queue);
7238        let pruned_failed: i64 = sqlx::query_scalar(&format!(
7239            r#"
7240            SELECT COALESCE(sum(pruned_failed_count), 0)::bigint
7241            FROM {schema}.queue_terminal_rollups
7242            WHERE queue = ANY($1)
7243            "#
7244        ))
7245        .bind(&queues)
7246        .fetch_one(pool)
7247        .await
7248        .map_err(map_sqlx_error)?;
7249        Ok(pruned_failed.max(0) as u64)
7250    }
7251
7252    async fn retry_job_tx<'a>(
7253        &self,
7254        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7255        job_id: i64,
7256    ) -> Result<Option<JobRow>, AwaError> {
7257        let schema = self.schema();
7258        let deleted_waiting: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
7259            r#"
7260            DELETE FROM {schema}.leases
7261            WHERE job_id = $1
7262              AND state = 'waiting_external'
7263            RETURNING
7264                ready_slot,
7265                ready_generation,
7266                job_id,
7267                queue,
7268                state,
7269                priority,
7270                attempt,
7271                run_lease,
7272                max_attempts,
7273                lane_seq,
7274                enqueue_shard,
7275                heartbeat_at,
7276                deadline_at,
7277                attempted_at,
7278                callback_id,
7279                callback_timeout_at
7280            "#
7281        ))
7282        .bind(job_id)
7283        .fetch_all(tx.as_mut())
7284        .await
7285        .map_err(map_sqlx_error)?;
7286
7287        if !deleted_waiting.is_empty() {
7288            let waiting = self
7289                .hydrate_deleted_leases_tx(tx, deleted_waiting)
7290                .await?
7291                .into_iter()
7292                .next()
7293                .expect("deleted waiting lease");
7294            let ready_payload = Self::payload_with_attempt_state(
7295                waiting.payload.clone(),
7296                waiting.progress.clone(),
7297            )?;
7298            let ready_row = ExistingReadyRow {
7299                attempt: 0,
7300                run_at: Utc::now(),
7301                attempted_at: None,
7302                ..waiting.clone().into_ready_row(Utc::now(), ready_payload)
7303            };
7304            self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(waiting.state))
7305                .await?;
7306            self.notify_queues_tx(tx, std::iter::once(waiting.queue.clone()))
7307                .await?;
7308            return Ok(Some(
7309                ReadyJobRow {
7310                    job_id: ready_row.job_id,
7311                    kind: ready_row.kind,
7312                    queue: ready_row.queue,
7313                    args: ready_row.args,
7314                    priority: ready_row.priority,
7315                    attempt: ready_row.attempt,
7316                    run_lease: ready_row.run_lease,
7317                    max_attempts: ready_row.max_attempts,
7318                    run_at: ready_row.run_at,
7319                    attempted_at: ready_row.attempted_at,
7320                    created_at: ready_row.created_at,
7321                    unique_key: ready_row.unique_key,
7322                    payload: ready_row.payload,
7323                }
7324                .into_job_row()?,
7325            ));
7326        }
7327
7328        let done_projection = done_row_projection("done", "ready");
7329        let ready_join = done_ready_join(schema, "done", "ready");
7330        let terminal: Option<DoneJobRow> = sqlx::query_as(&format!(
7331            r#"
7332            WITH deleted AS (
7333                DELETE FROM {schema}.done_entries
7334                WHERE (job_id, finalized_at) IN (
7335                SELECT job_id, finalized_at
7336                FROM {schema}.done_entries
7337                WHERE job_id = $1
7338                  AND state IN ('failed', 'cancelled')
7339                ORDER BY finalized_at DESC
7340                LIMIT 1
7341                FOR UPDATE SKIP LOCKED
7342            )
7343                RETURNING *
7344            )
7345            SELECT {done_projection}
7346            FROM deleted AS done
7347            {ready_join}
7348            "#
7349        ))
7350        .bind(job_id)
7351        .fetch_optional(tx.as_mut())
7352        .await
7353        .map_err(map_sqlx_error)?;
7354
7355        if let Some(terminal) = terminal {
7356            self.ensure_terminal_removed_receipt_closures_tx(tx, std::slice::from_ref(&terminal))
7357                .await?;
7358            // The DELETE FROM done_entries above removes one terminal row;
7359            // append a negative delta so exact counts stay in lockstep.
7360            self.decrement_live_terminal_counters_tx(
7361                tx,
7362                &Self::done_rows_to_counter_keys(std::slice::from_ref(&terminal)),
7363            )
7364            .await?;
7365            let ready_row = ExistingReadyRow {
7366                job_id: terminal.job_id,
7367                kind: terminal.kind,
7368                queue: terminal.queue.clone(),
7369                args: terminal.args,
7370                priority: terminal.priority,
7371                attempt: 0,
7372                run_lease: terminal.run_lease,
7373                max_attempts: terminal.max_attempts,
7374                run_at: Utc::now(),
7375                attempted_at: None,
7376                created_at: terminal.created_at,
7377                unique_key: terminal.unique_key,
7378                unique_states: terminal.unique_states,
7379                payload: terminal.payload,
7380            };
7381            self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(terminal.state))
7382                .await?;
7383            self.notify_queues_tx(tx, std::iter::once(terminal.queue.clone()))
7384                .await?;
7385            return Ok(Some(
7386                ReadyJobRow {
7387                    job_id: ready_row.job_id,
7388                    kind: ready_row.kind,
7389                    queue: ready_row.queue,
7390                    args: ready_row.args,
7391                    priority: ready_row.priority,
7392                    attempt: ready_row.attempt,
7393                    run_lease: ready_row.run_lease,
7394                    max_attempts: ready_row.max_attempts,
7395                    run_at: ready_row.run_at,
7396                    attempted_at: ready_row.attempted_at,
7397                    created_at: ready_row.created_at,
7398                    unique_key: ready_row.unique_key,
7399                    payload: ready_row.payload,
7400                }
7401                .into_job_row()?,
7402            ));
7403        }
7404
7405        Ok(None)
7406    }
7407
7408    pub async fn retry_job(&self, pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
7409        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
7410        let row = self.retry_job_tx(&mut tx, job_id).await?;
7411        tx.commit().await.map_err(map_sqlx_error)?;
7412        Ok(row)
7413    }
7414
7415    /// Retry jobs by id inside a single transaction. Duplicate ids are
7416    /// collapsed so each job is attempted at most once. Returns the
7417    /// retried rows together with the number of unique ids attempted:
7418    /// ids whose terminal row raced to another state or was pruned
7419    /// between the caller's scan and this call are skipped, so
7420    /// `attempted - retried.len()` is the count of unique ids that were
7421    /// requested but not retried.
7422    pub async fn retry_jobs_by_ids(
7423        &self,
7424        pool: &PgPool,
7425        ids: &[i64],
7426    ) -> Result<(Vec<JobRow>, u64), AwaError> {
7427        let unique_ids: BTreeSet<i64> = ids.iter().copied().collect();
7428        if unique_ids.is_empty() {
7429            return Ok((Vec::new(), 0));
7430        }
7431
7432        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
7433        let mut rows = Vec::with_capacity(unique_ids.len());
7434        for job_id in &unique_ids {
7435            if let Some(row) = self.retry_job_tx(&mut tx, *job_id).await? {
7436                rows.push(row);
7437            }
7438        }
7439        tx.commit().await.map_err(map_sqlx_error)?;
7440        Ok((rows, unique_ids.len() as u64))
7441    }
7442
7443    /// Write a `<outcome>` closure row for any matching open receipt.
7444    /// Idempotent: no-op if no row-local claim or compact claim-batch item
7445    /// exists for the `(job_id, run_lease)` pair, or if a closure already exists. Used
7446    /// by the admin cancel path to keep the receipt plane consistent
7447    /// with the job's new terminal state so rescue doesn't revive it.
7448    ///
7449    /// `FOR UPDATE` on the inner SELECT serialises the closure write
7450    /// against `ensure_running_leases_from_receipts_tx`
7451    /// (which also takes `FOR UPDATE` on the same claim evidence) and
7452    /// against concurrent rescue / re-close paths that might race the
7453    /// same `(job_id, run_lease)`. Without it, materialization could
7454    /// see the claim evidence, decide to materialize, and a concurrent
7455    /// admin cancel could write the closure between materialization's
7456    /// SELECT and the lease INSERT — leaving a `running` lease for a
7457    /// closed claim that admin cancel believes is fully shut down.
7458    async fn close_receipt_tx<'a>(
7459        &self,
7460        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7461        job_id: i64,
7462        run_lease: i64,
7463        outcome: &str,
7464    ) -> Result<(), AwaError> {
7465        self.close_receipt_pairs_tx(tx, &[(job_id, run_lease)], outcome)
7466            .await
7467    }
7468
7469    async fn close_receipt_pairs_tx<'a>(
7470        &self,
7471        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7472        pairs: &[(i64, i64)],
7473        outcome: &str,
7474    ) -> Result<(), AwaError> {
7475        if pairs.is_empty() || !self.lease_claim_receipts() {
7476            return Ok(());
7477        }
7478
7479        let unique_pairs: BTreeSet<(i64, i64)> = pairs.iter().copied().collect();
7480        let unique_jobs: Vec<(i64, i64)> = unique_pairs.iter().copied().collect();
7481        self.lock_receipt_attempts_tx(tx, &unique_jobs).await?;
7482
7483        let job_ids: Vec<i64> = unique_pairs.iter().map(|(job_id, _)| *job_id).collect();
7484        let run_leases: Vec<i64> = unique_pairs
7485            .iter()
7486            .map(|(_, run_lease)| *run_lease)
7487            .collect();
7488        let schema = self.schema();
7489        let closure_rel = format!("{schema}.lease_claim_closures");
7490        let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
7491        let closed_evidence =
7492            receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
7493
7494        sqlx::query(&format!(
7495            r#"
7496            WITH refs(job_id, run_lease) AS (
7497                SELECT * FROM unnest($1::bigint[], $2::bigint[])
7498            ),
7499            locked_claims AS (
7500                SELECT claims.claim_slot, claims.job_id, claims.run_lease
7501                FROM {schema}.lease_claims AS claims
7502                JOIN refs
7503                  ON refs.job_id = claims.job_id
7504                 AND refs.run_lease = claims.run_lease
7505                WHERE NOT {closed_evidence}
7506                FOR UPDATE OF claims
7507            ),
7508            locked_batch_claims AS (
7509                SELECT
7510                    claim_batches.claim_slot,
7511                    claim_batches.ready_slot,
7512                    claim_batches.ready_generation,
7513                    items.job_id,
7514                    items.run_lease,
7515                    items.receipt_id
7516                FROM {schema}.lease_claim_batches AS claim_batches
7517                CROSS JOIN LATERAL unnest(
7518                    claim_batches.job_ids,
7519                    claim_batches.run_leases,
7520                    claim_batches.receipt_ids
7521                ) AS items(job_id, run_lease, receipt_id)
7522                JOIN refs
7523                  ON refs.job_id = items.job_id
7524                 AND refs.run_lease = items.run_lease
7525                WHERE NOT EXISTS (
7526                      SELECT 1
7527                      FROM {schema}.lease_claim_closures AS closures
7528                      WHERE closures.claim_slot = claim_batches.claim_slot
7529                        AND closures.job_id = items.job_id
7530                        AND closures.run_lease = items.run_lease
7531                  )
7532                  AND NOT EXISTS (
7533                      SELECT 1
7534                      FROM {schema}.lease_claim_closure_batches AS closure_batches
7535                      WHERE closure_batches.claim_slot = claim_batches.claim_slot
7536                        AND closure_batches.receipt_ranges @> items.receipt_id
7537                  )
7538                  AND NOT EXISTS (
7539                      SELECT 1
7540                      FROM {schema}.done_entries AS done
7541                      WHERE done.job_id = items.job_id
7542                        AND done.run_lease = items.run_lease
7543                  )
7544                  AND NOT EXISTS (
7545                      SELECT 1
7546                      FROM {schema}.deferred_jobs AS deferred
7547                      WHERE deferred.job_id = items.job_id
7548                        AND deferred.run_lease = items.run_lease
7549                  )
7550                  AND NOT EXISTS (
7551                      SELECT 1
7552                      FROM {schema}.dlq_entries AS dlq
7553                      WHERE dlq.job_id = items.job_id
7554                        AND dlq.run_lease = items.run_lease
7555                  )
7556                FOR UPDATE OF claim_batches
7557            ),
7558            -- Row-sourced (deadline lease) claims keep their explicit
7559            -- closure: the queue/claim prune gates balance those against
7560            -- the lease_claims row via the closure JOIN. Compact
7561            -- batch-sourced claims have no lease_claims row to JOIN, so
7562            -- they are closed into lease_claim_closure_batches below.
7563            closing_claims AS (
7564                SELECT
7565                    locked_claims.claim_slot,
7566                    locked_claims.job_id,
7567                    locked_claims.run_lease,
7568                    $3 AS outcome,
7569                    clock_timestamp() AS closed_at
7570                FROM locked_claims
7571            ),
7572            inserted AS (
7573                INSERT INTO {schema}.lease_claim_closures (
7574                    claim_slot,
7575                    job_id,
7576                    run_lease,
7577                    outcome,
7578                    closed_at
7579                )
7580                SELECT
7581                    claim_slot,
7582                    job_id,
7583                    run_lease,
7584                    outcome,
7585                    closed_at
7586                FROM closing_claims
7587                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
7588                RETURNING claim_slot, job_id, run_lease, closed_at
7589            ),
7590            inserted_batches AS (
7591                INSERT INTO {schema}.lease_claim_closure_batches (
7592                    claim_slot,
7593                    ready_slot,
7594                    ready_generation,
7595                    outcome,
7596                    closed_count,
7597                    receipt_ids,
7598                    receipt_ranges,
7599                    closed_at
7600                )
7601                SELECT
7602                    locked_batch_claims.claim_slot,
7603                    locked_batch_claims.ready_slot,
7604                    locked_batch_claims.ready_generation,
7605                    $3,
7606                    count(*)::int,
7607                    array_agg(locked_batch_claims.receipt_id ORDER BY locked_batch_claims.receipt_id),
7608                    range_agg(int8range(locked_batch_claims.receipt_id, locked_batch_claims.receipt_id + 1, '[)') ORDER BY locked_batch_claims.receipt_id),
7609                    clock_timestamp()
7610                FROM locked_batch_claims
7611                GROUP BY
7612                    locked_batch_claims.claim_slot,
7613                    locked_batch_claims.ready_slot,
7614                    locked_batch_claims.ready_generation
7615                RETURNING claim_slot
7616            ),
7617            marked AS (
7618                UPDATE {schema}.lease_claims AS claims
7619                SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
7620                FROM inserted
7621                WHERE claims.claim_slot = inserted.claim_slot
7622                  AND claims.job_id = inserted.job_id
7623                  AND claims.run_lease = inserted.run_lease
7624                RETURNING claims.job_id
7625            )
7626            SELECT
7627                (SELECT count(*) FROM marked)
7628                + (SELECT count(*) FROM inserted_batches)
7629            "#
7630        ))
7631        .bind(&job_ids)
7632        .bind(&run_leases)
7633        .bind(outcome)
7634        .execute(tx.as_mut())
7635        .await
7636        .map_err(map_sqlx_error)?;
7637
7638        Ok(())
7639    }
7640
7641    /// Emit a `pg_notify('awa:cancel', ...)` inside the cancel
7642    /// transaction so any worker runtime currently executing this
7643    /// `(job_id, run_lease)` learns about the cancellation on commit
7644    /// and fires its in-flight cancel flag. Notifications are
7645    /// automatically discarded on rollback.
7646    async fn notify_cancellation_tx<'a>(
7647        &self,
7648        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7649        job_id: i64,
7650        run_lease: i64,
7651    ) -> Result<(), AwaError> {
7652        let payload = serde_json::json!({ "job_id": job_id, "run_lease": run_lease }).to_string();
7653        sqlx::query("SELECT pg_notify('awa:cancel', $1)")
7654            .bind(payload)
7655            .execute(tx.as_mut())
7656            .await
7657            .map_err(map_sqlx_error)?;
7658        Ok(())
7659    }
7660
7661    async fn cancel_job_tx<'a>(
7662        &self,
7663        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
7664        job_id: i64,
7665    ) -> Result<Option<CancelJobTxResult>, AwaError> {
7666        let schema = self.schema();
7667        let ready: Option<ReadyTransitionRow> = sqlx::query_as(&format!(
7668            r#"
7669            WITH target AS (
7670                SELECT ready.*
7671                FROM {schema}.ready_entries AS ready
7672                JOIN {schema}.queue_claim_heads AS claims
7673                  ON claims.queue = ready.queue
7674                 AND claims.priority = ready.priority
7675                 AND claims.enqueue_shard = ready.enqueue_shard
7676                WHERE ready.job_id = $1
7677                  AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
7678                  AND NOT EXISTS (
7679                      SELECT 1 FROM {schema}.ready_tombstones AS tomb
7680                      WHERE tomb.queue = ready.queue
7681                        AND tomb.priority = ready.priority
7682                        AND tomb.enqueue_shard = ready.enqueue_shard
7683                        AND tomb.lane_seq = ready.lane_seq
7684                        AND tomb.ready_slot = ready.ready_slot
7685                        AND tomb.ready_generation = ready.ready_generation
7686                  )
7687                ORDER BY ready.lane_seq DESC
7688                LIMIT 1
7689                FOR UPDATE OF ready SKIP LOCKED
7690            ),
7691            tombstone AS (
7692                INSERT INTO {schema}.ready_tombstones (
7693                    ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
7694                )
7695                SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
7696                FROM target
7697                ON CONFLICT DO NOTHING
7698            )
7699            SELECT
7700                ready_slot,
7701                ready_generation,
7702                job_id,
7703                kind,
7704                queue,
7705                args,
7706                priority,
7707                attempt,
7708                run_lease,
7709                max_attempts,
7710                lane_seq,
7711                enqueue_shard,
7712                run_at,
7713                attempted_at,
7714                created_at,
7715                unique_key,
7716                unique_states,
7717                COALESCE(payload, '{{}}'::jsonb) AS payload
7718            FROM target
7719            "#
7720        ))
7721        .bind(job_id)
7722        .fetch_optional(tx.as_mut())
7723        .await
7724        .map_err(map_sqlx_error)?;
7725
7726        if let Some(ready) = ready {
7727            let done =
7728                ready
7729                    .clone()
7730                    .into_done_row(JobState::Cancelled, Utc::now(), ready.payload.clone());
7731            self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Available))
7732                .await?;
7733            // If the cancelled lane was *exactly* at the claim head,
7734            // advance the head past it so the derived cursor-difference count
7735            // drops by 1 immediately.
7736            // When other unclaimed lanes still sit between claim_seq
7737            // and the cancelled lane_seq, leave claim_seq alone —
7738            // advancing would skip past those still-claimable rows.
7739            // Non-head deletes can leave a bounded dispatcher over-count until
7740            // later committed rows on the lane are claimed.
7741            let claim_cursor_advance = ClaimCursorAdvance {
7742                queue: ready.queue.clone(),
7743                priority: ready.priority,
7744                enqueue_shard: ready.enqueue_shard,
7745                next_seq: ready.lane_seq + 1,
7746                only_if_current: Some(ready.lane_seq),
7747            };
7748            return Ok(Some(CancelJobTxResult {
7749                row: done.into_job_row()?,
7750                claim_cursor_advance: Some(claim_cursor_advance),
7751            }));
7752        }
7753
7754        let deleted_lease: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
7755            r#"
7756            DELETE FROM {schema}.leases
7757            WHERE job_id = $1
7758              AND state IN ('running', 'waiting_external')
7759            RETURNING
7760                ready_slot,
7761                ready_generation,
7762                job_id,
7763                queue,
7764                state,
7765                priority,
7766                attempt,
7767                run_lease,
7768                max_attempts,
7769                lane_seq,
7770                enqueue_shard,
7771                heartbeat_at,
7772                deadline_at,
7773                attempted_at,
7774                callback_id,
7775                callback_timeout_at
7776            "#
7777        ))
7778        .bind(job_id)
7779        .fetch_all(tx.as_mut())
7780        .await
7781        .map_err(map_sqlx_error)?;
7782
7783        if !deleted_lease.is_empty() {
7784            let lease = self
7785                .hydrate_deleted_leases_tx(tx, deleted_lease)
7786                .await?
7787                .into_iter()
7788                .next()
7789                .expect("deleted running lease");
7790            let done_payload =
7791                Self::payload_with_attempt_state(lease.payload.clone(), lease.progress.clone())?;
7792            let done = lease
7793                .clone()
7794                .into_done_row(JobState::Cancelled, Utc::now(), done_payload);
7795            self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(lease.state))
7796                .await?;
7797            // Receipt-plane consistency: close any matching open
7798            // receipt so the ADR-023 anti-join no longer considers this
7799            // attempt live, and rescue doesn't try to revive it.
7800            self.close_receipt_tx(tx, lease.job_id, lease.run_lease, "cancelled")
7801                .await?;
7802            // Wake any worker currently executing this attempt.
7803            self.notify_cancellation_tx(tx, lease.job_id, lease.run_lease)
7804                .await?;
7805            return Ok(Some(CancelJobTxResult {
7806                row: done.into_job_row()?,
7807                claim_cursor_advance: None,
7808            }));
7809        }
7810
7811        // ADR-023 receipt-only cancel: the job may be running on a
7812        // receipt-backed short path that never materialized a `leases`
7813        // row. Find it by anti-joining lease_claims with
7814        // lease_claim_closures, cancel it by writing a closure and a
7815        // done row, and notify listening workers.
7816        if self.lease_claim_receipts() {
7817            let closure_rel = format!("{schema}.lease_claim_closures");
7818            let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
7819            let closed_evidence =
7820                receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
7821            type ReceiptCancelRow = (
7822                i32,
7823                i64,
7824                i32,
7825                i64,
7826                String,
7827                i16,
7828                i16,
7829                i16,
7830                i64,
7831                DateTime<Utc>,
7832                i64,
7833                bool,
7834            );
7835            let receipt: Option<ReceiptCancelRow> = sqlx::query_as(&format!(
7836                r#"
7837                WITH row_receipt AS (
7838                    SELECT
7839                        claims.claim_slot,
7840                        claims.run_lease,
7841                        claims.ready_slot,
7842                        claims.ready_generation,
7843                        claims.queue,
7844                        claims.priority,
7845                        claims.attempt,
7846                        claims.max_attempts,
7847                        claims.lane_seq,
7848                        claims.claimed_at,
7849                        claims.receipt_id,
7850                        false AS compact_batch
7851                    FROM {schema}.lease_claims AS claims
7852                    WHERE claims.job_id = $1
7853                      AND NOT {closed_evidence}
7854                    ORDER BY claims.run_lease DESC
7855                    LIMIT 1
7856                    FOR UPDATE OF claims SKIP LOCKED
7857                ),
7858                batch_receipt AS (
7859                    SELECT
7860                        claim_batches.claim_slot,
7861                        items.run_lease,
7862                        claim_batches.ready_slot,
7863                        claim_batches.ready_generation,
7864                        claim_batches.queue,
7865                        claim_batches.priority,
7866                        items.attempt,
7867                        items.max_attempts,
7868                        items.lane_seq,
7869                        claim_batches.claimed_at,
7870                        items.receipt_id,
7871                        true AS compact_batch
7872                    FROM {schema}.lease_claim_batches AS claim_batches
7873                    CROSS JOIN LATERAL unnest(
7874                        claim_batches.job_ids,
7875                        claim_batches.run_leases,
7876                        claim_batches.receipt_ids,
7877                        claim_batches.lane_seqs,
7878                        claim_batches.attempts,
7879                        claim_batches.max_attempts
7880                    ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
7881                    WHERE items.job_id = $1
7882                      AND NOT EXISTS (
7883                          SELECT 1
7884                          FROM {schema}.lease_claim_closures AS closures
7885                          WHERE closures.claim_slot = claim_batches.claim_slot
7886                            AND closures.job_id = items.job_id
7887                            AND closures.run_lease = items.run_lease
7888                      )
7889                      AND NOT EXISTS (
7890                          SELECT 1
7891                          FROM {schema}.lease_claim_closure_batches AS closure_batches
7892                          WHERE closure_batches.claim_slot = claim_batches.claim_slot
7893                            AND closure_batches.receipt_ranges @> items.receipt_id
7894                      )
7895                      AND NOT EXISTS (
7896                          SELECT 1
7897                          FROM {schema}.leases AS lease
7898                          WHERE lease.job_id = items.job_id
7899                            AND lease.run_lease = items.run_lease
7900                      )
7901                      AND NOT EXISTS (
7902                          SELECT 1 FROM {schema}.done_entries AS done
7903                          WHERE done.job_id = items.job_id
7904                            AND done.run_lease = items.run_lease
7905                      )
7906                      AND NOT EXISTS (
7907                          SELECT 1 FROM {schema}.deferred_jobs AS deferred
7908                          WHERE deferred.job_id = items.job_id
7909                            AND deferred.run_lease = items.run_lease
7910                      )
7911                      AND NOT EXISTS (
7912                          SELECT 1 FROM {schema}.dlq_entries AS dlq
7913                          WHERE dlq.job_id = items.job_id
7914                            AND dlq.run_lease = items.run_lease
7915                      )
7916                    ORDER BY items.run_lease DESC
7917                    LIMIT 1
7918                    FOR UPDATE OF claim_batches SKIP LOCKED
7919                )
7920                SELECT *
7921                FROM (
7922                    SELECT * FROM row_receipt
7923                    UNION ALL
7924                    SELECT * FROM batch_receipt
7925                ) AS receipt
7926                ORDER BY run_lease DESC
7927                LIMIT 1
7928                "#
7929            ))
7930            .bind(job_id)
7931            .fetch_optional(tx.as_mut())
7932            .await
7933            .map_err(map_sqlx_error)?;
7934
7935            if let Some((
7936                claim_slot,
7937                run_lease,
7938                ready_slot,
7939                ready_generation,
7940                queue,
7941                priority,
7942                attempt,
7943                max_attempts,
7944                lane_seq,
7945                claimed_at,
7946                receipt_id,
7947                compact_batch,
7948            )) = receipt
7949            {
7950                // Hydrate the ready row so we can synthesize the done
7951                // row with the original args/payload.
7952                let ready_match: Option<ReadyTransitionRow> = sqlx::query_as(&format!(
7953                    r#"
7954                    SELECT
7955                        ready_slot,
7956                        ready_generation,
7957                        job_id,
7958                        kind,
7959                        queue,
7960                        args,
7961                        priority,
7962                        attempt,
7963                        run_lease,
7964                        max_attempts,
7965                        lane_seq,
7966                        enqueue_shard,
7967                        run_at,
7968                        attempted_at,
7969                        created_at,
7970                        unique_key,
7971                        unique_states,
7972                        COALESCE(payload, '{{}}'::jsonb) AS payload
7973                    FROM {schema}.ready_entries
7974                    WHERE job_id = $1
7975                      AND ready_slot = $2
7976                      AND ready_generation = $3
7977                      AND queue = $4
7978                      AND lane_seq = $5
7979                    "#
7980                ))
7981                .bind(job_id)
7982                .bind(ready_slot)
7983                .bind(ready_generation)
7984                .bind(&queue)
7985                .bind(lane_seq)
7986                .fetch_optional(tx.as_mut())
7987                .await
7988                .map_err(map_sqlx_error)?;
7989
7990                let Some(ready) = ready_match else {
7991                    // Shouldn't happen — the claim references a ready
7992                    // row. Fall through to the deferred / not-found
7993                    // branches.
7994                    return Ok(None);
7995                };
7996
7997                let done = DoneJobRow {
7998                    ready_slot,
7999                    ready_generation,
8000                    job_id,
8001                    kind: ready.kind,
8002                    queue: queue.clone(),
8003                    args: ready.args,
8004                    state: JobState::Cancelled,
8005                    priority,
8006                    attempt,
8007                    run_lease,
8008                    max_attempts,
8009                    lane_seq,
8010                    enqueue_shard: ready.enqueue_shard,
8011                    run_at: ready.run_at,
8012                    attempted_at: Some(claimed_at),
8013                    finalized_at: Utc::now(),
8014                    created_at: ready.created_at,
8015                    unique_key: ready.unique_key,
8016                    unique_states: ready.unique_states,
8017                    payload: ready.payload,
8018                };
8019                self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Running))
8020                    .await?;
8021                // Write the closure into the same claim partition. A
8022                // compact batch-sourced claim has no lease_claims row, so
8023                // it closes into the batch ledger the queue prune gate
8024                // counts via compact_count; a deadline lease claim keeps
8025                // its explicit closure so the gate balances it against the
8026                // lease_claims row.
8027                if compact_batch {
8028                    sqlx::query(&format!(
8029                        r#"
8030                        INSERT INTO {schema}.lease_claim_closure_batches (
8031                            claim_slot,
8032                            ready_slot,
8033                            ready_generation,
8034                            outcome,
8035                            closed_count,
8036                            receipt_ids,
8037                            receipt_ranges,
8038                            closed_at
8039                        )
8040                        VALUES (
8041                            $1,
8042                            $2,
8043                            $3,
8044                            'cancelled',
8045                            1,
8046                            ARRAY[$4::bigint],
8047                            int8multirange(int8range($4::bigint, $4::bigint + 1, '[)')),
8048                            clock_timestamp()
8049                        )
8050                        "#
8051                    ))
8052                    .bind(claim_slot)
8053                    .bind(ready_slot)
8054                    .bind(ready_generation)
8055                    .bind(receipt_id)
8056                    .execute(tx.as_mut())
8057                    .await
8058                    .map_err(map_sqlx_error)?;
8059                } else {
8060                    sqlx::query(&format!(
8061                        r#"
8062                        WITH inserted AS (
8063                            INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
8064                            VALUES ($1, $2, $3, 'cancelled', clock_timestamp())
8065                            ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
8066                            RETURNING claim_slot, job_id, run_lease, closed_at
8067                        ),
8068                        marked AS (
8069                            UPDATE {schema}.lease_claims AS claims
8070                            SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
8071                            FROM inserted
8072                            WHERE claims.claim_slot = inserted.claim_slot
8073                              AND claims.job_id = inserted.job_id
8074                              AND claims.run_lease = inserted.run_lease
8075                            RETURNING claims.job_id
8076                        )
8077                        SELECT count(*) FROM marked
8078                        "#
8079                    ))
8080                    .bind(claim_slot)
8081                    .bind(job_id)
8082                    .bind(run_lease)
8083                    .execute(tx.as_mut())
8084                    .await
8085                    .map_err(map_sqlx_error)?;
8086                }
8087                // Defensive: between the leases DELETE at the top of
8088                // this function and the FOR UPDATE on claims above, a
8089                // concurrent `ensure_running_leases_from_receipts_tx`
8090                // can have materialized a `leases` row for this
8091                // (job_id, run_lease). Materialize and we both lock the
8092                // same claim evidence; whichever ran first commits, the
8093                // other replays under the new snapshot. If materialize
8094                // committed first, that lease is now an orphan pointing
8095                // at a job we're about to mark `cancelled`. Sweep it
8096                // defensively. If no race occurred this is a no-op.
8097                sqlx::query(&format!(
8098                    "DELETE FROM {schema}.leases WHERE job_id = $1 AND run_lease = $2"
8099                ))
8100                .bind(job_id)
8101                .bind(run_lease)
8102                .execute(tx.as_mut())
8103                .await
8104                .map_err(map_sqlx_error)?;
8105                self.notify_cancellation_tx(tx, job_id, run_lease).await?;
8106                return Ok(Some(CancelJobTxResult {
8107                    row: done.into_job_row()?,
8108                    claim_cursor_advance: None,
8109                }));
8110            }
8111        }
8112
8113        let deferred: Option<DeferredJobRow> = sqlx::query_as(&format!(
8114            r#"
8115            DELETE FROM {schema}.deferred_jobs
8116            WHERE job_id = $1
8117              AND state IN ('scheduled', 'retryable')
8118            RETURNING
8119                job_id,
8120                kind,
8121                queue,
8122                args,
8123                state,
8124                priority,
8125                attempt,
8126                run_lease,
8127                max_attempts,
8128                run_at,
8129                attempted_at,
8130                finalized_at,
8131                created_at,
8132                unique_key,
8133                unique_states,
8134                COALESCE(payload, '{{}}'::jsonb) AS payload
8135            "#
8136        ))
8137        .bind(job_id)
8138        .fetch_optional(tx.as_mut())
8139        .await
8140        .map_err(map_sqlx_error)?;
8141
8142        if let Some(deferred) = deferred {
8143            let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?;
8144            // A deferred-cancel never observed a claim, so it has no
8145            // shard assignment to inherit. The synthesized terminal row
8146            // is parked on shard 0 with a synthetic negative `lane_seq`
8147            // that keeps the `done_entries` PK
8148            // `(ready_slot, queue, priority, enqueue_shard, lane_seq)`
8149            // unique. `ensure_lane` registers shard 0 for the queue so
8150            // the lane row exists regardless of producer activity.
8151            self.ensure_lane(tx, &deferred.queue, deferred.priority, 0)
8152                .await?;
8153            let done = DoneJobRow {
8154                ready_slot,
8155                ready_generation,
8156                job_id: deferred.job_id,
8157                kind: deferred.kind,
8158                queue: deferred.queue.clone(),
8159                args: deferred.args,
8160                state: JobState::Cancelled,
8161                priority: deferred.priority,
8162                attempt: deferred.attempt,
8163                run_lease: deferred.run_lease,
8164                max_attempts: deferred.max_attempts,
8165                lane_seq: -deferred.job_id,
8166                enqueue_shard: 0,
8167                run_at: deferred.run_at,
8168                attempted_at: deferred.attempted_at,
8169                finalized_at: Utc::now(),
8170                created_at: deferred.created_at,
8171                unique_key: deferred.unique_key,
8172                unique_states: deferred.unique_states,
8173                payload: deferred.payload,
8174            };
8175            self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(deferred.state))
8176                .await?;
8177            return Ok(Some(CancelJobTxResult {
8178                row: done.into_job_row()?,
8179                claim_cursor_advance: None,
8180            }));
8181        }
8182
8183        Ok(None)
8184    }
8185
8186    pub async fn cancel_job(&self, pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
8187        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8188        let result = self.cancel_job_tx(&mut tx, job_id).await?;
8189        tx.commit().await.map_err(map_sqlx_error)?;
8190        if let Some(result) = result {
8191            if let Some(advance) = result.claim_cursor_advance.as_ref() {
8192                self.advance_claim_cursors(pool, std::slice::from_ref(advance))
8193                    .await;
8194            }
8195            Ok(Some(result.row))
8196        } else {
8197            Ok(None)
8198        }
8199    }
8200
8201    /// Transaction-scoped cancel used by [`crate::admin::cancel_tx`]. Runs the
8202    /// full cancellation on the caller's transaction and returns the cancelled
8203    /// row. Unlike [`Self::cancel_job`] this does NOT perform the post-commit
8204    /// claim-cursor advance, so the derived queue depth may briefly over-count
8205    /// by one until later committed rows on the lane are claimed.
8206    pub(crate) async fn cancel_job_in_tx<'a>(
8207        &self,
8208        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8209        job_id: i64,
8210    ) -> Result<Option<JobRow>, AwaError> {
8211        Ok(self
8212            .cancel_job_tx(tx, job_id)
8213            .await?
8214            .map(|result| result.row))
8215    }
8216
8217    pub async fn cancel_jobs_by_ids(
8218        &self,
8219        pool: &PgPool,
8220        ids: &[i64],
8221    ) -> Result<Vec<JobRow>, AwaError> {
8222        if ids.is_empty() {
8223            return Ok(Vec::new());
8224        }
8225
8226        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8227        let mut rows = Vec::with_capacity(ids.len());
8228        let mut claim_cursor_advances = Vec::new();
8229        for job_id in ids {
8230            if let Some(result) = self.cancel_job_tx(&mut tx, *job_id).await? {
8231                if let Some(advance) = result.claim_cursor_advance {
8232                    claim_cursor_advances.push(advance);
8233                }
8234                rows.push(result.row);
8235            }
8236        }
8237        tx.commit().await.map_err(map_sqlx_error)?;
8238        self.advance_claim_cursors(pool, &claim_cursor_advances)
8239            .await;
8240        Ok(rows)
8241    }
8242
8243    pub async fn set_priority(
8244        &self,
8245        pool: &PgPool,
8246        job_id: i64,
8247        priority: i16,
8248    ) -> Result<bool, AwaError> {
8249        if !(1..=4).contains(&priority) {
8250            return Err(AwaError::Validation(
8251                "priority must be between 1 and 4".to_string(),
8252            ));
8253        }
8254
8255        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8256        let result = self.set_priority_tx(&mut tx, job_id, priority).await?;
8257        tx.commit().await.map_err(map_sqlx_error)?;
8258        Ok(result)
8259    }
8260
8261    pub async fn set_priority_tx<'a>(
8262        &self,
8263        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8264        job_id: i64,
8265        priority: i16,
8266    ) -> Result<bool, AwaError> {
8267        if !(1..=4).contains(&priority) {
8268            return Err(AwaError::Validation(
8269                "priority must be between 1 and 4".to_string(),
8270            ));
8271        }
8272
8273        if self
8274            .update_deferred_batch_fields_tx(tx, job_id, None, Some(priority))
8275            .await?
8276        {
8277            return Ok(true);
8278        }
8279
8280        let result = self
8281            .move_ready_batch_fields_tx(tx, job_id, None, Some(priority))
8282            .await?;
8283        Ok(result.moved)
8284    }
8285
8286    pub async fn move_queue(
8287        &self,
8288        pool: &PgPool,
8289        job_id: i64,
8290        queue: &str,
8291        priority: Option<i16>,
8292    ) -> Result<bool, AwaError> {
8293        if queue.is_empty() || queue.len() > 200 {
8294            return Err(AwaError::Validation(
8295                "destination queue must be 1..=200 characters".to_string(),
8296            ));
8297        }
8298        if let Some(priority) = priority {
8299            if !(1..=4).contains(&priority) {
8300                return Err(AwaError::Validation(
8301                    "priority must be between 1 and 4".to_string(),
8302                ));
8303            }
8304        }
8305
8306        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8307        let result = self.move_queue_tx(&mut tx, job_id, queue, priority).await?;
8308        tx.commit().await.map_err(map_sqlx_error)?;
8309        Ok(result)
8310    }
8311
8312    pub async fn move_queue_tx<'a>(
8313        &self,
8314        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8315        job_id: i64,
8316        queue: &str,
8317        priority: Option<i16>,
8318    ) -> Result<bool, AwaError> {
8319        if queue.is_empty() || queue.len() > 200 {
8320            return Err(AwaError::Validation(
8321                "destination queue must be 1..=200 characters".to_string(),
8322            ));
8323        }
8324        if let Some(priority) = priority {
8325            if !(1..=4).contains(&priority) {
8326                return Err(AwaError::Validation(
8327                    "priority must be between 1 and 4".to_string(),
8328                ));
8329            }
8330        }
8331
8332        if self
8333            .update_deferred_batch_fields_tx(tx, job_id, Some(queue), priority)
8334            .await?
8335        {
8336            return Ok(true);
8337        }
8338
8339        let result = self
8340            .move_ready_batch_fields_tx(tx, job_id, Some(queue), priority)
8341            .await?;
8342        Ok(result.moved)
8343    }
8344
8345    async fn update_deferred_batch_fields_tx<'a>(
8346        &self,
8347        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8348        job_id: i64,
8349        queue: Option<&str>,
8350        priority: Option<i16>,
8351    ) -> Result<bool, AwaError> {
8352        let schema = self.schema();
8353        let row: Option<DeferredJobRow> = sqlx::query_as(&format!(
8354            r#"
8355            SELECT
8356                job_id,
8357                kind,
8358                queue,
8359                args,
8360                state,
8361                priority,
8362                attempt,
8363                run_lease,
8364                max_attempts,
8365                run_at,
8366                attempted_at,
8367                finalized_at,
8368                created_at,
8369                unique_key,
8370                unique_states,
8371                COALESCE(payload, '{{}}'::jsonb) AS payload
8372            FROM {schema}.deferred_jobs
8373            WHERE job_id = $1
8374              AND state = 'scheduled'
8375            FOR UPDATE SKIP LOCKED
8376            "#
8377        ))
8378        .bind(job_id)
8379        .fetch_optional(tx.as_mut())
8380        .await
8381        .map_err(map_sqlx_error)?;
8382
8383        let Some(row) = row else {
8384            return Ok(false);
8385        };
8386
8387        let old_queue = row.queue.clone();
8388        let old_priority = row.priority;
8389        let requested_queue = queue.unwrap_or(&old_queue);
8390        let old_logical_queue = self.logical_queue_name(&old_queue).to_string();
8391        let new_queue = if queue.is_some()
8392            && requested_queue != old_queue
8393            && requested_queue != old_logical_queue
8394        {
8395            self.queue_stripe_for_enqueue(requested_queue, &row.unique_key, row.job_id)
8396        } else {
8397            old_queue.clone()
8398        };
8399        let new_priority = priority.unwrap_or(old_priority);
8400        if new_queue == old_queue && new_priority == old_priority {
8401            return Ok(false);
8402        }
8403        let mut payload = RuntimePayload::from_json(row.payload)?;
8404        let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
8405            AwaError::Validation("queue storage payload metadata must be a JSON object".to_string())
8406        })?;
8407        if queue.is_some() {
8408            metadata
8409                .entry("_awa_original_queue".to_string())
8410                .or_insert_with(|| serde_json::Value::from(old_logical_queue));
8411        }
8412        if priority.is_some() {
8413            metadata
8414                .entry("_awa_original_priority".to_string())
8415                .or_insert_with(|| serde_json::Value::from(i64::from(old_priority)));
8416        }
8417
8418        sqlx::query(&format!(
8419            r#"
8420            UPDATE {schema}.deferred_jobs
8421            SET queue = $2,
8422                priority = $3,
8423                payload = $4
8424            WHERE job_id = $1
8425            "#
8426        ))
8427        .bind(job_id)
8428        .bind(new_queue)
8429        .bind(new_priority)
8430        .bind(storage_payload(&payload.into_json()))
8431        .execute(tx.as_mut())
8432        .await
8433        .map_err(map_sqlx_error)?;
8434        Ok(true)
8435    }
8436
8437    async fn move_ready_batch_fields_tx<'a>(
8438        &self,
8439        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8440        job_id: i64,
8441        queue: Option<&str>,
8442        priority: Option<i16>,
8443    ) -> Result<ReadyBatchMoveResult, AwaError> {
8444        let schema = self.schema();
8445        let ready: Option<ReadyTransitionRow> = sqlx::query_as(&format!(
8446            r#"
8447            WITH target AS (
8448                SELECT ready.*
8449                FROM {schema}.ready_entries AS ready
8450                JOIN {schema}.queue_claim_heads AS claims
8451                  ON claims.queue = ready.queue
8452                 AND claims.priority = ready.priority
8453                 AND claims.enqueue_shard = ready.enqueue_shard
8454                WHERE ready.job_id = $1
8455                  AND ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
8456                  AND NOT EXISTS (
8457                      SELECT 1 FROM {schema}.ready_tombstones AS tomb
8458                      WHERE tomb.queue = ready.queue
8459                        AND tomb.priority = ready.priority
8460                        AND tomb.enqueue_shard = ready.enqueue_shard
8461                        AND tomb.lane_seq = ready.lane_seq
8462                        AND tomb.ready_slot = ready.ready_slot
8463                        AND tomb.ready_generation = ready.ready_generation
8464                  )
8465                ORDER BY ready.lane_seq DESC
8466                LIMIT 1
8467                FOR UPDATE OF ready SKIP LOCKED
8468            )
8469            SELECT
8470                ready_slot,
8471                ready_generation,
8472                job_id,
8473                kind,
8474                queue,
8475                args,
8476                priority,
8477                attempt,
8478                run_lease,
8479                max_attempts,
8480                lane_seq,
8481                enqueue_shard,
8482                run_at,
8483                attempted_at,
8484                created_at,
8485                unique_key,
8486                unique_states,
8487                COALESCE(payload, '{{}}'::jsonb) AS payload
8488            FROM target
8489            "#
8490        ))
8491        .bind(job_id)
8492        .fetch_optional(tx.as_mut())
8493        .await
8494        .map_err(map_sqlx_error)?;
8495
8496        let Some(ready) = ready else {
8497            return Ok(ReadyBatchMoveResult { moved: false });
8498        };
8499
8500        let old_queue = ready.queue.clone();
8501        let old_priority = ready.priority;
8502        let requested_queue = queue.unwrap_or(&old_queue);
8503        let old_logical_queue = self.logical_queue_name(&old_queue).to_string();
8504        let new_queue = if queue.is_some()
8505            && requested_queue != old_queue
8506            && requested_queue != old_logical_queue
8507        {
8508            self.queue_stripe_for_enqueue(requested_queue, &ready.unique_key, ready.job_id)
8509        } else {
8510            old_queue.clone()
8511        };
8512        let new_priority = priority.unwrap_or(old_priority);
8513        if new_queue == old_queue && new_priority == old_priority {
8514            return Ok(ReadyBatchMoveResult { moved: false });
8515        }
8516        sqlx::query(&format!(
8517            r#"
8518            INSERT INTO {schema}.ready_tombstones (
8519                ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
8520            )
8521            VALUES ($1, $2, $3, $4, $5, $6, $7)
8522            ON CONFLICT DO NOTHING
8523            "#
8524        ))
8525        .bind(ready.ready_slot)
8526        .bind(ready.ready_generation)
8527        .bind(&ready.queue)
8528        .bind(ready.priority)
8529        .bind(ready.enqueue_shard)
8530        .bind(ready.lane_seq)
8531        .bind(ready.job_id)
8532        .execute(tx.as_mut())
8533        .await
8534        .map_err(map_sqlx_error)?;
8535
8536        let mut payload = RuntimePayload::from_json(ready.payload.clone())?;
8537        let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
8538            AwaError::Validation("queue storage payload metadata must be a JSON object".to_string())
8539        })?;
8540        if queue.is_some() {
8541            metadata
8542                .entry("_awa_original_queue".to_string())
8543                .or_insert_with(|| serde_json::Value::from(old_logical_queue));
8544        }
8545        if priority.is_some() {
8546            metadata
8547                .entry("_awa_original_priority".to_string())
8548                .or_insert_with(|| serde_json::Value::from(i64::from(old_priority)));
8549        }
8550
8551        let notify_queue = new_queue.clone();
8552        let ready_row = ready.into_existing_ready_row(new_queue, new_priority, payload.into_json());
8553        self.insert_existing_ready_rows_tx(tx, vec![ready_row], Some(JobState::Available))
8554            .await?;
8555        self.notify_queues_tx(tx, std::iter::once(notify_queue))
8556            .await?;
8557
8558        Ok(ReadyBatchMoveResult { moved: true })
8559    }
8560
8561    pub async fn age_waiting_priorities(
8562        &self,
8563        pool: &PgPool,
8564        aging_interval: Duration,
8565        limit: i64,
8566    ) -> Result<Vec<i64>, AwaError> {
8567        if limit <= 0 {
8568            return Ok(Vec::new());
8569        }
8570
8571        let cutoff = Utc::now()
8572            - TimeDelta::from_std(aging_interval)
8573                .map_err(|err| AwaError::Validation(format!("invalid aging interval: {err}")))?;
8574        let schema = self.schema();
8575        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8576
8577        let moved: Vec<ReadyTransitionRow> = sqlx::query_as(&format!(
8578            r#"
8579            WITH target AS (
8580                SELECT ready.*
8581                FROM {schema}.ready_entries AS ready
8582                JOIN {schema}.queue_claim_heads AS claims
8583                  ON claims.queue = ready.queue
8584                 AND claims.priority = ready.priority
8585                 AND claims.enqueue_shard = ready.enqueue_shard
8586                WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
8587                  AND ready.priority > 1
8588                  AND ready.run_at <= $1
8589                  AND NOT EXISTS (
8590                      SELECT 1 FROM {schema}.ready_tombstones AS tomb
8591                      WHERE tomb.queue = ready.queue
8592                        AND tomb.priority = ready.priority
8593                        AND tomb.enqueue_shard = ready.enqueue_shard
8594                        AND tomb.lane_seq = ready.lane_seq
8595                        AND tomb.ready_slot = ready.ready_slot
8596                        AND tomb.ready_generation = ready.ready_generation
8597                  )
8598                ORDER BY ready.run_at ASC, ready.lane_seq ASC
8599                LIMIT $2
8600                FOR UPDATE OF ready SKIP LOCKED
8601            ),
8602            tombstones AS (
8603                INSERT INTO {schema}.ready_tombstones (
8604                    ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
8605                )
8606                SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
8607                FROM target
8608                ON CONFLICT DO NOTHING
8609            )
8610            SELECT
8611                ready_slot,
8612                ready_generation,
8613                job_id,
8614                kind,
8615                queue,
8616                args,
8617                priority,
8618                attempt,
8619                run_lease,
8620                max_attempts,
8621                lane_seq,
8622                enqueue_shard,
8623                run_at,
8624                attempted_at,
8625                created_at,
8626                unique_key,
8627                unique_states,
8628                COALESCE(payload, '{{}}'::jsonb) AS payload
8629            FROM target
8630            "#
8631        ))
8632        .bind(cutoff)
8633        .bind(limit)
8634        .fetch_all(tx.as_mut())
8635        .await
8636        .map_err(map_sqlx_error)?;
8637
8638        if moved.is_empty() {
8639            tx.commit().await.map_err(map_sqlx_error)?;
8640            return Ok(Vec::new());
8641        }
8642
8643        let mut ids = Vec::with_capacity(moved.len());
8644        let mut queues = BTreeSet::new();
8645        let mut ready_rows = Vec::with_capacity(moved.len());
8646
8647        for row in moved {
8648            ids.push(row.job_id);
8649            queues.insert(row.queue.clone());
8650
8651            let mut payload = RuntimePayload::from_json(row.payload)?;
8652            let metadata = payload.metadata.as_object_mut().ok_or_else(|| {
8653                AwaError::Validation(
8654                    "queue storage payload metadata must be a JSON object".to_string(),
8655                )
8656            })?;
8657            metadata
8658                .entry("_awa_original_priority".to_string())
8659                .or_insert_with(|| serde_json::Value::from(i64::from(row.priority)));
8660
8661            ready_rows.push(ExistingReadyRow {
8662                job_id: row.job_id,
8663                kind: row.kind,
8664                queue: row.queue,
8665                args: row.args,
8666                priority: row.priority - 1,
8667                attempt: row.attempt,
8668                run_lease: row.run_lease,
8669                max_attempts: row.max_attempts,
8670                run_at: row.run_at,
8671                attempted_at: row.attempted_at,
8672                created_at: row.created_at,
8673                unique_key: row.unique_key,
8674                unique_states: row.unique_states,
8675                payload: payload.into_json(),
8676            });
8677        }
8678
8679        // Aging tombstones the source lane without moving its claim cursor,
8680        // then re-inserts on the target lane. The source lane can temporarily
8681        // over-count by `moved`; later claims advance over the tombstones.
8682        // The drift is bounded by aging rate × poll interval.
8683        self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(JobState::Available))
8684            .await?;
8685        self.notify_queues_tx(&mut tx, queues).await?;
8686        tx.commit().await.map_err(map_sqlx_error)?;
8687        Ok(ids)
8688    }
8689
8690    fn with_progress(
8691        payload: serde_json::Value,
8692        progress: Option<serde_json::Value>,
8693    ) -> Result<serde_json::Value, AwaError> {
8694        let mut payload = RuntimePayload::from_json(payload)?;
8695        payload.set_progress(progress);
8696        Ok(payload.into_json())
8697    }
8698
8699    async fn take_callback_result(
8700        &self,
8701        pool: &PgPool,
8702        job_id: i64,
8703        run_lease: i64,
8704    ) -> Result<serde_json::Value, AwaError> {
8705        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
8706        let mut row: Option<AttemptStateRow> = sqlx::query_as(&format!(
8707            r#"
8708            SELECT
8709                job_id,
8710                run_lease,
8711                progress,
8712                callback_filter,
8713                callback_on_complete,
8714                callback_on_fail,
8715                callback_transform,
8716                callback_result
8717            FROM {}
8718            WHERE job_id = $1
8719              AND run_lease = $2
8720            FOR UPDATE
8721            "#,
8722            self.attempt_state_table()
8723        ))
8724        .bind(job_id)
8725        .bind(run_lease)
8726        .fetch_optional(tx.as_mut())
8727        .await
8728        .map_err(map_sqlx_error)?;
8729
8730        let Some(mut row) = row.take() else {
8731            tx.commit().await.map_err(map_sqlx_error)?;
8732            return Ok(serde_json::Value::Null);
8733        };
8734
8735        let result = row
8736            .callback_result
8737            .take()
8738            .unwrap_or(serde_json::Value::Null);
8739
8740        if row.progress.is_none()
8741            && row.callback_filter.is_none()
8742            && row.callback_on_complete.is_none()
8743            && row.callback_on_fail.is_none()
8744            && row.callback_transform.is_none()
8745        {
8746            sqlx::query(&format!(
8747                "DELETE FROM {} WHERE job_id = $1 AND run_lease = $2",
8748                self.attempt_state_table()
8749            ))
8750            .bind(job_id)
8751            .bind(run_lease)
8752            .execute(tx.as_mut())
8753            .await
8754            .map_err(map_sqlx_error)?;
8755        } else {
8756            sqlx::query(&format!(
8757                "UPDATE {} SET callback_result = NULL, updated_at = clock_timestamp() WHERE job_id = $1 AND run_lease = $2",
8758                self.attempt_state_table()
8759            ))
8760            .bind(job_id)
8761            .bind(run_lease)
8762            .execute(tx.as_mut())
8763            .await
8764            .map_err(map_sqlx_error)?;
8765        }
8766
8767        tx.commit().await.map_err(map_sqlx_error)?;
8768        Ok(result)
8769    }
8770
8771    async fn backoff_at_tx<'a>(
8772        &self,
8773        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8774        attempt: i16,
8775        max_attempts: i16,
8776    ) -> Result<DateTime<Utc>, AwaError> {
8777        sqlx::query_scalar("SELECT clock_timestamp() + awa.backoff_duration($1, $2)")
8778            .bind(attempt)
8779            .bind(max_attempts)
8780            .fetch_one(tx.as_mut())
8781            .await
8782            .map_err(map_sqlx_error)
8783    }
8784
8785    async fn notify_queues_tx<'a>(
8786        &self,
8787        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8788        queues: impl IntoIterator<Item = String>,
8789    ) -> Result<(), AwaError> {
8790        // BTreeSet dedups so we never emit the same NOTIFY twice per
8791        // transaction. Multi-queue enqueues fold into a single round-trip via
8792        // `unnest($1::text[])` rather than one round-trip per channel.
8793        let channels: Vec<String> = queues
8794            .into_iter()
8795            .map(|queue| format!("awa:{}", self.logical_queue_name(&queue)))
8796            .collect::<BTreeSet<String>>()
8797            .into_iter()
8798            .collect();
8799        if channels.is_empty() {
8800            return Ok(());
8801        }
8802        sqlx::query("SELECT pg_notify(channel, '') FROM unnest($1::text[]) AS channel")
8803            .bind(&channels)
8804            .execute(tx.as_mut())
8805            .await
8806            .map_err(map_sqlx_error)?;
8807        Ok(())
8808    }
8809
8810    async fn lock_receipt_attempts_tx<'a>(
8811        &self,
8812        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8813        jobs: &[(i64, i64)],
8814    ) -> Result<(), AwaError> {
8815        if jobs.is_empty() {
8816            return Ok(());
8817        }
8818
8819        let unique_jobs: BTreeSet<(i64, i64)> = jobs.iter().copied().collect();
8820        let job_ids: Vec<i64> = unique_jobs.iter().map(|(job_id, _)| *job_id).collect();
8821        let run_leases: Vec<i64> = unique_jobs
8822            .iter()
8823            .map(|(_, run_lease)| *run_lease)
8824            .collect();
8825
8826        sqlx::query(
8827            r#"
8828            WITH locks AS MATERIALIZED (
8829                SELECT DISTINCT pg_catalog.hashtextextended(
8830                    format('awa.receipt.complete:%s:%s', job_id, run_lease),
8831                    0
8832                ) AS lock_key
8833                FROM unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease)
8834                ORDER BY lock_key
8835            )
8836            SELECT pg_catalog.pg_advisory_xact_lock(lock_key)
8837            FROM locks
8838            "#,
8839        )
8840        .bind(&job_ids)
8841        .bind(&run_leases)
8842        .execute(tx.as_mut())
8843        .await
8844        .map_err(map_sqlx_error)?;
8845
8846        Ok(())
8847    }
8848
8849    async fn ensure_running_leases_from_receipts_tx<'a>(
8850        &self,
8851        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
8852        jobs: &[(i64, i64)],
8853    ) -> Result<usize, AwaError> {
8854        if jobs.is_empty() {
8855            return Ok(0);
8856        }
8857
8858        self.lock_receipt_attempts_tx(tx, jobs).await?;
8859
8860        let schema = self.schema();
8861        let closure_rel = format!("{schema}.lease_claim_closures");
8862        let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
8863        let closed_evidence =
8864            receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
8865        let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
8866        let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
8867        let inserted: i64 = sqlx::query_scalar(&format!(
8868            r#"
8869            WITH inflight(job_id, run_lease) AS (
8870                SELECT * FROM unnest($1::bigint[], $2::bigint[])
8871            ),
8872            lease_ring AS (
8873                SELECT current_slot AS lease_slot, generation AS lease_generation
8874                FROM {schema}.lease_ring_state
8875                WHERE singleton = TRUE
8876            ),
8877            row_claim_refs AS (
8878                -- Source claim metadata directly from the partitioned
8879                -- lease_claims table anti-joined against every durable
8880                -- closure evidence shape.
8881                SELECT
8882                    claims.claim_slot,
8883                    claims.job_id,
8884                    claims.run_lease,
8885                    claims.ready_slot,
8886                    claims.ready_generation,
8887                    claims.queue,
8888                    claims.priority,
8889                    claims.attempt,
8890                    claims.max_attempts,
8891                    claims.lane_seq,
8892                    claims.enqueue_shard,
8893                    claims.claimed_at,
8894                    claims.deadline_at
8895                FROM {schema}.lease_claims AS claims
8896                JOIN inflight
8897                  ON inflight.job_id = claims.job_id
8898                 AND inflight.run_lease = claims.run_lease
8899                WHERE NOT {closed_evidence}
8900                FOR UPDATE OF claims
8901            ),
8902            batch_claim_refs AS (
8903                SELECT
8904                    claim_batches.claim_slot,
8905                    items.job_id,
8906                    items.run_lease,
8907                    claim_batches.ready_slot,
8908                    claim_batches.ready_generation,
8909                    claim_batches.queue,
8910                    claim_batches.priority,
8911                    items.attempt,
8912                    items.max_attempts,
8913                    items.lane_seq,
8914                    claim_batches.enqueue_shard,
8915                    claim_batches.claimed_at,
8916                    claim_batches.deadline_at
8917                FROM {schema}.lease_claim_batches AS claim_batches
8918                CROSS JOIN LATERAL unnest(
8919                    claim_batches.job_ids,
8920                    claim_batches.run_leases,
8921                    claim_batches.receipt_ids,
8922                    claim_batches.lane_seqs,
8923                    claim_batches.attempts,
8924                    claim_batches.max_attempts
8925                ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
8926                JOIN inflight
8927                  ON inflight.job_id = items.job_id
8928                 AND inflight.run_lease = items.run_lease
8929                WHERE NOT EXISTS (
8930                      SELECT 1
8931                      FROM {schema}.lease_claim_closures AS closures
8932                      WHERE closures.claim_slot = claim_batches.claim_slot
8933                        AND closures.job_id = items.job_id
8934                        AND closures.run_lease = items.run_lease
8935                  )
8936                  AND NOT EXISTS (
8937                      SELECT 1
8938                      FROM {schema}.lease_claim_closure_batches AS closure_batches
8939                      WHERE closure_batches.claim_slot = claim_batches.claim_slot
8940                        AND closure_batches.receipt_ranges @> items.receipt_id
8941                  )
8942                  AND NOT EXISTS (
8943                      SELECT 1
8944                      FROM {schema}.done_entries AS done
8945                      WHERE done.job_id = items.job_id
8946                        AND done.run_lease = items.run_lease
8947                  )
8948                  AND NOT EXISTS (
8949                      SELECT 1
8950                      FROM {schema}.deferred_jobs AS deferred
8951                      WHERE deferred.job_id = items.job_id
8952                        AND deferred.run_lease = items.run_lease
8953                  )
8954                  AND NOT EXISTS (
8955                      SELECT 1
8956                      FROM {schema}.dlq_entries AS dlq
8957                      WHERE dlq.job_id = items.job_id
8958                        AND dlq.run_lease = items.run_lease
8959                  )
8960                FOR UPDATE OF claim_batches
8961            ),
8962            claim_refs AS (
8963                SELECT * FROM row_claim_refs
8964                UNION ALL
8965                SELECT * FROM batch_claim_refs
8966            ),
8967            already_live AS (
8968                SELECT claim_refs.job_id, claim_refs.run_lease
8969                FROM claim_refs
8970                WHERE EXISTS (
8971                    SELECT 1
8972                    FROM {schema}.leases AS lease
8973                    WHERE lease.job_id = claim_refs.job_id
8974                      AND lease.run_lease = claim_refs.run_lease
8975                )
8976            ),
8977            inserted AS (
8978                INSERT INTO {schema}.leases (
8979                    lease_slot,
8980                    lease_generation,
8981                    ready_slot,
8982                    ready_generation,
8983                    job_id,
8984                    queue,
8985                    state,
8986                    priority,
8987                    attempt,
8988                    run_lease,
8989                    max_attempts,
8990                    lane_seq,
8991                    enqueue_shard,
8992                    heartbeat_at,
8993                    deadline_at,
8994                    attempted_at
8995                )
8996                SELECT
8997                    lease_ring.lease_slot,
8998                    lease_ring.lease_generation,
8999                    claim_refs.ready_slot,
9000                    claim_refs.ready_generation,
9001                    claim_refs.job_id,
9002                    claim_refs.queue,
9003                    'running'::awa.job_state,
9004                    claim_refs.priority,
9005                    claim_refs.attempt,
9006                    claim_refs.run_lease,
9007                    claim_refs.max_attempts,
9008                    claim_refs.lane_seq,
9009                    claim_refs.enqueue_shard,
9010                    clock_timestamp(),
9011                    -- Preserve the per-claim deadline so the lease-side
9012                    -- deadline rescue path picks up materialized claims
9013                    -- without an extra hop. NULL when receipts mode is
9014                    -- on with `deadline_duration = 0` (the short-job
9015                    -- shape that needs no deadline at all).
9016                    claim_refs.deadline_at,
9017                    claim_refs.claimed_at
9018                FROM claim_refs
9019                CROSS JOIN lease_ring
9020                WHERE NOT EXISTS (
9021                    SELECT 1
9022                    FROM {schema}.leases AS lease
9023                    WHERE lease.job_id = claim_refs.job_id
9024                      AND lease.run_lease = claim_refs.run_lease
9025                )
9026                RETURNING job_id, run_lease
9027            ),
9028            marked AS (
9029                UPDATE {schema}.lease_claims AS claims
9030                SET materialized_at = clock_timestamp()
9031                FROM (
9032                    SELECT job_id, run_lease FROM inserted
9033                    UNION
9034                    SELECT job_id, run_lease FROM already_live
9035                ) AS moved
9036                WHERE claims.job_id = moved.job_id
9037                  AND claims.run_lease = moved.run_lease
9038                RETURNING claims.job_id
9039            )
9040            SELECT count(*)::bigint
9041            FROM (
9042                SELECT job_id, run_lease FROM inserted
9043                UNION
9044                SELECT job_id, run_lease FROM already_live
9045            ) AS moved
9046            "#
9047        ))
9048        .bind(&job_ids)
9049        .bind(&run_leases)
9050        .fetch_one(tx.as_mut())
9051        .await
9052        .map_err(map_sqlx_error)?;
9053        Ok(inserted as usize)
9054    }
9055
9056    async fn ensure_mutable_running_attempt_tx<'a>(
9057        &self,
9058        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9059        job_id: i64,
9060        run_lease: i64,
9061    ) -> Result<(), AwaError> {
9062        if self.lease_claim_receipts() {
9063            self.ensure_running_leases_from_receipts_tx(tx, &[(job_id, run_lease)])
9064                .await?;
9065        }
9066        Ok(())
9067    }
9068
9069    async fn upsert_attempt_state_from_receipts_tx<'a>(
9070        &self,
9071        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9072        jobs: &[(i64, i64)],
9073    ) -> Result<usize, AwaError> {
9074        if jobs.is_empty() {
9075            return Ok(0);
9076        }
9077
9078        self.lock_receipt_attempts_tx(tx, jobs).await?;
9079
9080        let schema = self.schema();
9081        let closure_rel = format!("{schema}.lease_claim_closures");
9082        let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
9083        let closed_evidence =
9084            receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
9085        let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
9086        let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
9087        let updated: i64 = sqlx::query_scalar(&format!(
9088            r#"
9089            WITH inflight(job_id, run_lease) AS (
9090                SELECT * FROM unnest($1::bigint[], $2::bigint[])
9091            ),
9092            row_claim_refs AS (
9093                -- Source open-claim identity from lease_claims
9094                -- anti-joined against every durable closure evidence
9095                -- shape.
9096                SELECT claims.job_id, claims.run_lease
9097                FROM {schema}.lease_claims AS claims
9098                JOIN inflight
9099                  ON inflight.job_id = claims.job_id
9100                 AND inflight.run_lease = claims.run_lease
9101                WHERE NOT {closed_evidence}
9102                FOR UPDATE OF claims
9103            ),
9104            batch_claim_refs AS (
9105                SELECT items.job_id, items.run_lease
9106                FROM {schema}.lease_claim_batches AS claim_batches
9107                CROSS JOIN LATERAL unnest(
9108                    claim_batches.job_ids,
9109                    claim_batches.run_leases,
9110                    claim_batches.receipt_ids
9111                ) AS items(job_id, run_lease, receipt_id)
9112                JOIN inflight
9113                  ON inflight.job_id = items.job_id
9114                 AND inflight.run_lease = items.run_lease
9115                WHERE NOT EXISTS (
9116                      SELECT 1
9117                      FROM {schema}.lease_claim_closures AS closures
9118                      WHERE closures.claim_slot = claim_batches.claim_slot
9119                        AND closures.job_id = items.job_id
9120                        AND closures.run_lease = items.run_lease
9121                  )
9122                  AND NOT EXISTS (
9123                      SELECT 1
9124                      FROM {schema}.lease_claim_closure_batches AS closure_batches
9125                      WHERE closure_batches.claim_slot = claim_batches.claim_slot
9126                        AND closure_batches.receipt_ranges @> items.receipt_id
9127                  )
9128                  AND NOT EXISTS (
9129                      SELECT 1 FROM {schema}.done_entries AS done
9130                      WHERE done.job_id = items.job_id
9131                        AND done.run_lease = items.run_lease
9132                  )
9133                  AND NOT EXISTS (
9134                      SELECT 1 FROM {schema}.deferred_jobs AS deferred
9135                      WHERE deferred.job_id = items.job_id
9136                        AND deferred.run_lease = items.run_lease
9137                  )
9138                  AND NOT EXISTS (
9139                      SELECT 1 FROM {schema}.dlq_entries AS dlq
9140                      WHERE dlq.job_id = items.job_id
9141                        AND dlq.run_lease = items.run_lease
9142                  )
9143                FOR UPDATE OF claim_batches
9144            ),
9145            claim_refs AS (
9146                SELECT * FROM row_claim_refs
9147                UNION ALL
9148                SELECT * FROM batch_claim_refs
9149            ),
9150            upserted AS (
9151                INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at)
9152                SELECT claim_refs.job_id, claim_refs.run_lease, clock_timestamp(), clock_timestamp()
9153                FROM claim_refs
9154                ON CONFLICT (job_id, run_lease)
9155                DO UPDATE SET
9156                    heartbeat_at = clock_timestamp(),
9157                    updated_at = clock_timestamp()
9158                RETURNING job_id, run_lease
9159            ),
9160            marked AS (
9161                UPDATE {schema}.lease_claims AS claims
9162                SET materialized_at = COALESCE(claims.materialized_at, clock_timestamp())
9163                FROM row_claim_refs
9164                WHERE claims.job_id = row_claim_refs.job_id
9165                  AND claims.run_lease = row_claim_refs.run_lease
9166                RETURNING claims.job_id
9167            )
9168            SELECT count(*)::bigint FROM upserted
9169            "#
9170        ))
9171        .bind(&job_ids)
9172        .bind(&run_leases)
9173        .fetch_one(tx.as_mut())
9174        .await
9175        .map_err(map_sqlx_error)?;
9176        Ok(updated as usize)
9177    }
9178
9179    async fn upsert_attempt_state_progress_from_receipts_tx<'a>(
9180        &self,
9181        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9182        jobs: &[(i64, i64, serde_json::Value)],
9183    ) -> Result<usize, AwaError> {
9184        if jobs.is_empty() {
9185            return Ok(0);
9186        }
9187
9188        let receipt_jobs: Vec<(i64, i64)> = jobs
9189            .iter()
9190            .map(|(job_id, run_lease, _)| (*job_id, *run_lease))
9191            .collect();
9192        self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?;
9193
9194        let schema = self.schema();
9195        let closure_rel = format!("{schema}.lease_claim_closures");
9196        let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
9197        let closed_evidence =
9198            receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
9199        let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _, _)| *job_id).collect();
9200        let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect();
9201        let progress: Vec<serde_json::Value> = jobs
9202            .iter()
9203            .map(|(_, _, progress)| progress.clone())
9204            .collect();
9205        let updated: i64 = sqlx::query_scalar(&format!(
9206            r#"
9207            WITH inflight(job_id, run_lease, progress) AS (
9208                SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::jsonb[])
9209            ),
9210            row_claim_refs AS (
9211                -- Same anti-join pattern as the heartbeat-only path
9212                -- above.
9213                SELECT claims.job_id, claims.run_lease, inflight.progress
9214                FROM {schema}.lease_claims AS claims
9215                JOIN inflight
9216                  ON inflight.job_id = claims.job_id
9217                 AND inflight.run_lease = claims.run_lease
9218                WHERE NOT {closed_evidence}
9219                FOR UPDATE OF claims
9220            ),
9221            batch_claim_refs AS (
9222                SELECT items.job_id, items.run_lease, inflight.progress
9223                FROM {schema}.lease_claim_batches AS claim_batches
9224                CROSS JOIN LATERAL unnest(
9225                    claim_batches.job_ids,
9226                    claim_batches.run_leases,
9227                    claim_batches.receipt_ids
9228                ) AS items(job_id, run_lease, receipt_id)
9229                JOIN inflight
9230                  ON inflight.job_id = items.job_id
9231                 AND inflight.run_lease = items.run_lease
9232                WHERE NOT EXISTS (
9233                      SELECT 1
9234                      FROM {schema}.lease_claim_closures AS closures
9235                      WHERE closures.claim_slot = claim_batches.claim_slot
9236                        AND closures.job_id = items.job_id
9237                        AND closures.run_lease = items.run_lease
9238                  )
9239                  AND NOT EXISTS (
9240                      SELECT 1
9241                      FROM {schema}.lease_claim_closure_batches AS closure_batches
9242                      WHERE closure_batches.claim_slot = claim_batches.claim_slot
9243                        AND closure_batches.receipt_ranges @> items.receipt_id
9244                  )
9245                  AND NOT EXISTS (
9246                      SELECT 1 FROM {schema}.done_entries AS done
9247                      WHERE done.job_id = items.job_id
9248                        AND done.run_lease = items.run_lease
9249                  )
9250                  AND NOT EXISTS (
9251                      SELECT 1 FROM {schema}.deferred_jobs AS deferred
9252                      WHERE deferred.job_id = items.job_id
9253                        AND deferred.run_lease = items.run_lease
9254                  )
9255                  AND NOT EXISTS (
9256                      SELECT 1 FROM {schema}.dlq_entries AS dlq
9257                      WHERE dlq.job_id = items.job_id
9258                        AND dlq.run_lease = items.run_lease
9259                  )
9260                FOR UPDATE OF claim_batches
9261            ),
9262            claim_refs AS (
9263                SELECT * FROM row_claim_refs
9264                UNION ALL
9265                SELECT * FROM batch_claim_refs
9266            ),
9267            upserted AS (
9268                INSERT INTO {schema}.attempt_state (
9269                    job_id,
9270                    run_lease,
9271                    heartbeat_at,
9272                    progress,
9273                    updated_at
9274                )
9275                SELECT
9276                    claim_refs.job_id,
9277                    claim_refs.run_lease,
9278                    clock_timestamp(),
9279                    claim_refs.progress,
9280                    clock_timestamp()
9281                FROM claim_refs
9282                ON CONFLICT (job_id, run_lease)
9283                DO UPDATE SET
9284                    heartbeat_at = clock_timestamp(),
9285                    progress = EXCLUDED.progress,
9286                    updated_at = clock_timestamp()
9287                RETURNING job_id, run_lease
9288            ),
9289            marked AS (
9290                UPDATE {schema}.lease_claims AS claims
9291                SET materialized_at = COALESCE(claims.materialized_at, clock_timestamp())
9292                FROM row_claim_refs
9293                WHERE claims.job_id = row_claim_refs.job_id
9294                  AND claims.run_lease = row_claim_refs.run_lease
9295                RETURNING claims.job_id
9296            )
9297            SELECT count(*)::bigint FROM upserted
9298            "#
9299        ))
9300        .bind(&job_ids)
9301        .bind(&run_leases)
9302        .bind(&progress)
9303        .fetch_one(tx.as_mut())
9304        .await
9305        .map_err(map_sqlx_error)?;
9306        Ok(updated as usize)
9307    }
9308
9309    async fn hydrate_deleted_leases_tx<'a>(
9310        &self,
9311        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9312        deleted: Vec<DeletedLeaseRow>,
9313    ) -> Result<Vec<LeaseTransitionRow>, AwaError> {
9314        if deleted.is_empty() {
9315            return Ok(Vec::new());
9316        }
9317
9318        let schema = self.schema();
9319        let ready_slots: Vec<i32> = deleted.iter().map(|row| row.ready_slot).collect();
9320        let ready_generations: Vec<i64> = deleted.iter().map(|row| row.ready_generation).collect();
9321        let queues: Vec<String> = deleted.iter().map(|row| row.queue.clone()).collect();
9322        let enqueue_shards: Vec<i16> = deleted.iter().map(|row| row.enqueue_shard).collect();
9323        let lane_seqs: Vec<i64> = deleted.iter().map(|row| row.lane_seq).collect();
9324        let job_ids: Vec<i64> = deleted.iter().map(|row| row.job_id).collect();
9325        let run_leases: Vec<i64> = deleted.iter().map(|row| row.run_lease).collect();
9326
9327        let ready_rows: Vec<ReadySnapshotRow> = sqlx::query_as(&format!(
9328            r#"
9329            WITH refs(ready_slot, ready_generation, queue, enqueue_shard, lane_seq, job_id) AS (
9330                SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::bigint[], $6::bigint[])
9331            )
9332            SELECT
9333                ready.ready_slot,
9334                ready.ready_generation,
9335                ready.job_id,
9336                ready.kind,
9337                ready.queue,
9338                ready.args,
9339                ready.lane_seq,
9340                ready.enqueue_shard,
9341                ready.run_at,
9342                ready.created_at,
9343                ready.unique_key,
9344                ready.unique_states,
9345                COALESCE(ready.payload, '{{}}'::jsonb) AS payload
9346            FROM refs
9347            JOIN {schema}.ready_entries AS ready
9348              ON ready.ready_slot = refs.ready_slot
9349             AND ready.ready_generation = refs.ready_generation
9350             AND ready.queue = refs.queue
9351             AND ready.enqueue_shard = refs.enqueue_shard
9352             AND ready.lane_seq = refs.lane_seq
9353             AND ready.job_id = refs.job_id
9354            "#
9355        ))
9356        .bind(&ready_slots)
9357        .bind(&ready_generations)
9358        .bind(&queues)
9359        .bind(&enqueue_shards)
9360        .bind(&lane_seqs)
9361        .bind(&job_ids)
9362        .fetch_all(tx.as_mut())
9363        .await
9364        .map_err(map_sqlx_error)?;
9365
9366        let attempt_rows: Vec<AttemptStateRow> = sqlx::query_as(&format!(
9367            r#"
9368            WITH refs(job_id, run_lease) AS (
9369                SELECT * FROM unnest($1::bigint[], $2::bigint[])
9370            )
9371            DELETE FROM {schema}.attempt_state AS attempt
9372            USING refs
9373            WHERE attempt.job_id = refs.job_id
9374              AND attempt.run_lease = refs.run_lease
9375            RETURNING
9376                attempt.job_id,
9377                attempt.run_lease,
9378                attempt.progress,
9379                attempt.callback_filter,
9380                attempt.callback_on_complete,
9381                attempt.callback_on_fail,
9382                attempt.callback_transform,
9383                attempt.callback_result
9384            "#
9385        ))
9386        .bind(&job_ids)
9387        .bind(&run_leases)
9388        .fetch_all(tx.as_mut())
9389        .await
9390        .map_err(map_sqlx_error)?;
9391
9392        // Hydrate runs as part of every rescue path that DELETE'd
9393        // a leases row (heartbeat / deadline / callback timeout
9394        // rescue, plus admin cancel of running attempts). For
9395        // receipt-backed attempts those leases came from
9396        // `ensure_running_leases_from_receipts_tx`, which leaves a
9397        // `lease_claims` row behind with `materialized_at` set. The
9398        // rescue itself closes the lease but never wrote a closure
9399        // for the original receipt, so the claim sat "open" until
9400        // partition prune — `load_job` and any
9401        // `lease_claims`-aware count then double-counted the
9402        // attempt as `running` even after it had moved to
9403        // retryable / failed / completed. Write the closure here so
9404        // the receipt plane mirrors the lease plane: when the lease
9405        // is gone, the receipt is gone too.
9406        sqlx::query(&format!(
9407            r#"
9408            WITH refs(job_id, run_lease) AS (
9409                SELECT * FROM unnest($1::bigint[], $2::bigint[])
9410            ),
9411            row_claims AS (
9412                SELECT claims.claim_slot, claims.job_id, claims.run_lease,
9413                       'rescue', clock_timestamp()
9414                FROM {schema}.lease_claims AS claims
9415                JOIN refs
9416                  ON refs.job_id = claims.job_id
9417                 AND refs.run_lease = claims.run_lease
9418            ),
9419            -- Compact batch-sourced claims (including those materialized
9420            -- into `leases` without a lease_claims row) have no row to
9421            -- JOIN, so they close into the batch ledger the queue prune
9422            -- gate counts via compact_count. The double-close guards
9423            -- below keep a re-hydrated rescued receipt from writing a
9424            -- second batch closure.
9425            batch_claims AS (
9426                SELECT
9427                    claim_batches.claim_slot,
9428                    claim_batches.ready_slot,
9429                    claim_batches.ready_generation,
9430                    items.job_id,
9431                    items.run_lease,
9432                    items.receipt_id
9433                FROM {schema}.lease_claim_batches AS claim_batches
9434                CROSS JOIN LATERAL unnest(
9435                    claim_batches.job_ids,
9436                    claim_batches.run_leases,
9437                    claim_batches.receipt_ids
9438                ) AS items(job_id, run_lease, receipt_id)
9439                JOIN refs
9440                  ON refs.job_id = items.job_id
9441                 AND refs.run_lease = items.run_lease
9442                WHERE NOT EXISTS (
9443                      SELECT 1
9444                      FROM {schema}.lease_claim_closures AS closures
9445                      WHERE closures.claim_slot = claim_batches.claim_slot
9446                        AND closures.job_id = items.job_id
9447                        AND closures.run_lease = items.run_lease
9448                  )
9449                  AND NOT EXISTS (
9450                      SELECT 1
9451                      FROM {schema}.lease_claim_closure_batches AS closure_batches
9452                      WHERE closure_batches.claim_slot = claim_batches.claim_slot
9453                        AND closure_batches.receipt_ranges @> items.receipt_id
9454                  )
9455            ),
9456            inserted AS (
9457                INSERT INTO {schema}.lease_claim_closures
9458                    (claim_slot, job_id, run_lease, outcome, closed_at)
9459                SELECT * FROM row_claims
9460                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
9461                RETURNING claim_slot, job_id, run_lease, closed_at
9462            ),
9463            inserted_batches AS (
9464                INSERT INTO {schema}.lease_claim_closure_batches (
9465                    claim_slot,
9466                    ready_slot,
9467                    ready_generation,
9468                    outcome,
9469                    closed_count,
9470                    receipt_ids,
9471                    receipt_ranges,
9472                    closed_at
9473                )
9474                SELECT
9475                    batch_claims.claim_slot,
9476                    batch_claims.ready_slot,
9477                    batch_claims.ready_generation,
9478                    'rescue',
9479                    count(*)::int,
9480                    array_agg(batch_claims.receipt_id ORDER BY batch_claims.receipt_id),
9481                    range_agg(int8range(batch_claims.receipt_id, batch_claims.receipt_id + 1, '[)') ORDER BY batch_claims.receipt_id),
9482                    clock_timestamp()
9483                FROM batch_claims
9484                GROUP BY
9485                    batch_claims.claim_slot,
9486                    batch_claims.ready_slot,
9487                    batch_claims.ready_generation
9488                RETURNING claim_slot
9489            ),
9490            marked AS (
9491                UPDATE {schema}.lease_claims AS claims
9492                SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
9493                FROM inserted
9494                WHERE claims.claim_slot = inserted.claim_slot
9495                  AND claims.job_id = inserted.job_id
9496                  AND claims.run_lease = inserted.run_lease
9497                RETURNING claims.job_id
9498            )
9499            SELECT
9500                (SELECT count(*) FROM marked)
9501                + (SELECT count(*) FROM inserted_batches)
9502            "#
9503        ))
9504        .bind(&job_ids)
9505        .bind(&run_leases)
9506        .execute(tx.as_mut())
9507        .await
9508        .map_err(map_sqlx_error)?;
9509
9510        let ready_map: BTreeMap<(i32, i64, String, i16, i64, i64), ReadySnapshotRow> = ready_rows
9511            .into_iter()
9512            .map(|row| {
9513                (
9514                    (
9515                        row.ready_slot,
9516                        row.ready_generation,
9517                        row.queue.clone(),
9518                        row.enqueue_shard,
9519                        row.lane_seq,
9520                        row.job_id,
9521                    ),
9522                    row,
9523                )
9524            })
9525            .collect();
9526
9527        let attempt_map: BTreeMap<(i64, i64), AttemptStateRow> = attempt_rows
9528            .into_iter()
9529            .map(|row| ((row.job_id, row.run_lease), row))
9530            .collect();
9531
9532        let mut hydrated = Vec::with_capacity(deleted.len());
9533        for deleted_row in deleted {
9534            let ready = ready_map
9535                .get(&(
9536                    deleted_row.ready_slot,
9537                    deleted_row.ready_generation,
9538                    deleted_row.queue.clone(),
9539                    deleted_row.enqueue_shard,
9540                    deleted_row.lane_seq,
9541                    deleted_row.job_id,
9542                ))
9543                .ok_or_else(|| {
9544                    AwaError::Validation(format!(
9545                        "queue storage ready row missing for deleted lease job {} run_lease {}",
9546                        deleted_row.job_id, deleted_row.run_lease
9547                    ))
9548                })?;
9549            let attempt = attempt_map.get(&(deleted_row.job_id, deleted_row.run_lease));
9550
9551            hydrated.push(LeaseTransitionRow {
9552                ready_slot: deleted_row.ready_slot,
9553                ready_generation: deleted_row.ready_generation,
9554                job_id: deleted_row.job_id,
9555                kind: ready.kind.clone(),
9556                queue: ready.queue.clone(),
9557                args: ready.args.clone(),
9558                state: deleted_row.state,
9559                priority: deleted_row.priority,
9560                attempt: deleted_row.attempt,
9561                run_lease: deleted_row.run_lease,
9562                max_attempts: deleted_row.max_attempts,
9563                lane_seq: deleted_row.lane_seq,
9564                enqueue_shard: deleted_row.enqueue_shard,
9565                run_at: ready.run_at,
9566                attempted_at: deleted_row.attempted_at,
9567                created_at: ready.created_at,
9568                unique_key: ready.unique_key.clone(),
9569                unique_states: ready.unique_states.clone(),
9570                payload: ready.payload.clone(),
9571                progress: attempt.and_then(|row| row.progress.clone()),
9572            });
9573        }
9574
9575        Ok(hydrated)
9576    }
9577
9578    async fn close_open_receipt_claim_tx<'a>(
9579        &self,
9580        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9581        job_id: i64,
9582        run_lease: i64,
9583        outcome: &str,
9584    ) -> Result<Option<LeaseTransitionRow>, AwaError> {
9585        if !self.lease_claim_receipts() {
9586            return Ok(None);
9587        }
9588
9589        self.lock_receipt_attempts_tx(tx, &[(job_id, run_lease)])
9590            .await?;
9591
9592        let schema = self.schema();
9593        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
9594            r#"
9595            WITH row_target AS (
9596                -- Target is the open claim identified from the
9597                -- partitioned lease_claims table anti-joined against
9598                -- durable closure evidence.
9599                SELECT
9600                    claims.claim_slot,
9601                    NULL::bigint AS batch_id,
9602                    claims.receipt_id,
9603                    claims.ready_slot,
9604                    claims.ready_generation,
9605                    claims.job_id,
9606                    claims.queue,
9607                    'running'::awa.job_state AS state,
9608                    claims.priority,
9609                    claims.attempt,
9610                    claims.run_lease,
9611                    claims.max_attempts,
9612                    claims.lane_seq,
9613                    claims.enqueue_shard,
9614                    claims.claimed_at AS attempted_at
9615                FROM {schema}.lease_claims AS claims
9616                WHERE claims.job_id = $1
9617                  AND claims.run_lease = $2
9618                  AND claims.closed_at IS NULL
9619                  AND NOT EXISTS (
9620                      SELECT 1 FROM {schema}.lease_claim_closures AS closures
9621                      WHERE closures.claim_slot = claims.claim_slot
9622                        AND closures.job_id = claims.job_id
9623                        AND closures.run_lease = claims.run_lease
9624                  )
9625                  AND NOT EXISTS (
9626                      SELECT 1
9627                      FROM {schema}.lease_claim_closure_batches AS closure_batches
9628                      WHERE closure_batches.receipt_ranges @> claims.receipt_id
9629                  )
9630                  AND NOT EXISTS (
9631                      SELECT 1 FROM {schema}.done_entries AS done
9632                      WHERE done.job_id = claims.job_id
9633                        AND done.run_lease = claims.run_lease
9634                  )
9635                  AND NOT EXISTS (
9636                      SELECT 1 FROM {schema}.deferred_jobs AS deferred
9637                      WHERE deferred.job_id = claims.job_id
9638                        AND deferred.run_lease = claims.run_lease
9639                  )
9640                  AND NOT EXISTS (
9641                      SELECT 1 FROM {schema}.dlq_entries AS dlq
9642                      WHERE dlq.job_id = claims.job_id
9643                        AND dlq.run_lease = claims.run_lease
9644                  )
9645                  AND NOT EXISTS (
9646                      SELECT 1 FROM {schema}.leases AS lease
9647                      WHERE lease.job_id = claims.job_id
9648                        AND lease.run_lease = claims.run_lease
9649                  )
9650                FOR UPDATE OF claims
9651            ),
9652            batch_target AS (
9653                SELECT
9654                    claim_batches.claim_slot,
9655                    claim_batches.batch_id,
9656                    items.receipt_id,
9657                    claim_batches.ready_slot,
9658                    claim_batches.ready_generation,
9659                    items.job_id,
9660                    claim_batches.queue,
9661                    'running'::awa.job_state AS state,
9662                    claim_batches.priority,
9663                    items.attempt,
9664                    items.run_lease,
9665                    items.max_attempts,
9666                    items.lane_seq,
9667                    claim_batches.enqueue_shard,
9668                    claim_batches.claimed_at AS attempted_at
9669                FROM {schema}.lease_claim_batches AS claim_batches
9670                CROSS JOIN LATERAL unnest(
9671                    claim_batches.job_ids,
9672                    claim_batches.run_leases,
9673                    claim_batches.receipt_ids,
9674                    claim_batches.lane_seqs,
9675                    claim_batches.attempts,
9676                    claim_batches.max_attempts
9677                ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
9678                WHERE items.job_id = $1
9679                  AND items.run_lease = $2
9680                  AND NOT EXISTS (
9681                      SELECT 1 FROM {schema}.lease_claim_closures AS closures
9682                      WHERE closures.claim_slot = claim_batches.claim_slot
9683                        AND closures.job_id = items.job_id
9684                        AND closures.run_lease = items.run_lease
9685                  )
9686                  AND NOT EXISTS (
9687                      SELECT 1
9688                      FROM {schema}.lease_claim_closure_batches AS closure_batches
9689                      WHERE closure_batches.claim_slot = claim_batches.claim_slot
9690                        AND closure_batches.receipt_ranges @> items.receipt_id
9691                  )
9692                  AND NOT EXISTS (
9693                      SELECT 1 FROM {schema}.done_entries AS done
9694                      WHERE done.job_id = items.job_id
9695                        AND done.run_lease = items.run_lease
9696                  )
9697                  AND NOT EXISTS (
9698                      SELECT 1 FROM {schema}.deferred_jobs AS deferred
9699                      WHERE deferred.job_id = items.job_id
9700                        AND deferred.run_lease = items.run_lease
9701                  )
9702                  AND NOT EXISTS (
9703                      SELECT 1 FROM {schema}.dlq_entries AS dlq
9704                      WHERE dlq.job_id = items.job_id
9705                        AND dlq.run_lease = items.run_lease
9706                  )
9707                  AND NOT EXISTS (
9708                      SELECT 1 FROM {schema}.leases AS lease
9709                      WHERE lease.job_id = items.job_id
9710                        AND lease.run_lease = items.run_lease
9711                  )
9712                FOR UPDATE OF claim_batches
9713            ),
9714            target AS (
9715                SELECT * FROM row_target
9716                UNION ALL
9717                SELECT * FROM batch_target
9718                LIMIT 1
9719            ),
9720            -- A row target (batch_id IS NULL) is a deadline lease claim:
9721            -- close it explicitly so the prune gates balance it against
9722            -- the lease_claims row. A batch target has no lease_claims
9723            -- row to JOIN, so it closes into the batch ledger below.
9724            inserted AS (
9725                INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
9726                SELECT target.claim_slot, target.job_id, target.run_lease, $3, clock_timestamp()
9727                FROM target
9728                WHERE target.batch_id IS NULL
9729                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
9730                RETURNING claim_slot, job_id, run_lease, closed_at
9731            ),
9732            inserted_batches AS (
9733                INSERT INTO {schema}.lease_claim_closure_batches (
9734                    claim_slot,
9735                    ready_slot,
9736                    ready_generation,
9737                    outcome,
9738                    closed_count,
9739                    receipt_ids,
9740                    receipt_ranges,
9741                    closed_at
9742                )
9743                SELECT
9744                    target.claim_slot,
9745                    target.ready_slot,
9746                    target.ready_generation,
9747                    $3,
9748                    count(*)::int,
9749                    array_agg(target.receipt_id ORDER BY target.receipt_id),
9750                    range_agg(int8range(target.receipt_id, target.receipt_id + 1, '[)') ORDER BY target.receipt_id),
9751                    clock_timestamp()
9752                FROM target
9753                WHERE target.batch_id IS NOT NULL
9754                GROUP BY target.claim_slot, target.ready_slot, target.ready_generation
9755                RETURNING claim_slot
9756            ),
9757            marked AS (
9758                UPDATE {schema}.lease_claims AS claims
9759                SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
9760                FROM inserted
9761                WHERE claims.claim_slot = inserted.claim_slot
9762                  AND claims.job_id = inserted.job_id
9763                  AND claims.run_lease = inserted.run_lease
9764                RETURNING claims.job_id
9765            ),
9766            -- The batch ledger has no job_id/run_lease column to RETURN,
9767            -- so derive the closed batch target from `target` gated on
9768            -- the batch insert having fired (target is LIMIT 1).
9769            closed_target AS (
9770                SELECT claim_slot, job_id, run_lease FROM inserted
9771                UNION ALL
9772                SELECT target.claim_slot, target.job_id, target.run_lease
9773                FROM target
9774                WHERE target.batch_id IS NOT NULL
9775                  AND EXISTS (SELECT 1 FROM inserted_batches)
9776            )
9777            SELECT
9778                target.ready_slot,
9779                target.ready_generation,
9780                target.job_id,
9781                target.queue,
9782                target.state,
9783                target.priority,
9784                target.attempt,
9785                target.run_lease,
9786                target.max_attempts,
9787                target.lane_seq,
9788                target.enqueue_shard,
9789                NULL::timestamptz AS heartbeat_at,
9790                NULL::timestamptz AS deadline_at,
9791                target.attempted_at,
9792                NULL::uuid AS callback_id,
9793                NULL::timestamptz AS callback_timeout_at
9794            FROM target
9795            JOIN closed_target
9796              ON closed_target.claim_slot = target.claim_slot
9797             AND closed_target.job_id = target.job_id
9798             AND closed_target.run_lease = target.run_lease
9799            "#
9800        ))
9801        .bind(job_id)
9802        .bind(run_lease)
9803        .bind(outcome)
9804        .fetch_all(tx.as_mut())
9805        .await
9806        .map_err(map_sqlx_error)?;
9807
9808        if deleted.is_empty() {
9809            return Ok(None);
9810        }
9811
9812        let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
9813        Ok(moved.into_iter().next())
9814    }
9815
9816    async fn take_running_attempt_tx<'a>(
9817        &self,
9818        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9819        job_id: i64,
9820        run_lease: i64,
9821        receipt_outcome: &str,
9822    ) -> Result<Option<LeaseTransitionRow>, AwaError> {
9823        if let Some(moved) = self
9824            .close_open_receipt_claim_tx(tx, job_id, run_lease, receipt_outcome)
9825            .await?
9826        {
9827            return Ok(Some(moved));
9828        }
9829
9830        let schema = self.schema();
9831        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
9832            r#"
9833            DELETE FROM {schema}.leases
9834            WHERE job_id = $1
9835              AND run_lease = $2
9836              AND state = 'running'
9837            RETURNING
9838                ready_slot,
9839                ready_generation,
9840                job_id,
9841                queue,
9842                state,
9843                priority,
9844                attempt,
9845                run_lease,
9846                max_attempts,
9847                lane_seq,
9848                enqueue_shard,
9849                heartbeat_at,
9850                deadline_at,
9851                attempted_at,
9852                callback_id,
9853                callback_timeout_at
9854            "#
9855        ))
9856        .bind(job_id)
9857        .bind(run_lease)
9858        .fetch_all(tx.as_mut())
9859        .await
9860        .map_err(map_sqlx_error)?;
9861
9862        if deleted.is_empty() {
9863            return Ok(None);
9864        }
9865
9866        let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
9867        Ok(moved.into_iter().next())
9868    }
9869
9870    async fn rescue_stale_receipt_claims_tx<'a>(
9871        &self,
9872        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9873        cutoff: DateTime<Utc>,
9874    ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
9875        let mut rescued = Vec::new();
9876        let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT;
9877        let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?;
9878
9879        if let Some(slot) = preferred_slot {
9880            let mut slot_rescued = self
9881                .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining)
9882                .await?;
9883            remaining = remaining.saturating_sub(slot_rescued.len() as i64);
9884            rescued.append(&mut slot_rescued);
9885            if remaining == 0 {
9886                return Ok(rescued);
9887            }
9888        }
9889
9890        for slot in 0..self.claim_slot_count() {
9891            let slot = slot as i32;
9892            if Some(slot) == preferred_slot {
9893                continue;
9894            }
9895
9896            let mut slot_rescued = self
9897                .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining)
9898                .await?;
9899            remaining = remaining.saturating_sub(slot_rescued.len() as i64);
9900            rescued.append(&mut slot_rescued);
9901
9902            if remaining == 0 {
9903                break;
9904            }
9905        }
9906
9907        Ok(rescued)
9908    }
9909
9910    async fn rescue_stale_receipt_claims_for_slot_tx<'a>(
9911        &self,
9912        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
9913        slot: i32,
9914        cutoff: DateTime<Utc>,
9915        rescue_limit: i64,
9916    ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
9917        let schema = self.schema();
9918        let claim_child = claim_child_name(schema, slot as usize);
9919        let claim_batch_child = claim_batch_child_name(schema, slot as usize);
9920        let closure_child = closure_child_name(schema, slot as usize);
9921        let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize);
9922        let rescued: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
9923            r#"
9924            WITH cursor_row AS (
9925                SELECT
9926                    rescue_cursor_claimed_at,
9927                    rescue_cursor_job_id,
9928                    rescue_cursor_run_lease
9929                FROM {schema}.claim_ring_slots
9930                WHERE slot = $1
9931                FOR UPDATE
9932            ),
9933            claim_source AS MATERIALIZED (
9934                SELECT
9935                    claims.claim_slot,
9936                    NULL::bigint AS batch_id,
9937                    claims.ready_slot,
9938                    claims.ready_generation,
9939                    claims.job_id,
9940                    claims.queue,
9941                    claims.priority,
9942                    claims.attempt,
9943                    claims.run_lease,
9944                    claims.max_attempts,
9945                    claims.lane_seq,
9946                    claims.enqueue_shard,
9947                    claims.receipt_id,
9948                    claims.claimed_at,
9949                    claims.closed_at,
9950                    false AS compact_batch
9951                FROM {claim_child} AS claims
9952                WHERE claims.claim_slot = $1
9953                UNION ALL
9954                SELECT
9955                    claim_batches.claim_slot,
9956                    claim_batches.batch_id,
9957                    claim_batches.ready_slot,
9958                    claim_batches.ready_generation,
9959                    items.job_id,
9960                    claim_batches.queue,
9961                    claim_batches.priority,
9962                    items.attempt,
9963                    items.run_lease,
9964                    items.max_attempts,
9965                    items.lane_seq,
9966                    claim_batches.enqueue_shard,
9967                    items.receipt_id,
9968                    claim_batches.claimed_at,
9969                    NULL::timestamptz AS closed_at,
9970                    true AS compact_batch
9971                FROM {claim_batch_child} AS claim_batches
9972                CROSS JOIN LATERAL unnest(
9973                    claim_batches.job_ids,
9974                    claim_batches.run_leases,
9975                    claim_batches.receipt_ids,
9976                    claim_batches.lane_seqs,
9977                    claim_batches.attempts,
9978                    claim_batches.max_attempts
9979                ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
9980                WHERE claim_batches.claim_slot = $1
9981            ),
9982            after_cursor AS MATERIALIZED (
9983                SELECT
9984                    claim_source.claim_slot,
9985                    claim_source.job_id,
9986                    claim_source.run_lease,
9987                    row_number() OVER (
9988                        ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
9989                    ) AS rn
9990                FROM claim_source
9991                CROSS JOIN cursor_row
9992                WHERE claim_source.claimed_at < $2
9993                  AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease)
9994                      > (
9995                          cursor_row.rescue_cursor_claimed_at,
9996                          cursor_row.rescue_cursor_job_id,
9997                          cursor_row.rescue_cursor_run_lease
9998                        )
9999                ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
10000                LIMIT $3
10001            ),
10002            after_count AS (
10003                SELECT count(*)::bigint AS count FROM after_cursor
10004            ),
10005            before_cursor AS MATERIALIZED (
10006                SELECT
10007                    claim_source.claim_slot,
10008                    claim_source.job_id,
10009                    claim_source.run_lease,
10010                    after_count.count + row_number() OVER (
10011                        ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
10012                    ) AS rn
10013                FROM claim_source
10014                CROSS JOIN cursor_row
10015                CROSS JOIN after_count
10016                WHERE after_count.count < $3
10017                  AND claim_source.claimed_at < $2
10018                  AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease)
10019                      <= (
10020                          cursor_row.rescue_cursor_claimed_at,
10021                          cursor_row.rescue_cursor_job_id,
10022                          cursor_row.rescue_cursor_run_lease
10023                        )
10024                ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease
10025                LIMIT (SELECT GREATEST($3 - count, 0) FROM after_count)
10026            ),
10027            candidate_keys AS MATERIALIZED (
10028                SELECT claim_slot, job_id, run_lease, rn FROM after_cursor
10029                UNION ALL
10030                SELECT claim_slot, job_id, run_lease, rn FROM before_cursor
10031            ),
10032            candidates AS MATERIALIZED (
10033                SELECT
10034                    claim_source.claim_slot,
10035                    claim_source.batch_id,
10036                    claim_source.ready_slot,
10037                    claim_source.ready_generation,
10038                    claim_source.job_id,
10039                    claim_source.queue,
10040                    claim_source.priority,
10041                    claim_source.attempt,
10042                    claim_source.run_lease,
10043                    claim_source.max_attempts,
10044                    claim_source.lane_seq,
10045                    claim_source.enqueue_shard,
10046                    claim_source.receipt_id,
10047                    claim_source.claimed_at,
10048                    claim_source.closed_at,
10049                    claim_source.compact_batch,
10050                    COALESCE(attempt.heartbeat_at, claim_source.claimed_at) < $2 AS is_stale,
10051                    (
10052                        claim_source.closed_at IS NOT NULL
10053                        OR EXISTS (
10054                            SELECT 1 FROM {closure_child} AS closures
10055                            WHERE closures.claim_slot = claim_source.claim_slot
10056                              AND closures.job_id = claim_source.job_id
10057                              AND closures.run_lease = claim_source.run_lease
10058                        )
10059                        OR EXISTS (
10060                            SELECT 1
10061                            FROM {closure_batch_child} AS closure_batches
10062                            WHERE closure_batches.receipt_ranges @> claim_source.receipt_id
10063                        )
10064                        OR EXISTS (
10065                            SELECT 1 FROM {schema}.done_entries AS done
10066                            WHERE done.job_id = claim_source.job_id
10067                              AND done.run_lease = claim_source.run_lease
10068                        )
10069                        OR EXISTS (
10070                            SELECT 1 FROM {schema}.deferred_jobs AS deferred
10071                            WHERE deferred.job_id = claim_source.job_id
10072                              AND deferred.run_lease = claim_source.run_lease
10073                        )
10074                        OR EXISTS (
10075                            SELECT 1 FROM {schema}.dlq_entries AS dlq
10076                            WHERE dlq.job_id = claim_source.job_id
10077                              AND dlq.run_lease = claim_source.run_lease
10078                        )
10079                    ) AS is_closed,
10080                    EXISTS (
10081                        SELECT 1 FROM {schema}.leases AS lease
10082                        WHERE lease.job_id = claim_source.job_id
10083                          AND lease.run_lease = claim_source.run_lease
10084                    ) AS is_lease_managed,
10085                    candidate_keys.rn
10086                FROM candidate_keys
10087                JOIN claim_source
10088                  ON claim_source.claim_slot = candidate_keys.claim_slot
10089                 AND claim_source.job_id = candidate_keys.job_id
10090                 AND claim_source.run_lease = candidate_keys.run_lease
10091                LEFT JOIN {schema}.attempt_state AS attempt
10092                  ON attempt.job_id = claim_source.job_id
10093                 AND attempt.run_lease = claim_source.run_lease
10094            ),
10095            stale_candidates AS (
10096                SELECT candidates.*
10097                FROM candidates
10098                WHERE NOT is_closed
10099                  AND NOT is_lease_managed
10100                  AND is_stale
10101                ORDER BY rn
10102                LIMIT $4
10103            ),
10104            stale_row_locked AS (
10105                SELECT stale_candidates.*
10106                FROM stale_candidates
10107                JOIN {claim_child} AS claims
10108                  ON claims.claim_slot = stale_candidates.claim_slot
10109                 AND claims.job_id = stale_candidates.job_id
10110                 AND claims.run_lease = stale_candidates.run_lease
10111                WHERE NOT stale_candidates.compact_batch
10112                  AND claims.closed_at IS NULL
10113                  AND NOT EXISTS (
10114                      SELECT 1 FROM {closure_child} AS closures
10115                      WHERE closures.claim_slot = claims.claim_slot
10116                        AND closures.job_id = claims.job_id
10117                        AND closures.run_lease = claims.run_lease
10118                  )
10119                  AND NOT EXISTS (
10120                      SELECT 1
10121                      FROM {closure_batch_child} AS closure_batches
10122                      WHERE closure_batches.receipt_ranges @> claims.receipt_id
10123                  )
10124                  AND NOT EXISTS (
10125                      SELECT 1 FROM {schema}.done_entries AS done
10126                      WHERE done.job_id = claims.job_id
10127                        AND done.run_lease = claims.run_lease
10128                  )
10129                  AND NOT EXISTS (
10130                      SELECT 1 FROM {schema}.deferred_jobs AS deferred
10131                      WHERE deferred.job_id = claims.job_id
10132                        AND deferred.run_lease = claims.run_lease
10133                  )
10134                  AND NOT EXISTS (
10135                      SELECT 1 FROM {schema}.dlq_entries AS dlq
10136                      WHERE dlq.job_id = claims.job_id
10137                        AND dlq.run_lease = claims.run_lease
10138                  )
10139                  AND pg_catalog.pg_try_advisory_xact_lock(
10140                      pg_catalog.hashtextextended(
10141                          format('awa.receipt.complete:%s:%s', claims.job_id, claims.run_lease),
10142                          0
10143                      )
10144                  )
10145                  -- A claim that already materialized into `leases` is
10146                  -- on the lease-side heartbeat-rescue path (see
10147                  -- `rescue_stale_heartbeats`). Rescuing it again here
10148                  -- would write a second closure for an attempt the
10149                  -- runtime is still tracking via its lease row.
10150                  AND NOT EXISTS (
10151                      SELECT 1 FROM {schema}.leases AS lease
10152	                      WHERE lease.job_id = claims.job_id
10153	                        AND lease.run_lease = claims.run_lease
10154	                  )
10155                FOR UPDATE OF claims SKIP LOCKED
10156            ),
10157            stale_batch_locked AS (
10158                SELECT stale_candidates.*
10159                FROM stale_candidates
10160                JOIN {claim_batch_child} AS claim_batches
10161                  ON claim_batches.claim_slot = stale_candidates.claim_slot
10162                 AND claim_batches.batch_id = stale_candidates.batch_id
10163                WHERE stale_candidates.compact_batch
10164                  AND NOT EXISTS (
10165                      SELECT 1 FROM {closure_child} AS closures
10166                      WHERE closures.claim_slot = stale_candidates.claim_slot
10167                        AND closures.job_id = stale_candidates.job_id
10168                        AND closures.run_lease = stale_candidates.run_lease
10169                  )
10170                  AND NOT EXISTS (
10171                      SELECT 1
10172                      FROM {closure_batch_child} AS closure_batches
10173                      WHERE closure_batches.receipt_ranges @> stale_candidates.receipt_id
10174                  )
10175                  AND NOT EXISTS (
10176                      SELECT 1 FROM {schema}.done_entries AS done
10177                      WHERE done.job_id = stale_candidates.job_id
10178                        AND done.run_lease = stale_candidates.run_lease
10179                  )
10180                  AND NOT EXISTS (
10181                      SELECT 1 FROM {schema}.deferred_jobs AS deferred
10182                      WHERE deferred.job_id = stale_candidates.job_id
10183                        AND deferred.run_lease = stale_candidates.run_lease
10184                  )
10185                  AND NOT EXISTS (
10186                      SELECT 1 FROM {schema}.dlq_entries AS dlq
10187                      WHERE dlq.job_id = stale_candidates.job_id
10188                        AND dlq.run_lease = stale_candidates.run_lease
10189                  )
10190                  AND pg_catalog.pg_try_advisory_xact_lock(
10191                      pg_catalog.hashtextextended(
10192                          format('awa.receipt.complete:%s:%s', stale_candidates.job_id, stale_candidates.run_lease),
10193                          0
10194                      )
10195                  )
10196                  AND NOT EXISTS (
10197                      SELECT 1 FROM {schema}.leases AS lease
10198                      WHERE lease.job_id = stale_candidates.job_id
10199                        AND lease.run_lease = stale_candidates.run_lease
10200                  )
10201                FOR UPDATE OF claim_batches SKIP LOCKED
10202            ),
10203            stale_locked AS (
10204                SELECT * FROM stale_row_locked
10205                UNION ALL
10206                SELECT * FROM stale_batch_locked
10207            ),
10208            -- Row-sourced (deadline lease) claims close explicitly so the
10209            -- prune gates balance them against the lease_claims row.
10210            -- Compact batch-sourced claims have no lease_claims row to
10211            -- JOIN, so they close into the batch ledger the queue prune
10212            -- gate counts via compact_count.
10213            inserted AS (
10214                INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
10215                SELECT stale_locked.claim_slot, stale_locked.job_id, stale_locked.run_lease, 'rescued', clock_timestamp()
10216                FROM stale_locked
10217                WHERE NOT stale_locked.compact_batch
10218                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
10219                RETURNING claim_slot, job_id, run_lease, closed_at
10220            ),
10221            inserted_batches AS (
10222                INSERT INTO {schema}.lease_claim_closure_batches (
10223                    claim_slot,
10224                    ready_slot,
10225                    ready_generation,
10226                    outcome,
10227                    closed_count,
10228                    receipt_ids,
10229                    receipt_ranges,
10230                    closed_at
10231                )
10232                SELECT
10233                    stale_locked.claim_slot,
10234                    stale_locked.ready_slot,
10235                    stale_locked.ready_generation,
10236                    'rescued',
10237                    count(*)::int,
10238                    array_agg(stale_locked.receipt_id ORDER BY stale_locked.receipt_id),
10239                    range_agg(int8range(stale_locked.receipt_id, stale_locked.receipt_id + 1, '[)') ORDER BY stale_locked.receipt_id),
10240                    clock_timestamp()
10241                FROM stale_locked
10242                WHERE stale_locked.compact_batch
10243                GROUP BY
10244                    stale_locked.claim_slot,
10245                    stale_locked.ready_slot,
10246                    stale_locked.ready_generation
10247                RETURNING claim_slot
10248            ),
10249            closed_locked AS (
10250                SELECT claim_slot, job_id, run_lease FROM inserted
10251                UNION ALL
10252                SELECT stale_locked.claim_slot, stale_locked.job_id, stale_locked.run_lease
10253                FROM stale_locked
10254                WHERE stale_locked.compact_batch
10255                  AND EXISTS (SELECT 1 FROM inserted_batches)
10256            ),
10257            marked AS (
10258                UPDATE {schema}.lease_claims AS claims
10259                SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
10260                FROM inserted
10261                WHERE claims.claim_slot = inserted.claim_slot
10262                  AND claims.job_id = inserted.job_id
10263                  AND claims.run_lease = inserted.run_lease
10264                RETURNING claims.job_id
10265            ),
10266            annotated AS (
10267                SELECT
10268                    candidates.*,
10269                    (
10270                        candidates.is_closed
10271                        OR candidates.is_lease_managed
10272                        OR NOT candidates.is_stale
10273                        OR EXISTS (
10274                            SELECT 1 FROM closed_locked
10275                            WHERE closed_locked.claim_slot = candidates.claim_slot
10276                              AND closed_locked.job_id = candidates.job_id
10277                              AND closed_locked.run_lease = candidates.run_lease
10278                        )
10279                    ) AS advanceable
10280                FROM candidates
10281            ),
10282            bounded AS (
10283                SELECT
10284                    annotated.*,
10285                    min(CASE WHEN NOT annotated.advanceable THEN annotated.rn END) OVER () AS first_blocked_rn
10286                FROM annotated
10287            ),
10288            advance_target AS (
10289                SELECT claimed_at, job_id, run_lease
10290                FROM bounded
10291                WHERE first_blocked_rn IS NULL OR rn < first_blocked_rn
10292                ORDER BY rn DESC
10293                LIMIT 1
10294            ),
10295            advance_cursor AS (
10296                UPDATE {schema}.claim_ring_slots AS slots
10297                SET rescue_cursor_claimed_at = advance_target.claimed_at,
10298                    rescue_cursor_job_id = advance_target.job_id,
10299                    rescue_cursor_run_lease = advance_target.run_lease
10300                FROM advance_target
10301                WHERE slots.slot = $1
10302                RETURNING slots.slot
10303            ),
10304            cursor_advance AS (
10305                SELECT count(*) FROM advance_cursor
10306            )
10307            SELECT
10308                stale_locked.ready_slot,
10309                stale_locked.ready_generation,
10310                stale_locked.job_id,
10311                stale_locked.queue,
10312                'running'::awa.job_state AS state,
10313                stale_locked.priority,
10314                stale_locked.attempt,
10315                stale_locked.run_lease,
10316                stale_locked.max_attempts,
10317                stale_locked.lane_seq,
10318                stale_locked.enqueue_shard,
10319                stale_locked.claimed_at AS attempted_at
10320            FROM stale_locked
10321            JOIN closed_locked
10322              ON closed_locked.claim_slot = stale_locked.claim_slot
10323             AND closed_locked.job_id = stale_locked.job_id
10324             AND closed_locked.run_lease = stale_locked.run_lease
10325            CROSS JOIN cursor_advance
10326            "#
10327        ))
10328        .bind(slot)
10329        .bind(cutoff)
10330        .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT)
10331        .bind(rescue_limit)
10332        .fetch_all(tx.as_mut())
10333        .await
10334        .map_err(map_sqlx_error)?;
10335        Ok(rescued)
10336    }
10337
10338    /// Receipt-side counterpart to `rescue_expired_deadlines`: scans
10339    /// `lease_claims` for rows whose per-claim `deadline_at` has
10340    /// passed but which still don't have a closure or a materialized
10341    /// lease row. Each match gets a `'deadline_expired'` closure
10342    /// written and is returned for the maintenance caller to convert
10343    /// into a deferred / DLQ row, exactly as the lease-side path does.
10344    ///
10345    /// The two anti-joins mirror `rescue_stale_receipt_claims_tx`'s
10346    /// disambiguation: a claim that has already materialized into
10347    /// `leases` is on the lease-side deadline-rescue path, and
10348    /// rescuing it here would double-close it.
10349    ///
10350    /// Unlike the lease-side scan, receipt deadline rescue walks each
10351    /// claim partition through a tiny cursor ordered by deadline. That
10352    /// keeps a long MVCC horizon from making every maintenance tick
10353    /// re-prove old successful receipt completions until claim prune
10354    /// can truncate the partition.
10355    async fn rescue_expired_receipt_deadlines_tx<'a>(
10356        &self,
10357        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
10358    ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
10359        let mut rescued = Vec::new();
10360        let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT;
10361        let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?;
10362
10363        if let Some(slot) = preferred_slot {
10364            let mut slot_rescued = self
10365                .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining)
10366                .await?;
10367            remaining = remaining.saturating_sub(slot_rescued.len() as i64);
10368            rescued.append(&mut slot_rescued);
10369            if remaining == 0 {
10370                return Ok(rescued);
10371            }
10372        }
10373
10374        for slot in 0..self.claim_slot_count() {
10375            let slot = slot as i32;
10376            if Some(slot) == preferred_slot {
10377                continue;
10378            }
10379
10380            let mut slot_rescued = self
10381                .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining)
10382                .await?;
10383            remaining = remaining.saturating_sub(slot_rescued.len() as i64);
10384            rescued.append(&mut slot_rescued);
10385
10386            if remaining == 0 {
10387                break;
10388            }
10389        }
10390
10391        Ok(rescued)
10392    }
10393
10394    async fn oldest_initialized_claim_slot_tx<'a>(
10395        &self,
10396        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
10397    ) -> Result<Option<i32>, AwaError> {
10398        let schema = self.schema();
10399        let preferred_slot = sqlx::query_as::<_, (i32, i64, i32)>(&format!(
10400            r#"
10401            SELECT current_slot, generation, slot_count
10402            FROM {schema}.claim_ring_state
10403            WHERE singleton = TRUE
10404            "#
10405        ))
10406        .fetch_optional(tx.as_mut())
10407        .await
10408        .map_err(map_sqlx_error)?
10409        .and_then(|(current_slot, generation, slot_count)| {
10410            oldest_initialized_ring_slot(current_slot, generation, slot_count)
10411                .map(|(slot, _generation)| slot)
10412                .filter(|slot| *slot >= 0 && (*slot as usize) < self.claim_slot_count())
10413        });
10414
10415        Ok(preferred_slot)
10416    }
10417
10418    async fn rescue_expired_receipt_deadlines_for_slot_tx<'a>(
10419        &self,
10420        tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
10421        slot: i32,
10422        rescue_limit: i64,
10423    ) -> Result<Vec<DeletedLeaseRow>, AwaError> {
10424        if rescue_limit <= 0 {
10425            return Ok(Vec::new());
10426        }
10427
10428        let schema = self.schema();
10429        let claim_child = claim_child_name(schema, slot as usize);
10430        let closure_child = closure_child_name(schema, slot as usize);
10431        let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize);
10432        let closed_evidence =
10433            receipt_closed_evidence_sql(schema, &closure_child, &closure_batch_child, "claims");
10434        let rescued: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
10435            r#"
10436            WITH cursor_row AS (
10437                SELECT
10438                    deadline_cursor_deadline_at,
10439                    deadline_cursor_job_id,
10440                    deadline_cursor_run_lease
10441                FROM {schema}.claim_ring_slots
10442                WHERE slot = $1
10443                FOR UPDATE
10444            ),
10445            after_cursor AS MATERIALIZED (
10446                SELECT
10447                    claims.claim_slot,
10448                    claims.job_id,
10449                    claims.run_lease,
10450                    row_number() OVER (
10451                        ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10452                    ) AS rn
10453                FROM {claim_child} AS claims
10454                CROSS JOIN cursor_row
10455                WHERE claims.claim_slot = $1
10456                  AND claims.deadline_at IS NOT NULL
10457                  AND claims.deadline_at < clock_timestamp()
10458                  AND (claims.deadline_at, claims.job_id, claims.run_lease)
10459                      > (
10460                          cursor_row.deadline_cursor_deadline_at,
10461                          cursor_row.deadline_cursor_job_id,
10462                          cursor_row.deadline_cursor_run_lease
10463                        )
10464                ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10465                LIMIT $2
10466            ),
10467            after_count AS (
10468                SELECT count(*)::bigint AS count FROM after_cursor
10469            ),
10470            before_cursor AS MATERIALIZED (
10471                SELECT
10472                    claims.claim_slot,
10473                    claims.job_id,
10474                    claims.run_lease,
10475                    after_count.count + row_number() OVER (
10476                        ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10477                    ) AS rn
10478                FROM {claim_child} AS claims
10479                CROSS JOIN cursor_row
10480                CROSS JOIN after_count
10481                WHERE after_count.count < $2
10482                  AND claims.claim_slot = $1
10483                  AND claims.deadline_at IS NOT NULL
10484                  AND claims.deadline_at < clock_timestamp()
10485                  AND (claims.deadline_at, claims.job_id, claims.run_lease)
10486                      <= (
10487                          cursor_row.deadline_cursor_deadline_at,
10488                          cursor_row.deadline_cursor_job_id,
10489                          cursor_row.deadline_cursor_run_lease
10490                        )
10491                ORDER BY claims.deadline_at, claims.job_id, claims.run_lease
10492                LIMIT (SELECT GREATEST($2 - count, 0) FROM after_count)
10493            ),
10494            candidate_keys AS MATERIALIZED (
10495                SELECT claim_slot, job_id, run_lease, rn FROM after_cursor
10496                UNION ALL
10497                SELECT claim_slot, job_id, run_lease, rn FROM before_cursor
10498            ),
10499            candidates AS MATERIALIZED (
10500                SELECT
10501                    claims.claim_slot,
10502                    claims.ready_slot,
10503                    claims.ready_generation,
10504                    claims.job_id,
10505                    claims.queue,
10506                    'running'::awa.job_state AS state,
10507                    claims.priority,
10508                    claims.attempt,
10509                    claims.run_lease,
10510                    claims.max_attempts,
10511                    claims.lane_seq,
10512                    claims.enqueue_shard,
10513                    claims.claimed_at,
10514                    claims.deadline_at,
10515                    claims.deadline_at < clock_timestamp() AS is_expired,
10516                    {closed_evidence} AS is_closed,
10517                    EXISTS (
10518                        SELECT 1 FROM {schema}.leases AS lease
10519                        WHERE lease.job_id = claims.job_id
10520                          AND lease.run_lease = claims.run_lease
10521                    ) AS is_lease_managed,
10522                    candidate_keys.rn
10523                FROM candidate_keys
10524                JOIN {claim_child} AS claims
10525                  ON claims.claim_slot = candidate_keys.claim_slot
10526                 AND claims.job_id = candidate_keys.job_id
10527                 AND claims.run_lease = candidate_keys.run_lease
10528            ),
10529            expired_candidates AS (
10530                SELECT candidates.*
10531                FROM candidates
10532                WHERE NOT is_closed
10533                  AND NOT is_lease_managed
10534                  AND is_expired
10535                ORDER BY rn
10536                LIMIT $3
10537            ),
10538            expired_locked AS (
10539                SELECT expired_candidates.*
10540                FROM expired_candidates
10541                JOIN {claim_child} AS claims
10542                  ON claims.claim_slot = expired_candidates.claim_slot
10543                 AND claims.job_id = expired_candidates.job_id
10544                 AND claims.run_lease = expired_candidates.run_lease
10545                WHERE NOT {closed_evidence}
10546                  AND claims.deadline_at < clock_timestamp()
10547                  AND pg_catalog.pg_try_advisory_xact_lock(
10548                      pg_catalog.hashtextextended(
10549                          format('awa.receipt.complete:%s:%s', claims.job_id, claims.run_lease),
10550                          0
10551                      )
10552                  )
10553                  AND NOT EXISTS (
10554                      SELECT 1 FROM {schema}.leases AS lease
10555                      WHERE lease.job_id = claims.job_id
10556                        AND lease.run_lease = claims.run_lease
10557                  )
10558                FOR UPDATE OF claims SKIP LOCKED
10559            ),
10560            inserted AS (
10561                INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at)
10562                SELECT
10563                    expired_locked.claim_slot,
10564                    expired_locked.job_id,
10565                    expired_locked.run_lease,
10566                    'deadline_expired',
10567                    clock_timestamp()
10568                FROM expired_locked
10569                ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING
10570                RETURNING claim_slot, job_id, run_lease, closed_at
10571            ),
10572            marked AS (
10573                UPDATE {schema}.lease_claims AS claims
10574                SET closed_at = COALESCE(claims.closed_at, inserted.closed_at)
10575                FROM inserted
10576                WHERE claims.claim_slot = inserted.claim_slot
10577                  AND claims.job_id = inserted.job_id
10578                  AND claims.run_lease = inserted.run_lease
10579                RETURNING claims.job_id
10580            ),
10581            annotated AS (
10582                SELECT
10583                    candidates.*,
10584                    (
10585                        candidates.is_closed
10586                        OR candidates.is_lease_managed
10587                        OR EXISTS (
10588                            SELECT 1 FROM inserted
10589                            WHERE inserted.claim_slot = candidates.claim_slot
10590                              AND inserted.job_id = candidates.job_id
10591                              AND inserted.run_lease = candidates.run_lease
10592                        )
10593                    ) AS advanceable
10594                FROM candidates
10595            ),
10596            bounded AS (
10597                SELECT
10598                    annotated.*,
10599                    min(CASE WHEN NOT annotated.advanceable THEN annotated.rn END) OVER () AS first_blocked_rn
10600                FROM annotated
10601            ),
10602            advance_target AS (
10603                SELECT deadline_at, job_id, run_lease
10604                FROM bounded
10605                WHERE first_blocked_rn IS NULL OR rn < first_blocked_rn
10606                ORDER BY rn DESC
10607                LIMIT 1
10608            ),
10609            advance_cursor AS (
10610                UPDATE {schema}.claim_ring_slots AS slots
10611                SET deadline_cursor_deadline_at = advance_target.deadline_at,
10612                    deadline_cursor_job_id = advance_target.job_id,
10613                    deadline_cursor_run_lease = advance_target.run_lease
10614                FROM advance_target
10615                WHERE slots.slot = $1
10616                RETURNING slots.slot
10617            ),
10618            cursor_advance AS (
10619                SELECT count(*) FROM advance_cursor
10620            )
10621            SELECT
10622                expired_locked.ready_slot,
10623                expired_locked.ready_generation,
10624                expired_locked.job_id,
10625                expired_locked.queue,
10626                expired_locked.state,
10627                expired_locked.priority,
10628                expired_locked.attempt,
10629                expired_locked.run_lease,
10630                expired_locked.max_attempts,
10631                expired_locked.lane_seq,
10632                expired_locked.enqueue_shard,
10633                expired_locked.claimed_at AS attempted_at
10634            FROM expired_locked
10635            JOIN inserted
10636              ON inserted.claim_slot = expired_locked.claim_slot
10637             AND inserted.job_id = expired_locked.job_id
10638             AND inserted.run_lease = expired_locked.run_lease
10639            CROSS JOIN cursor_advance
10640            "#
10641        ))
10642        .bind(slot)
10643        .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT)
10644        .bind(rescue_limit)
10645        .fetch_all(tx.as_mut())
10646        .await
10647        .map_err(map_sqlx_error)?;
10648        Ok(rescued)
10649    }
10650
10651    pub async fn load_job(&self, pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
10652        let schema = self.schema();
10653        let closure_rel = format!("{schema}.lease_claim_closures");
10654        let closure_batch_rel = format!("{schema}.lease_claim_closure_batches");
10655        let closed_evidence =
10656            receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims");
10657        let mut candidates = Vec::new();
10658
10659        let ready_rows: Vec<ReadyJobRow> = sqlx::query_as(&format!(
10660            r#"
10661            SELECT
10662                job_id,
10663                kind,
10664                queue,
10665                args,
10666                priority,
10667                attempt,
10668                run_lease,
10669                max_attempts,
10670                run_at,
10671                attempted_at,
10672                created_at,
10673                unique_key,
10674                unique_states,
10675                COALESCE(payload, '{{}}'::jsonb) AS payload
10676            FROM {schema}.ready_entries
10677            WHERE job_id = $1
10678            ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC
10679            "#,
10680        ))
10681        .bind(job_id)
10682        .fetch_all(pool)
10683        .await
10684        .map_err(map_sqlx_error)?;
10685        for row in ready_rows {
10686            candidates.push(row.into_job_row()?);
10687        }
10688
10689        let deferred_rows: Vec<DeferredJobRow> = sqlx::query_as(&format!(
10690            r#"
10691            SELECT
10692                job_id,
10693                kind,
10694                queue,
10695                args,
10696                state,
10697                priority,
10698                attempt,
10699                run_lease,
10700                max_attempts,
10701                run_at,
10702                attempted_at,
10703                finalized_at,
10704                created_at,
10705                unique_key,
10706                unique_states,
10707                COALESCE(payload, '{{}}'::jsonb) AS payload
10708            FROM {schema}.deferred_jobs
10709            WHERE job_id = $1
10710            "#,
10711        ))
10712        .bind(job_id)
10713        .fetch_all(pool)
10714        .await
10715        .map_err(map_sqlx_error)?;
10716        for row in deferred_rows {
10717            candidates.push(row.into_job_row()?);
10718        }
10719
10720        let lease_rows: Vec<LeaseJobRow> = sqlx::query_as(&format!(
10721            r#"
10722            SELECT
10723                lease.ready_slot,
10724                lease.ready_generation,
10725                lease.job_id,
10726                ready.kind,
10727                ready.queue,
10728                ready.args,
10729                lease.state,
10730                lease.priority,
10731                lease.attempt,
10732                lease.run_lease,
10733                lease.max_attempts,
10734                lease.lane_seq,
10735                ready.run_at,
10736                COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at,
10737                lease.deadline_at,
10738                lease.attempted_at,
10739                NULL::timestamptz AS finalized_at,
10740                ready.created_at,
10741                ready.unique_key,
10742                ready.unique_states,
10743                lease.callback_id,
10744                lease.callback_timeout_at,
10745                attempt.callback_filter,
10746                attempt.callback_on_complete,
10747                attempt.callback_on_fail,
10748                attempt.callback_transform,
10749                COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
10750                attempt.progress,
10751                attempt.callback_result
10752            FROM {schema}.leases AS lease
10753            JOIN {schema}.ready_entries AS ready
10754              ON ready.ready_slot = lease.ready_slot
10755             AND ready.ready_generation = lease.ready_generation
10756             AND ready.queue = lease.queue
10757             AND ready.priority = lease.priority
10758             AND ready.enqueue_shard = lease.enqueue_shard
10759             AND ready.lane_seq = lease.lane_seq
10760            LEFT JOIN {schema}.attempt_state AS attempt
10761              ON attempt.job_id = lease.job_id
10762             AND attempt.run_lease = lease.run_lease
10763            WHERE lease.job_id = $1
10764            ORDER BY lease.run_lease DESC
10765            "#,
10766        ))
10767        .bind(job_id)
10768        .fetch_all(pool)
10769        .await
10770        .map_err(map_sqlx_error)?;
10771        for row in lease_rows {
10772            candidates.push(row.into_job_row()?);
10773        }
10774
10775        // Report receipt-backed attempts as running by anti-joining
10776        // lease_claims against every durable closure evidence shape.
10777        let lease_claim_rows: Vec<LeaseJobRow> = sqlx::query_as(&format!(
10778            r#"
10779            SELECT
10780                claims.ready_slot,
10781                claims.ready_generation,
10782                claims.job_id,
10783                ready.kind,
10784                ready.queue,
10785                ready.args,
10786                'running'::awa.job_state AS state,
10787                claims.priority,
10788                claims.attempt,
10789                claims.run_lease,
10790                claims.max_attempts,
10791                claims.lane_seq,
10792                ready.run_at,
10793                attempt.heartbeat_at,
10794                claims.deadline_at,
10795                claims.claimed_at AS attempted_at,
10796                NULL::timestamptz AS finalized_at,
10797                ready.created_at,
10798                ready.unique_key,
10799                ready.unique_states,
10800                NULL::uuid AS callback_id,
10801                NULL::timestamptz AS callback_timeout_at,
10802                attempt.callback_filter,
10803                attempt.callback_on_complete,
10804                attempt.callback_on_fail,
10805                attempt.callback_transform,
10806                COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
10807                attempt.progress,
10808                attempt.callback_result
10809            FROM {schema}.lease_claims AS claims
10810            JOIN {schema}.ready_entries AS ready
10811              ON ready.ready_slot = claims.ready_slot
10812             AND ready.ready_generation = claims.ready_generation
10813             AND ready.queue = claims.queue
10814             AND ready.priority = claims.priority
10815             AND ready.enqueue_shard = claims.enqueue_shard
10816             AND ready.lane_seq = claims.lane_seq
10817            LEFT JOIN {schema}.attempt_state AS attempt
10818              ON attempt.job_id = claims.job_id
10819             AND attempt.run_lease = claims.run_lease
10820            WHERE claims.job_id = $1
10821              AND NOT {closed_evidence}
10822              -- Exclude claims that have already been materialized into
10823              -- leases — the lease-backed branch above already reports
10824              -- those.
10825              AND NOT EXISTS (
10826                  SELECT 1 FROM {schema}.leases AS lease
10827                  WHERE lease.job_id = claims.job_id
10828                    AND lease.run_lease = claims.run_lease
10829              )
10830            ORDER BY claims.run_lease DESC
10831            "#,
10832        ))
10833        .bind(job_id)
10834        .fetch_all(pool)
10835        .await
10836        .map_err(map_sqlx_error)?;
10837        for row in lease_claim_rows {
10838            candidates.push(row.into_job_row()?);
10839        }
10840
10841        // Zero-deadline receipt claims are stored as compact batches. Expand
10842        // them only for this admin read, and report still-open items as
10843        // running until durable closure, terminal, or materialized-lease
10844        // evidence supersedes the claim.
10845        let lease_claim_batch_rows: Vec<LeaseJobRow> = sqlx::query_as(&format!(
10846            r#"
10847            SELECT
10848                claim_batches.ready_slot,
10849                claim_batches.ready_generation,
10850                items.job_id,
10851                ready.kind,
10852                ready.queue,
10853                ready.args,
10854                'running'::awa.job_state AS state,
10855                claim_batches.priority,
10856                items.attempt,
10857                items.run_lease,
10858                items.max_attempts,
10859                items.lane_seq,
10860                ready.run_at,
10861                attempt.heartbeat_at,
10862                claim_batches.deadline_at,
10863                claim_batches.claimed_at AS attempted_at,
10864                NULL::timestamptz AS finalized_at,
10865                ready.created_at,
10866                ready.unique_key,
10867                ready.unique_states,
10868                NULL::uuid AS callback_id,
10869                NULL::timestamptz AS callback_timeout_at,
10870                attempt.callback_filter,
10871                attempt.callback_on_complete,
10872                attempt.callback_on_fail,
10873                attempt.callback_transform,
10874                COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
10875                attempt.progress,
10876                attempt.callback_result
10877            FROM {schema}.lease_claim_batches AS claim_batches
10878            CROSS JOIN LATERAL unnest(
10879                claim_batches.job_ids,
10880                claim_batches.run_leases,
10881                claim_batches.receipt_ids,
10882                claim_batches.lane_seqs,
10883                claim_batches.attempts,
10884                claim_batches.max_attempts
10885            ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts)
10886            JOIN {schema}.ready_entries AS ready
10887              ON ready.ready_slot = claim_batches.ready_slot
10888             AND ready.ready_generation = claim_batches.ready_generation
10889             AND ready.queue = claim_batches.queue
10890             AND ready.enqueue_shard = claim_batches.enqueue_shard
10891             AND ready.lane_seq = items.lane_seq
10892             AND ready.job_id = items.job_id
10893            LEFT JOIN {schema}.attempt_state AS attempt
10894              ON attempt.job_id = items.job_id
10895             AND attempt.run_lease = items.run_lease
10896            WHERE items.job_id = $1
10897              AND NOT EXISTS (
10898                  SELECT 1
10899                  FROM {schema}.lease_claim_closures AS closures
10900                  WHERE closures.claim_slot = claim_batches.claim_slot
10901                    AND closures.job_id = items.job_id
10902                    AND closures.run_lease = items.run_lease
10903              )
10904              AND NOT EXISTS (
10905                  SELECT 1
10906                  FROM {schema}.lease_claim_closure_batches AS closure_batches
10907                  WHERE closure_batches.claim_slot = claim_batches.claim_slot
10908                    AND closure_batches.receipt_ranges @> items.receipt_id
10909              )
10910              AND NOT EXISTS (
10911                  SELECT 1
10912                  FROM {schema}.leases AS lease
10913                  WHERE lease.job_id = items.job_id
10914                    AND lease.run_lease = items.run_lease
10915              )
10916              AND NOT EXISTS (
10917                  SELECT 1 FROM {schema}.done_entries AS done
10918                  WHERE done.job_id = items.job_id
10919                    AND done.run_lease = items.run_lease
10920              )
10921              AND NOT EXISTS (
10922                  SELECT 1 FROM {schema}.deferred_jobs AS deferred
10923                  WHERE deferred.job_id = items.job_id
10924                    AND deferred.run_lease = items.run_lease
10925              )
10926              AND NOT EXISTS (
10927                  SELECT 1 FROM {schema}.dlq_entries AS dlq
10928                  WHERE dlq.job_id = items.job_id
10929                    AND dlq.run_lease = items.run_lease
10930              )
10931            ORDER BY items.run_lease DESC
10932            "#,
10933        ))
10934        .bind(job_id)
10935        .fetch_all(pool)
10936        .await
10937        .map_err(map_sqlx_error)?;
10938        for row in lease_claim_batch_rows {
10939            candidates.push(row.into_job_row()?);
10940        }
10941
10942        let done_rows: Vec<DoneJobRow> = sqlx::query_as(&format!(
10943            r#"
10944            SELECT
10945                ready_slot,
10946                ready_generation,
10947                job_id,
10948                kind,
10949                queue,
10950                args,
10951                state,
10952                priority,
10953                attempt,
10954                run_lease,
10955                max_attempts,
10956                lane_seq,
10957                enqueue_shard,
10958                run_at,
10959                attempted_at,
10960                finalized_at,
10961                created_at,
10962                unique_key,
10963                unique_states,
10964                payload
10965            FROM {schema}.terminal_jobs AS done
10966            WHERE done.job_id = $1
10967            ORDER BY done.run_lease DESC, done.finalized_at DESC
10968            "#,
10969        ))
10970        .bind(job_id)
10971        .fetch_all(pool)
10972        .await
10973        .map_err(map_sqlx_error)?;
10974        for row in done_rows {
10975            candidates.push(row.into_job_row()?);
10976        }
10977
10978        let dlq_rows: Vec<DlqJobRow> = sqlx::query_as(&format!(
10979            r#"
10980            SELECT
10981                job_id,
10982                kind,
10983                queue,
10984                args,
10985                state,
10986                priority,
10987                attempt,
10988                run_lease,
10989                max_attempts,
10990                run_at,
10991                attempted_at,
10992                finalized_at,
10993                created_at,
10994                unique_key,
10995                unique_states,
10996                COALESCE(payload, '{{}}'::jsonb) AS payload,
10997                dlq_reason,
10998                dlq_at,
10999                original_run_lease
11000            FROM {schema}.dlq_entries
11001            WHERE job_id = $1
11002            ORDER BY dlq_at DESC
11003            "#,
11004        ))
11005        .bind(job_id)
11006        .fetch_all(pool)
11007        .await
11008        .map_err(map_sqlx_error)?;
11009        for row in dlq_rows {
11010            candidates.push(row.into_job_row()?);
11011        }
11012
11013        Ok(candidates.into_iter().max_by_key(|job| {
11014            (
11015                job.run_lease,
11016                transition_timestamp(job),
11017                state_rank(job.state),
11018            )
11019        }))
11020    }
11021
11022    pub async fn register_callback(
11023        &self,
11024        pool: &PgPool,
11025        job_id: i64,
11026        run_lease: i64,
11027        timeout: Duration,
11028    ) -> Result<Uuid, AwaError> {
11029        let callback_id = Uuid::new_v4();
11030        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11031        self.ensure_mutable_running_attempt_tx(&mut tx, job_id, run_lease)
11032            .await?;
11033        let updated = sqlx::query(&format!(
11034            r#"
11035            UPDATE {}
11036            SET callback_id = $2,
11037                callback_timeout_at = clock_timestamp() + make_interval(secs => $3)
11038            WHERE job_id = $1
11039              AND state = 'running'
11040              AND run_lease = $4
11041            "#,
11042            self.leases_table()
11043        ))
11044        .bind(job_id)
11045        .bind(callback_id)
11046        .bind(timeout.as_secs_f64())
11047        .bind(run_lease)
11048        .execute(tx.as_mut())
11049        .await
11050        .map_err(map_sqlx_error)?;
11051
11052        if updated.rows_affected() == 0 {
11053            tx.rollback().await.map_err(map_sqlx_error)?;
11054            return Err(AwaError::Validation("job is not in running state".into()));
11055        }
11056
11057        sqlx::query(&format!(
11058            r#"
11059            UPDATE {}
11060            SET callback_filter = NULL,
11061                callback_on_complete = NULL,
11062                callback_on_fail = NULL,
11063                callback_transform = NULL,
11064                updated_at = clock_timestamp()
11065            WHERE job_id = $1
11066              AND run_lease = $2
11067            "#,
11068            self.attempt_state_table()
11069        ))
11070        .bind(job_id)
11071        .bind(run_lease)
11072        .execute(tx.as_mut())
11073        .await
11074        .map_err(map_sqlx_error)?;
11075
11076        sqlx::query(&format!(
11077            r#"
11078            DELETE FROM {}
11079            WHERE job_id = $1
11080              AND run_lease = $2
11081              AND progress IS NULL
11082              AND callback_result IS NULL
11083              AND callback_filter IS NULL
11084              AND callback_on_complete IS NULL
11085              AND callback_on_fail IS NULL
11086              AND callback_transform IS NULL
11087            "#,
11088            self.attempt_state_table()
11089        ))
11090        .bind(job_id)
11091        .bind(run_lease)
11092        .execute(tx.as_mut())
11093        .await
11094        .map_err(map_sqlx_error)?;
11095
11096        tx.commit().await.map_err(map_sqlx_error)?;
11097        Ok(callback_id)
11098    }
11099
11100    pub async fn register_callback_with_config(
11101        &self,
11102        pool: &PgPool,
11103        job_id: i64,
11104        run_lease: i64,
11105        timeout: Duration,
11106        config: &CallbackConfig,
11107    ) -> Result<Uuid, AwaError> {
11108        if config.is_empty() {
11109            return self
11110                .register_callback(pool, job_id, run_lease, timeout)
11111                .await;
11112        }
11113
11114        #[cfg(feature = "cel")]
11115        {
11116            for (name, expr) in [
11117                ("filter", &config.filter),
11118                ("on_complete", &config.on_complete),
11119                ("on_fail", &config.on_fail),
11120                ("transform", &config.transform),
11121            ] {
11122                if let Some(src) = expr {
11123                    let program = cel::Program::compile(src).map_err(|e| {
11124                        AwaError::Validation(format!("invalid CEL expression for {name}: {e}"))
11125                    })?;
11126                    let references = program.references();
11127                    let bad_vars: Vec<String> = references
11128                        .variables()
11129                        .into_iter()
11130                        .filter(|v| *v != "payload")
11131                        .map(str::to_string)
11132                        .collect();
11133                    if !bad_vars.is_empty() {
11134                        return Err(AwaError::Validation(format!(
11135                            "CEL expression for {name} references undeclared variable(s): {}; only 'payload' is available",
11136                            bad_vars.join(", ")
11137                        )));
11138                    }
11139                }
11140            }
11141        }
11142
11143        #[cfg(not(feature = "cel"))]
11144        {
11145            if !config.is_empty() {
11146                return Err(AwaError::Validation(
11147                    "CEL expressions require the 'cel' feature".into(),
11148                ));
11149            }
11150        }
11151
11152        let callback_id = Uuid::new_v4();
11153        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11154        self.ensure_mutable_running_attempt_tx(&mut tx, job_id, run_lease)
11155            .await?;
11156        let updated = sqlx::query(&format!(
11157            r#"
11158            UPDATE {}
11159            SET callback_id = $2,
11160                callback_timeout_at = clock_timestamp() + make_interval(secs => $3)
11161            WHERE job_id = $1
11162              AND state = 'running'
11163              AND run_lease = $4
11164            "#,
11165            self.leases_table()
11166        ))
11167        .bind(job_id)
11168        .bind(callback_id)
11169        .bind(timeout.as_secs_f64())
11170        .bind(run_lease)
11171        .execute(tx.as_mut())
11172        .await
11173        .map_err(map_sqlx_error)?;
11174
11175        if updated.rows_affected() == 0 {
11176            tx.rollback().await.map_err(map_sqlx_error)?;
11177            return Err(AwaError::Validation("job is not in running state".into()));
11178        }
11179
11180        sqlx::query(&format!(
11181            r#"
11182            INSERT INTO {} (
11183                job_id,
11184                run_lease,
11185                callback_filter,
11186                callback_on_complete,
11187                callback_on_fail,
11188                callback_transform,
11189                updated_at
11190            )
11191            VALUES ($1, $2, $3, $4, $5, $6, clock_timestamp())
11192            ON CONFLICT (job_id, run_lease)
11193            DO UPDATE SET
11194                callback_filter = EXCLUDED.callback_filter,
11195                callback_on_complete = EXCLUDED.callback_on_complete,
11196                callback_on_fail = EXCLUDED.callback_on_fail,
11197                callback_transform = EXCLUDED.callback_transform,
11198                updated_at = clock_timestamp()
11199            "#,
11200            self.attempt_state_table()
11201        ))
11202        .bind(job_id)
11203        .bind(run_lease)
11204        .bind(&config.filter)
11205        .bind(&config.on_complete)
11206        .bind(&config.on_fail)
11207        .bind(&config.transform)
11208        .execute(tx.as_mut())
11209        .await
11210        .map_err(map_sqlx_error)?;
11211
11212        tx.commit().await.map_err(map_sqlx_error)?;
11213        Ok(callback_id)
11214    }
11215
11216    pub async fn cancel_callback(
11217        &self,
11218        pool: &PgPool,
11219        job_id: i64,
11220        run_lease: i64,
11221    ) -> Result<bool, AwaError> {
11222        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11223        let result = sqlx::query(&format!(
11224            r#"
11225            UPDATE {}
11226            SET callback_id = NULL,
11227                callback_timeout_at = NULL
11228            WHERE job_id = $1
11229              AND callback_id IS NOT NULL
11230              AND state = 'running'
11231              AND run_lease = $2
11232            "#,
11233            self.leases_table()
11234        ))
11235        .bind(job_id)
11236        .bind(run_lease)
11237        .execute(tx.as_mut())
11238        .await
11239        .map_err(map_sqlx_error)?;
11240        if result.rows_affected() == 0 {
11241            tx.rollback().await.map_err(map_sqlx_error)?;
11242            return Ok(false);
11243        }
11244
11245        sqlx::query(&format!(
11246            r#"
11247            UPDATE {}
11248            SET callback_filter = NULL,
11249                callback_on_complete = NULL,
11250                callback_on_fail = NULL,
11251                callback_transform = NULL,
11252                updated_at = clock_timestamp()
11253            WHERE job_id = $1
11254              AND run_lease = $2
11255            "#,
11256            self.attempt_state_table()
11257        ))
11258        .bind(job_id)
11259        .bind(run_lease)
11260        .execute(tx.as_mut())
11261        .await
11262        .map_err(map_sqlx_error)?;
11263
11264        sqlx::query(&format!(
11265            r#"
11266            DELETE FROM {}
11267            WHERE job_id = $1
11268              AND run_lease = $2
11269              AND progress IS NULL
11270              AND callback_result IS NULL
11271              AND callback_filter IS NULL
11272              AND callback_on_complete IS NULL
11273              AND callback_on_fail IS NULL
11274              AND callback_transform IS NULL
11275            "#,
11276            self.attempt_state_table()
11277        ))
11278        .bind(job_id)
11279        .bind(run_lease)
11280        .execute(tx.as_mut())
11281        .await
11282        .map_err(map_sqlx_error)?;
11283
11284        tx.commit().await.map_err(map_sqlx_error)?;
11285        Ok(true)
11286    }
11287
11288    /// Load the currently-active lease row for `job_id` (running or
11289    /// waiting_external) inside a caller-owned transaction. Used by ADR-029
11290    /// helpers that need the post-park snapshot — including the
11291    /// `callback_id` and `callback_timeout_at` written by
11292    /// `register_callback()` — without leaving the transaction that just
11293    /// performed the parking UPDATE.
11294    pub async fn load_active_lease_in_tx(
11295        &self,
11296        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11297        job_id: i64,
11298        run_lease: i64,
11299    ) -> Result<Option<JobRow>, AwaError> {
11300        let schema = self.schema();
11301        let row: Option<LeaseJobRow> = sqlx::query_as(&format!(
11302            r#"
11303            SELECT
11304                lease.ready_slot,
11305                lease.ready_generation,
11306                lease.job_id,
11307                ready.kind,
11308                ready.queue,
11309                ready.args,
11310                lease.state,
11311                lease.priority,
11312                lease.attempt,
11313                lease.run_lease,
11314                lease.max_attempts,
11315                lease.lane_seq,
11316                ready.run_at,
11317                COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at,
11318                lease.deadline_at,
11319                lease.attempted_at,
11320                NULL::timestamptz AS finalized_at,
11321                ready.created_at,
11322                ready.unique_key,
11323                ready.unique_states,
11324                lease.callback_id,
11325                lease.callback_timeout_at,
11326                attempt.callback_filter,
11327                attempt.callback_on_complete,
11328                attempt.callback_on_fail,
11329                attempt.callback_transform,
11330                COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
11331                attempt.progress,
11332                attempt.callback_result
11333            FROM {schema}.leases AS lease
11334            JOIN {schema}.ready_entries AS ready
11335              ON ready.ready_slot = lease.ready_slot
11336             AND ready.ready_generation = lease.ready_generation
11337             AND ready.queue = lease.queue
11338             AND ready.priority = lease.priority
11339             AND ready.enqueue_shard = lease.enqueue_shard
11340             AND ready.lane_seq = lease.lane_seq
11341            LEFT JOIN {schema}.attempt_state AS attempt
11342              ON attempt.job_id = lease.job_id
11343             AND attempt.run_lease = lease.run_lease
11344            WHERE lease.job_id = $1
11345              AND lease.run_lease = $2
11346            "#,
11347        ))
11348        .bind(job_id)
11349        .bind(run_lease)
11350        .fetch_optional(tx.as_mut())
11351        .await
11352        .map_err(map_sqlx_error)?;
11353        row.map(LeaseJobRow::into_job_row).transpose()
11354    }
11355
11356    pub async fn enter_callback_wait(
11357        &self,
11358        pool: &PgPool,
11359        job_id: i64,
11360        run_lease: i64,
11361        callback_id: Uuid,
11362    ) -> Result<bool, AwaError> {
11363        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11364        let entered = self
11365            .enter_callback_wait_in_tx(&mut tx, job_id, run_lease, callback_id)
11366            .await?;
11367        tx.commit().await.map_err(map_sqlx_error)?;
11368        Ok(entered)
11369    }
11370
11371    /// Transaction-aware variant of [`Self::enter_callback_wait`] (ADR-029).
11372    /// Returns whether the row transitioned to `waiting_external`.
11373    pub async fn enter_callback_wait_in_tx(
11374        &self,
11375        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11376        job_id: i64,
11377        run_lease: i64,
11378        callback_id: Uuid,
11379    ) -> Result<bool, AwaError> {
11380        let result = sqlx::query(&format!(
11381            r#"
11382            UPDATE {}
11383            SET state = 'waiting_external',
11384                heartbeat_at = NULL,
11385                deadline_at = NULL
11386            WHERE job_id = $1
11387              AND state = 'running'
11388              AND run_lease = $2
11389              AND callback_id = $3
11390            "#,
11391            self.leases_table()
11392        ))
11393        .bind(job_id)
11394        .bind(run_lease)
11395        .bind(callback_id)
11396        .execute(tx.as_mut())
11397        .await
11398        .map_err(map_sqlx_error)?;
11399        Ok(result.rows_affected() > 0)
11400    }
11401
11402    pub async fn check_callback_state(
11403        &self,
11404        pool: &PgPool,
11405        job_id: i64,
11406        callback_id: Uuid,
11407    ) -> Result<CallbackPollResult, AwaError> {
11408        let row: Option<(JobState, Option<Uuid>, i64, Option<serde_json::Value>)> =
11409            sqlx::query_as(&format!(
11410                r#"
11411                SELECT
11412                    lease.state,
11413                    lease.callback_id,
11414                    lease.run_lease,
11415                    attempt.callback_result
11416                FROM {} AS lease
11417                LEFT JOIN {} AS attempt
11418                  ON attempt.job_id = lease.job_id
11419                 AND attempt.run_lease = lease.run_lease
11420                WHERE lease.job_id = $1
11421                ORDER BY lease.run_lease DESC
11422                LIMIT 1
11423                "#,
11424                self.leases_table(),
11425                self.attempt_state_table()
11426            ))
11427            .bind(job_id)
11428            .fetch_optional(pool)
11429            .await
11430            .map_err(map_sqlx_error)?;
11431
11432        match row {
11433            Some((JobState::Running, None, run_lease, Some(_))) => {
11434                let result = self.take_callback_result(pool, job_id, run_lease).await?;
11435                Ok(CallbackPollResult::Resolved(result))
11436            }
11437            Some((state, Some(current_callback_id), _, _))
11438                if current_callback_id != callback_id =>
11439            {
11440                Ok(CallbackPollResult::Stale {
11441                    token: callback_id,
11442                    current: current_callback_id,
11443                    state,
11444                })
11445            }
11446            Some((JobState::WaitingExternal, Some(current), _, _)) if current == callback_id => {
11447                Ok(CallbackPollResult::Pending)
11448            }
11449            Some((state, _, _, _)) => Ok(CallbackPollResult::UnexpectedState {
11450                token: callback_id,
11451                state,
11452            }),
11453            None => {
11454                if let Some(job) = self.load_job(pool, job_id).await? {
11455                    Ok(CallbackPollResult::UnexpectedState {
11456                        token: callback_id,
11457                        state: job.state,
11458                    })
11459                } else {
11460                    Ok(CallbackPollResult::NotFound)
11461                }
11462            }
11463        }
11464    }
11465
11466    pub async fn callback_job(
11467        &self,
11468        pool: &PgPool,
11469        callback_id: Uuid,
11470        run_lease: Option<i64>,
11471    ) -> Result<Option<JobRow>, AwaError> {
11472        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11473        let result = self
11474            .callback_job_in_tx(&mut tx, callback_id, run_lease, false)
11475            .await?;
11476        tx.commit().await.map_err(map_sqlx_error)?;
11477        Ok(result)
11478    }
11479
11480    /// Transaction-aware variant of [`Self::callback_job`] (ADR-029).
11481    /// When `for_update` is `true` the join's lease row is locked with
11482    /// `FOR UPDATE OF lease`, mirroring the canonical `resolve_callback`
11483    /// lookup that takes a row lock on `awa.jobs_hot` before evaluating
11484    /// the callback policy and committing the resulting transition.
11485    pub async fn callback_job_in_tx(
11486        &self,
11487        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11488        callback_id: Uuid,
11489        run_lease: Option<i64>,
11490        for_update: bool,
11491    ) -> Result<Option<JobRow>, AwaError> {
11492        let lock_clause = if for_update {
11493            "FOR UPDATE OF lease"
11494        } else {
11495            ""
11496        };
11497        let row: Option<LeaseJobRow> = sqlx::query_as(&format!(
11498            r#"
11499            SELECT
11500                lease.ready_slot,
11501                lease.ready_generation,
11502                lease.job_id,
11503                ready.kind,
11504                ready.queue,
11505                ready.args,
11506                lease.state,
11507                lease.priority,
11508                lease.attempt,
11509                lease.run_lease,
11510                lease.max_attempts,
11511                lease.lane_seq,
11512                ready.run_at,
11513                COALESCE(attempt.heartbeat_at, lease.heartbeat_at) AS heartbeat_at,
11514                lease.deadline_at,
11515                lease.attempted_at,
11516                NULL::timestamptz AS finalized_at,
11517                ready.created_at,
11518                ready.unique_key,
11519                ready.unique_states,
11520                lease.callback_id,
11521                lease.callback_timeout_at,
11522                attempt.callback_filter,
11523                attempt.callback_on_complete,
11524                attempt.callback_on_fail,
11525                attempt.callback_transform,
11526                COALESCE(ready.payload, '{{}}'::jsonb) AS payload,
11527                attempt.progress,
11528                attempt.callback_result
11529            FROM {} AS lease
11530            JOIN {schema}.ready_entries AS ready
11531              ON ready.ready_slot = lease.ready_slot
11532             AND ready.ready_generation = lease.ready_generation
11533             AND ready.queue = lease.queue
11534             AND ready.priority = lease.priority
11535             AND ready.enqueue_shard = lease.enqueue_shard
11536             AND ready.lane_seq = lease.lane_seq
11537            LEFT JOIN {schema}.attempt_state AS attempt
11538              ON attempt.job_id = lease.job_id
11539             AND attempt.run_lease = lease.run_lease
11540            WHERE lease.callback_id = $1
11541              AND lease.state IN ('waiting_external', 'running')
11542              AND ($2::bigint IS NULL OR lease.run_lease = $2)
11543            ORDER BY lease.run_lease DESC
11544            LIMIT 1
11545            {lock_clause}
11546            "#,
11547            self.leases_table(),
11548            schema = self.schema(),
11549        ))
11550        .bind(callback_id)
11551        .bind(run_lease)
11552        .fetch_optional(tx.as_mut())
11553        .await
11554        .map_err(map_sqlx_error)?;
11555
11556        row.map(LeaseJobRow::into_job_row).transpose()
11557    }
11558
11559    #[tracing::instrument(skip(self, pool, payload), name = "queue_storage.complete_external")]
11560    pub async fn complete_external(
11561        &self,
11562        pool: &PgPool,
11563        callback_id: Uuid,
11564        payload: Option<serde_json::Value>,
11565        run_lease: Option<i64>,
11566        resume: bool,
11567    ) -> Result<JobRow, AwaError> {
11568        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11569        let result = self
11570            .complete_external_in_tx(&mut tx, callback_id, payload, run_lease, resume)
11571            .await?;
11572        tx.commit().await.map_err(map_sqlx_error)?;
11573        Ok(result)
11574    }
11575
11576    /// Transaction-aware variant of [`Self::complete_external`] (ADR-029).
11577    /// The non-resume path returns the post-completion `JobRow` directly
11578    /// from the `done_row` insert. The resume path returns the parked-row
11579    /// snapshot reloaded inside the same transaction via
11580    /// [`Self::load_active_lease_in_tx`] — i.e. it does not leave the
11581    /// caller's transaction to refresh state.
11582    pub async fn complete_external_in_tx(
11583        &self,
11584        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11585        callback_id: Uuid,
11586        payload: Option<serde_json::Value>,
11587        run_lease: Option<i64>,
11588        resume: bool,
11589    ) -> Result<JobRow, AwaError> {
11590        if resume {
11591            let resumed: Option<(i64, i64)> = sqlx::query_as(&format!(
11592                r#"
11593                UPDATE {}
11594                SET state = 'running',
11595                    callback_id = NULL,
11596                    callback_timeout_at = NULL,
11597                    heartbeat_at = clock_timestamp()
11598                WHERE callback_id = $1
11599                  AND state IN ('waiting_external', 'running')
11600                  AND ($2::bigint IS NULL OR run_lease = $2)
11601                RETURNING job_id, run_lease
11602                "#,
11603                self.leases_table()
11604            ))
11605            .bind(callback_id)
11606            .bind(run_lease)
11607            .fetch_optional(tx.as_mut())
11608            .await
11609            .map_err(map_sqlx_error)?;
11610
11611            let Some((job_id, resumed_run_lease)) = resumed else {
11612                return Err(AwaError::CallbackNotFound {
11613                    callback_id: callback_id.to_string(),
11614                });
11615            };
11616
11617            sqlx::query(&format!(
11618                r#"
11619                INSERT INTO {} (
11620                    job_id,
11621                    run_lease,
11622                    callback_filter,
11623                    callback_on_complete,
11624                    callback_on_fail,
11625                    callback_transform,
11626                    callback_result,
11627                    updated_at
11628                )
11629                VALUES ($1, $2, NULL, NULL, NULL, NULL, $3, clock_timestamp())
11630                ON CONFLICT (job_id, run_lease)
11631                DO UPDATE SET
11632                    callback_filter = NULL,
11633                    callback_on_complete = NULL,
11634                    callback_on_fail = NULL,
11635                    callback_transform = NULL,
11636                    callback_result = EXCLUDED.callback_result,
11637                    updated_at = clock_timestamp()
11638                "#,
11639                self.attempt_state_table()
11640            ))
11641            .bind(job_id)
11642            .bind(resumed_run_lease)
11643            .bind(payload.unwrap_or(serde_json::Value::Null))
11644            .execute(tx.as_mut())
11645            .await
11646            .map_err(map_sqlx_error)?;
11647
11648            return self
11649                .load_active_lease_in_tx(tx, job_id, resumed_run_lease)
11650                .await?
11651                .ok_or(AwaError::CallbackNotFound {
11652                    callback_id: callback_id.to_string(),
11653                });
11654        }
11655
11656        let schema = self.schema();
11657        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
11658            r#"
11659            DELETE FROM {schema}.leases
11660            WHERE callback_id = $1
11661              AND state IN ('waiting_external', 'running')
11662              AND ($2::bigint IS NULL OR run_lease = $2)
11663            RETURNING
11664                ready_slot,
11665                ready_generation,
11666                job_id,
11667                queue,
11668                state,
11669                priority,
11670                attempt,
11671                run_lease,
11672                max_attempts,
11673                lane_seq,
11674                enqueue_shard,
11675                heartbeat_at,
11676                deadline_at,
11677                attempted_at,
11678                callback_id,
11679                callback_timeout_at
11680            "#
11681        ))
11682        .bind(callback_id)
11683        .bind(run_lease)
11684        .fetch_all(tx.as_mut())
11685        .await
11686        .map_err(map_sqlx_error)?;
11687
11688        if deleted.is_empty() {
11689            return Err(AwaError::CallbackNotFound {
11690                callback_id: callback_id.to_string(),
11691            });
11692        }
11693
11694        let completed_pairs: Vec<(i64, i64)> = deleted
11695            .iter()
11696            .map(|row| (row.job_id, row.run_lease))
11697            .collect();
11698        self.close_receipt_pairs_tx(tx, &completed_pairs, "completed")
11699            .await?;
11700
11701        let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
11702        let moved = moved.into_iter().next().expect("deleted callback lease");
11703
11704        let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
11705            moved.payload.clone(),
11706            moved.progress.clone(),
11707        )?)?;
11708        payload.set_progress(None);
11709        let done_row =
11710            moved
11711                .clone()
11712                .into_done_row(JobState::Completed, Utc::now(), payload.into_json());
11713        self.insert_done_rows_tx(tx, std::slice::from_ref(&done_row), Some(moved.state))
11714            .await?;
11715        done_row.into_job_row()
11716    }
11717
11718    pub async fn fail_external(
11719        &self,
11720        pool: &PgPool,
11721        callback_id: Uuid,
11722        error: &str,
11723        run_lease: Option<i64>,
11724    ) -> Result<JobRow, AwaError> {
11725        self.fail_external_with_error_entry(
11726            pool,
11727            callback_id,
11728            serde_json::json!({ "error": error }),
11729            run_lease,
11730        )
11731        .await
11732    }
11733
11734    pub async fn fail_external_with_error_entry(
11735        &self,
11736        pool: &PgPool,
11737        callback_id: Uuid,
11738        error_entry: serde_json::Value,
11739        run_lease: Option<i64>,
11740    ) -> Result<JobRow, AwaError> {
11741        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11742        let result = self
11743            .fail_external_with_error_entry_in_tx(&mut tx, callback_id, error_entry, run_lease)
11744            .await?;
11745        tx.commit().await.map_err(map_sqlx_error)?;
11746        Ok(result)
11747    }
11748
11749    /// Transaction-aware variant of [`Self::fail_external_with_error_entry`]
11750    /// (ADR-029).
11751    pub async fn fail_external_with_error_entry_in_tx(
11752        &self,
11753        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11754        callback_id: Uuid,
11755        error_entry: serde_json::Value,
11756        run_lease: Option<i64>,
11757    ) -> Result<JobRow, AwaError> {
11758        let schema = self.schema();
11759        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
11760            r#"
11761            DELETE FROM {schema}.leases
11762            WHERE callback_id = $1
11763              AND state IN ('waiting_external', 'running')
11764              AND ($2::bigint IS NULL OR run_lease = $2)
11765            RETURNING
11766                ready_slot,
11767                ready_generation,
11768                job_id,
11769                queue,
11770                state,
11771                priority,
11772                attempt,
11773                run_lease,
11774                max_attempts,
11775                lane_seq,
11776                enqueue_shard,
11777                heartbeat_at,
11778                deadline_at,
11779                attempted_at,
11780                callback_id,
11781                callback_timeout_at
11782            "#
11783        ))
11784        .bind(callback_id)
11785        .bind(run_lease)
11786        .fetch_all(tx.as_mut())
11787        .await
11788        .map_err(map_sqlx_error)?;
11789
11790        if deleted.is_empty() {
11791            return Err(AwaError::CallbackNotFound {
11792                callback_id: callback_id.to_string(),
11793            });
11794        }
11795
11796        let failed_pairs: Vec<(i64, i64)> = deleted
11797            .iter()
11798            .map(|row| (row.job_id, row.run_lease))
11799            .collect();
11800        self.close_receipt_pairs_tx(tx, &failed_pairs, "failed")
11801            .await?;
11802
11803        let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
11804        let moved = moved.into_iter().next().expect("deleted callback lease");
11805
11806        let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
11807            moved.payload.clone(),
11808            moved.progress.clone(),
11809        )?)?;
11810        let mut error_entry = match error_entry {
11811            serde_json::Value::Object(map) => serde_json::Value::Object(map),
11812            other => serde_json::json!({ "error": other }),
11813        };
11814        let error_obj = error_entry
11815            .as_object_mut()
11816            .ok_or_else(|| AwaError::Validation("callback error entry must be an object".into()))?;
11817        error_obj
11818            .entry("attempt".to_string())
11819            .or_insert_with(|| serde_json::Value::from(i64::from(moved.attempt)));
11820        error_obj
11821            .entry("at".to_string())
11822            .or_insert_with(|| serde_json::Value::String(Utc::now().to_rfc3339()));
11823        error_obj
11824            .entry("terminal".to_string())
11825            .or_insert(serde_json::Value::Bool(true));
11826        payload.push_error(error_entry);
11827        let done_row =
11828            moved
11829                .clone()
11830                .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
11831        self.insert_done_rows_tx(tx, std::slice::from_ref(&done_row), Some(moved.state))
11832            .await?;
11833        done_row.into_job_row()
11834    }
11835
11836    pub async fn retry_external(
11837        &self,
11838        pool: &PgPool,
11839        callback_id: Uuid,
11840        run_lease: Option<i64>,
11841    ) -> Result<JobRow, AwaError> {
11842        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11843        let result = self
11844            .retry_external_in_tx(&mut tx, callback_id, run_lease)
11845            .await?;
11846        tx.commit().await.map_err(map_sqlx_error)?;
11847        Ok(result)
11848    }
11849
11850    /// Transaction-aware variant of [`Self::retry_external`] (ADR-029).
11851    pub async fn retry_external_in_tx(
11852        &self,
11853        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
11854        callback_id: Uuid,
11855        run_lease: Option<i64>,
11856    ) -> Result<JobRow, AwaError> {
11857        let schema = self.schema();
11858        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
11859            r#"
11860            DELETE FROM {schema}.leases
11861            WHERE callback_id = $1
11862              AND state = 'waiting_external'
11863              AND ($2::bigint IS NULL OR run_lease = $2)
11864            RETURNING
11865                ready_slot,
11866                ready_generation,
11867                job_id,
11868                queue,
11869                state,
11870                priority,
11871                attempt,
11872                run_lease,
11873                max_attempts,
11874                lane_seq,
11875                enqueue_shard,
11876                heartbeat_at,
11877                deadline_at,
11878                attempted_at,
11879                callback_id,
11880                callback_timeout_at
11881            "#
11882        ))
11883        .bind(callback_id)
11884        .bind(run_lease)
11885        .fetch_all(tx.as_mut())
11886        .await
11887        .map_err(map_sqlx_error)?;
11888
11889        if deleted.is_empty() {
11890            return Err(AwaError::CallbackNotFound {
11891                callback_id: callback_id.to_string(),
11892            });
11893        }
11894
11895        let retryable_pairs: Vec<(i64, i64)> = deleted
11896            .iter()
11897            .map(|row| (row.job_id, row.run_lease))
11898            .collect();
11899        self.close_receipt_pairs_tx(tx, &retryable_pairs, "retryable")
11900            .await?;
11901
11902        let moved = self.hydrate_deleted_leases_tx(tx, deleted).await?;
11903        let moved = moved.into_iter().next().expect("deleted callback lease");
11904
11905        let ready_payload =
11906            Self::payload_with_attempt_state(moved.payload.clone(), moved.progress.clone())?;
11907
11908        let ready_row = ExistingReadyRow {
11909            attempt: 0,
11910            run_at: Utc::now(),
11911            ..moved.clone().into_ready_row(Utc::now(), ready_payload)
11912        };
11913        self.insert_existing_ready_rows_tx(tx, vec![ready_row.clone()], Some(moved.state))
11914            .await?;
11915        self.notify_queues_tx(tx, std::iter::once(moved.queue.clone()))
11916            .await?;
11917        ReadyJobRow {
11918            job_id: ready_row.job_id,
11919            kind: ready_row.kind,
11920            queue: ready_row.queue,
11921            args: ready_row.args,
11922            priority: ready_row.priority,
11923            attempt: ready_row.attempt,
11924            run_lease: ready_row.run_lease,
11925            max_attempts: ready_row.max_attempts,
11926            run_at: ready_row.run_at,
11927            attempted_at: ready_row.attempted_at,
11928            created_at: ready_row.created_at,
11929            unique_key: ready_row.unique_key,
11930            payload: ready_row.payload,
11931        }
11932        .into_job_row()
11933    }
11934
11935    pub async fn heartbeat_callback(
11936        &self,
11937        pool: &PgPool,
11938        callback_id: Uuid,
11939        timeout: Duration,
11940    ) -> Result<JobRow, AwaError> {
11941        let updated: Option<(i64, i64)> = sqlx::query_as(&format!(
11942            r#"
11943            UPDATE {}
11944            SET callback_timeout_at = clock_timestamp() + make_interval(secs => $2)
11945            WHERE callback_id = $1
11946              AND state = 'waiting_external'
11947            RETURNING job_id, run_lease
11948            "#,
11949            self.leases_table()
11950        ))
11951        .bind(callback_id)
11952        .bind(timeout.as_secs_f64())
11953        .fetch_optional(pool)
11954        .await
11955        .map_err(map_sqlx_error)?;
11956
11957        let Some((job_id, _run_lease)) = updated else {
11958            return Err(AwaError::CallbackNotFound {
11959                callback_id: callback_id.to_string(),
11960            });
11961        };
11962
11963        self.load_job(pool, job_id)
11964            .await?
11965            .ok_or(AwaError::CallbackNotFound {
11966                callback_id: callback_id.to_string(),
11967            })
11968    }
11969
11970    pub async fn flush_progress(
11971        &self,
11972        pool: &PgPool,
11973        job_id: i64,
11974        run_lease: i64,
11975        progress: serde_json::Value,
11976    ) -> Result<(), AwaError> {
11977        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
11978        if self.lease_claim_receipts() {
11979            self.upsert_attempt_state_progress_from_receipts_tx(
11980                &mut tx,
11981                &[(job_id, run_lease, progress.clone())],
11982            )
11983            .await?;
11984        }
11985        sqlx::query(&format!(
11986            r#"
11987            INSERT INTO {} (job_id, run_lease, progress, updated_at)
11988            SELECT lease.job_id, lease.run_lease, $3, clock_timestamp()
11989            FROM {} AS lease
11990            WHERE lease.job_id = $1
11991              AND lease.run_lease = $2
11992              AND lease.state IN ('running', 'waiting_external')
11993            ON CONFLICT (job_id, run_lease)
11994            DO UPDATE SET
11995                progress = EXCLUDED.progress,
11996                updated_at = clock_timestamp()
11997            "#,
11998            self.attempt_state_table(),
11999            self.leases_table()
12000        ))
12001        .bind(job_id)
12002        .bind(run_lease)
12003        .bind(progress)
12004        .execute(tx.as_mut())
12005        .await
12006        .map_err(map_sqlx_error)?;
12007        tx.commit().await.map_err(map_sqlx_error)?;
12008        Ok(())
12009    }
12010
12011    pub async fn heartbeat_batch(
12012        &self,
12013        pool: &PgPool,
12014        jobs: &[(i64, i64)],
12015    ) -> Result<usize, AwaError> {
12016        if jobs.is_empty() {
12017            return Ok(0);
12018        }
12019
12020        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12021        let mut updated = 0_usize;
12022        if self.lease_claim_receipts() {
12023            // #169 B1: in receipts mode, attempt_state is the
12024            // authoritative heartbeat home. Skip the `UPDATE leases SET
12025            // heartbeat_at = ...` entirely — that write was the
12026            // dominant per-heartbeat non-HOT update on the partitioned
12027            // `leases` table (every state-indexed partial index pays 2
12028            // dead entries per write under sustained churn). Compat
12029            // reads in `LeaseJobRow` SELECTs and the `awa.jobs` view
12030            // COALESCE attempt_state.heartbeat_at first.
12031            updated += self
12032                .upsert_attempt_state_from_receipts_tx(&mut tx, jobs)
12033                .await?;
12034        } else {
12035            // Legacy non-receipts mode (custom schemas with
12036            // `lease_claim_receipts=FALSE`): the `leases.heartbeat_at`
12037            // write is still the heartbeat home, since there is no
12038            // upsert_attempt_state path firing.
12039            let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _)| *job_id).collect();
12040            let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease)| *run_lease).collect();
12041            let result = sqlx::query(&format!(
12042                r#"
12043                WITH inflight AS (
12044                    SELECT * FROM unnest($1::bigint[], $2::bigint[]) AS v(job_id, run_lease)
12045                )
12046                UPDATE {table}
12047                SET heartbeat_at = clock_timestamp()
12048                FROM inflight
12049                WHERE {table}.job_id = inflight.job_id
12050                  AND {table}.run_lease = inflight.run_lease
12051                  AND {table}.state = 'running'
12052                "#,
12053                table = self.leases_table(),
12054            ))
12055            .bind(&job_ids)
12056            .bind(&run_leases)
12057            .execute(tx.as_mut())
12058            .await
12059            .map_err(map_sqlx_error)?;
12060            updated += result.rows_affected() as usize;
12061        }
12062        tx.commit().await.map_err(map_sqlx_error)?;
12063        Ok(updated)
12064    }
12065
12066    pub async fn heartbeat_progress_batch(
12067        &self,
12068        pool: &PgPool,
12069        jobs: &[(i64, i64, serde_json::Value)],
12070    ) -> Result<usize, AwaError> {
12071        if jobs.is_empty() {
12072            return Ok(0);
12073        }
12074
12075        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12076        let updated = if self.lease_claim_receipts() {
12077            // #169 B1: receipts mode is the only supported shape for
12078            // the default `awa` schema. attempt_state already carries
12079            // heartbeat_at + progress, so the `UPDATE leases SET
12080            // heartbeat_at` + nested `INSERT INTO attempt_state` CTE
12081            // is collapsed into the single attempt_state upsert. The
12082            // upsert sources open-claim identity from `lease_claims`
12083            // anti-joined against durable closure evidence so it picks
12084            // up every open claim regardless of whether
12085            // `materialize_claims` has fanned it out to `leases` yet.
12086            self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs)
12087                .await?
12088        } else {
12089            // Legacy non-receipts mode keeps the old CTE shape that
12090            // updates leases.heartbeat_at + leases.progress
12091            // (via attempt_state upsert) in a single round-trip.
12092            let schema = self.schema();
12093            let job_ids: Vec<i64> = jobs.iter().map(|(job_id, _, _)| *job_id).collect();
12094            let run_leases: Vec<i64> = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect();
12095            let progress: Vec<serde_json::Value> =
12096                jobs.iter().map(|(_, _, value)| value.clone()).collect();
12097            let lease_updated: i64 = sqlx::query_scalar(&format!(
12098                r#"
12099                WITH inflight AS (
12100                    SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::jsonb[]) AS v(job_id, run_lease, progress)
12101                ),
12102                updated AS (
12103                    UPDATE {table} AS lease
12104                    SET heartbeat_at = clock_timestamp()
12105                    FROM inflight
12106                    WHERE lease.job_id = inflight.job_id
12107                      AND lease.run_lease = inflight.run_lease
12108                      AND lease.state = 'running'
12109                    RETURNING lease.job_id, lease.run_lease, inflight.progress
12110                ),
12111                upsert_attempt AS (
12112                    INSERT INTO {schema}.attempt_state (job_id, run_lease, progress, updated_at)
12113                    SELECT job_id, run_lease, progress, clock_timestamp()
12114                    FROM updated
12115                    ON CONFLICT (job_id, run_lease)
12116                    DO UPDATE SET
12117                        progress = EXCLUDED.progress,
12118                        updated_at = clock_timestamp()
12119                )
12120                SELECT count(*)::bigint FROM updated
12121                "#,
12122                table = self.leases_table(),
12123            ))
12124            .bind(&job_ids)
12125            .bind(&run_leases)
12126            .bind(&progress)
12127            .fetch_one(tx.as_mut())
12128            .await
12129            .map_err(map_sqlx_error)?;
12130            lease_updated as usize
12131        };
12132        tx.commit().await.map_err(map_sqlx_error)?;
12133        Ok(updated)
12134    }
12135
12136    pub async fn retry_after(
12137        &self,
12138        pool: &PgPool,
12139        job_id: i64,
12140        run_lease: i64,
12141        retry_after: Duration,
12142        progress: Option<serde_json::Value>,
12143    ) -> Result<Option<JobRow>, AwaError> {
12144        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12145        let result = self
12146            .retry_after_in_tx(&mut tx, job_id, run_lease, retry_after, progress)
12147            .await?;
12148        tx.commit().await.map_err(map_sqlx_error)?;
12149        Ok(result)
12150    }
12151
12152    /// Transaction-aware variant of [`Self::retry_after`]. Caller owns the
12153    /// transaction lifecycle so the move can commit atomically alongside
12154    /// follow-up `INSERT`s (ADR-029).
12155    pub async fn retry_after_in_tx(
12156        &self,
12157        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12158        job_id: i64,
12159        run_lease: i64,
12160        retry_after: Duration,
12161        progress: Option<serde_json::Value>,
12162    ) -> Result<Option<JobRow>, AwaError> {
12163        let Some(moved) = self
12164            .take_running_attempt_tx(tx, job_id, run_lease, "retryable")
12165            .await?
12166        else {
12167            return Ok(None);
12168        };
12169        let now = self.current_timestamp_tx(tx).await?;
12170
12171        let payload =
12172            Self::with_progress(moved.payload.clone(), progress.or(moved.progress.clone()))?;
12173        let deferred = moved.clone().into_deferred_row(
12174            JobState::Retryable,
12175            now + TimeDelta::from_std(retry_after).map_err(|err| {
12176                AwaError::Validation(format!("invalid retry_after duration: {err}"))
12177            })?,
12178            Some(now),
12179            payload,
12180        );
12181        self.insert_deferred_rows_tx(tx, vec![deferred.clone()], Some(moved.state))
12182            .await?;
12183        Ok(Some(deferred.into_job_row()?))
12184    }
12185
12186    pub async fn snooze(
12187        &self,
12188        pool: &PgPool,
12189        job_id: i64,
12190        run_lease: i64,
12191        snooze_for: Duration,
12192        progress: Option<serde_json::Value>,
12193    ) -> Result<Option<JobRow>, AwaError> {
12194        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12195        let Some(moved) = self
12196            .take_running_attempt_tx(&mut tx, job_id, run_lease, "scheduled")
12197            .await?
12198        else {
12199            tx.commit().await.map_err(map_sqlx_error)?;
12200            return Ok(None);
12201        };
12202        let now = self.current_timestamp_tx(&mut tx).await?;
12203
12204        let payload =
12205            Self::with_progress(moved.payload.clone(), progress.or(moved.progress.clone()))?;
12206        let mut deferred = moved.clone().into_deferred_row(
12207            JobState::Scheduled,
12208            now + TimeDelta::from_std(snooze_for)
12209                .map_err(|err| AwaError::Validation(format!("invalid snooze duration: {err}")))?,
12210            None,
12211            payload,
12212        );
12213        deferred.attempt = deferred.attempt.saturating_sub(1);
12214        self.insert_deferred_rows_tx(&mut tx, vec![deferred.clone()], Some(moved.state))
12215            .await?;
12216        tx.commit().await.map_err(map_sqlx_error)?;
12217        Ok(Some(deferred.into_job_row()?))
12218    }
12219
12220    /// Insert a canonical attempt's transition-time deferred successor,
12221    /// degrading to a cancelled terminal row if a newer duplicate already
12222    /// owns the unique claim. Canonical deletion releases any claim before
12223    /// this call, so the deferred insert must acquire from `None` even when
12224    /// the old `running` state was inside the mask.
12225    pub(crate) async fn insert_migrated_deferred_or_cancel_duplicate_tx(
12226        &self,
12227        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12228        row: JobRow,
12229        unique_states: Option<String>,
12230        duplicate_error: &str,
12231    ) -> Result<JobRow, AwaError> {
12232        let payload = Self::payload_from_parts(row.metadata, row.tags, row.errors, row.progress)?;
12233        let deferred = DeferredJobRow {
12234            job_id: row.id,
12235            kind: row.kind,
12236            queue: row.queue,
12237            args: row.args,
12238            state: row.state,
12239            priority: row.priority,
12240            attempt: row.attempt,
12241            run_lease: row.run_lease,
12242            max_attempts: row.max_attempts,
12243            run_at: row.run_at,
12244            attempted_at: row.attempted_at,
12245            finalized_at: row.finalized_at,
12246            created_at: row.created_at,
12247            unique_key: row.unique_key,
12248            unique_states,
12249            payload,
12250        };
12251
12252        {
12253            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12254            match self
12255                .insert_deferred_rows_tx(&mut savepoint, vec![deferred.clone()], None)
12256                .await
12257            {
12258                Ok(_) => {
12259                    savepoint.commit().await.map_err(map_sqlx_error)?;
12260                    return deferred.into_job_row();
12261                }
12262                Err(AwaError::UniqueConflict { .. }) => {
12263                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12264                }
12265                Err(err) => return Err(err),
12266            }
12267        }
12268
12269        let mut cancelled_payload = RuntimePayload::from_json(deferred.payload.clone())?;
12270        cancelled_payload.push_error(lifecycle_error(duplicate_error, deferred.attempt, true));
12271        let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?;
12272        self.ensure_lane(tx, &deferred.queue, deferred.priority, 0)
12273            .await?;
12274        let mut done = DoneJobRow {
12275            ready_slot,
12276            ready_generation,
12277            job_id: deferred.job_id,
12278            kind: deferred.kind,
12279            queue: deferred.queue,
12280            args: deferred.args,
12281            state: JobState::Cancelled,
12282            priority: deferred.priority,
12283            attempt: deferred.attempt,
12284            run_lease: deferred.run_lease,
12285            max_attempts: deferred.max_attempts,
12286            // The successor never entered a ready lane. Mirror deferred-job
12287            // cancellation with a synthetic negative sequence on shard 0.
12288            lane_seq: -deferred.job_id,
12289            enqueue_shard: 0,
12290            run_at: deferred.run_at,
12291            attempted_at: deferred.attempted_at,
12292            finalized_at: self.current_timestamp_tx(tx).await?,
12293            created_at: deferred.created_at,
12294            unique_key: deferred.unique_key,
12295            unique_states: deferred.unique_states,
12296            payload: cancelled_payload.into_json(),
12297        };
12298
12299        {
12300            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12301            match self
12302                .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), None)
12303                .await
12304            {
12305                Ok(_) => {
12306                    savepoint.commit().await.map_err(map_sqlx_error)?;
12307                    return done.into_job_row();
12308                }
12309                Err(AwaError::UniqueConflict { .. }) => {
12310                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12311                }
12312                Err(err) => return Err(err),
12313            }
12314        }
12315
12316        // If `cancelled` itself claims uniqueness, the newer holder still
12317        // wins. The cancelled evidence carries no claim to release, so strip
12318        // the uniqueness fields exactly like rescue duplicate cancellation.
12319        done.unique_key = None;
12320        done.unique_states = None;
12321        self.insert_done_rows_tx(tx, std::slice::from_ref(&done), None)
12322            .await?;
12323        done.into_job_row()
12324    }
12325
12326    /// Rescue-path fallback for unique-claim conflicts: insert the rescued
12327    /// attempt's deferred successor, degrading to a cancelled terminal row
12328    /// when the job's unique claim is held by a newer duplicate — the claim
12329    /// holder wins, and one poisoned row must not abort the whole batched
12330    /// rescue transaction. The insert attempt runs inside a savepoint so the
12331    /// conflict leaves the outer transaction usable.
12332    async fn insert_rescued_deferred_or_cancel_tx(
12333        &self,
12334        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12335        row: LeaseTransitionRow,
12336        deferred: DeferredJobRow,
12337        duplicate_error: &str,
12338    ) -> Result<JobRow, AwaError> {
12339        let old_state = row.state;
12340        {
12341            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12342            match self
12343                .insert_deferred_rows_tx(&mut savepoint, vec![deferred.clone()], Some(old_state))
12344                .await
12345            {
12346                Ok(_) => {
12347                    savepoint.commit().await.map_err(map_sqlx_error)?;
12348                    return deferred.into_job_row();
12349                }
12350                Err(AwaError::UniqueConflict { .. }) => {
12351                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12352                }
12353                Err(err) => return Err(err),
12354            }
12355        }
12356        self.cancel_rescued_duplicate_tx(tx, row, deferred.payload, duplicate_error)
12357            .await
12358    }
12359
12360    /// Same fallback for rescue arms that write a terminal row directly
12361    /// (callback timeout at max attempts): a `failed` insert that conflicts
12362    /// with a newer duplicate's claim becomes a cancellation instead.
12363    async fn insert_rescued_done_or_cancel_tx(
12364        &self,
12365        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12366        row: LeaseTransitionRow,
12367        done: DoneJobRow,
12368        duplicate_error: &str,
12369    ) -> Result<JobRow, AwaError> {
12370        let old_state = row.state;
12371        {
12372            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12373            match self
12374                .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state))
12375                .await
12376            {
12377                Ok(_) => {
12378                    savepoint.commit().await.map_err(map_sqlx_error)?;
12379                    return done.into_job_row();
12380                }
12381                Err(AwaError::UniqueConflict { .. }) => {
12382                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12383                }
12384                Err(err) => return Err(err),
12385            }
12386        }
12387        self.cancel_rescued_duplicate_tx(tx, row, done.payload, duplicate_error)
12388            .await
12389    }
12390
12391    /// Write the cancelled terminal row for a rescue that lost its unique
12392    /// claim to a newer duplicate. `cancelled` normally sits outside the
12393    /// claiming mask so the insert is conflict-free; a mask that claims
12394    /// `cancelled` as well conflicts again, and then the row is written with
12395    /// its unique key stripped — the attempt held no claim (that is what
12396    /// made it conflict), so there is nothing to release, and losing the key
12397    /// on this terminal row beats aborting rescue for the whole cluster.
12398    async fn cancel_rescued_duplicate_tx(
12399        &self,
12400        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12401        row: LeaseTransitionRow,
12402        base_payload: serde_json::Value,
12403        duplicate_error: &str,
12404    ) -> Result<JobRow, AwaError> {
12405        let old_state = row.state;
12406        let mut payload = RuntimePayload::from_json(base_payload)?;
12407        payload.push_error(lifecycle_error(duplicate_error, row.attempt, true));
12408        let done = row.into_done_row(JobState::Cancelled, Utc::now(), payload.into_json());
12409
12410        {
12411            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12412            match self
12413                .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state))
12414                .await
12415            {
12416                Ok(_) => {
12417                    savepoint.commit().await.map_err(map_sqlx_error)?;
12418                    return done.into_job_row();
12419                }
12420                Err(AwaError::UniqueConflict { .. }) => {
12421                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12422                }
12423                Err(err) => return Err(err),
12424            }
12425        }
12426
12427        let mut stripped = done;
12428        stripped.unique_key = None;
12429        stripped.unique_states = None;
12430        self.insert_done_rows_tx(tx, std::slice::from_ref(&stripped), Some(old_state))
12431            .await?;
12432        stripped.into_job_row()
12433    }
12434
12435    pub async fn cancel_running(
12436        &self,
12437        pool: &PgPool,
12438        job_id: i64,
12439        run_lease: i64,
12440        reason: &str,
12441        progress: Option<serde_json::Value>,
12442    ) -> Result<Option<JobRow>, AwaError> {
12443        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12444        let result = self
12445            .cancel_running_in_tx(&mut tx, job_id, run_lease, reason, progress)
12446            .await?;
12447        tx.commit().await.map_err(map_sqlx_error)?;
12448        Ok(result)
12449    }
12450
12451    /// Transaction-aware variant of [`Self::cancel_running`] (ADR-029).
12452    pub async fn cancel_running_in_tx(
12453        &self,
12454        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12455        job_id: i64,
12456        run_lease: i64,
12457        reason: &str,
12458        progress: Option<serde_json::Value>,
12459    ) -> Result<Option<JobRow>, AwaError> {
12460        let Some(moved) = self
12461            .take_running_attempt_tx(tx, job_id, run_lease, "cancelled")
12462            .await?
12463        else {
12464            return Ok(None);
12465        };
12466
12467        let mut payload = RuntimePayload::from_json(Self::with_progress(
12468            moved.payload.clone(),
12469            progress.or(moved.progress.clone()),
12470        )?)?;
12471        payload.push_error(lifecycle_error(
12472            format!("cancelled: {reason}"),
12473            moved.attempt,
12474            false,
12475        ));
12476        let done =
12477            moved
12478                .clone()
12479                .into_done_row(JobState::Cancelled, Utc::now(), payload.into_json());
12480        self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12481            .await?;
12482        Ok(Some(done.into_job_row()?))
12483    }
12484
12485    pub async fn fail_terminal(
12486        &self,
12487        pool: &PgPool,
12488        job_id: i64,
12489        run_lease: i64,
12490        error: &str,
12491        progress: Option<serde_json::Value>,
12492    ) -> Result<Option<JobRow>, AwaError> {
12493        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12494        let result = self
12495            .fail_terminal_in_tx(&mut tx, job_id, run_lease, error, progress)
12496            .await?;
12497        tx.commit().await.map_err(map_sqlx_error)?;
12498        Ok(result)
12499    }
12500
12501    /// Transaction-aware variant of [`Self::fail_terminal`] (ADR-029).
12502    pub async fn fail_terminal_in_tx(
12503        &self,
12504        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12505        job_id: i64,
12506        run_lease: i64,
12507        error: &str,
12508        progress: Option<serde_json::Value>,
12509    ) -> Result<Option<JobRow>, AwaError> {
12510        let Some(moved) = self
12511            .take_running_attempt_tx(tx, job_id, run_lease, "failed")
12512            .await?
12513        else {
12514            return Ok(None);
12515        };
12516
12517        let mut payload = RuntimePayload::from_json(Self::with_progress(
12518            moved.payload.clone(),
12519            progress.or(moved.progress.clone()),
12520        )?)?;
12521        payload.push_error(lifecycle_error(error, moved.attempt, true));
12522        let done = moved
12523            .clone()
12524            .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
12525        self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12526            .await?;
12527        Ok(Some(done.into_job_row()?))
12528    }
12529
12530    pub async fn fail_to_dlq(
12531        &self,
12532        pool: &PgPool,
12533        job_id: i64,
12534        run_lease: i64,
12535        dlq_reason: &str,
12536        error: &str,
12537        progress: Option<serde_json::Value>,
12538    ) -> Result<Option<JobRow>, AwaError> {
12539        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12540        let result = self
12541            .fail_to_dlq_in_tx(&mut tx, job_id, run_lease, dlq_reason, error, progress)
12542            .await?;
12543        tx.commit().await.map_err(map_sqlx_error)?;
12544        Ok(result)
12545    }
12546
12547    /// Transaction-aware variant of [`Self::fail_to_dlq`] (ADR-029).
12548    pub async fn fail_to_dlq_in_tx(
12549        &self,
12550        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12551        job_id: i64,
12552        run_lease: i64,
12553        dlq_reason: &str,
12554        error: &str,
12555        progress: Option<serde_json::Value>,
12556    ) -> Result<Option<JobRow>, AwaError> {
12557        let Some(moved) = self
12558            .take_running_attempt_tx(tx, job_id, run_lease, "dlq")
12559            .await?
12560        else {
12561            return Ok(None);
12562        };
12563
12564        let finalized_at = Utc::now();
12565        let dlq_at = finalized_at;
12566        let mut payload = RuntimePayload::from_json(Self::with_progress(
12567            moved.payload.clone(),
12568            progress.or(moved.progress.clone()),
12569        )?)?;
12570        payload.push_error(lifecycle_error(error, moved.attempt, true));
12571        let dlq_row = moved.clone().into_dlq_row(
12572            finalized_at,
12573            payload.into_json(),
12574            dlq_reason.to_string(),
12575            dlq_at,
12576        );
12577        self.insert_dlq_rows_tx(tx, std::slice::from_ref(&dlq_row), Some(moved.state))
12578            .await?;
12579        Ok(Some(dlq_row.into_job_row()?))
12580    }
12581
12582    pub async fn move_failed_to_dlq(
12583        &self,
12584        pool: &PgPool,
12585        job_id: i64,
12586        dlq_reason: &str,
12587    ) -> Result<Option<JobRow>, AwaError> {
12588        let schema = self.schema();
12589        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12590        let done_projection = done_row_projection("done", "ready");
12591        let ready_join = done_ready_join(schema, "done", "ready");
12592        let moved: Option<DoneJobRow> = sqlx::query_as(&format!(
12593            r#"
12594            WITH deleted AS (
12595                DELETE FROM {schema}.done_entries
12596                WHERE (job_id, finalized_at) IN (
12597                    SELECT job_id, finalized_at
12598                    FROM {schema}.done_entries
12599                    WHERE job_id = $1
12600                      AND state = 'failed'
12601                    ORDER BY finalized_at DESC
12602                    LIMIT 1
12603                    FOR UPDATE SKIP LOCKED
12604                )
12605                RETURNING *
12606            )
12607            SELECT {done_projection}
12608            FROM deleted AS done
12609            {ready_join}
12610            "#
12611        ))
12612        .bind(job_id)
12613        .fetch_optional(tx.as_mut())
12614        .await
12615        .map_err(map_sqlx_error)?;
12616
12617        let Some(moved) = moved else {
12618            tx.commit().await.map_err(map_sqlx_error)?;
12619            return Ok(None);
12620        };
12621
12622        self.ensure_terminal_removed_receipt_closures_tx(&mut tx, std::slice::from_ref(&moved))
12623            .await?;
12624        // DLQ inserts are outside done_entries, so the source terminal DELETE
12625        // appends a negative delta.
12626        self.decrement_live_terminal_counters_tx(
12627            &mut tx,
12628            &Self::done_rows_to_counter_keys(std::slice::from_ref(&moved)),
12629        )
12630        .await?;
12631        let dlq_row = moved
12632            .clone()
12633            .into_dlq_row(dlq_reason.to_string(), Utc::now());
12634        self.insert_dlq_rows_tx(&mut tx, std::slice::from_ref(&dlq_row), Some(moved.state))
12635            .await?;
12636        tx.commit().await.map_err(map_sqlx_error)?;
12637        Ok(Some(dlq_row.into_job_row()?))
12638    }
12639
12640    #[tracing::instrument(
12641        skip(self, pool, dlq_reason),
12642        fields(kind = ?kind, queue = ?queue),
12643        name = "queue_storage.bulk_move_failed_to_dlq"
12644    )]
12645    pub async fn bulk_move_failed_to_dlq(
12646        &self,
12647        pool: &PgPool,
12648        kind: Option<&str>,
12649        queue: Option<&str>,
12650        dlq_reason: &str,
12651    ) -> Result<u64, AwaError> {
12652        let schema = self.schema();
12653        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12654        let done_projection = done_row_projection("done", "ready");
12655        let ready_join = done_ready_join(schema, "done", "ready");
12656        let moved: Vec<DoneJobRow> = sqlx::query_as(&format!(
12657            r#"
12658            WITH deleted AS (
12659                DELETE FROM {schema}.done_entries
12660                WHERE state = 'failed'
12661                  AND ($1::text IS NULL OR kind = $1)
12662                  AND ($2::text IS NULL OR queue = $2)
12663                RETURNING *
12664            )
12665            SELECT {done_projection}
12666            FROM deleted AS done
12667            {ready_join}
12668            "#
12669        ))
12670        .bind(kind)
12671        .bind(queue)
12672        .fetch_all(tx.as_mut())
12673        .await
12674        .map_err(map_sqlx_error)?;
12675
12676        if moved.is_empty() {
12677            tx.commit().await.map_err(map_sqlx_error)?;
12678            return Ok(0);
12679        }
12680
12681        self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &moved)
12682            .await?;
12683        // Bulk DLQ move deletes from done_entries. Append negative deltas by
12684        // the per-group magnitudes of the moved rows.
12685        self.decrement_live_terminal_counters_tx(&mut tx, &Self::done_rows_to_counter_keys(&moved))
12686            .await?;
12687        let dlq_at = Utc::now();
12688        let rows: Vec<DlqJobRow> = moved
12689            .into_iter()
12690            .map(|row| row.into_dlq_row(dlq_reason.to_string(), dlq_at))
12691            .collect();
12692        self.insert_dlq_rows_tx(&mut tx, &rows, Some(JobState::Failed))
12693            .await?;
12694        tx.commit().await.map_err(map_sqlx_error)?;
12695        Ok(rows.len() as u64)
12696    }
12697
12698    pub async fn retry_from_dlq(
12699        &self,
12700        pool: &PgPool,
12701        job_id: i64,
12702        opts: &RetryFromDlqOpts,
12703    ) -> Result<Option<JobRow>, AwaError> {
12704        let schema = self.schema();
12705        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12706        let moved: Option<DlqJobRow> = sqlx::query_as(&format!(
12707            r#"
12708            DELETE FROM {schema}.dlq_entries
12709            WHERE job_id = $1
12710            RETURNING
12711                job_id,
12712                kind,
12713                queue,
12714                args,
12715                state,
12716                priority,
12717                attempt,
12718                run_lease,
12719                max_attempts,
12720                run_at,
12721                attempted_at,
12722                finalized_at,
12723                created_at,
12724                unique_key,
12725                unique_states,
12726                COALESCE(payload, '{{}}'::jsonb) AS payload,
12727                dlq_reason,
12728                dlq_at,
12729                original_run_lease
12730            "#
12731        ))
12732        .bind(job_id)
12733        .fetch_optional(tx.as_mut())
12734        .await
12735        .map_err(map_sqlx_error)?;
12736
12737        let Some(moved) = moved else {
12738            tx.commit().await.map_err(map_sqlx_error)?;
12739            return Ok(None);
12740        };
12741
12742        let queue = opts.queue.clone().unwrap_or_else(|| moved.queue.clone());
12743        let priority = opts.priority.unwrap_or(moved.priority);
12744        let mut payload = RuntimePayload::from_json(moved.payload.clone())?;
12745        payload.set_progress(None);
12746        let payload = payload.into_json();
12747
12748        if let Some(run_at) = opts.run_at.filter(|run_at| *run_at > Utc::now()) {
12749            let deferred = moved.into_retry_deferred_row(queue, priority, run_at, payload);
12750            self.insert_deferred_rows_tx(&mut tx, vec![deferred.clone()], Some(JobState::Failed))
12751                .await?;
12752            tx.commit().await.map_err(map_sqlx_error)?;
12753            return Ok(Some(deferred.into_job_row()?));
12754        }
12755
12756        let ready = moved.into_retry_ready_row(queue.clone(), priority, Utc::now(), payload);
12757        self.insert_existing_ready_rows_tx(&mut tx, vec![ready.clone()], Some(JobState::Failed))
12758            .await?;
12759        self.notify_queues_tx(&mut tx, std::iter::once(queue))
12760            .await?;
12761        tx.commit().await.map_err(map_sqlx_error)?;
12762        Ok(Some(
12763            ReadyJobRow {
12764                job_id: ready.job_id,
12765                kind: ready.kind,
12766                queue: ready.queue,
12767                args: ready.args,
12768                priority: ready.priority,
12769                attempt: ready.attempt,
12770                run_lease: ready.run_lease,
12771                max_attempts: ready.max_attempts,
12772                run_at: ready.run_at,
12773                attempted_at: ready.attempted_at,
12774                created_at: ready.created_at,
12775                unique_key: ready.unique_key,
12776                payload: ready.payload,
12777            }
12778            .into_job_row()?,
12779        ))
12780    }
12781
12782    #[tracing::instrument(
12783        skip(self, pool, filter),
12784        fields(kind = ?filter.kind, queue = ?filter.queue, tag = ?filter.tag),
12785        name = "queue_storage.bulk_retry_from_dlq"
12786    )]
12787    pub async fn bulk_retry_from_dlq(
12788        &self,
12789        pool: &PgPool,
12790        filter: &ListDlqFilter,
12791    ) -> Result<u64, AwaError> {
12792        let schema = self.schema();
12793        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12794        let moved: Vec<DlqJobRow> = sqlx::query_as(&format!(
12795            r#"
12796            DELETE FROM {schema}.dlq_entries
12797            WHERE ($1::text IS NULL OR kind = $1)
12798              AND ($2::text IS NULL OR queue = $2)
12799              AND ($3::text IS NULL OR payload -> 'tags' ? $3)
12800              AND (
12801                  ($4::bigint IS NULL AND $5::timestamptz IS NULL)
12802                  OR ($4::bigint IS NOT NULL AND $5::timestamptz IS NULL AND job_id < $4)
12803                  OR ($4::bigint IS NULL AND $5::timestamptz IS NOT NULL AND dlq_at < $5)
12804                  OR (
12805                      $4::bigint IS NOT NULL
12806                      AND $5::timestamptz IS NOT NULL
12807                      AND (dlq_at, job_id) < ($5, $4)
12808                  )
12809              )
12810            RETURNING
12811                job_id,
12812                kind,
12813                queue,
12814                args,
12815                state,
12816                priority,
12817                attempt,
12818                run_lease,
12819                max_attempts,
12820                run_at,
12821                attempted_at,
12822                finalized_at,
12823                created_at,
12824                unique_key,
12825                unique_states,
12826                COALESCE(payload, '{{}}'::jsonb) AS payload,
12827                dlq_reason,
12828                dlq_at,
12829                original_run_lease
12830            "#
12831        ))
12832        .bind(&filter.kind)
12833        .bind(&filter.queue)
12834        .bind(&filter.tag)
12835        .bind(filter.before_id)
12836        .bind(filter.before_dlq_at)
12837        .fetch_all(tx.as_mut())
12838        .await
12839        .map_err(map_sqlx_error)?;
12840
12841        if moved.is_empty() {
12842            tx.commit().await.map_err(map_sqlx_error)?;
12843            return Ok(0);
12844        }
12845
12846        let run_at = Utc::now();
12847        let mut queues = BTreeSet::new();
12848        let mut ready_rows = Vec::with_capacity(moved.len());
12849        for moved_row in moved {
12850            let queue = moved_row.queue.clone();
12851            let priority = moved_row.priority;
12852            queues.insert(queue.clone());
12853            let mut payload = RuntimePayload::from_json(moved_row.payload.clone())?;
12854            payload.set_progress(None);
12855            ready_rows.push(moved_row.into_retry_ready_row(
12856                queue,
12857                priority,
12858                run_at,
12859                payload.into_json(),
12860            ));
12861        }
12862
12863        let revived = ready_rows.len() as u64;
12864        self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(JobState::Failed))
12865            .await?;
12866        self.notify_queues_tx(&mut tx, queues).await?;
12867        tx.commit().await.map_err(map_sqlx_error)?;
12868        Ok(revived)
12869    }
12870
12871    pub async fn discard_failed_by_kind(&self, pool: &PgPool, kind: &str) -> Result<u64, AwaError> {
12872        let schema = self.schema();
12873        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12874
12875        let done_projection = done_row_projection("done", "ready");
12876        let ready_join = done_ready_join(schema, "done", "ready");
12877        let deleted_done: Vec<DoneJobRow> = sqlx::query_as(&format!(
12878            r#"
12879            WITH deleted AS (
12880                DELETE FROM {schema}.done_entries
12881                WHERE kind = $1
12882                  AND state = 'failed'
12883                RETURNING *
12884            )
12885            SELECT {done_projection}
12886            FROM deleted AS done
12887            {ready_join}
12888            "#
12889        ))
12890        .bind(kind)
12891        .fetch_all(tx.as_mut())
12892        .await
12893        .map_err(map_sqlx_error)?;
12894
12895        let deleted_dlq: Vec<DlqJobRow> = sqlx::query_as(&format!(
12896            r#"
12897            DELETE FROM {schema}.dlq_entries
12898            WHERE kind = $1
12899            RETURNING
12900                job_id,
12901                kind,
12902                queue,
12903                args,
12904                state,
12905                priority,
12906                attempt,
12907                run_lease,
12908                max_attempts,
12909                run_at,
12910                attempted_at,
12911                finalized_at,
12912                created_at,
12913                unique_key,
12914                unique_states,
12915                COALESCE(payload, '{{}}'::jsonb) AS payload,
12916                dlq_reason,
12917                dlq_at,
12918                original_run_lease
12919            "#
12920        ))
12921        .bind(kind)
12922        .fetch_all(tx.as_mut())
12923        .await
12924        .map_err(map_sqlx_error)?;
12925
12926        // Discard removes terminal `failed` rows from done_entries (and
12927        // `failed` DLQ rows from dlq_entries — exact terminal counts are keyed
12928        // on done_entries only). Append negative deltas by the per-group
12929        // magnitudes of `deleted_done`.
12930        self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &deleted_done)
12931            .await?;
12932        self.decrement_live_terminal_counters_tx(
12933            &mut tx,
12934            &Self::done_rows_to_counter_keys(&deleted_done),
12935        )
12936        .await?;
12937
12938        for row in &deleted_done {
12939            self.sync_unique_claim(
12940                &mut tx,
12941                row.job_id,
12942                &row.unique_key,
12943                row.unique_states.as_deref(),
12944                Some(row.state),
12945                None,
12946            )
12947            .await?;
12948        }
12949
12950        for row in &deleted_dlq {
12951            self.sync_unique_claim(
12952                &mut tx,
12953                row.job_id,
12954                &row.unique_key,
12955                row.unique_states.as_deref(),
12956                Some(row.state),
12957                None,
12958            )
12959            .await?;
12960        }
12961
12962        tx.commit().await.map_err(map_sqlx_error)?;
12963        Ok((deleted_done.len() + deleted_dlq.len()) as u64)
12964    }
12965
12966    pub async fn fail_retryable(
12967        &self,
12968        pool: &PgPool,
12969        job_id: i64,
12970        run_lease: i64,
12971        error: &str,
12972        progress: Option<serde_json::Value>,
12973    ) -> Result<Option<JobRow>, AwaError> {
12974        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12975        let result = self
12976            .fail_retryable_in_tx(&mut tx, job_id, run_lease, error, progress)
12977            .await?;
12978        tx.commit().await.map_err(map_sqlx_error)?;
12979        Ok(result)
12980    }
12981
12982    /// Transaction-aware variant of [`Self::fail_retryable`] (ADR-029).
12983    pub async fn fail_retryable_in_tx(
12984        &self,
12985        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12986        job_id: i64,
12987        run_lease: i64,
12988        error: &str,
12989        progress: Option<serde_json::Value>,
12990    ) -> Result<Option<JobRow>, AwaError> {
12991        let Some(moved) = self
12992            .take_running_attempt_tx(tx, job_id, run_lease, "retryable")
12993            .await?
12994        else {
12995            return Ok(None);
12996        };
12997
12998        let mut payload = RuntimePayload::from_json(Self::with_progress(
12999            moved.payload.clone(),
13000            progress.or(moved.progress.clone()),
13001        )?)?;
13002        let exhausted = moved.attempt >= moved.max_attempts;
13003        payload.push_error(lifecycle_error(error, moved.attempt, exhausted));
13004
13005        if exhausted {
13006            let done =
13007                moved
13008                    .clone()
13009                    .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
13010            self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
13011                .await?;
13012            return Ok(Some(done.into_job_row()?));
13013        }
13014
13015        let deferred = moved.clone().into_deferred_row(
13016            JobState::Retryable,
13017            self.backoff_at_tx(tx, moved.attempt, moved.max_attempts)
13018                .await?,
13019            Some(Utc::now()),
13020            payload.into_json(),
13021        );
13022        self.insert_deferred_rows_tx(tx, vec![deferred.clone()], Some(moved.state))
13023            .await?;
13024        Ok(Some(deferred.into_job_row()?))
13025    }
13026
13027    #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_stale_heartbeats")]
13028    pub async fn rescue_stale_heartbeats(
13029        &self,
13030        pool: &PgPool,
13031        staleness: Duration,
13032    ) -> Result<Vec<JobRow>, AwaError> {
13033        let schema = self.schema();
13034        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13035        let cutoff = Utc::now()
13036            - TimeDelta::from_std(staleness)
13037                .map_err(|err| AwaError::Validation(format!("invalid staleness: {err}")))?;
13038        // #169 B1: the staleness predicate prefers
13039        // `attempt_state.heartbeat_at` (receipts-mode source of truth,
13040        // where heartbeat_batch writes go) and falls back to
13041        // `leases.heartbeat_at` (legacy non-receipts source). Two
13042        // cases this covers:
13043        //
13044        //   * Receipts mode, claim materialized into `leases` via
13045        //     callback registration / progress upsert / equivalent.
13046        //     The leases row was written once at materialize time and
13047        //     its `heartbeat_at` is stale by definition. The fresh
13048        //     value lives on `attempt_state.heartbeat_at`. COALESCE
13049        //     picks `attempt` so a healthy worker isn't falsely
13050        //     rescued — and a dead worker IS rescued, which the
13051        //     receipt-side path can't do (its anti-join below
13052        //     excludes materialized leases to avoid double-closure).
13053        //   * Legacy non-receipts mode: attempt_state.heartbeat_at is
13054        //     never written; COALESCE falls back to
13055        //     leases.heartbeat_at — same shape as pre-B1.
13056        //
13057        // The dropped `idx_state_hb` (v025) doesn't hurt this scan:
13058        // the planner can satisfy the `state='running'` prefix via the
13059        // surviving `(state, deadline_at)` or
13060        // `(state, callback_timeout_at)` indexes, followed by a heap
13061        // recheck of the COALESCE. Bounded by running-lease count and
13062        // called at 30s cadence — cheap.
13063        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
13064            r#"
13065            DELETE FROM {schema}.leases AS target
13066            WHERE (target.job_id, target.run_lease) IN (
13067                SELECT lease.job_id, lease.run_lease
13068                FROM {schema}.leases AS lease
13069                LEFT JOIN {schema}.attempt_state AS attempt
13070                  ON attempt.job_id = lease.job_id
13071                 AND attempt.run_lease = lease.run_lease
13072                WHERE lease.state = 'running'
13073                  AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1
13074                ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC
13075                LIMIT 500
13076                FOR UPDATE OF lease SKIP LOCKED
13077            )
13078            RETURNING
13079                ready_slot,
13080                ready_generation,
13081                job_id,
13082                queue,
13083                state,
13084                priority,
13085                attempt,
13086                run_lease,
13087                max_attempts,
13088                lane_seq,
13089                enqueue_shard,
13090                heartbeat_at,
13091                deadline_at,
13092                attempted_at,
13093                callback_id,
13094                callback_timeout_at
13095            "#
13096        ))
13097        .bind(cutoff)
13098        .fetch_all(tx.as_mut())
13099        .await
13100        .map_err(map_sqlx_error)?;
13101
13102        let rescued_receipts = if self.lease_claim_receipts() {
13103            self.rescue_stale_receipt_claims_tx(&mut tx, cutoff).await?
13104        } else {
13105            Vec::new()
13106        };
13107
13108        if deleted.is_empty() && rescued_receipts.is_empty() {
13109            tx.commit().await.map_err(map_sqlx_error)?;
13110            return Ok(Vec::new());
13111        }
13112
13113        let moved_leases = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13114        let moved_receipts = self
13115            .hydrate_deleted_leases_tx(&mut tx, rescued_receipts)
13116            .await?;
13117
13118        let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len());
13119        for row in moved_leases {
13120            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13121                row.payload.clone(),
13122                row.progress.clone(),
13123            )?)?;
13124            payload.push_error(lifecycle_error(
13125                "heartbeat stale: worker presumed dead",
13126                row.attempt,
13127                false,
13128            ));
13129            let deferred = row.clone().into_deferred_row(
13130                JobState::Retryable,
13131                self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13132                    .await?,
13133                Some(Utc::now()),
13134                payload.into_json(),
13135            );
13136            let rescued_row = self
13137                .insert_rescued_deferred_or_cancel_tx(
13138                    &mut tx,
13139                    row,
13140                    deferred,
13141                    "rescued as duplicate: unique claim held by a newer job",
13142                )
13143                .await?;
13144            rescued.push(rescued_row);
13145        }
13146        for row in moved_receipts {
13147            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13148                row.payload.clone(),
13149                row.progress.clone(),
13150            )?)?;
13151            payload.push_error(lifecycle_error(
13152                "receipt claim stale: worker presumed dead",
13153                row.attempt,
13154                false,
13155            ));
13156            let deferred = row.clone().into_deferred_row(
13157                JobState::Retryable,
13158                self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13159                    .await?,
13160                Some(Utc::now()),
13161                payload.into_json(),
13162            );
13163            let rescued_row = self
13164                .insert_rescued_deferred_or_cancel_tx(
13165                    &mut tx,
13166                    row,
13167                    deferred,
13168                    "rescued as duplicate: unique claim held by a newer job",
13169                )
13170                .await?;
13171            rescued.push(rescued_row);
13172        }
13173        tx.commit().await.map_err(map_sqlx_error)?;
13174        Ok(rescued)
13175    }
13176
13177    #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_expired_deadlines")]
13178    pub async fn rescue_expired_deadlines(&self, pool: &PgPool) -> Result<Vec<JobRow>, AwaError> {
13179        let schema = self.schema();
13180        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13181        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
13182            r#"
13183            DELETE FROM {schema}.leases
13184            WHERE job_id IN (
13185                SELECT job_id
13186                FROM {schema}.leases
13187                WHERE state = 'running'
13188                  AND deadline_at IS NOT NULL
13189                  AND deadline_at < clock_timestamp()
13190                ORDER BY deadline_at ASC
13191                LIMIT 500
13192                FOR UPDATE SKIP LOCKED
13193            )
13194            RETURNING
13195                ready_slot,
13196                ready_generation,
13197                job_id,
13198                queue,
13199                state,
13200                priority,
13201                attempt,
13202                run_lease,
13203                max_attempts,
13204                lane_seq,
13205                enqueue_shard,
13206                heartbeat_at,
13207                deadline_at,
13208                attempted_at,
13209                callback_id,
13210                callback_timeout_at
13211            "#
13212        ))
13213        .fetch_all(tx.as_mut())
13214        .await
13215        .map_err(map_sqlx_error)?;
13216
13217        // Receipts-mode short-path claims hold their deadline on
13218        // `lease_claims.deadline_at` rather than on a `leases` row, so
13219        // the receipt-plane needs its own scan; merge both populations
13220        // into one `moved` set so the maintenance caller observes a
13221        // single rescue batch per tick.
13222        let receipt_deleted = if self.lease_claim_receipts() {
13223            self.rescue_expired_receipt_deadlines_tx(&mut tx).await?
13224        } else {
13225            Vec::new()
13226        };
13227
13228        if deleted.is_empty() && receipt_deleted.is_empty() {
13229            tx.commit().await.map_err(map_sqlx_error)?;
13230            return Ok(Vec::new());
13231        }
13232
13233        let mut moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13234        moved.extend(
13235            self.hydrate_deleted_leases_tx(&mut tx, receipt_deleted)
13236                .await?,
13237        );
13238
13239        let mut rescued = Vec::with_capacity(moved.len());
13240        for row in moved {
13241            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13242                row.payload.clone(),
13243                row.progress.clone(),
13244            )?)?;
13245            payload.push_error(lifecycle_error(
13246                "hard deadline exceeded",
13247                row.attempt,
13248                false,
13249            ));
13250            let deferred = row.clone().into_deferred_row(
13251                JobState::Retryable,
13252                self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13253                    .await?,
13254                Some(Utc::now()),
13255                payload.into_json(),
13256            );
13257            let rescued_row = self
13258                .insert_rescued_deferred_or_cancel_tx(
13259                    &mut tx,
13260                    row,
13261                    deferred,
13262                    "rescued as duplicate: unique claim held by a newer job",
13263                )
13264                .await?;
13265            rescued.push(rescued_row);
13266        }
13267        tx.commit().await.map_err(map_sqlx_error)?;
13268        Ok(rescued)
13269    }
13270
13271    #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_expired_callbacks")]
13272    pub async fn rescue_expired_callbacks(&self, pool: &PgPool) -> Result<Vec<JobRow>, AwaError> {
13273        let schema = self.schema();
13274        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13275        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
13276            r#"
13277            DELETE FROM {schema}.leases
13278            WHERE job_id IN (
13279                SELECT job_id
13280                FROM {schema}.leases
13281                WHERE state = 'waiting_external'
13282                  AND callback_timeout_at IS NOT NULL
13283                  AND callback_timeout_at < clock_timestamp()
13284                ORDER BY callback_timeout_at ASC
13285                LIMIT 500
13286                FOR UPDATE SKIP LOCKED
13287            )
13288            RETURNING
13289                ready_slot,
13290                ready_generation,
13291                job_id,
13292                queue,
13293                state,
13294                priority,
13295                attempt,
13296                run_lease,
13297                max_attempts,
13298                lane_seq,
13299                enqueue_shard,
13300                heartbeat_at,
13301                deadline_at,
13302                attempted_at,
13303                callback_id,
13304                callback_timeout_at
13305            "#
13306        ))
13307        .fetch_all(tx.as_mut())
13308        .await
13309        .map_err(map_sqlx_error)?;
13310
13311        if deleted.is_empty() {
13312            tx.commit().await.map_err(map_sqlx_error)?;
13313            return Ok(Vec::new());
13314        }
13315
13316        let moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13317
13318        let mut rescued = Vec::with_capacity(moved.len());
13319        for row in moved {
13320            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13321                row.payload.clone(),
13322                row.progress.clone(),
13323            )?)?;
13324            let exhausted = row.attempt >= row.max_attempts;
13325            payload.push_error(lifecycle_error(
13326                "callback timed out",
13327                row.attempt,
13328                exhausted,
13329            ));
13330            if exhausted {
13331                let done =
13332                    row.clone()
13333                        .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
13334                let rescued_row = self
13335                    .insert_rescued_done_or_cancel_tx(
13336                        &mut tx,
13337                        row,
13338                        done,
13339                        "rescued as duplicate: unique claim held by a newer job",
13340                    )
13341                    .await?;
13342                rescued.push(rescued_row);
13343            } else {
13344                let deferred = row.clone().into_deferred_row(
13345                    JobState::Retryable,
13346                    self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13347                        .await?,
13348                    Some(Utc::now()),
13349                    payload.into_json(),
13350                );
13351                let rescued_row = self
13352                    .insert_rescued_deferred_or_cancel_tx(
13353                        &mut tx,
13354                        row,
13355                        deferred,
13356                        "rescued as duplicate: unique claim held by a newer job",
13357                    )
13358                    .await?;
13359                rescued.push(rescued_row);
13360            }
13361        }
13362        tx.commit().await.map_err(map_sqlx_error)?;
13363        Ok(rescued)
13364    }
13365
13366    pub async fn promote_due(
13367        &self,
13368        pool: &PgPool,
13369        state: JobState,
13370        batch_size: i64,
13371    ) -> Result<usize, AwaError> {
13372        if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 {
13373            return Ok(0);
13374        }
13375
13376        let schema = self.schema();
13377        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13378        let moved: Vec<DeferredJobRow> = sqlx::query_as(&format!(
13379            r#"
13380            DELETE FROM {schema}.deferred_jobs
13381            WHERE job_id IN (
13382                SELECT job_id
13383                FROM {schema}.deferred_jobs
13384                WHERE state = $1
13385                  AND run_at <= clock_timestamp()
13386                  AND NOT EXISTS (
13387                      SELECT 1 FROM awa.queue_meta
13388                      WHERE queue = {schema}.deferred_jobs.queue AND paused = TRUE
13389                  )
13390                ORDER BY run_at ASC, priority ASC, job_id ASC
13391                LIMIT $2
13392                FOR UPDATE SKIP LOCKED
13393            )
13394            RETURNING
13395                job_id,
13396                kind,
13397                queue,
13398                args,
13399                state,
13400                priority,
13401                attempt,
13402                run_lease,
13403                max_attempts,
13404                run_at,
13405                attempted_at,
13406                finalized_at,
13407                created_at,
13408                unique_key,
13409                unique_states,
13410                COALESCE(payload, '{{}}'::jsonb) AS payload
13411            "#
13412        ))
13413        .bind(state)
13414        .bind(batch_size)
13415        .fetch_all(tx.as_mut())
13416        .await
13417        .map_err(map_sqlx_error)?;
13418
13419        if moved.is_empty() {
13420            tx.commit().await.map_err(map_sqlx_error)?;
13421            return Ok(0);
13422        }
13423
13424        let ready_rows: Vec<ExistingReadyRow> = moved
13425            .iter()
13426            .cloned()
13427            .map(|row| ExistingReadyRow {
13428                job_id: row.job_id,
13429                kind: row.kind,
13430                queue: row.queue,
13431                args: row.args,
13432                priority: row.priority,
13433                attempt: row.attempt,
13434                run_lease: row.run_lease,
13435                max_attempts: row.max_attempts,
13436                run_at: Utc::now(),
13437                attempted_at: row.attempted_at,
13438                created_at: row.created_at,
13439                unique_key: row.unique_key,
13440                unique_states: row.unique_states,
13441                payload: row.payload,
13442            })
13443            .collect();
13444        let queues = ready_rows
13445            .iter()
13446            .map(|row| row.queue.clone())
13447            .collect::<Vec<_>>();
13448        self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(state))
13449            .await?;
13450        self.notify_queues_tx(&mut tx, queues).await?;
13451        tx.commit().await.map_err(map_sqlx_error)?;
13452        Ok(moved.len())
13453    }
13454
13455    async fn relation_has_rows_tx(
13456        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
13457        relation: &str,
13458    ) -> Result<bool, AwaError> {
13459        sqlx::query_scalar(&format!("SELECT EXISTS (SELECT 1 FROM {relation} LIMIT 1)"))
13460            .fetch_one(tx.as_mut())
13461            .await
13462            .map_err(map_sqlx_error)
13463    }
13464
13465    /// Returns `true` if any of `relations` has at least one row, checking in
13466    /// order and stopping at the first non-empty one. Used by the ring prune
13467    /// paths to prove an oldest sealed slot is empty before taking
13468    /// `ACCESS EXCLUSIVE` and issuing a destructive `TRUNCATE` on it.
13469    async fn any_relation_has_rows_tx(
13470        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
13471        relations: &[&str],
13472    ) -> Result<bool, AwaError> {
13473        for relation in relations {
13474            if Self::relation_has_rows_tx(tx, relation).await? {
13475                return Ok(true);
13476            }
13477        }
13478        Ok(false)
13479    }
13480
13481    #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate")]
13482    pub async fn rotate(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
13483        let schema = self.schema();
13484        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13485
13486        let state: (i32, i64, i32) = sqlx::query_as(&format!(
13487            r#"
13488            SELECT current_slot, generation, slot_count
13489            FROM {schema}.queue_ring_state
13490            WHERE singleton = TRUE
13491            FOR UPDATE
13492            "#
13493        ))
13494        .fetch_one(tx.as_mut())
13495        .await
13496        .map_err(map_sqlx_error)?;
13497
13498        let next_slot = (state.0 + 1).rem_euclid(state.2);
13499        let ready_busy =
13500            Self::relation_has_rows_tx(&mut tx, &ready_child_name(schema, next_slot as usize))
13501                .await?;
13502        let claim_attempt_batch_busy = Self::relation_has_rows_tx(
13503            &mut tx,
13504            &ready_claim_attempt_batch_child_name(schema, next_slot as usize),
13505        )
13506        .await?;
13507        let done_busy =
13508            Self::relation_has_rows_tx(&mut tx, &done_child_name(schema, next_slot as usize))
13509                .await?;
13510        let tombstone_busy = Self::relation_has_rows_tx(
13511            &mut tx,
13512            &ready_tombstone_child_name(schema, next_slot as usize),
13513        )
13514        .await?;
13515        let segment_busy = Self::relation_has_rows_tx(
13516            &mut tx,
13517            &ready_segment_child_name(schema, next_slot as usize),
13518        )
13519        .await?;
13520        let receipt_batch_busy = Self::relation_has_rows_tx(
13521            &mut tx,
13522            &receipt_completion_batch_child_name(schema, next_slot as usize),
13523        )
13524        .await?;
13525        let receipt_tombstone_busy = Self::relation_has_rows_tx(
13526            &mut tx,
13527            &receipt_completion_tombstone_child_name(schema, next_slot as usize),
13528        )
13529        .await?;
13530        let terminal_delta_busy = Self::relation_has_rows_tx(
13531            &mut tx,
13532            &terminal_delta_child_name(schema, next_slot as usize),
13533        )
13534        .await?;
13535
13536        if ready_busy
13537            || claim_attempt_batch_busy
13538            || done_busy
13539            || tombstone_busy
13540            || segment_busy
13541            || receipt_batch_busy
13542            || receipt_tombstone_busy
13543            || terminal_delta_busy
13544        {
13545            tx.commit().await.map_err(map_sqlx_error)?;
13546            return Ok(RotateOutcome::SkippedBusy {
13547                slot: next_slot,
13548                busy: BusyCounts {
13549                    queue_ready: busy_indicator(ready_busy),
13550                    queue_claim_attempt_batches: busy_indicator(claim_attempt_batch_busy),
13551                    queue_done: busy_indicator(done_busy),
13552                    queue_tombstones: busy_indicator(tombstone_busy),
13553                    queue_ready_segments: busy_indicator(segment_busy),
13554                    queue_receipt_completion_batches: busy_indicator(receipt_batch_busy),
13555                    queue_receipt_completion_tombstones: busy_indicator(receipt_tombstone_busy),
13556                    queue_terminal_deltas: busy_indicator(terminal_delta_busy),
13557                    ..Default::default()
13558                },
13559            });
13560        }
13561
13562        let next_generation = state.1 + 1;
13563
13564        sqlx::query(&format!(
13565            r#"
13566            UPDATE {schema}.queue_ring_state
13567            SET current_slot = $1,
13568                generation = $2
13569            WHERE singleton = TRUE
13570            "#
13571        ))
13572        .bind(next_slot)
13573        .bind(next_generation)
13574        .execute(tx.as_mut())
13575        .await
13576        .map_err(map_sqlx_error)?;
13577
13578        sqlx::query(&format!(
13579            r#"
13580            UPDATE {schema}.queue_ring_slots
13581            SET generation = $2
13582            WHERE slot = $1
13583            "#
13584        ))
13585        .bind(next_slot)
13586        .bind(next_generation)
13587        .execute(tx.as_mut())
13588        .await
13589        .map_err(map_sqlx_error)?;
13590
13591        tx.commit().await.map_err(map_sqlx_error)?;
13592        Ok(RotateOutcome::Rotated {
13593            slot: next_slot,
13594            generation: next_generation,
13595        })
13596    }
13597
13598    #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_leases")]
13599    pub async fn rotate_leases(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
13600        let schema = self.schema();
13601        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13602
13603        // FOR UPDATE serialises with prune_oldest_leases and parallel
13604        // rotators. Without it two rotators can both pass the busy-check,
13605        // both compute the same next_slot, and the loser's CAS update
13606        // wastes work. `RotateLeasesPlan` in
13607        // `correctness/storage/AwaStorageLockOrder.tla` requires this
13608        // lock as the first acquired resource for the rotation tx.
13609        let state: (i32, i64, i32) = sqlx::query_as(&format!(
13610            r#"
13611            SELECT current_slot, generation, slot_count
13612            FROM {schema}.lease_ring_state
13613            WHERE singleton = TRUE
13614            FOR UPDATE
13615            "#
13616        ))
13617        .fetch_one(tx.as_mut())
13618        .await
13619        .map_err(map_sqlx_error)?;
13620
13621        let next_slot = (state.0 + 1).rem_euclid(state.2);
13622        let lease_busy =
13623            Self::relation_has_rows_tx(&mut tx, &lease_child_name(schema, next_slot as usize))
13624                .await?;
13625
13626        if lease_busy {
13627            tx.commit().await.map_err(map_sqlx_error)?;
13628            return Ok(RotateOutcome::SkippedBusy {
13629                slot: next_slot,
13630                busy: BusyCounts {
13631                    leases: busy_indicator(lease_busy),
13632                    ..Default::default()
13633                },
13634            });
13635        }
13636
13637        let next_generation = state.1 + 1;
13638
13639        let rotated = sqlx::query(&format!(
13640            r#"
13641            UPDATE {schema}.lease_ring_state
13642            SET current_slot = $1,
13643                generation = $2
13644            WHERE singleton = TRUE
13645              AND current_slot = $3
13646              AND generation = $4
13647            "#
13648        ))
13649        .bind(next_slot)
13650        .bind(next_generation)
13651        .bind(state.0)
13652        .bind(state.1)
13653        .execute(tx.as_mut())
13654        .await
13655        .map_err(map_sqlx_error)?;
13656
13657        if rotated.rows_affected() == 0 {
13658            // Another rotator beat us to the CAS. Report the bounded
13659            // presence evidence we sampled before giving up.
13660            tx.commit().await.map_err(map_sqlx_error)?;
13661            return Ok(RotateOutcome::SkippedBusy {
13662                slot: next_slot,
13663                busy: BusyCounts {
13664                    leases: busy_indicator(lease_busy),
13665                    ..Default::default()
13666                },
13667            });
13668        }
13669
13670        tx.commit().await.map_err(map_sqlx_error)?;
13671        Ok(RotateOutcome::Rotated {
13672            slot: next_slot,
13673            generation: next_generation,
13674        })
13675    }
13676
13677    /// Rebuild terminal counters from scratch by truncating folded live counts
13678    /// and pending deltas, then re-aggregating terminal history. Run this when:
13679    ///
13680    /// - upgrading from a pre-#290 fleet, where in-flight binaries wrote
13681    ///   terminal history without maintaining the counter,
13682    /// - after any incident that may have left the counter inconsistent
13683    ///   with terminal history, or
13684    /// - as a routine drift-recovery step before relying on counter-fed
13685    ///   reads for billing-grade accuracy.
13686    ///
13687    /// The rebuild is wrapped in an advisory lock so concurrent writers
13688    /// don't interleave new inserts mid-rebuild. The lock key is scoped
13689    /// to the queue-storage schema, so other schemas / other operations
13690    /// are unaffected. Writers that hit the lock will block briefly
13691    /// rather than fail.
13692    ///
13693    /// **Operator note:** this is best run on a quiesced fleet (workers
13694    /// paused or fully drained). Concurrent inserts during the rebuild
13695    /// will block on the lock; long-held locks can stall the fleet. The
13696    /// rebuild itself is O(rows in `{schema}.done_entries`). Compact receipt
13697    /// completions remain counted directly from retained
13698    /// `receipt_completion_batches` until queue prune folds them into
13699    /// permanent rollups.
13700    #[tracing::instrument(skip(self, pool), name = "queue_storage.rebuild_terminal_counters")]
13701    pub async fn rebuild_terminal_counters(&self, pool: &PgPool) -> Result<i64, AwaError> {
13702        let schema = self.schema();
13703        let rebuild_lock_name = format!("awa.queue_storage.rebuild_terminal_counters:{schema}");
13704        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13705
13706        sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
13707            .bind(&rebuild_lock_name)
13708            .execute(tx.as_mut())
13709            .await
13710            .map_err(map_sqlx_error)?;
13711
13712        sqlx::query(&format!(
13713            "LOCK TABLE {schema}.done_entries, \
13714             {schema}.receipt_completion_batches, \
13715             {schema}.receipt_completion_tombstones, \
13716             {schema}.queue_terminal_count_deltas, \
13717             {schema}.queue_terminal_live_counts \
13718             IN ACCESS EXCLUSIVE MODE"
13719        ))
13720        .execute(tx.as_mut())
13721        .await
13722        .map_err(map_sqlx_error)?;
13723
13724        sqlx::query(&format!(
13725            "TRUNCATE TABLE {schema}.queue_terminal_live_counts, {schema}.queue_terminal_count_deltas"
13726        ))
13727        .execute(tx.as_mut())
13728        .await
13729        .map_err(map_sqlx_error)?;
13730
13731        let inserted: i64 = sqlx::query_scalar(&format!(
13732            r#"
13733            WITH inserted AS (
13734                INSERT INTO {schema}.queue_terminal_live_counts AS counts (
13735                    ready_slot, queue, priority, enqueue_shard, counter_bucket, live_terminal_count
13736                )
13737                SELECT
13738                    ready_slot,
13739                    queue,
13740                    priority,
13741                    enqueue_shard,
13742                    mod(
13743                        mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint)
13744                            + {TERMINAL_COUNTER_BUCKETS}::bigint,
13745                        {TERMINAL_COUNTER_BUCKETS}::bigint
13746                    )::smallint AS counter_bucket,
13747                    count(*)::bigint
13748                FROM {schema}.done_entries
13749                GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket
13750                RETURNING 1
13751            )
13752            SELECT COALESCE(count(*), 0)::bigint FROM inserted
13753            "#
13754        ))
13755        .fetch_one(tx.as_mut())
13756        .await
13757        .map_err(map_sqlx_error)?;
13758
13759        // Flip the trust marker. From this point the read path
13760        // (queue_counts_exact) uses the counter for done-entry terminal rows;
13761        // before this call, it falls back to scanning terminal_jobs.
13762        sqlx::query(&format!(
13763            r#"
13764            UPDATE {schema}.queue_ring_state
13765            SET terminal_counter_trusted_at = now()
13766            WHERE singleton = TRUE
13767            "#
13768        ))
13769        .execute(tx.as_mut())
13770        .await
13771        .map_err(map_sqlx_error)?;
13772
13773        tx.commit().await.map_err(map_sqlx_error)?;
13774        Ok(inserted)
13775    }
13776
13777    /// Check whether the live terminal counter has been marked trusted
13778    /// for exact reads. Returns `true` for fresh installs (the trust
13779    /// marker is auto-set by `prepare_schema` when `done_entries` is
13780    /// empty) and after a successful
13781    /// [`Self::rebuild_terminal_counters`]; returns `false` on an
13782    /// existing install that has not yet been rebuilt under the new
13783    /// counter-aware code path.
13784    ///
13785    /// `queue_counts_exact` consults this to decide between the
13786    /// counter-fed fast path and the scan-based fallback. Single-row
13787    /// PK fetch; negligible cost per call.
13788    pub async fn terminal_counter_trusted(&self, pool: &PgPool) -> Result<bool, AwaError> {
13789        let schema = self.schema();
13790        let trusted: Option<bool> = sqlx::query_scalar(&format!(
13791            "SELECT terminal_counter_trusted_at IS NOT NULL \
13792             FROM {schema}.queue_ring_state WHERE singleton = TRUE"
13793        ))
13794        .fetch_optional(pool)
13795        .await
13796        .map_err(map_sqlx_error)?;
13797        // Missing row = pre-#290 schema that hasn't run the new
13798        // prepare_schema yet. Treat as untrusted; the scan path is
13799        // still correct.
13800        Ok(trusted.unwrap_or(false))
13801    }
13802
13803    /// Fold append-only `done_entries` terminal-count deltas into
13804    /// `queue_terminal_live_counts` for sealed queue slots.
13805    ///
13806    /// `done_entries` terminal inserts and deletes append signed rows into
13807    /// `queue_terminal_count_deltas` instead of updating the live counter on
13808    /// the user-facing hot path. Compact receipt completions are counted from
13809    /// retained batch rows directly. Exact reads sum folded live counts plus
13810    /// pending deltas and retained compact batches, so this rollup can run
13811    /// asynchronously. It intentionally
13812    /// skips the current queue slot and any slot with active leases or open
13813    /// receipt claims; that keeps rollup away from the segment receiving hot
13814    /// completions and avoids racing future terminal deltas for the same
13815    /// ready generation. Rollup also stands down while another backend pins the
13816    /// MVCC horizon, leaving the append-only deltas visible to exact reads
13817    /// until the horizon clears.
13818    #[tracing::instrument(skip(self, pool), name = "queue_storage.rollup_terminal_count_deltas")]
13819    pub async fn rollup_terminal_count_deltas(
13820        &self,
13821        pool: &PgPool,
13822        max_slots: usize,
13823    ) -> Result<TerminalDeltaRollupOutcome, AwaError> {
13824        if max_slots == 0 {
13825            return Ok(TerminalDeltaRollupOutcome::default());
13826        }
13827
13828        let schema = self.schema();
13829        let current_slot: i32 = sqlx::query_scalar(&format!(
13830            r#"
13831            SELECT current_slot
13832            FROM {schema}.queue_ring_state
13833            WHERE singleton = TRUE
13834            "#
13835        ))
13836        .fetch_one(pool)
13837        .await
13838        .map_err(map_sqlx_error)?;
13839
13840        let slots = self
13841            .terminal_delta_rollup_candidates(pool, current_slot)
13842            .await?;
13843
13844        let mut outcome = TerminalDeltaRollupOutcome::default();
13845        for (slot, generation) in slots {
13846            if outcome.rolled_slots >= max_slots {
13847                break;
13848            }
13849            match self
13850                .rollup_terminal_count_delta_slot(pool, slot, generation)
13851                .await?
13852            {
13853                TerminalDeltaSlotRollup::Empty => {}
13854                TerminalDeltaSlotRollup::Rolled {
13855                    delta_rows,
13856                    grouped_keys,
13857                } => {
13858                    outcome.rolled_slots += 1;
13859                    outcome.delta_rows += delta_rows;
13860                    outcome.grouped_keys += grouped_keys;
13861                }
13862                TerminalDeltaSlotRollup::SkippedActive => {
13863                    outcome.skipped_active_slots += 1;
13864                }
13865                TerminalDeltaSlotRollup::SkippedMvccPinned => {
13866                    outcome.skipped_mvcc_pinned = true;
13867                    break;
13868                }
13869                TerminalDeltaSlotRollup::Blocked => {
13870                    outcome.blocked_slots += 1;
13871                }
13872            }
13873        }
13874
13875        Ok(outcome)
13876    }
13877
13878    async fn terminal_delta_rollup_mvcc_horizon_pinned_tx(
13879        tx: &mut sqlx::Transaction<'_, Postgres>,
13880    ) -> Result<bool, AwaError> {
13881        let pinned: bool = sqlx::query_scalar(
13882            r#"
13883            SELECT EXISTS (
13884                SELECT 1
13885                FROM pg_stat_activity
13886                WHERE datname = current_database()
13887                  AND pid <> pg_backend_pid()
13888                  AND backend_type = 'client backend'
13889                  AND (
13890                      backend_xmin IS NOT NULL
13891                      OR (
13892                          backend_xid IS NOT NULL
13893                          AND state LIKE 'idle in transaction%'
13894                      )
13895                  )
13896            )
13897            "#,
13898        )
13899        .fetch_one(tx.as_mut())
13900        .await
13901        .map_err(map_sqlx_error)?;
13902        Ok(pinned)
13903    }
13904
13905    async fn terminal_delta_rollup_candidates(
13906        &self,
13907        pool: &PgPool,
13908        current_slot: i32,
13909    ) -> Result<Vec<(i32, i64)>, AwaError> {
13910        let schema = self.schema();
13911        let sealed_slots: Vec<(i32, i64)> = sqlx::query_as(&format!(
13912            r#"
13913            SELECT slot, generation
13914            FROM {schema}.queue_ring_slots
13915            WHERE generation >= 0
13916              AND slot <> $1
13917            ORDER BY generation ASC, slot ASC
13918            "#
13919        ))
13920        .bind(current_slot)
13921        .fetch_all(pool)
13922        .await
13923        .map_err(map_sqlx_error)?;
13924
13925        let mut pending_slots = Vec::new();
13926        for (slot, generation) in sealed_slots {
13927            let Ok(slot_index) = usize::try_from(slot) else {
13928                continue;
13929            };
13930            let delta_child = terminal_delta_child_name(schema, slot_index);
13931            let has_pending: bool = sqlx::query_scalar(&format!(
13932                r#"
13933                SELECT EXISTS (
13934                    SELECT 1
13935                    FROM {delta_child}
13936                    WHERE ready_generation = $1
13937                    LIMIT 1
13938                )
13939                "#
13940            ))
13941            .bind(generation)
13942            .fetch_one(pool)
13943            .await
13944            .map_err(map_sqlx_error)?;
13945            if has_pending {
13946                pending_slots.push((slot, generation));
13947            }
13948        }
13949
13950        Ok(pending_slots)
13951    }
13952
13953    async fn rollup_terminal_count_delta_slot(
13954        &self,
13955        pool: &PgPool,
13956        slot: i32,
13957        generation: i64,
13958    ) -> Result<TerminalDeltaSlotRollup, AwaError> {
13959        let schema = self.schema();
13960        let delta_child = terminal_delta_child_name(schema, slot as usize);
13961        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13962
13963        let current_slot: i32 = sqlx::query_scalar(&format!(
13964            r#"
13965            SELECT current_slot
13966            FROM {schema}.queue_ring_state
13967            WHERE singleton = TRUE
13968            FOR UPDATE
13969            "#
13970        ))
13971        .fetch_one(tx.as_mut())
13972        .await
13973        .map_err(map_sqlx_error)?;
13974
13975        let slot_generation: Option<i64> = sqlx::query_scalar(&format!(
13976            r#"
13977            SELECT generation
13978            FROM {schema}.queue_ring_slots
13979            WHERE slot = $1
13980            FOR UPDATE
13981            "#
13982        ))
13983        .bind(slot)
13984        .fetch_optional(tx.as_mut())
13985        .await
13986        .map_err(map_sqlx_error)?;
13987
13988        let Some(slot_generation) = slot_generation else {
13989            tx.commit().await.map_err(map_sqlx_error)?;
13990            return Ok(TerminalDeltaSlotRollup::Empty);
13991        };
13992
13993        if current_slot == slot {
13994            tx.commit().await.map_err(map_sqlx_error)?;
13995            return Ok(TerminalDeltaSlotRollup::SkippedActive);
13996        }
13997
13998        if slot_generation != generation {
13999            tx.commit().await.map_err(map_sqlx_error)?;
14000            return Ok(TerminalDeltaSlotRollup::Empty);
14001        }
14002
14003        if Self::terminal_delta_rollup_mvcc_horizon_pinned_tx(&mut tx).await? {
14004            tx.commit().await.map_err(map_sqlx_error)?;
14005            return Ok(TerminalDeltaSlotRollup::SkippedMvccPinned);
14006        }
14007
14008        let ready_child = ready_child_name(schema, slot as usize);
14009        let pending_ready: bool = sqlx::query_scalar(&format!(
14010            r#"
14011            WITH claim_cursors AS MATERIALIZED (
14012                SELECT
14013                    queue,
14014                    priority,
14015                    enqueue_shard,
14016                    {schema}.sequence_next_value(seq_name) AS claim_seq
14017                FROM {schema}.queue_claim_heads
14018            )
14019            SELECT EXISTS (
14020                SELECT 1
14021                FROM claim_cursors AS claims
14022                CROSS JOIN LATERAL (
14023                    SELECT 1
14024                    FROM {ready_child} AS ready
14025                    WHERE ready.ready_generation = $1
14026                      AND ready.queue = claims.queue
14027                      AND ready.priority = claims.priority
14028                      AND ready.enqueue_shard = claims.enqueue_shard
14029                      AND ready.lane_seq >= claims.claim_seq
14030                    LIMIT 1
14031                ) AS pending_ready
14032                LIMIT 1
14033            )
14034            "#
14035        ))
14036        .bind(generation)
14037        .fetch_one(tx.as_mut())
14038        .await
14039        .map_err(map_sqlx_error)?;
14040
14041        if pending_ready {
14042            tx.commit().await.map_err(map_sqlx_error)?;
14043            return Ok(TerminalDeltaSlotRollup::SkippedActive);
14044        }
14045
14046        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14047
14048        let lock_delta = sqlx::query(&format!(
14049            "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE"
14050        ))
14051        .execute(tx.as_mut())
14052        .await;
14053
14054        match lock_delta {
14055            Ok(_) => {}
14056            Err(err) if is_lock_contention_error(&err) => {
14057                let _ = tx.rollback().await;
14058                return Ok(TerminalDeltaSlotRollup::Blocked);
14059            }
14060            Err(err) => {
14061                let _ = tx.rollback().await;
14062                return Err(map_sqlx_error(err));
14063            }
14064        }
14065
14066        let active_leases =
14067            queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
14068
14069        if active_leases {
14070            tx.commit().await.map_err(map_sqlx_error)?;
14071            return Ok(TerminalDeltaSlotRollup::SkippedActive);
14072        }
14073
14074        let unclosed_claim_refs =
14075            queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
14076
14077        if unclosed_claim_refs {
14078            tx.commit().await.map_err(map_sqlx_error)?;
14079            return Ok(TerminalDeltaSlotRollup::SkippedActive);
14080        }
14081
14082        let delta_rows: i64 = sqlx::query_scalar(&format!(
14083            r#"
14084            SELECT count(*)::bigint
14085            FROM {delta_child}
14086            WHERE ready_generation = $1
14087            "#
14088        ))
14089        .bind(generation)
14090        .fetch_one(tx.as_mut())
14091        .await
14092        .map_err(map_sqlx_error)?;
14093
14094        if delta_rows == 0 {
14095            tx.commit().await.map_err(map_sqlx_error)?;
14096            return Ok(TerminalDeltaSlotRollup::Empty);
14097        }
14098
14099        let grouped_keys: i64 = sqlx::query_scalar(&format!(
14100            r#"
14101            WITH grouped AS MATERIALIZED (
14102                SELECT
14103                    ready_slot,
14104                    queue,
14105                    priority,
14106                    enqueue_shard,
14107                    counter_bucket,
14108                    SUM(terminal_delta)::bigint AS delta
14109                FROM {delta_child}
14110                WHERE ready_generation = $1
14111                GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket
14112                HAVING SUM(terminal_delta) <> 0
14113            ),
14114            updated AS (
14115                UPDATE {schema}.queue_terminal_live_counts AS counts
14116                SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta)
14117                FROM grouped
14118                WHERE counts.ready_slot = grouped.ready_slot
14119                  AND counts.queue = grouped.queue
14120                  AND counts.priority = grouped.priority
14121                  AND counts.enqueue_shard = grouped.enqueue_shard
14122                  AND counts.counter_bucket = grouped.counter_bucket
14123                RETURNING
14124                    counts.ready_slot,
14125                    counts.queue,
14126                    counts.priority,
14127                    counts.enqueue_shard,
14128                    counts.counter_bucket
14129            ),
14130            inserted AS (
14131                INSERT INTO {schema}.queue_terminal_live_counts AS counts (
14132                    ready_slot,
14133                    queue,
14134                    priority,
14135                    enqueue_shard,
14136                    counter_bucket,
14137                    live_terminal_count
14138                )
14139                SELECT
14140                    grouped.ready_slot,
14141                    grouped.queue,
14142                    grouped.priority,
14143                    grouped.enqueue_shard,
14144                    grouped.counter_bucket,
14145                    grouped.delta
14146                FROM grouped
14147                WHERE grouped.delta > 0
14148                  AND NOT EXISTS (
14149                      SELECT 1
14150                      FROM updated
14151                      WHERE updated.ready_slot = grouped.ready_slot
14152                        AND updated.queue = grouped.queue
14153                        AND updated.priority = grouped.priority
14154                        AND updated.enqueue_shard = grouped.enqueue_shard
14155                        AND updated.counter_bucket = grouped.counter_bucket
14156                  )
14157                ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket
14158                ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE
14159                SET live_terminal_count =
14160                    GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count)
14161                RETURNING 1
14162            )
14163            SELECT count(*)::bigint FROM grouped
14164            "#
14165        ))
14166        .bind(generation)
14167        .fetch_one(tx.as_mut())
14168        .await
14169        .map_err(map_sqlx_error)?;
14170
14171        let truncate_delta = sqlx::query(&format!("TRUNCATE TABLE {delta_child}"))
14172            .execute(tx.as_mut())
14173            .await;
14174
14175        match truncate_delta {
14176            Ok(_) => {
14177                tx.commit().await.map_err(map_sqlx_error)?;
14178                Ok(TerminalDeltaSlotRollup::Rolled {
14179                    delta_rows,
14180                    grouped_keys,
14181                })
14182            }
14183            Err(err) if is_lock_contention_error(&err) => {
14184                let _ = tx.rollback().await;
14185                Ok(TerminalDeltaSlotRollup::Blocked)
14186            }
14187            Err(err) => {
14188                let _ = tx.rollback().await;
14189                Err(map_sqlx_error(err))
14190            }
14191        }
14192    }
14193
14194    /// Prune the oldest sealed queue ring slot.
14195    ///
14196    /// `failed_retention` is the floor for `failed` terminal rows: rows
14197    /// with `finalized_at` inside the floor are re-homed into the live
14198    /// `done_entries` segment (still retryable) instead of being
14199    /// truncated with the slot. `Duration::ZERO` disables the floor and
14200    /// prunes every terminal row. Granularity is whole seconds,
14201    /// matching DLQ retention. The floor applies to `failed` only —
14202    /// `cancelled` is an explicit operator action and is always pruned.
14203    #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest")]
14204    pub async fn prune_oldest(
14205        &self,
14206        pool: &PgPool,
14207        failed_retention: Duration,
14208    ) -> Result<PruneOutcome, AwaError> {
14209        let schema = self.schema();
14210        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14211
14212        let state: (i32,) = sqlx::query_as(&format!(
14213            r#"
14214            SELECT current_slot
14215            FROM {schema}.queue_ring_state
14216            WHERE singleton = TRUE
14217            FOR UPDATE
14218            "#
14219        ))
14220        .fetch_one(tx.as_mut())
14221        .await
14222        .map_err(map_sqlx_error)?;
14223
14224        let target: Option<(i32, i64)> = sqlx::query_as(&format!(
14225            r#"
14226            SELECT slot, generation
14227            FROM {schema}.queue_ring_slots
14228            WHERE generation >= 0
14229              AND slot <> $1
14230            ORDER BY generation ASC, slot ASC
14231            LIMIT 1
14232            FOR UPDATE
14233            "#
14234        ))
14235        .bind(state.0)
14236        .fetch_optional(tx.as_mut())
14237        .await
14238        .map_err(map_sqlx_error)?;
14239
14240        let Some((slot, generation)) = target else {
14241            tx.commit().await.map_err(map_sqlx_error)?;
14242            return Ok(PruneOutcome::Noop);
14243        };
14244
14245        let ready_child = ready_child_name(schema, slot as usize);
14246        let claim_attempt_child = ready_claim_attempt_batch_child_name(schema, slot as usize);
14247        let done_child = done_child_name(schema, slot as usize);
14248        let tomb_child = ready_tombstone_child_name(schema, slot as usize);
14249        let segment_child = ready_segment_child_name(schema, slot as usize);
14250        let receipt_batch_child = receipt_completion_batch_child_name(schema, slot as usize);
14251        let receipt_tomb_child = receipt_completion_tombstone_child_name(schema, slot as usize);
14252        let delta_child = terminal_delta_child_name(schema, slot as usize);
14253
14254        // Idle-prune churn guard: on an idle queue the rotation cursor keeps
14255        // sealing empty slots, so without this gate the prune would re-take
14256        // ACCESS EXCLUSIVE and re-issue TRUNCATE on an already-empty sealed
14257        // slot every maintenance tick — continuous relfilenode churn and
14258        // WAL-flushing commits with zero jobs in the system. That idle DDL is
14259        // near-free on local-SSD Postgres but saturates the IO path on
14260        // disaggregated-storage engines such as AlloyDB, inflating p99 latency
14261        // for co-located application queries. If every child partition this
14262        // prune would truncate is already empty there is nothing to reclaim,
14263        // so commit and report Noop instead of truncating. The slot is sealed
14264        // (slot <> current_slot) and its ring-state and slot rows are held FOR
14265        // UPDATE above, so the emptiness snapshot cannot race a writer.
14266        let slot_has_reclaimable_rows = Self::any_relation_has_rows_tx(
14267            &mut tx,
14268            &[
14269                &ready_child,
14270                &claim_attempt_child,
14271                &done_child,
14272                &tomb_child,
14273                &segment_child,
14274                &receipt_batch_child,
14275                &receipt_tomb_child,
14276                &delta_child,
14277            ],
14278        )
14279        .await?;
14280        if !slot_has_reclaimable_rows {
14281            // The destructive prune path also deletes this slot's
14282            // queue_terminal_live_counts rows alongside the TRUNCATE. Those
14283            // rows are a materialized per-slot aggregate of done_entries, so an
14284            // empty done partition normally means none exist. Counter drift
14285            // (possible during a rolling upgrade from a pre-counter binary) can
14286            // nonetheless leave orphan rows for a slot whose ring children are
14287            // already empty. Preserve that cleanup on the skip path so avoiding
14288            // the (no-op) TRUNCATE does not strand counter rows and let the
14289            // exact per-queue counters drift indefinitely. This is a cheap
14290            // 0-row DELETE in the steady state and needs no ACCESS EXCLUSIVE.
14291            sqlx::query(&format!(
14292                "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1"
14293            ))
14294            .bind(slot)
14295            .execute(tx.as_mut())
14296            .await
14297            .map_err(map_sqlx_error)?;
14298            tx.commit().await.map_err(map_sqlx_error)?;
14299            return Ok(PruneOutcome::Noop);
14300        }
14301
14302        // Beyond the emptiness gate above, queue prune has three liveness
14303        // skip gates. Check them before taking ACCESS EXCLUSIVE on the
14304        // ready/terminal child tables so a known-busy slot does not block hot
14305        // claim reads behind a doomed truncate attempt. These are only MVCC
14306        // snapshots; the same gates are checked again after the exclusive
14307        // locks.
14308        let active_leases =
14309            queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
14310
14311        if active_leases {
14312            tx.commit().await.map_err(map_sqlx_error)?;
14313            return Ok(PruneOutcome::SkippedActive {
14314                slot,
14315                reason: SkipReason::QueueActiveLeases,
14316                count: busy_indicator(active_leases),
14317            });
14318        }
14319
14320        let pending =
14321            queue_prune_has_pending_ready_tx(&mut tx, schema, &ready_child, generation).await?;
14322
14323        if pending {
14324            tx.commit().await.map_err(map_sqlx_error)?;
14325            return Ok(PruneOutcome::SkippedActive {
14326                slot,
14327                reason: SkipReason::QueuePendingReady,
14328                count: busy_indicator(pending),
14329            });
14330        }
14331
14332        let unclosed_claim_refs =
14333            queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
14334
14335        if unclosed_claim_refs {
14336            tx.commit().await.map_err(map_sqlx_error)?;
14337            return Ok(PruneOutcome::SkippedActive {
14338                slot,
14339                reason: SkipReason::QueueUnclosedClaimRefs,
14340                count: busy_indicator(unclosed_claim_refs),
14341            });
14342        }
14343
14344        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14345
14346        let lock_tables = sqlx::query(&format!(
14347            "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE"
14348        ))
14349        .execute(tx.as_mut())
14350        .await;
14351
14352        if let Err(err) = lock_tables {
14353            let _ = tx.rollback().await;
14354            if is_lock_contention_error(&err) {
14355                return Ok(PruneOutcome::Blocked { slot });
14356            }
14357            return Err(map_sqlx_error(err));
14358        }
14359
14360        // The pre-lock gates are only an MVCC snapshot. If an in-flight
14361        // producer or claimer committed while the exclusive partition
14362        // lock waited, these post-lock gates see it before TRUNCATE.
14363        let active_leases =
14364            queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
14365        if active_leases {
14366            tx.commit().await.map_err(map_sqlx_error)?;
14367            return Ok(PruneOutcome::SkippedActive {
14368                slot,
14369                reason: SkipReason::QueueActiveLeases,
14370                count: busy_indicator(active_leases),
14371            });
14372        }
14373
14374        let pending =
14375            queue_prune_has_pending_ready_tx(&mut tx, schema, &ready_child, generation).await?;
14376        if pending {
14377            tx.commit().await.map_err(map_sqlx_error)?;
14378            return Ok(PruneOutcome::SkippedActive {
14379                slot,
14380                reason: SkipReason::QueuePendingReady,
14381                count: busy_indicator(pending),
14382            });
14383        }
14384
14385        let unclosed_claim_refs =
14386            queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
14387        if unclosed_claim_refs {
14388            tx.commit().await.map_err(map_sqlx_error)?;
14389            return Ok(PruneOutcome::SkippedActive {
14390                slot,
14391                reason: SkipReason::QueueUnclosedClaimRefs,
14392                count: busy_indicator(unclosed_claim_refs),
14393            });
14394        }
14395
14396        let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX);
14397        let retention_floor = failed_retention_secs > 0;
14398
14399        // #337: failed terminal rows still inside the retention floor
14400        // are re-homed into the live done segment before the TRUNCATE
14401        // so they stay retryable. The queue_ring_state row is held FOR
14402        // UPDATE above, which serializes against rotate — the current
14403        // slot cannot move while carried rows are re-inserted into it.
14404        let carried_failed_rows = if retention_floor {
14405            let (current_generation,): (i64,) = sqlx::query_as(&format!(
14406                "SELECT generation FROM {schema}.queue_ring_slots WHERE slot = $1"
14407            ))
14408            .bind(state.0)
14409            .fetch_one(tx.as_mut())
14410            .await
14411            .map_err(map_sqlx_error)?;
14412
14413            // Carried rows are widened to the self-contained synthetic
14414            // done-row shape (the COALESCE chain mirrors
14415            // `done_row_projection`) because the ready rows they would
14416            // otherwise hydrate from are truncated with this slot.
14417            // Deliberately no ON CONFLICT: lane_seq comes from
14418            // never-resetting sequences, so a PK collision means
14419            // corrupted terminal state and must abort the prune loudly
14420            // rather than silently drop a terminal fact.
14421            let carried = sqlx::query(&format!(
14422                r#"
14423                INSERT INTO {schema}.done_entries (
14424                    ready_slot, ready_generation, job_id, kind, queue,
14425                    args, state, priority, attempt, run_lease,
14426                    max_attempts, lane_seq, enqueue_shard, run_at,
14427                    attempted_at, finalized_at, created_at, unique_key,
14428                    unique_states, payload
14429                )
14430                SELECT
14431                    $2,
14432                    $3,
14433                    done.job_id,
14434                    done.kind,
14435                    done.queue,
14436                    COALESCE(done.args, ready.args, '{{}}'::jsonb),
14437                    done.state,
14438                    done.priority,
14439                    done.attempt,
14440                    done.run_lease,
14441                    COALESCE(done.max_attempts, ready.max_attempts, 25::smallint),
14442                    done.lane_seq,
14443                    done.enqueue_shard,
14444                    COALESCE(done.run_at, ready.run_at, done.finalized_at),
14445                    COALESCE(done.attempted_at, ready.attempted_at),
14446                    done.finalized_at,
14447                    COALESCE(done.created_at, ready.created_at, done.finalized_at),
14448                    COALESCE(done.unique_key, ready.unique_key),
14449                    COALESCE(done.unique_states, ready.unique_states),
14450                    COALESCE(done.payload, ready.payload, '{{}}'::jsonb)
14451                FROM {done_child} AS done
14452                LEFT JOIN {ready_child} AS ready
14453                  ON ready.ready_slot = done.ready_slot
14454                 AND ready.ready_generation = done.ready_generation
14455                 AND ready.queue = done.queue
14456                 AND ready.priority = done.priority
14457                 AND ready.enqueue_shard = done.enqueue_shard
14458                 AND ready.lane_seq = done.lane_seq
14459                WHERE done.state = 'failed'
14460                  AND done.finalized_at >= now() - make_interval(secs => $1::bigint)
14461                "#
14462            ))
14463            .bind(failed_retention_secs)
14464            .bind(state.0)
14465            .bind(current_generation)
14466            .execute(tx.as_mut())
14467            .await
14468            .map_err(map_sqlx_error)?
14469            .rows_affected();
14470
14471            if carried > 0 {
14472                // This slot's live counter rows and delta segment are
14473                // wiped below, so the carried rows' exact-count
14474                // evidence moves with them: re-append positive deltas
14475                // under the current slot, grouped exactly like the
14476                // completion path's delta append.
14477                sqlx::query(&format!(
14478                    r#"
14479                    INSERT INTO {schema}.queue_terminal_count_deltas (
14480                        ready_slot,
14481                        ready_generation,
14482                        queue,
14483                        priority,
14484                        enqueue_shard,
14485                        counter_bucket,
14486                        terminal_delta
14487                    )
14488                    SELECT
14489                        $2,
14490                        $3,
14491                        queue,
14492                        priority,
14493                        enqueue_shard,
14494                        mod(
14495                            mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint)
14496                                + {TERMINAL_COUNTER_BUCKETS}::bigint,
14497                            {TERMINAL_COUNTER_BUCKETS}::bigint
14498                        )::smallint AS counter_bucket,
14499                        count(*)::bigint
14500                    FROM {done_child}
14501                    WHERE state = 'failed'
14502                      AND finalized_at >= now() - make_interval(secs => $1::bigint)
14503                    GROUP BY queue, priority, enqueue_shard, counter_bucket
14504                    ORDER BY queue, priority, enqueue_shard, counter_bucket
14505                    "#
14506                ))
14507                .bind(failed_retention_secs)
14508                .bind(state.0)
14509                .bind(current_generation)
14510                .execute(tx.as_mut())
14511                .await
14512                .map_err(map_sqlx_error)?;
14513            }
14514            carried
14515        } else {
14516            0
14517        };
14518
14519        // #290: scan the about-to-be-truncated partition for the rollup
14520        // fold. The rollup column is *permanent* state, so we MUST fold
14521        // from ground truth (the `{done_child}` partition itself), not
14522        // from `queue_terminal_live_counts` — the counter can drift
14523        // briefly during a rolling upgrade from a pre-counter binary
14524        // and folding drift into the rollup would bake it in forever.
14525        // The counter rows for this slot are still cleaned up after the
14526        // truncate; reads from the counter (queue_counts_exact in #305)
14527        // may transiently disagree with the rollup until the operator
14528        // runs `awa storage rebuild-terminal-counters`, but the rollup
14529        // itself stays authoritative. See PR #304 reviewer finding
14530        // "High: A1 can persist stale counter state before the
14531        // read-switch PR" for the trade-off.
14532        //
14533        // Carried failed rows are excluded from both columns: they are
14534        // still live in `done_entries`, so folding them into the
14535        // permanent rollup would double-count them.
14536        let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = sqlx::query_as(&format!(
14537            r#"
14538            WITH done_counts AS (
14539                SELECT
14540                    queue,
14541                    priority,
14542                    (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed,
14543                    (count(*) FILTER (
14544                        WHERE state = 'failed'
14545                          AND (
14546                              $2::boolean IS FALSE
14547                              OR finalized_at < now() - make_interval(secs => $1::bigint)
14548                          )
14549                    ))::bigint AS failed
14550                FROM {done_child}
14551                GROUP BY queue, priority
14552            ),
14553            -- Scan the entire partition being truncated (all generations),
14554            -- mirroring done_counts' whole-child scan above. The rotate guard
14555            -- keeps a partition to one generation at prune time, so this is the
14556            -- same set today; counting the whole child keeps the rollup fold
14557            -- conservative even if that invariant is ever weakened (a stray
14558            -- generation's batches would otherwise be truncated without being
14559            -- folded into the rollup -> terminal-count undercount).
14560            batch_rows AS (
14561                SELECT
14562                    batch.ready_generation,
14563                    batch.queue,
14564                    batch.priority,
14565                    item.job_id,
14566                    item.run_lease
14567                FROM {receipt_batch_child} AS batch
14568                CROSS JOIN LATERAL unnest(
14569                    batch.job_ids,
14570                    batch.run_leases
14571                ) AS item(job_id, run_lease)
14572            ),
14573            batch_counts AS (
14574                SELECT
14575                    batch_rows.queue,
14576                    batch_rows.priority,
14577                    count(*)::bigint AS completed
14578                FROM batch_rows
14579                WHERE NOT EXISTS (
14580                    SELECT 1
14581                    FROM {receipt_tomb_child} AS tomb
14582                    WHERE tomb.ready_generation = batch_rows.ready_generation
14583                      AND tomb.job_id = batch_rows.job_id
14584                      AND tomb.run_lease = batch_rows.run_lease
14585                )
14586                GROUP BY batch_rows.queue, batch_rows.priority
14587            ),
14588            keys AS (
14589                SELECT queue, priority FROM done_counts
14590                UNION
14591                SELECT queue, priority FROM batch_counts
14592            )
14593            SELECT
14594                keys.queue,
14595                keys.priority,
14596                (
14597                    COALESCE(done_counts.completed, 0)
14598                    + COALESCE(batch_counts.completed, 0)
14599                )::bigint AS pruned_completed_count,
14600                COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count
14601            FROM keys
14602            LEFT JOIN done_counts USING (queue, priority)
14603            LEFT JOIN batch_counts USING (queue, priority)
14604            "#
14605        ))
14606        .bind(failed_retention_secs)
14607        .bind(retention_floor)
14608        .fetch_all(tx.as_mut())
14609        .await
14610        .map_err(map_sqlx_error)?;
14611
14612        let truncate = sqlx::query(&format!(
14613            "TRUNCATE TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child}"
14614        ))
14615        .execute(tx.as_mut())
14616        .await;
14617
14618        match truncate {
14619            Ok(_) => {
14620                if !pruned_terminal_counts.is_empty() {
14621                    self.adjust_terminal_rollups_batch(&mut tx, pruned_terminal_counts.into_iter())
14622                        .await?;
14623                }
14624
14625                // #290: the live counter rows for this slot are about to
14626                // be orphans (their underlying done_entries rows just got
14627                // truncated). Delete them in the same transaction. The
14628                // rollup fold above already captured ground-truth from
14629                // the partition scan; this just cleans up the counter
14630                // index entries so a future insert into a re-rotated
14631                // slot starts from zero.
14632                sqlx::query(&format!(
14633                    "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1"
14634                ))
14635                .bind(slot)
14636                .execute(tx.as_mut())
14637                .await
14638                .map_err(map_sqlx_error)?;
14639                tx.commit().await.map_err(map_sqlx_error)?;
14640                if carried_failed_rows > 0 {
14641                    tracing::info!(
14642                        slot,
14643                        carried_failed_rows,
14644                        "Carried failed terminal rows inside the retention floor to the live done segment"
14645                    );
14646                }
14647                Ok(PruneOutcome::Pruned {
14648                    slot,
14649                    carried_failed_rows,
14650                })
14651            }
14652            Err(err) if is_lock_contention_error(&err) => {
14653                let _ = tx.rollback().await;
14654                Ok(PruneOutcome::Blocked { slot })
14655            }
14656            Err(err) => {
14657                let _ = tx.rollback().await;
14658                Err(map_sqlx_error(err))
14659            }
14660        }
14661    }
14662
14663    #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest_leases")]
14664    pub async fn prune_oldest_leases(&self, pool: &PgPool) -> Result<PruneOutcome, AwaError> {
14665        let schema = self.schema();
14666        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14667
14668        // `PruneLeasesPlan` in
14669        // `correctness/storage/AwaStorageLockOrder.tla` requires the
14670        // sequence `lease_ring_state FOR UPDATE` →
14671        // `lease_ring_slots[slot] FOR UPDATE` → `ACCESS EXCLUSIVE` on
14672        // the child. The child lock is bounded by a short
14673        // transaction-local lock_timeout because pure NOWAIT can starve
14674        // under continuous parent-partition readers, while an unbounded
14675        // wait would put maintenance at the head of the relation-lock
14676        // queue indefinitely. Without these locks a concurrent
14677        // rotator can flip the cursor under the prune's liveness check
14678        // (current_slot recheck races a CAS update) and prune what
14679        // should be the active partition.
14680        let state: (i32, i64, i32) = sqlx::query_as(&format!(
14681            r#"
14682            SELECT current_slot, generation, slot_count
14683            FROM {schema}.lease_ring_state
14684            WHERE singleton = TRUE
14685            FOR UPDATE
14686            "#
14687        ))
14688        .fetch_one(tx.as_mut())
14689        .await
14690        .map_err(map_sqlx_error)?;
14691
14692        let Some((slot, _generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2)
14693        else {
14694            tx.commit().await.map_err(map_sqlx_error)?;
14695            return Ok(PruneOutcome::Noop);
14696        };
14697
14698        let slot_locked: Option<i32> = sqlx::query_scalar(&format!(
14699            r#"
14700            SELECT slot FROM {schema}.lease_ring_slots
14701            WHERE slot = $1
14702            FOR UPDATE
14703            "#
14704        ))
14705        .bind(slot)
14706        .fetch_optional(tx.as_mut())
14707        .await
14708        .map_err(map_sqlx_error)?;
14709
14710        if slot_locked.is_none() {
14711            tx.commit().await.map_err(map_sqlx_error)?;
14712            return Ok(PruneOutcome::Noop);
14713        }
14714
14715        let lease_child = lease_child_name(schema, slot as usize);
14716
14717        // Idle-prune churn guard (see `prune_oldest`): skip the destructive
14718        // ACCESS EXCLUSIVE + TRUNCATE when the sealed lease slot is already
14719        // empty, turning a per-tick truncate of an empty slot into a Noop and
14720        // eliminating the relfilenode/WAL churn that disaggregated-storage
14721        // engines such as AlloyDB amplify into application p99 latency. The
14722        // slot is sealed and its ring-state and slot rows are held FOR UPDATE
14723        // above, so the emptiness snapshot cannot race a writer.
14724        if !Self::relation_has_rows_tx(&mut tx, &lease_child).await? {
14725            tx.commit().await.map_err(map_sqlx_error)?;
14726            return Ok(PruneOutcome::Noop);
14727        }
14728
14729        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14730
14731        let lock_table = sqlx::query(&format!(
14732            "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE"
14733        ))
14734        .execute(tx.as_mut())
14735        .await;
14736
14737        if let Err(err) = lock_table {
14738            let _ = tx.rollback().await;
14739            if is_lock_contention_error(&err) {
14740                return Ok(PruneOutcome::Blocked { slot });
14741            }
14742            return Err(map_sqlx_error(err));
14743        }
14744
14745        let current_slot: i32 = sqlx::query_scalar(&format!(
14746            r#"
14747            SELECT current_slot
14748            FROM {schema}.lease_ring_state
14749            WHERE singleton = TRUE
14750            "#
14751        ))
14752        .fetch_one(tx.as_mut())
14753        .await
14754        .map_err(map_sqlx_error)?;
14755
14756        if current_slot == slot {
14757            tx.commit().await.map_err(map_sqlx_error)?;
14758            return Ok(PruneOutcome::SkippedActive {
14759                slot,
14760                reason: SkipReason::LeaseCurrent,
14761                count: 0,
14762            });
14763        }
14764
14765        let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?;
14766
14767        if active_leases {
14768            tx.commit().await.map_err(map_sqlx_error)?;
14769            return Ok(PruneOutcome::SkippedActive {
14770                slot,
14771                reason: SkipReason::LeaseActive,
14772                count: busy_indicator(active_leases),
14773            });
14774        }
14775
14776        let truncate = sqlx::query(&format!("TRUNCATE TABLE {lease_child}"))
14777            .execute(tx.as_mut())
14778            .await;
14779
14780        match truncate {
14781            Ok(_) => {
14782                tx.commit().await.map_err(map_sqlx_error)?;
14783                Ok(PruneOutcome::Pruned {
14784                    slot,
14785                    carried_failed_rows: 0,
14786                })
14787            }
14788            Err(err) if is_lock_contention_error(&err) => {
14789                let _ = tx.rollback().await;
14790                Ok(PruneOutcome::Blocked { slot })
14791            }
14792            Err(err) => {
14793                let _ = tx.rollback().await;
14794                Err(map_sqlx_error(err))
14795            }
14796        }
14797    }
14798
14799    pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> {
14800        sqlx::query(&format!("VACUUM {}", self.leases_table()))
14801            .execute(pool)
14802            .await
14803            .map_err(map_sqlx_error)?;
14804        Ok(())
14805    }
14806
14807    /// ADR-023 claim-ring rotation. Parallel of `rotate_leases`.
14808    ///
14809    /// Advances `claim_ring_state.current_slot` via compare-and-swap. Before
14810    /// flipping the cursor the target partition must be drained: the
14811    /// `lease_claims_<next>`, `lease_claim_batches_<next>`,
14812    /// `lease_claim_closures_<next>`, and
14813    /// `lease_claim_closure_batches_<next>` child tables
14814    /// must be empty. This is what the `rotate → prune → rotate` ring
14815    /// invariant requires — we only hand out a slot to new claims when a
14816    /// prior prune has truncated it.
14817    #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_claims")]
14818    pub async fn rotate_claims(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
14819        let schema = self.schema();
14820        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14821
14822        let state: (i32, i64, i32) = sqlx::query_as(&format!(
14823            r#"
14824            SELECT current_slot, generation, slot_count
14825            FROM {schema}.claim_ring_state
14826            WHERE singleton = TRUE
14827            FOR UPDATE
14828            "#
14829        ))
14830        .fetch_one(tx.as_mut())
14831        .await
14832        .map_err(map_sqlx_error)?;
14833
14834        let next_slot = (state.0 + 1).rem_euclid(state.2);
14835
14836        // Busy check: all children of the incoming slot must be empty.
14837        // A non-empty `lease_claims_<next>` means the previous lap's
14838        // prune hasn't run (or didn't complete); rotating anyway would
14839        // mix fresh claims with legacy rows and defeat the point of
14840        // partitioning. Non-empty closure children mean prune fell behind
14841        // on receipt-closure evidence specifically.
14842        let claim_busy =
14843            Self::relation_has_rows_tx(&mut tx, &claim_child_name(schema, next_slot as usize))
14844                .await?;
14845        let claim_batch_busy = Self::relation_has_rows_tx(
14846            &mut tx,
14847            &claim_batch_child_name(schema, next_slot as usize),
14848        )
14849        .await?;
14850        let closure_busy =
14851            Self::relation_has_rows_tx(&mut tx, &closure_child_name(schema, next_slot as usize))
14852                .await?;
14853        let closure_batch_busy = Self::relation_has_rows_tx(
14854            &mut tx,
14855            &claim_closure_batch_child_name(schema, next_slot as usize),
14856        )
14857        .await?;
14858
14859        if claim_busy || claim_batch_busy || closure_busy || closure_batch_busy {
14860            tx.commit().await.map_err(map_sqlx_error)?;
14861            return Ok(RotateOutcome::SkippedBusy {
14862                slot: next_slot,
14863                busy: BusyCounts {
14864                    claims: busy_indicator(claim_busy || claim_batch_busy),
14865                    closures: busy_indicator(closure_busy),
14866                    closure_batches: busy_indicator(closure_batch_busy),
14867                    ..Default::default()
14868                },
14869            });
14870        }
14871
14872        let next_generation = state.1 + 1;
14873
14874        let rotated = sqlx::query(&format!(
14875            r#"
14876            UPDATE {schema}.claim_ring_state
14877            SET current_slot = $1,
14878                generation = $2
14879            WHERE singleton = TRUE
14880              AND current_slot = $3
14881              AND generation = $4
14882            "#
14883        ))
14884        .bind(next_slot)
14885        .bind(next_generation)
14886        .bind(state.0)
14887        .bind(state.1)
14888        .execute(tx.as_mut())
14889        .await
14890        .map_err(map_sqlx_error)?;
14891
14892        if rotated.rows_affected() == 0 {
14893            // Lost the CAS race. Report the bounded presence evidence we
14894            // sampled before giving up.
14895            tx.commit().await.map_err(map_sqlx_error)?;
14896            return Ok(RotateOutcome::SkippedBusy {
14897                slot: next_slot,
14898                busy: BusyCounts {
14899                    claims: busy_indicator(claim_busy || claim_batch_busy),
14900                    closures: busy_indicator(closure_busy),
14901                    closure_batches: busy_indicator(closure_batch_busy),
14902                    ..Default::default()
14903                },
14904            });
14905        }
14906
14907        tx.commit().await.map_err(map_sqlx_error)?;
14908        Ok(RotateOutcome::Rotated {
14909            slot: next_slot,
14910            generation: next_generation,
14911        })
14912    }
14913
14914    /// ADR-023 claim-ring prune. Parallel of `prune_oldest_leases`.
14915    ///
14916    /// Reclaims the oldest initialized (sealed) claim-ring slot by
14917    /// `TRUNCATE`-ing its `lease_claims_<slot>`,
14918    /// `lease_claim_batches_<slot>`, `lease_claim_closures_<slot>`, and
14919    /// `lease_claim_closure_batches_<slot>` children. Takes the full ADR-023
14920    /// lock sequence:
14921    ///
14922    /// 1. `FOR UPDATE` on `claim_ring_state` (serialises with rotate).
14923    /// 2. `FOR UPDATE` on the target `claim_ring_slots` row.
14924    /// 3. Proves every claim in the sealed partition has closure evidence.
14925    /// 4. `LOCK TABLE ACCESS EXCLUSIVE` on all claim-ring children with a
14926    ///    short transaction-local `lock_timeout` (serialises with in-flight
14927    ///    claim/complete/rescue writers; gives up under sustained
14928    ///    contention).
14929    /// 5. Rechecks the slot is not the current one and that every
14930    ///    claim still has closure evidence.
14931    /// 6. `TRUNCATE` all claim-ring children in a single statement.
14932    ///
14933    /// The "every claim has closure evidence" precondition is what ADR-023
14934    /// calls `PartitionTruncateSafety`. If an open claim remains in the
14935    /// partition, prune returns `SkippedActive` and the claim has to
14936    /// drain by normal completion or be rescued by
14937    /// `rescue_stale_receipt_claims_tx` before prune will try again.
14938    #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest_claims")]
14939    pub async fn prune_oldest_claims(&self, pool: &PgPool) -> Result<PruneOutcome, AwaError> {
14940        let schema = self.schema();
14941        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14942
14943        let state: (i32, i64, i32) = sqlx::query_as(&format!(
14944            r#"
14945            SELECT current_slot, generation, slot_count
14946            FROM {schema}.claim_ring_state
14947            WHERE singleton = TRUE
14948            FOR UPDATE
14949            "#
14950        ))
14951        .fetch_one(tx.as_mut())
14952        .await
14953        .map_err(map_sqlx_error)?;
14954
14955        let Some((slot, _generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2)
14956        else {
14957            tx.commit().await.map_err(map_sqlx_error)?;
14958            return Ok(PruneOutcome::Noop);
14959        };
14960
14961        // Lock the slot row so concurrent rotate/prune observe the same
14962        // state machine transition.
14963        let slot_locked: Option<i32> = sqlx::query_scalar(&format!(
14964            r#"
14965            SELECT slot FROM {schema}.claim_ring_slots
14966            WHERE slot = $1
14967            FOR UPDATE
14968            "#
14969        ))
14970        .bind(slot)
14971        .fetch_optional(tx.as_mut())
14972        .await
14973        .map_err(map_sqlx_error)?;
14974
14975        if slot_locked.is_none() {
14976            tx.commit().await.map_err(map_sqlx_error)?;
14977            return Ok(PruneOutcome::Noop);
14978        }
14979
14980        let claim_child = claim_child_name(schema, slot as usize);
14981        let claim_batch_child = claim_batch_child_name(schema, slot as usize);
14982        let closure_child = closure_child_name(schema, slot as usize);
14983        let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize);
14984
14985        // Idle-prune churn guard (see `prune_oldest`): if every claim-ring
14986        // child partition for this sealed slot is already empty there is
14987        // nothing to reclaim, so commit a Noop instead of re-truncating an
14988        // empty slot every maintenance tick, eliminating the relfilenode/WAL
14989        // churn that disaggregated-storage engines such as AlloyDB amplify
14990        // into application p99 latency. The slot is sealed and its ring-state
14991        // and slot rows are held FOR UPDATE above, so the emptiness snapshot
14992        // cannot race a writer.
14993        let slot_has_reclaimable_rows = Self::any_relation_has_rows_tx(
14994            &mut tx,
14995            &[
14996                &claim_child,
14997                &claim_batch_child,
14998                &closure_child,
14999                &closure_batch_child,
15000            ],
15001        )
15002        .await?;
15003        if !slot_has_reclaimable_rows {
15004            tx.commit().await.map_err(map_sqlx_error)?;
15005            return Ok(PruneOutcome::Noop);
15006        }
15007
15008        // Before taking ACCESS EXCLUSIVE on the child partitions, prove
15009        // the sealed slot is actually reclaimable. This optimistic proof
15010        // avoids blocking claim/complete traffic behind a prune attempt
15011        // that is already known to be doomed. A claimer that read the
15012        // old current slot before rotation may still commit while the
15013        // exclusive lock waits, so the open-claim proof is repeated once
15014        // the child locks are held.
15015        let open_claims = claim_prune_has_open_claims_tx(
15016            &mut tx,
15017            schema,
15018            &claim_child,
15019            &claim_batch_child,
15020            &closure_child,
15021            &closure_batch_child,
15022        )
15023        .await?;
15024
15025        if open_claims {
15026            tx.commit().await.map_err(map_sqlx_error)?;
15027            return Ok(PruneOutcome::SkippedActive {
15028                slot,
15029                reason: SkipReason::ClaimOpen,
15030                count: busy_indicator(open_claims),
15031            });
15032        }
15033
15034        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
15035
15036        let lock_tables = sqlx::query(&format!(
15037            "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE"
15038        ))
15039        .execute(tx.as_mut())
15040        .await;
15041
15042        if let Err(err) = lock_tables {
15043            let _ = tx.rollback().await;
15044            if is_lock_contention_error(&err) {
15045                return Ok(PruneOutcome::Blocked { slot });
15046            }
15047            return Err(map_sqlx_error(err));
15048        }
15049
15050        // After taking ACCESS EXCLUSIVE, recheck that the slot is still
15051        // not the current one. The ring-state lock should already make
15052        // this stable, but keeping the explicit gate documents the
15053        // truncate precondition.
15054        let current_slot: i32 = sqlx::query_scalar(&format!(
15055            r#"
15056            SELECT current_slot FROM {schema}.claim_ring_state WHERE singleton = TRUE
15057            "#
15058        ))
15059        .fetch_one(tx.as_mut())
15060        .await
15061        .map_err(map_sqlx_error)?;
15062
15063        if current_slot == slot {
15064            tx.commit().await.map_err(map_sqlx_error)?;
15065            return Ok(PruneOutcome::SkippedActive {
15066                slot,
15067                reason: SkipReason::ClaimCurrent,
15068                count: 0,
15069            });
15070        }
15071
15072        let open_claims = claim_prune_has_open_claims_tx(
15073            &mut tx,
15074            schema,
15075            &claim_child,
15076            &claim_batch_child,
15077            &closure_child,
15078            &closure_batch_child,
15079        )
15080        .await?;
15081        if open_claims {
15082            tx.commit().await.map_err(map_sqlx_error)?;
15083            return Ok(PruneOutcome::SkippedActive {
15084                slot,
15085                reason: SkipReason::ClaimOpen,
15086                count: busy_indicator(open_claims),
15087            });
15088        }
15089
15090        let truncate = sqlx::query(&format!(
15091            "TRUNCATE TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child}"
15092        ))
15093        .execute(tx.as_mut())
15094        .await;
15095
15096        match truncate {
15097            Ok(_) => {
15098                sqlx::query(&format!(
15099                    r#"
15100                    UPDATE {schema}.claim_ring_slots
15101                    SET rescue_cursor_claimed_at = '-infinity'::timestamptz,
15102                        rescue_cursor_job_id = 0,
15103                        rescue_cursor_run_lease = 0,
15104                        deadline_cursor_deadline_at = '-infinity'::timestamptz,
15105                        deadline_cursor_job_id = 0,
15106                        deadline_cursor_run_lease = 0
15107                    WHERE slot = $1
15108                    "#
15109                ))
15110                .bind(slot)
15111                .execute(tx.as_mut())
15112                .await
15113                .map_err(map_sqlx_error)?;
15114                tx.commit().await.map_err(map_sqlx_error)?;
15115                Ok(PruneOutcome::Pruned {
15116                    slot,
15117                    carried_failed_rows: 0,
15118                })
15119            }
15120            Err(err) if is_lock_contention_error(&err) => {
15121                let _ = tx.rollback().await;
15122                Ok(PruneOutcome::Blocked { slot })
15123            }
15124            Err(err) => {
15125                let _ = tx.rollback().await;
15126                Err(map_sqlx_error(err))
15127            }
15128        }
15129    }
15130
15131    fn job_id_sequence(&self) -> String {
15132        format!("{}.job_id_seq", self.schema())
15133    }
15134
15135    fn leases_table(&self) -> String {
15136        format!("{}.{}", self.schema(), self.leases_relname())
15137    }
15138
15139    fn attempt_state_table(&self) -> String {
15140        format!("{}.{}", self.schema(), self.attempt_state_relname())
15141    }
15142}