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    /// Rescue-path fallback for unique-claim conflicts: insert the rescued
12221    /// attempt's deferred successor, degrading to a cancelled terminal row
12222    /// when the job's unique claim is held by a newer duplicate — the claim
12223    /// holder wins, and one poisoned row must not abort the whole batched
12224    /// rescue transaction. The insert attempt runs inside a savepoint so the
12225    /// conflict leaves the outer transaction usable.
12226    async fn insert_rescued_deferred_or_cancel_tx(
12227        &self,
12228        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12229        row: LeaseTransitionRow,
12230        deferred: DeferredJobRow,
12231        duplicate_error: &str,
12232    ) -> Result<JobRow, AwaError> {
12233        let old_state = row.state;
12234        {
12235            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12236            match self
12237                .insert_deferred_rows_tx(&mut savepoint, vec![deferred.clone()], Some(old_state))
12238                .await
12239            {
12240                Ok(_) => {
12241                    savepoint.commit().await.map_err(map_sqlx_error)?;
12242                    return deferred.into_job_row();
12243                }
12244                Err(AwaError::UniqueConflict { .. }) => {
12245                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12246                }
12247                Err(err) => return Err(err),
12248            }
12249        }
12250        self.cancel_rescued_duplicate_tx(tx, row, deferred.payload, duplicate_error)
12251            .await
12252    }
12253
12254    /// Same fallback for rescue arms that write a terminal row directly
12255    /// (callback timeout at max attempts): a `failed` insert that conflicts
12256    /// with a newer duplicate's claim becomes a cancellation instead.
12257    async fn insert_rescued_done_or_cancel_tx(
12258        &self,
12259        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12260        row: LeaseTransitionRow,
12261        done: DoneJobRow,
12262        duplicate_error: &str,
12263    ) -> Result<JobRow, AwaError> {
12264        let old_state = row.state;
12265        {
12266            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12267            match self
12268                .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state))
12269                .await
12270            {
12271                Ok(_) => {
12272                    savepoint.commit().await.map_err(map_sqlx_error)?;
12273                    return done.into_job_row();
12274                }
12275                Err(AwaError::UniqueConflict { .. }) => {
12276                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12277                }
12278                Err(err) => return Err(err),
12279            }
12280        }
12281        self.cancel_rescued_duplicate_tx(tx, row, done.payload, duplicate_error)
12282            .await
12283    }
12284
12285    /// Write the cancelled terminal row for a rescue that lost its unique
12286    /// claim to a newer duplicate. `cancelled` normally sits outside the
12287    /// claiming mask so the insert is conflict-free; a mask that claims
12288    /// `cancelled` as well conflicts again, and then the row is written with
12289    /// its unique key stripped — the attempt held no claim (that is what
12290    /// made it conflict), so there is nothing to release, and losing the key
12291    /// on this terminal row beats aborting rescue for the whole cluster.
12292    async fn cancel_rescued_duplicate_tx(
12293        &self,
12294        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12295        row: LeaseTransitionRow,
12296        base_payload: serde_json::Value,
12297        duplicate_error: &str,
12298    ) -> Result<JobRow, AwaError> {
12299        let old_state = row.state;
12300        let mut payload = RuntimePayload::from_json(base_payload)?;
12301        payload.push_error(lifecycle_error(duplicate_error, row.attempt, true));
12302        let done = row.into_done_row(JobState::Cancelled, Utc::now(), payload.into_json());
12303
12304        {
12305            let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?;
12306            match self
12307                .insert_done_rows_tx(&mut savepoint, std::slice::from_ref(&done), Some(old_state))
12308                .await
12309            {
12310                Ok(_) => {
12311                    savepoint.commit().await.map_err(map_sqlx_error)?;
12312                    return done.into_job_row();
12313                }
12314                Err(AwaError::UniqueConflict { .. }) => {
12315                    savepoint.rollback().await.map_err(map_sqlx_error)?;
12316                }
12317                Err(err) => return Err(err),
12318            }
12319        }
12320
12321        let mut stripped = done;
12322        stripped.unique_key = None;
12323        stripped.unique_states = None;
12324        self.insert_done_rows_tx(tx, std::slice::from_ref(&stripped), Some(old_state))
12325            .await?;
12326        stripped.into_job_row()
12327    }
12328
12329    pub async fn cancel_running(
12330        &self,
12331        pool: &PgPool,
12332        job_id: i64,
12333        run_lease: i64,
12334        reason: &str,
12335        progress: Option<serde_json::Value>,
12336    ) -> Result<Option<JobRow>, AwaError> {
12337        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12338        let result = self
12339            .cancel_running_in_tx(&mut tx, job_id, run_lease, reason, progress)
12340            .await?;
12341        tx.commit().await.map_err(map_sqlx_error)?;
12342        Ok(result)
12343    }
12344
12345    /// Transaction-aware variant of [`Self::cancel_running`] (ADR-029).
12346    pub async fn cancel_running_in_tx(
12347        &self,
12348        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12349        job_id: i64,
12350        run_lease: i64,
12351        reason: &str,
12352        progress: Option<serde_json::Value>,
12353    ) -> Result<Option<JobRow>, AwaError> {
12354        let Some(moved) = self
12355            .take_running_attempt_tx(tx, job_id, run_lease, "cancelled")
12356            .await?
12357        else {
12358            return Ok(None);
12359        };
12360
12361        let mut payload = RuntimePayload::from_json(Self::with_progress(
12362            moved.payload.clone(),
12363            progress.or(moved.progress.clone()),
12364        )?)?;
12365        payload.push_error(lifecycle_error(
12366            format!("cancelled: {reason}"),
12367            moved.attempt,
12368            false,
12369        ));
12370        let done =
12371            moved
12372                .clone()
12373                .into_done_row(JobState::Cancelled, Utc::now(), payload.into_json());
12374        self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12375            .await?;
12376        Ok(Some(done.into_job_row()?))
12377    }
12378
12379    pub async fn fail_terminal(
12380        &self,
12381        pool: &PgPool,
12382        job_id: i64,
12383        run_lease: i64,
12384        error: &str,
12385        progress: Option<serde_json::Value>,
12386    ) -> Result<Option<JobRow>, AwaError> {
12387        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12388        let result = self
12389            .fail_terminal_in_tx(&mut tx, job_id, run_lease, error, progress)
12390            .await?;
12391        tx.commit().await.map_err(map_sqlx_error)?;
12392        Ok(result)
12393    }
12394
12395    /// Transaction-aware variant of [`Self::fail_terminal`] (ADR-029).
12396    pub async fn fail_terminal_in_tx(
12397        &self,
12398        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12399        job_id: i64,
12400        run_lease: i64,
12401        error: &str,
12402        progress: Option<serde_json::Value>,
12403    ) -> Result<Option<JobRow>, AwaError> {
12404        let Some(moved) = self
12405            .take_running_attempt_tx(tx, job_id, run_lease, "failed")
12406            .await?
12407        else {
12408            return Ok(None);
12409        };
12410
12411        let mut payload = RuntimePayload::from_json(Self::with_progress(
12412            moved.payload.clone(),
12413            progress.or(moved.progress.clone()),
12414        )?)?;
12415        payload.push_error(lifecycle_error(error, moved.attempt, true));
12416        let done = moved
12417            .clone()
12418            .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
12419        self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12420            .await?;
12421        Ok(Some(done.into_job_row()?))
12422    }
12423
12424    pub async fn fail_to_dlq(
12425        &self,
12426        pool: &PgPool,
12427        job_id: i64,
12428        run_lease: i64,
12429        dlq_reason: &str,
12430        error: &str,
12431        progress: Option<serde_json::Value>,
12432    ) -> Result<Option<JobRow>, AwaError> {
12433        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12434        let result = self
12435            .fail_to_dlq_in_tx(&mut tx, job_id, run_lease, dlq_reason, error, progress)
12436            .await?;
12437        tx.commit().await.map_err(map_sqlx_error)?;
12438        Ok(result)
12439    }
12440
12441    /// Transaction-aware variant of [`Self::fail_to_dlq`] (ADR-029).
12442    pub async fn fail_to_dlq_in_tx(
12443        &self,
12444        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12445        job_id: i64,
12446        run_lease: i64,
12447        dlq_reason: &str,
12448        error: &str,
12449        progress: Option<serde_json::Value>,
12450    ) -> Result<Option<JobRow>, AwaError> {
12451        let Some(moved) = self
12452            .take_running_attempt_tx(tx, job_id, run_lease, "dlq")
12453            .await?
12454        else {
12455            return Ok(None);
12456        };
12457
12458        let finalized_at = Utc::now();
12459        let dlq_at = finalized_at;
12460        let mut payload = RuntimePayload::from_json(Self::with_progress(
12461            moved.payload.clone(),
12462            progress.or(moved.progress.clone()),
12463        )?)?;
12464        payload.push_error(lifecycle_error(error, moved.attempt, true));
12465        let dlq_row = moved.clone().into_dlq_row(
12466            finalized_at,
12467            payload.into_json(),
12468            dlq_reason.to_string(),
12469            dlq_at,
12470        );
12471        self.insert_dlq_rows_tx(tx, std::slice::from_ref(&dlq_row), Some(moved.state))
12472            .await?;
12473        Ok(Some(dlq_row.into_job_row()?))
12474    }
12475
12476    pub async fn move_failed_to_dlq(
12477        &self,
12478        pool: &PgPool,
12479        job_id: i64,
12480        dlq_reason: &str,
12481    ) -> Result<Option<JobRow>, AwaError> {
12482        let schema = self.schema();
12483        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12484        let done_projection = done_row_projection("done", "ready");
12485        let ready_join = done_ready_join(schema, "done", "ready");
12486        let moved: Option<DoneJobRow> = sqlx::query_as(&format!(
12487            r#"
12488            WITH deleted AS (
12489                DELETE FROM {schema}.done_entries
12490                WHERE (job_id, finalized_at) IN (
12491                    SELECT job_id, finalized_at
12492                    FROM {schema}.done_entries
12493                    WHERE job_id = $1
12494                      AND state = 'failed'
12495                    ORDER BY finalized_at DESC
12496                    LIMIT 1
12497                    FOR UPDATE SKIP LOCKED
12498                )
12499                RETURNING *
12500            )
12501            SELECT {done_projection}
12502            FROM deleted AS done
12503            {ready_join}
12504            "#
12505        ))
12506        .bind(job_id)
12507        .fetch_optional(tx.as_mut())
12508        .await
12509        .map_err(map_sqlx_error)?;
12510
12511        let Some(moved) = moved else {
12512            tx.commit().await.map_err(map_sqlx_error)?;
12513            return Ok(None);
12514        };
12515
12516        self.ensure_terminal_removed_receipt_closures_tx(&mut tx, std::slice::from_ref(&moved))
12517            .await?;
12518        // DLQ inserts are outside done_entries, so the source terminal DELETE
12519        // appends a negative delta.
12520        self.decrement_live_terminal_counters_tx(
12521            &mut tx,
12522            &Self::done_rows_to_counter_keys(std::slice::from_ref(&moved)),
12523        )
12524        .await?;
12525        let dlq_row = moved
12526            .clone()
12527            .into_dlq_row(dlq_reason.to_string(), Utc::now());
12528        self.insert_dlq_rows_tx(&mut tx, std::slice::from_ref(&dlq_row), Some(moved.state))
12529            .await?;
12530        tx.commit().await.map_err(map_sqlx_error)?;
12531        Ok(Some(dlq_row.into_job_row()?))
12532    }
12533
12534    #[tracing::instrument(
12535        skip(self, pool, dlq_reason),
12536        fields(kind = ?kind, queue = ?queue),
12537        name = "queue_storage.bulk_move_failed_to_dlq"
12538    )]
12539    pub async fn bulk_move_failed_to_dlq(
12540        &self,
12541        pool: &PgPool,
12542        kind: Option<&str>,
12543        queue: Option<&str>,
12544        dlq_reason: &str,
12545    ) -> Result<u64, AwaError> {
12546        let schema = self.schema();
12547        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12548        let done_projection = done_row_projection("done", "ready");
12549        let ready_join = done_ready_join(schema, "done", "ready");
12550        let moved: Vec<DoneJobRow> = sqlx::query_as(&format!(
12551            r#"
12552            WITH deleted AS (
12553                DELETE FROM {schema}.done_entries
12554                WHERE state = 'failed'
12555                  AND ($1::text IS NULL OR kind = $1)
12556                  AND ($2::text IS NULL OR queue = $2)
12557                RETURNING *
12558            )
12559            SELECT {done_projection}
12560            FROM deleted AS done
12561            {ready_join}
12562            "#
12563        ))
12564        .bind(kind)
12565        .bind(queue)
12566        .fetch_all(tx.as_mut())
12567        .await
12568        .map_err(map_sqlx_error)?;
12569
12570        if moved.is_empty() {
12571            tx.commit().await.map_err(map_sqlx_error)?;
12572            return Ok(0);
12573        }
12574
12575        self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &moved)
12576            .await?;
12577        // Bulk DLQ move deletes from done_entries. Append negative deltas by
12578        // the per-group magnitudes of the moved rows.
12579        self.decrement_live_terminal_counters_tx(&mut tx, &Self::done_rows_to_counter_keys(&moved))
12580            .await?;
12581        let dlq_at = Utc::now();
12582        let rows: Vec<DlqJobRow> = moved
12583            .into_iter()
12584            .map(|row| row.into_dlq_row(dlq_reason.to_string(), dlq_at))
12585            .collect();
12586        self.insert_dlq_rows_tx(&mut tx, &rows, Some(JobState::Failed))
12587            .await?;
12588        tx.commit().await.map_err(map_sqlx_error)?;
12589        Ok(rows.len() as u64)
12590    }
12591
12592    pub async fn retry_from_dlq(
12593        &self,
12594        pool: &PgPool,
12595        job_id: i64,
12596        opts: &RetryFromDlqOpts,
12597    ) -> Result<Option<JobRow>, AwaError> {
12598        let schema = self.schema();
12599        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12600        let moved: Option<DlqJobRow> = sqlx::query_as(&format!(
12601            r#"
12602            DELETE FROM {schema}.dlq_entries
12603            WHERE job_id = $1
12604            RETURNING
12605                job_id,
12606                kind,
12607                queue,
12608                args,
12609                state,
12610                priority,
12611                attempt,
12612                run_lease,
12613                max_attempts,
12614                run_at,
12615                attempted_at,
12616                finalized_at,
12617                created_at,
12618                unique_key,
12619                unique_states,
12620                COALESCE(payload, '{{}}'::jsonb) AS payload,
12621                dlq_reason,
12622                dlq_at,
12623                original_run_lease
12624            "#
12625        ))
12626        .bind(job_id)
12627        .fetch_optional(tx.as_mut())
12628        .await
12629        .map_err(map_sqlx_error)?;
12630
12631        let Some(moved) = moved else {
12632            tx.commit().await.map_err(map_sqlx_error)?;
12633            return Ok(None);
12634        };
12635
12636        let queue = opts.queue.clone().unwrap_or_else(|| moved.queue.clone());
12637        let priority = opts.priority.unwrap_or(moved.priority);
12638        let mut payload = RuntimePayload::from_json(moved.payload.clone())?;
12639        payload.set_progress(None);
12640        let payload = payload.into_json();
12641
12642        if let Some(run_at) = opts.run_at.filter(|run_at| *run_at > Utc::now()) {
12643            let deferred = moved.into_retry_deferred_row(queue, priority, run_at, payload);
12644            self.insert_deferred_rows_tx(&mut tx, vec![deferred.clone()], Some(JobState::Failed))
12645                .await?;
12646            tx.commit().await.map_err(map_sqlx_error)?;
12647            return Ok(Some(deferred.into_job_row()?));
12648        }
12649
12650        let ready = moved.into_retry_ready_row(queue.clone(), priority, Utc::now(), payload);
12651        self.insert_existing_ready_rows_tx(&mut tx, vec![ready.clone()], Some(JobState::Failed))
12652            .await?;
12653        self.notify_queues_tx(&mut tx, std::iter::once(queue))
12654            .await?;
12655        tx.commit().await.map_err(map_sqlx_error)?;
12656        Ok(Some(
12657            ReadyJobRow {
12658                job_id: ready.job_id,
12659                kind: ready.kind,
12660                queue: ready.queue,
12661                args: ready.args,
12662                priority: ready.priority,
12663                attempt: ready.attempt,
12664                run_lease: ready.run_lease,
12665                max_attempts: ready.max_attempts,
12666                run_at: ready.run_at,
12667                attempted_at: ready.attempted_at,
12668                created_at: ready.created_at,
12669                unique_key: ready.unique_key,
12670                payload: ready.payload,
12671            }
12672            .into_job_row()?,
12673        ))
12674    }
12675
12676    #[tracing::instrument(
12677        skip(self, pool, filter),
12678        fields(kind = ?filter.kind, queue = ?filter.queue, tag = ?filter.tag),
12679        name = "queue_storage.bulk_retry_from_dlq"
12680    )]
12681    pub async fn bulk_retry_from_dlq(
12682        &self,
12683        pool: &PgPool,
12684        filter: &ListDlqFilter,
12685    ) -> Result<u64, AwaError> {
12686        let schema = self.schema();
12687        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12688        let moved: Vec<DlqJobRow> = sqlx::query_as(&format!(
12689            r#"
12690            DELETE FROM {schema}.dlq_entries
12691            WHERE ($1::text IS NULL OR kind = $1)
12692              AND ($2::text IS NULL OR queue = $2)
12693              AND ($3::text IS NULL OR payload -> 'tags' ? $3)
12694              AND (
12695                  ($4::bigint IS NULL AND $5::timestamptz IS NULL)
12696                  OR ($4::bigint IS NOT NULL AND $5::timestamptz IS NULL AND job_id < $4)
12697                  OR ($4::bigint IS NULL AND $5::timestamptz IS NOT NULL AND dlq_at < $5)
12698                  OR (
12699                      $4::bigint IS NOT NULL
12700                      AND $5::timestamptz IS NOT NULL
12701                      AND (dlq_at, job_id) < ($5, $4)
12702                  )
12703              )
12704            RETURNING
12705                job_id,
12706                kind,
12707                queue,
12708                args,
12709                state,
12710                priority,
12711                attempt,
12712                run_lease,
12713                max_attempts,
12714                run_at,
12715                attempted_at,
12716                finalized_at,
12717                created_at,
12718                unique_key,
12719                unique_states,
12720                COALESCE(payload, '{{}}'::jsonb) AS payload,
12721                dlq_reason,
12722                dlq_at,
12723                original_run_lease
12724            "#
12725        ))
12726        .bind(&filter.kind)
12727        .bind(&filter.queue)
12728        .bind(&filter.tag)
12729        .bind(filter.before_id)
12730        .bind(filter.before_dlq_at)
12731        .fetch_all(tx.as_mut())
12732        .await
12733        .map_err(map_sqlx_error)?;
12734
12735        if moved.is_empty() {
12736            tx.commit().await.map_err(map_sqlx_error)?;
12737            return Ok(0);
12738        }
12739
12740        let run_at = Utc::now();
12741        let mut queues = BTreeSet::new();
12742        let mut ready_rows = Vec::with_capacity(moved.len());
12743        for moved_row in moved {
12744            let queue = moved_row.queue.clone();
12745            let priority = moved_row.priority;
12746            queues.insert(queue.clone());
12747            let mut payload = RuntimePayload::from_json(moved_row.payload.clone())?;
12748            payload.set_progress(None);
12749            ready_rows.push(moved_row.into_retry_ready_row(
12750                queue,
12751                priority,
12752                run_at,
12753                payload.into_json(),
12754            ));
12755        }
12756
12757        let revived = ready_rows.len() as u64;
12758        self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(JobState::Failed))
12759            .await?;
12760        self.notify_queues_tx(&mut tx, queues).await?;
12761        tx.commit().await.map_err(map_sqlx_error)?;
12762        Ok(revived)
12763    }
12764
12765    pub async fn discard_failed_by_kind(&self, pool: &PgPool, kind: &str) -> Result<u64, AwaError> {
12766        let schema = self.schema();
12767        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12768
12769        let done_projection = done_row_projection("done", "ready");
12770        let ready_join = done_ready_join(schema, "done", "ready");
12771        let deleted_done: Vec<DoneJobRow> = sqlx::query_as(&format!(
12772            r#"
12773            WITH deleted AS (
12774                DELETE FROM {schema}.done_entries
12775                WHERE kind = $1
12776                  AND state = 'failed'
12777                RETURNING *
12778            )
12779            SELECT {done_projection}
12780            FROM deleted AS done
12781            {ready_join}
12782            "#
12783        ))
12784        .bind(kind)
12785        .fetch_all(tx.as_mut())
12786        .await
12787        .map_err(map_sqlx_error)?;
12788
12789        let deleted_dlq: Vec<DlqJobRow> = sqlx::query_as(&format!(
12790            r#"
12791            DELETE FROM {schema}.dlq_entries
12792            WHERE kind = $1
12793            RETURNING
12794                job_id,
12795                kind,
12796                queue,
12797                args,
12798                state,
12799                priority,
12800                attempt,
12801                run_lease,
12802                max_attempts,
12803                run_at,
12804                attempted_at,
12805                finalized_at,
12806                created_at,
12807                unique_key,
12808                unique_states,
12809                COALESCE(payload, '{{}}'::jsonb) AS payload,
12810                dlq_reason,
12811                dlq_at,
12812                original_run_lease
12813            "#
12814        ))
12815        .bind(kind)
12816        .fetch_all(tx.as_mut())
12817        .await
12818        .map_err(map_sqlx_error)?;
12819
12820        // Discard removes terminal `failed` rows from done_entries (and
12821        // `failed` DLQ rows from dlq_entries — exact terminal counts are keyed
12822        // on done_entries only). Append negative deltas by the per-group
12823        // magnitudes of `deleted_done`.
12824        self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &deleted_done)
12825            .await?;
12826        self.decrement_live_terminal_counters_tx(
12827            &mut tx,
12828            &Self::done_rows_to_counter_keys(&deleted_done),
12829        )
12830        .await?;
12831
12832        for row in &deleted_done {
12833            self.sync_unique_claim(
12834                &mut tx,
12835                row.job_id,
12836                &row.unique_key,
12837                row.unique_states.as_deref(),
12838                Some(row.state),
12839                None,
12840            )
12841            .await?;
12842        }
12843
12844        for row in &deleted_dlq {
12845            self.sync_unique_claim(
12846                &mut tx,
12847                row.job_id,
12848                &row.unique_key,
12849                row.unique_states.as_deref(),
12850                Some(row.state),
12851                None,
12852            )
12853            .await?;
12854        }
12855
12856        tx.commit().await.map_err(map_sqlx_error)?;
12857        Ok((deleted_done.len() + deleted_dlq.len()) as u64)
12858    }
12859
12860    pub async fn fail_retryable(
12861        &self,
12862        pool: &PgPool,
12863        job_id: i64,
12864        run_lease: i64,
12865        error: &str,
12866        progress: Option<serde_json::Value>,
12867    ) -> Result<Option<JobRow>, AwaError> {
12868        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12869        let result = self
12870            .fail_retryable_in_tx(&mut tx, job_id, run_lease, error, progress)
12871            .await?;
12872        tx.commit().await.map_err(map_sqlx_error)?;
12873        Ok(result)
12874    }
12875
12876    /// Transaction-aware variant of [`Self::fail_retryable`] (ADR-029).
12877    pub async fn fail_retryable_in_tx(
12878        &self,
12879        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
12880        job_id: i64,
12881        run_lease: i64,
12882        error: &str,
12883        progress: Option<serde_json::Value>,
12884    ) -> Result<Option<JobRow>, AwaError> {
12885        let Some(moved) = self
12886            .take_running_attempt_tx(tx, job_id, run_lease, "retryable")
12887            .await?
12888        else {
12889            return Ok(None);
12890        };
12891
12892        let mut payload = RuntimePayload::from_json(Self::with_progress(
12893            moved.payload.clone(),
12894            progress.or(moved.progress.clone()),
12895        )?)?;
12896        let exhausted = moved.attempt >= moved.max_attempts;
12897        payload.push_error(lifecycle_error(error, moved.attempt, exhausted));
12898
12899        if exhausted {
12900            let done =
12901                moved
12902                    .clone()
12903                    .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
12904            self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(moved.state))
12905                .await?;
12906            return Ok(Some(done.into_job_row()?));
12907        }
12908
12909        let deferred = moved.clone().into_deferred_row(
12910            JobState::Retryable,
12911            self.backoff_at_tx(tx, moved.attempt, moved.max_attempts)
12912                .await?,
12913            Some(Utc::now()),
12914            payload.into_json(),
12915        );
12916        self.insert_deferred_rows_tx(tx, vec![deferred.clone()], Some(moved.state))
12917            .await?;
12918        Ok(Some(deferred.into_job_row()?))
12919    }
12920
12921    #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_stale_heartbeats")]
12922    pub async fn rescue_stale_heartbeats(
12923        &self,
12924        pool: &PgPool,
12925        staleness: Duration,
12926    ) -> Result<Vec<JobRow>, AwaError> {
12927        let schema = self.schema();
12928        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
12929        let cutoff = Utc::now()
12930            - TimeDelta::from_std(staleness)
12931                .map_err(|err| AwaError::Validation(format!("invalid staleness: {err}")))?;
12932        // #169 B1: the staleness predicate prefers
12933        // `attempt_state.heartbeat_at` (receipts-mode source of truth,
12934        // where heartbeat_batch writes go) and falls back to
12935        // `leases.heartbeat_at` (legacy non-receipts source). Two
12936        // cases this covers:
12937        //
12938        //   * Receipts mode, claim materialized into `leases` via
12939        //     callback registration / progress upsert / equivalent.
12940        //     The leases row was written once at materialize time and
12941        //     its `heartbeat_at` is stale by definition. The fresh
12942        //     value lives on `attempt_state.heartbeat_at`. COALESCE
12943        //     picks `attempt` so a healthy worker isn't falsely
12944        //     rescued — and a dead worker IS rescued, which the
12945        //     receipt-side path can't do (its anti-join below
12946        //     excludes materialized leases to avoid double-closure).
12947        //   * Legacy non-receipts mode: attempt_state.heartbeat_at is
12948        //     never written; COALESCE falls back to
12949        //     leases.heartbeat_at — same shape as pre-B1.
12950        //
12951        // The dropped `idx_state_hb` (v025) doesn't hurt this scan:
12952        // the planner can satisfy the `state='running'` prefix via the
12953        // surviving `(state, deadline_at)` or
12954        // `(state, callback_timeout_at)` indexes, followed by a heap
12955        // recheck of the COALESCE. Bounded by running-lease count and
12956        // called at 30s cadence — cheap.
12957        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
12958            r#"
12959            DELETE FROM {schema}.leases AS target
12960            WHERE (target.job_id, target.run_lease) IN (
12961                SELECT lease.job_id, lease.run_lease
12962                FROM {schema}.leases AS lease
12963                LEFT JOIN {schema}.attempt_state AS attempt
12964                  ON attempt.job_id = lease.job_id
12965                 AND attempt.run_lease = lease.run_lease
12966                WHERE lease.state = 'running'
12967                  AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1
12968                ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC
12969                LIMIT 500
12970                FOR UPDATE OF lease SKIP LOCKED
12971            )
12972            RETURNING
12973                ready_slot,
12974                ready_generation,
12975                job_id,
12976                queue,
12977                state,
12978                priority,
12979                attempt,
12980                run_lease,
12981                max_attempts,
12982                lane_seq,
12983                enqueue_shard,
12984                heartbeat_at,
12985                deadline_at,
12986                attempted_at,
12987                callback_id,
12988                callback_timeout_at
12989            "#
12990        ))
12991        .bind(cutoff)
12992        .fetch_all(tx.as_mut())
12993        .await
12994        .map_err(map_sqlx_error)?;
12995
12996        let rescued_receipts = if self.lease_claim_receipts() {
12997            self.rescue_stale_receipt_claims_tx(&mut tx, cutoff).await?
12998        } else {
12999            Vec::new()
13000        };
13001
13002        if deleted.is_empty() && rescued_receipts.is_empty() {
13003            tx.commit().await.map_err(map_sqlx_error)?;
13004            return Ok(Vec::new());
13005        }
13006
13007        let moved_leases = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13008        let moved_receipts = self
13009            .hydrate_deleted_leases_tx(&mut tx, rescued_receipts)
13010            .await?;
13011
13012        let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len());
13013        for row in moved_leases {
13014            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13015                row.payload.clone(),
13016                row.progress.clone(),
13017            )?)?;
13018            payload.push_error(lifecycle_error(
13019                "heartbeat stale: worker presumed dead",
13020                row.attempt,
13021                false,
13022            ));
13023            let deferred = row.clone().into_deferred_row(
13024                JobState::Retryable,
13025                self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13026                    .await?,
13027                Some(Utc::now()),
13028                payload.into_json(),
13029            );
13030            let rescued_row = self
13031                .insert_rescued_deferred_or_cancel_tx(
13032                    &mut tx,
13033                    row,
13034                    deferred,
13035                    "rescued as duplicate: unique claim held by a newer job",
13036                )
13037                .await?;
13038            rescued.push(rescued_row);
13039        }
13040        for row in moved_receipts {
13041            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13042                row.payload.clone(),
13043                row.progress.clone(),
13044            )?)?;
13045            payload.push_error(lifecycle_error(
13046                "receipt claim stale: worker presumed dead",
13047                row.attempt,
13048                false,
13049            ));
13050            let deferred = row.clone().into_deferred_row(
13051                JobState::Retryable,
13052                self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13053                    .await?,
13054                Some(Utc::now()),
13055                payload.into_json(),
13056            );
13057            let rescued_row = self
13058                .insert_rescued_deferred_or_cancel_tx(
13059                    &mut tx,
13060                    row,
13061                    deferred,
13062                    "rescued as duplicate: unique claim held by a newer job",
13063                )
13064                .await?;
13065            rescued.push(rescued_row);
13066        }
13067        tx.commit().await.map_err(map_sqlx_error)?;
13068        Ok(rescued)
13069    }
13070
13071    #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_expired_deadlines")]
13072    pub async fn rescue_expired_deadlines(&self, pool: &PgPool) -> Result<Vec<JobRow>, AwaError> {
13073        let schema = self.schema();
13074        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13075        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
13076            r#"
13077            DELETE FROM {schema}.leases
13078            WHERE job_id IN (
13079                SELECT job_id
13080                FROM {schema}.leases
13081                WHERE state = 'running'
13082                  AND deadline_at IS NOT NULL
13083                  AND deadline_at < clock_timestamp()
13084                ORDER BY deadline_at ASC
13085                LIMIT 500
13086                FOR UPDATE SKIP LOCKED
13087            )
13088            RETURNING
13089                ready_slot,
13090                ready_generation,
13091                job_id,
13092                queue,
13093                state,
13094                priority,
13095                attempt,
13096                run_lease,
13097                max_attempts,
13098                lane_seq,
13099                enqueue_shard,
13100                heartbeat_at,
13101                deadline_at,
13102                attempted_at,
13103                callback_id,
13104                callback_timeout_at
13105            "#
13106        ))
13107        .fetch_all(tx.as_mut())
13108        .await
13109        .map_err(map_sqlx_error)?;
13110
13111        // Receipts-mode short-path claims hold their deadline on
13112        // `lease_claims.deadline_at` rather than on a `leases` row, so
13113        // the receipt-plane needs its own scan; merge both populations
13114        // into one `moved` set so the maintenance caller observes a
13115        // single rescue batch per tick.
13116        let receipt_deleted = if self.lease_claim_receipts() {
13117            self.rescue_expired_receipt_deadlines_tx(&mut tx).await?
13118        } else {
13119            Vec::new()
13120        };
13121
13122        if deleted.is_empty() && receipt_deleted.is_empty() {
13123            tx.commit().await.map_err(map_sqlx_error)?;
13124            return Ok(Vec::new());
13125        }
13126
13127        let mut moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13128        moved.extend(
13129            self.hydrate_deleted_leases_tx(&mut tx, receipt_deleted)
13130                .await?,
13131        );
13132
13133        let mut rescued = Vec::with_capacity(moved.len());
13134        for row in moved {
13135            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13136                row.payload.clone(),
13137                row.progress.clone(),
13138            )?)?;
13139            payload.push_error(lifecycle_error(
13140                "hard deadline exceeded",
13141                row.attempt,
13142                false,
13143            ));
13144            let deferred = row.clone().into_deferred_row(
13145                JobState::Retryable,
13146                self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13147                    .await?,
13148                Some(Utc::now()),
13149                payload.into_json(),
13150            );
13151            let rescued_row = self
13152                .insert_rescued_deferred_or_cancel_tx(
13153                    &mut tx,
13154                    row,
13155                    deferred,
13156                    "rescued as duplicate: unique claim held by a newer job",
13157                )
13158                .await?;
13159            rescued.push(rescued_row);
13160        }
13161        tx.commit().await.map_err(map_sqlx_error)?;
13162        Ok(rescued)
13163    }
13164
13165    #[tracing::instrument(skip(self, pool), name = "queue_storage.rescue_expired_callbacks")]
13166    pub async fn rescue_expired_callbacks(&self, pool: &PgPool) -> Result<Vec<JobRow>, AwaError> {
13167        let schema = self.schema();
13168        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13169        let deleted: Vec<DeletedLeaseRow> = sqlx::query_as(&format!(
13170            r#"
13171            DELETE FROM {schema}.leases
13172            WHERE job_id IN (
13173                SELECT job_id
13174                FROM {schema}.leases
13175                WHERE state = 'waiting_external'
13176                  AND callback_timeout_at IS NOT NULL
13177                  AND callback_timeout_at < clock_timestamp()
13178                ORDER BY callback_timeout_at ASC
13179                LIMIT 500
13180                FOR UPDATE SKIP LOCKED
13181            )
13182            RETURNING
13183                ready_slot,
13184                ready_generation,
13185                job_id,
13186                queue,
13187                state,
13188                priority,
13189                attempt,
13190                run_lease,
13191                max_attempts,
13192                lane_seq,
13193                enqueue_shard,
13194                heartbeat_at,
13195                deadline_at,
13196                attempted_at,
13197                callback_id,
13198                callback_timeout_at
13199            "#
13200        ))
13201        .fetch_all(tx.as_mut())
13202        .await
13203        .map_err(map_sqlx_error)?;
13204
13205        if deleted.is_empty() {
13206            tx.commit().await.map_err(map_sqlx_error)?;
13207            return Ok(Vec::new());
13208        }
13209
13210        let moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?;
13211
13212        let mut rescued = Vec::with_capacity(moved.len());
13213        for row in moved {
13214            let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state(
13215                row.payload.clone(),
13216                row.progress.clone(),
13217            )?)?;
13218            let exhausted = row.attempt >= row.max_attempts;
13219            payload.push_error(lifecycle_error(
13220                "callback timed out",
13221                row.attempt,
13222                exhausted,
13223            ));
13224            if exhausted {
13225                let done =
13226                    row.clone()
13227                        .into_done_row(JobState::Failed, Utc::now(), payload.into_json());
13228                let rescued_row = self
13229                    .insert_rescued_done_or_cancel_tx(
13230                        &mut tx,
13231                        row,
13232                        done,
13233                        "rescued as duplicate: unique claim held by a newer job",
13234                    )
13235                    .await?;
13236                rescued.push(rescued_row);
13237            } else {
13238                let deferred = row.clone().into_deferred_row(
13239                    JobState::Retryable,
13240                    self.backoff_at_tx(&mut tx, row.attempt, row.max_attempts)
13241                        .await?,
13242                    Some(Utc::now()),
13243                    payload.into_json(),
13244                );
13245                let rescued_row = self
13246                    .insert_rescued_deferred_or_cancel_tx(
13247                        &mut tx,
13248                        row,
13249                        deferred,
13250                        "rescued as duplicate: unique claim held by a newer job",
13251                    )
13252                    .await?;
13253                rescued.push(rescued_row);
13254            }
13255        }
13256        tx.commit().await.map_err(map_sqlx_error)?;
13257        Ok(rescued)
13258    }
13259
13260    pub async fn promote_due(
13261        &self,
13262        pool: &PgPool,
13263        state: JobState,
13264        batch_size: i64,
13265    ) -> Result<usize, AwaError> {
13266        if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 {
13267            return Ok(0);
13268        }
13269
13270        let schema = self.schema();
13271        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13272        let moved: Vec<DeferredJobRow> = sqlx::query_as(&format!(
13273            r#"
13274            DELETE FROM {schema}.deferred_jobs
13275            WHERE job_id IN (
13276                SELECT job_id
13277                FROM {schema}.deferred_jobs
13278                WHERE state = $1
13279                  AND run_at <= clock_timestamp()
13280                  AND NOT EXISTS (
13281                      SELECT 1 FROM awa.queue_meta
13282                      WHERE queue = {schema}.deferred_jobs.queue AND paused = TRUE
13283                  )
13284                ORDER BY run_at ASC, priority ASC, job_id ASC
13285                LIMIT $2
13286                FOR UPDATE SKIP LOCKED
13287            )
13288            RETURNING
13289                job_id,
13290                kind,
13291                queue,
13292                args,
13293                state,
13294                priority,
13295                attempt,
13296                run_lease,
13297                max_attempts,
13298                run_at,
13299                attempted_at,
13300                finalized_at,
13301                created_at,
13302                unique_key,
13303                unique_states,
13304                COALESCE(payload, '{{}}'::jsonb) AS payload
13305            "#
13306        ))
13307        .bind(state)
13308        .bind(batch_size)
13309        .fetch_all(tx.as_mut())
13310        .await
13311        .map_err(map_sqlx_error)?;
13312
13313        if moved.is_empty() {
13314            tx.commit().await.map_err(map_sqlx_error)?;
13315            return Ok(0);
13316        }
13317
13318        let ready_rows: Vec<ExistingReadyRow> = moved
13319            .iter()
13320            .cloned()
13321            .map(|row| ExistingReadyRow {
13322                job_id: row.job_id,
13323                kind: row.kind,
13324                queue: row.queue,
13325                args: row.args,
13326                priority: row.priority,
13327                attempt: row.attempt,
13328                run_lease: row.run_lease,
13329                max_attempts: row.max_attempts,
13330                run_at: Utc::now(),
13331                attempted_at: row.attempted_at,
13332                created_at: row.created_at,
13333                unique_key: row.unique_key,
13334                unique_states: row.unique_states,
13335                payload: row.payload,
13336            })
13337            .collect();
13338        let queues = ready_rows
13339            .iter()
13340            .map(|row| row.queue.clone())
13341            .collect::<Vec<_>>();
13342        self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(state))
13343            .await?;
13344        self.notify_queues_tx(&mut tx, queues).await?;
13345        tx.commit().await.map_err(map_sqlx_error)?;
13346        Ok(moved.len())
13347    }
13348
13349    async fn relation_has_rows_tx(
13350        tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
13351        relation: &str,
13352    ) -> Result<bool, AwaError> {
13353        sqlx::query_scalar(&format!("SELECT EXISTS (SELECT 1 FROM {relation} LIMIT 1)"))
13354            .fetch_one(tx.as_mut())
13355            .await
13356            .map_err(map_sqlx_error)
13357    }
13358
13359    #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate")]
13360    pub async fn rotate(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
13361        let schema = self.schema();
13362        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13363
13364        let state: (i32, i64, i32) = sqlx::query_as(&format!(
13365            r#"
13366            SELECT current_slot, generation, slot_count
13367            FROM {schema}.queue_ring_state
13368            WHERE singleton = TRUE
13369            FOR UPDATE
13370            "#
13371        ))
13372        .fetch_one(tx.as_mut())
13373        .await
13374        .map_err(map_sqlx_error)?;
13375
13376        let next_slot = (state.0 + 1).rem_euclid(state.2);
13377        let ready_busy =
13378            Self::relation_has_rows_tx(&mut tx, &ready_child_name(schema, next_slot as usize))
13379                .await?;
13380        let claim_attempt_batch_busy = Self::relation_has_rows_tx(
13381            &mut tx,
13382            &ready_claim_attempt_batch_child_name(schema, next_slot as usize),
13383        )
13384        .await?;
13385        let done_busy =
13386            Self::relation_has_rows_tx(&mut tx, &done_child_name(schema, next_slot as usize))
13387                .await?;
13388        let tombstone_busy = Self::relation_has_rows_tx(
13389            &mut tx,
13390            &ready_tombstone_child_name(schema, next_slot as usize),
13391        )
13392        .await?;
13393        let segment_busy = Self::relation_has_rows_tx(
13394            &mut tx,
13395            &ready_segment_child_name(schema, next_slot as usize),
13396        )
13397        .await?;
13398        let receipt_batch_busy = Self::relation_has_rows_tx(
13399            &mut tx,
13400            &receipt_completion_batch_child_name(schema, next_slot as usize),
13401        )
13402        .await?;
13403        let receipt_tombstone_busy = Self::relation_has_rows_tx(
13404            &mut tx,
13405            &receipt_completion_tombstone_child_name(schema, next_slot as usize),
13406        )
13407        .await?;
13408        let terminal_delta_busy = Self::relation_has_rows_tx(
13409            &mut tx,
13410            &terminal_delta_child_name(schema, next_slot as usize),
13411        )
13412        .await?;
13413
13414        if ready_busy
13415            || claim_attempt_batch_busy
13416            || done_busy
13417            || tombstone_busy
13418            || segment_busy
13419            || receipt_batch_busy
13420            || receipt_tombstone_busy
13421            || terminal_delta_busy
13422        {
13423            tx.commit().await.map_err(map_sqlx_error)?;
13424            return Ok(RotateOutcome::SkippedBusy {
13425                slot: next_slot,
13426                busy: BusyCounts {
13427                    queue_ready: busy_indicator(ready_busy),
13428                    queue_claim_attempt_batches: busy_indicator(claim_attempt_batch_busy),
13429                    queue_done: busy_indicator(done_busy),
13430                    queue_tombstones: busy_indicator(tombstone_busy),
13431                    queue_ready_segments: busy_indicator(segment_busy),
13432                    queue_receipt_completion_batches: busy_indicator(receipt_batch_busy),
13433                    queue_receipt_completion_tombstones: busy_indicator(receipt_tombstone_busy),
13434                    queue_terminal_deltas: busy_indicator(terminal_delta_busy),
13435                    ..Default::default()
13436                },
13437            });
13438        }
13439
13440        let next_generation = state.1 + 1;
13441
13442        sqlx::query(&format!(
13443            r#"
13444            UPDATE {schema}.queue_ring_state
13445            SET current_slot = $1,
13446                generation = $2
13447            WHERE singleton = TRUE
13448            "#
13449        ))
13450        .bind(next_slot)
13451        .bind(next_generation)
13452        .execute(tx.as_mut())
13453        .await
13454        .map_err(map_sqlx_error)?;
13455
13456        sqlx::query(&format!(
13457            r#"
13458            UPDATE {schema}.queue_ring_slots
13459            SET generation = $2
13460            WHERE slot = $1
13461            "#
13462        ))
13463        .bind(next_slot)
13464        .bind(next_generation)
13465        .execute(tx.as_mut())
13466        .await
13467        .map_err(map_sqlx_error)?;
13468
13469        tx.commit().await.map_err(map_sqlx_error)?;
13470        Ok(RotateOutcome::Rotated {
13471            slot: next_slot,
13472            generation: next_generation,
13473        })
13474    }
13475
13476    #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_leases")]
13477    pub async fn rotate_leases(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
13478        let schema = self.schema();
13479        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13480
13481        // FOR UPDATE serialises with prune_oldest_leases and parallel
13482        // rotators. Without it two rotators can both pass the busy-check,
13483        // both compute the same next_slot, and the loser's CAS update
13484        // wastes work. `RotateLeasesPlan` in
13485        // `correctness/storage/AwaStorageLockOrder.tla` requires this
13486        // lock as the first acquired resource for the rotation tx.
13487        let state: (i32, i64, i32) = sqlx::query_as(&format!(
13488            r#"
13489            SELECT current_slot, generation, slot_count
13490            FROM {schema}.lease_ring_state
13491            WHERE singleton = TRUE
13492            FOR UPDATE
13493            "#
13494        ))
13495        .fetch_one(tx.as_mut())
13496        .await
13497        .map_err(map_sqlx_error)?;
13498
13499        let next_slot = (state.0 + 1).rem_euclid(state.2);
13500        let lease_busy =
13501            Self::relation_has_rows_tx(&mut tx, &lease_child_name(schema, next_slot as usize))
13502                .await?;
13503
13504        if lease_busy {
13505            tx.commit().await.map_err(map_sqlx_error)?;
13506            return Ok(RotateOutcome::SkippedBusy {
13507                slot: next_slot,
13508                busy: BusyCounts {
13509                    leases: busy_indicator(lease_busy),
13510                    ..Default::default()
13511                },
13512            });
13513        }
13514
13515        let next_generation = state.1 + 1;
13516
13517        let rotated = sqlx::query(&format!(
13518            r#"
13519            UPDATE {schema}.lease_ring_state
13520            SET current_slot = $1,
13521                generation = $2
13522            WHERE singleton = TRUE
13523              AND current_slot = $3
13524              AND generation = $4
13525            "#
13526        ))
13527        .bind(next_slot)
13528        .bind(next_generation)
13529        .bind(state.0)
13530        .bind(state.1)
13531        .execute(tx.as_mut())
13532        .await
13533        .map_err(map_sqlx_error)?;
13534
13535        if rotated.rows_affected() == 0 {
13536            // Another rotator beat us to the CAS. Report the bounded
13537            // presence evidence we sampled before giving up.
13538            tx.commit().await.map_err(map_sqlx_error)?;
13539            return Ok(RotateOutcome::SkippedBusy {
13540                slot: next_slot,
13541                busy: BusyCounts {
13542                    leases: busy_indicator(lease_busy),
13543                    ..Default::default()
13544                },
13545            });
13546        }
13547
13548        tx.commit().await.map_err(map_sqlx_error)?;
13549        Ok(RotateOutcome::Rotated {
13550            slot: next_slot,
13551            generation: next_generation,
13552        })
13553    }
13554
13555    /// Rebuild terminal counters from scratch by truncating folded live counts
13556    /// and pending deltas, then re-aggregating terminal history. Run this when:
13557    ///
13558    /// - upgrading from a pre-#290 fleet, where in-flight binaries wrote
13559    ///   terminal history without maintaining the counter,
13560    /// - after any incident that may have left the counter inconsistent
13561    ///   with terminal history, or
13562    /// - as a routine drift-recovery step before relying on counter-fed
13563    ///   reads for billing-grade accuracy.
13564    ///
13565    /// The rebuild is wrapped in an advisory lock so concurrent writers
13566    /// don't interleave new inserts mid-rebuild. The lock key is scoped
13567    /// to the queue-storage schema, so other schemas / other operations
13568    /// are unaffected. Writers that hit the lock will block briefly
13569    /// rather than fail.
13570    ///
13571    /// **Operator note:** this is best run on a quiesced fleet (workers
13572    /// paused or fully drained). Concurrent inserts during the rebuild
13573    /// will block on the lock; long-held locks can stall the fleet. The
13574    /// rebuild itself is O(rows in `{schema}.done_entries`). Compact receipt
13575    /// completions remain counted directly from retained
13576    /// `receipt_completion_batches` until queue prune folds them into
13577    /// permanent rollups.
13578    #[tracing::instrument(skip(self, pool), name = "queue_storage.rebuild_terminal_counters")]
13579    pub async fn rebuild_terminal_counters(&self, pool: &PgPool) -> Result<i64, AwaError> {
13580        let schema = self.schema();
13581        let rebuild_lock_name = format!("awa.queue_storage.rebuild_terminal_counters:{schema}");
13582        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13583
13584        sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))")
13585            .bind(&rebuild_lock_name)
13586            .execute(tx.as_mut())
13587            .await
13588            .map_err(map_sqlx_error)?;
13589
13590        sqlx::query(&format!(
13591            "LOCK TABLE {schema}.done_entries, \
13592             {schema}.receipt_completion_batches, \
13593             {schema}.receipt_completion_tombstones, \
13594             {schema}.queue_terminal_count_deltas, \
13595             {schema}.queue_terminal_live_counts \
13596             IN ACCESS EXCLUSIVE MODE"
13597        ))
13598        .execute(tx.as_mut())
13599        .await
13600        .map_err(map_sqlx_error)?;
13601
13602        sqlx::query(&format!(
13603            "TRUNCATE TABLE {schema}.queue_terminal_live_counts, {schema}.queue_terminal_count_deltas"
13604        ))
13605        .execute(tx.as_mut())
13606        .await
13607        .map_err(map_sqlx_error)?;
13608
13609        let inserted: i64 = sqlx::query_scalar(&format!(
13610            r#"
13611            WITH inserted AS (
13612                INSERT INTO {schema}.queue_terminal_live_counts AS counts (
13613                    ready_slot, queue, priority, enqueue_shard, counter_bucket, live_terminal_count
13614                )
13615                SELECT
13616                    ready_slot,
13617                    queue,
13618                    priority,
13619                    enqueue_shard,
13620                    mod(
13621                        mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint)
13622                            + {TERMINAL_COUNTER_BUCKETS}::bigint,
13623                        {TERMINAL_COUNTER_BUCKETS}::bigint
13624                    )::smallint AS counter_bucket,
13625                    count(*)::bigint
13626                FROM {schema}.done_entries
13627                GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket
13628                RETURNING 1
13629            )
13630            SELECT COALESCE(count(*), 0)::bigint FROM inserted
13631            "#
13632        ))
13633        .fetch_one(tx.as_mut())
13634        .await
13635        .map_err(map_sqlx_error)?;
13636
13637        // Flip the trust marker. From this point the read path
13638        // (queue_counts_exact) uses the counter for done-entry terminal rows;
13639        // before this call, it falls back to scanning terminal_jobs.
13640        sqlx::query(&format!(
13641            r#"
13642            UPDATE {schema}.queue_ring_state
13643            SET terminal_counter_trusted_at = now()
13644            WHERE singleton = TRUE
13645            "#
13646        ))
13647        .execute(tx.as_mut())
13648        .await
13649        .map_err(map_sqlx_error)?;
13650
13651        tx.commit().await.map_err(map_sqlx_error)?;
13652        Ok(inserted)
13653    }
13654
13655    /// Check whether the live terminal counter has been marked trusted
13656    /// for exact reads. Returns `true` for fresh installs (the trust
13657    /// marker is auto-set by `prepare_schema` when `done_entries` is
13658    /// empty) and after a successful
13659    /// [`Self::rebuild_terminal_counters`]; returns `false` on an
13660    /// existing install that has not yet been rebuilt under the new
13661    /// counter-aware code path.
13662    ///
13663    /// `queue_counts_exact` consults this to decide between the
13664    /// counter-fed fast path and the scan-based fallback. Single-row
13665    /// PK fetch; negligible cost per call.
13666    pub async fn terminal_counter_trusted(&self, pool: &PgPool) -> Result<bool, AwaError> {
13667        let schema = self.schema();
13668        let trusted: Option<bool> = sqlx::query_scalar(&format!(
13669            "SELECT terminal_counter_trusted_at IS NOT NULL \
13670             FROM {schema}.queue_ring_state WHERE singleton = TRUE"
13671        ))
13672        .fetch_optional(pool)
13673        .await
13674        .map_err(map_sqlx_error)?;
13675        // Missing row = pre-#290 schema that hasn't run the new
13676        // prepare_schema yet. Treat as untrusted; the scan path is
13677        // still correct.
13678        Ok(trusted.unwrap_or(false))
13679    }
13680
13681    /// Fold append-only `done_entries` terminal-count deltas into
13682    /// `queue_terminal_live_counts` for sealed queue slots.
13683    ///
13684    /// `done_entries` terminal inserts and deletes append signed rows into
13685    /// `queue_terminal_count_deltas` instead of updating the live counter on
13686    /// the user-facing hot path. Compact receipt completions are counted from
13687    /// retained batch rows directly. Exact reads sum folded live counts plus
13688    /// pending deltas and retained compact batches, so this rollup can run
13689    /// asynchronously. It intentionally
13690    /// skips the current queue slot and any slot with active leases or open
13691    /// receipt claims; that keeps rollup away from the segment receiving hot
13692    /// completions and avoids racing future terminal deltas for the same
13693    /// ready generation. Rollup also stands down while another backend pins the
13694    /// MVCC horizon, leaving the append-only deltas visible to exact reads
13695    /// until the horizon clears.
13696    #[tracing::instrument(skip(self, pool), name = "queue_storage.rollup_terminal_count_deltas")]
13697    pub async fn rollup_terminal_count_deltas(
13698        &self,
13699        pool: &PgPool,
13700        max_slots: usize,
13701    ) -> Result<TerminalDeltaRollupOutcome, AwaError> {
13702        if max_slots == 0 {
13703            return Ok(TerminalDeltaRollupOutcome::default());
13704        }
13705
13706        let schema = self.schema();
13707        let current_slot: i32 = sqlx::query_scalar(&format!(
13708            r#"
13709            SELECT current_slot
13710            FROM {schema}.queue_ring_state
13711            WHERE singleton = TRUE
13712            "#
13713        ))
13714        .fetch_one(pool)
13715        .await
13716        .map_err(map_sqlx_error)?;
13717
13718        let slots = self
13719            .terminal_delta_rollup_candidates(pool, current_slot)
13720            .await?;
13721
13722        let mut outcome = TerminalDeltaRollupOutcome::default();
13723        for (slot, generation) in slots {
13724            if outcome.rolled_slots >= max_slots {
13725                break;
13726            }
13727            match self
13728                .rollup_terminal_count_delta_slot(pool, slot, generation)
13729                .await?
13730            {
13731                TerminalDeltaSlotRollup::Empty => {}
13732                TerminalDeltaSlotRollup::Rolled {
13733                    delta_rows,
13734                    grouped_keys,
13735                } => {
13736                    outcome.rolled_slots += 1;
13737                    outcome.delta_rows += delta_rows;
13738                    outcome.grouped_keys += grouped_keys;
13739                }
13740                TerminalDeltaSlotRollup::SkippedActive => {
13741                    outcome.skipped_active_slots += 1;
13742                }
13743                TerminalDeltaSlotRollup::SkippedMvccPinned => {
13744                    outcome.skipped_mvcc_pinned = true;
13745                    break;
13746                }
13747                TerminalDeltaSlotRollup::Blocked => {
13748                    outcome.blocked_slots += 1;
13749                }
13750            }
13751        }
13752
13753        Ok(outcome)
13754    }
13755
13756    async fn terminal_delta_rollup_mvcc_horizon_pinned_tx(
13757        tx: &mut sqlx::Transaction<'_, Postgres>,
13758    ) -> Result<bool, AwaError> {
13759        let pinned: bool = sqlx::query_scalar(
13760            r#"
13761            SELECT EXISTS (
13762                SELECT 1
13763                FROM pg_stat_activity
13764                WHERE datname = current_database()
13765                  AND pid <> pg_backend_pid()
13766                  AND backend_type = 'client backend'
13767                  AND (
13768                      backend_xmin IS NOT NULL
13769                      OR (
13770                          backend_xid IS NOT NULL
13771                          AND state LIKE 'idle in transaction%'
13772                      )
13773                  )
13774            )
13775            "#,
13776        )
13777        .fetch_one(tx.as_mut())
13778        .await
13779        .map_err(map_sqlx_error)?;
13780        Ok(pinned)
13781    }
13782
13783    async fn terminal_delta_rollup_candidates(
13784        &self,
13785        pool: &PgPool,
13786        current_slot: i32,
13787    ) -> Result<Vec<(i32, i64)>, AwaError> {
13788        let schema = self.schema();
13789        let sealed_slots: Vec<(i32, i64)> = sqlx::query_as(&format!(
13790            r#"
13791            SELECT slot, generation
13792            FROM {schema}.queue_ring_slots
13793            WHERE generation >= 0
13794              AND slot <> $1
13795            ORDER BY generation ASC, slot ASC
13796            "#
13797        ))
13798        .bind(current_slot)
13799        .fetch_all(pool)
13800        .await
13801        .map_err(map_sqlx_error)?;
13802
13803        let mut pending_slots = Vec::new();
13804        for (slot, generation) in sealed_slots {
13805            let Ok(slot_index) = usize::try_from(slot) else {
13806                continue;
13807            };
13808            let delta_child = terminal_delta_child_name(schema, slot_index);
13809            let has_pending: bool = sqlx::query_scalar(&format!(
13810                r#"
13811                SELECT EXISTS (
13812                    SELECT 1
13813                    FROM {delta_child}
13814                    WHERE ready_generation = $1
13815                    LIMIT 1
13816                )
13817                "#
13818            ))
13819            .bind(generation)
13820            .fetch_one(pool)
13821            .await
13822            .map_err(map_sqlx_error)?;
13823            if has_pending {
13824                pending_slots.push((slot, generation));
13825            }
13826        }
13827
13828        Ok(pending_slots)
13829    }
13830
13831    async fn rollup_terminal_count_delta_slot(
13832        &self,
13833        pool: &PgPool,
13834        slot: i32,
13835        generation: i64,
13836    ) -> Result<TerminalDeltaSlotRollup, AwaError> {
13837        let schema = self.schema();
13838        let delta_child = terminal_delta_child_name(schema, slot as usize);
13839        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
13840
13841        let current_slot: i32 = sqlx::query_scalar(&format!(
13842            r#"
13843            SELECT current_slot
13844            FROM {schema}.queue_ring_state
13845            WHERE singleton = TRUE
13846            FOR UPDATE
13847            "#
13848        ))
13849        .fetch_one(tx.as_mut())
13850        .await
13851        .map_err(map_sqlx_error)?;
13852
13853        let slot_generation: Option<i64> = sqlx::query_scalar(&format!(
13854            r#"
13855            SELECT generation
13856            FROM {schema}.queue_ring_slots
13857            WHERE slot = $1
13858            FOR UPDATE
13859            "#
13860        ))
13861        .bind(slot)
13862        .fetch_optional(tx.as_mut())
13863        .await
13864        .map_err(map_sqlx_error)?;
13865
13866        let Some(slot_generation) = slot_generation else {
13867            tx.commit().await.map_err(map_sqlx_error)?;
13868            return Ok(TerminalDeltaSlotRollup::Empty);
13869        };
13870
13871        if current_slot == slot {
13872            tx.commit().await.map_err(map_sqlx_error)?;
13873            return Ok(TerminalDeltaSlotRollup::SkippedActive);
13874        }
13875
13876        if slot_generation != generation {
13877            tx.commit().await.map_err(map_sqlx_error)?;
13878            return Ok(TerminalDeltaSlotRollup::Empty);
13879        }
13880
13881        if Self::terminal_delta_rollup_mvcc_horizon_pinned_tx(&mut tx).await? {
13882            tx.commit().await.map_err(map_sqlx_error)?;
13883            return Ok(TerminalDeltaSlotRollup::SkippedMvccPinned);
13884        }
13885
13886        let ready_child = ready_child_name(schema, slot as usize);
13887        let pending_ready: bool = sqlx::query_scalar(&format!(
13888            r#"
13889            WITH claim_cursors AS MATERIALIZED (
13890                SELECT
13891                    queue,
13892                    priority,
13893                    enqueue_shard,
13894                    {schema}.sequence_next_value(seq_name) AS claim_seq
13895                FROM {schema}.queue_claim_heads
13896            )
13897            SELECT EXISTS (
13898                SELECT 1
13899                FROM claim_cursors AS claims
13900                CROSS JOIN LATERAL (
13901                    SELECT 1
13902                    FROM {ready_child} AS ready
13903                    WHERE ready.ready_generation = $1
13904                      AND ready.queue = claims.queue
13905                      AND ready.priority = claims.priority
13906                      AND ready.enqueue_shard = claims.enqueue_shard
13907                      AND ready.lane_seq >= claims.claim_seq
13908                    LIMIT 1
13909                ) AS pending_ready
13910                LIMIT 1
13911            )
13912            "#
13913        ))
13914        .bind(generation)
13915        .fetch_one(tx.as_mut())
13916        .await
13917        .map_err(map_sqlx_error)?;
13918
13919        if pending_ready {
13920            tx.commit().await.map_err(map_sqlx_error)?;
13921            return Ok(TerminalDeltaSlotRollup::SkippedActive);
13922        }
13923
13924        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
13925
13926        let lock_delta = sqlx::query(&format!(
13927            "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE"
13928        ))
13929        .execute(tx.as_mut())
13930        .await;
13931
13932        match lock_delta {
13933            Ok(_) => {}
13934            Err(err) if is_lock_contention_error(&err) => {
13935                let _ = tx.rollback().await;
13936                return Ok(TerminalDeltaSlotRollup::Blocked);
13937            }
13938            Err(err) => {
13939                let _ = tx.rollback().await;
13940                return Err(map_sqlx_error(err));
13941            }
13942        }
13943
13944        let active_leases =
13945            queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
13946
13947        if active_leases {
13948            tx.commit().await.map_err(map_sqlx_error)?;
13949            return Ok(TerminalDeltaSlotRollup::SkippedActive);
13950        }
13951
13952        let unclosed_claim_refs =
13953            queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
13954
13955        if unclosed_claim_refs {
13956            tx.commit().await.map_err(map_sqlx_error)?;
13957            return Ok(TerminalDeltaSlotRollup::SkippedActive);
13958        }
13959
13960        let delta_rows: i64 = sqlx::query_scalar(&format!(
13961            r#"
13962            SELECT count(*)::bigint
13963            FROM {delta_child}
13964            WHERE ready_generation = $1
13965            "#
13966        ))
13967        .bind(generation)
13968        .fetch_one(tx.as_mut())
13969        .await
13970        .map_err(map_sqlx_error)?;
13971
13972        if delta_rows == 0 {
13973            tx.commit().await.map_err(map_sqlx_error)?;
13974            return Ok(TerminalDeltaSlotRollup::Empty);
13975        }
13976
13977        let grouped_keys: i64 = sqlx::query_scalar(&format!(
13978            r#"
13979            WITH grouped AS MATERIALIZED (
13980                SELECT
13981                    ready_slot,
13982                    queue,
13983                    priority,
13984                    enqueue_shard,
13985                    counter_bucket,
13986                    SUM(terminal_delta)::bigint AS delta
13987                FROM {delta_child}
13988                WHERE ready_generation = $1
13989                GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket
13990                HAVING SUM(terminal_delta) <> 0
13991            ),
13992            updated AS (
13993                UPDATE {schema}.queue_terminal_live_counts AS counts
13994                SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta)
13995                FROM grouped
13996                WHERE counts.ready_slot = grouped.ready_slot
13997                  AND counts.queue = grouped.queue
13998                  AND counts.priority = grouped.priority
13999                  AND counts.enqueue_shard = grouped.enqueue_shard
14000                  AND counts.counter_bucket = grouped.counter_bucket
14001                RETURNING
14002                    counts.ready_slot,
14003                    counts.queue,
14004                    counts.priority,
14005                    counts.enqueue_shard,
14006                    counts.counter_bucket
14007            ),
14008            inserted AS (
14009                INSERT INTO {schema}.queue_terminal_live_counts AS counts (
14010                    ready_slot,
14011                    queue,
14012                    priority,
14013                    enqueue_shard,
14014                    counter_bucket,
14015                    live_terminal_count
14016                )
14017                SELECT
14018                    grouped.ready_slot,
14019                    grouped.queue,
14020                    grouped.priority,
14021                    grouped.enqueue_shard,
14022                    grouped.counter_bucket,
14023                    grouped.delta
14024                FROM grouped
14025                WHERE grouped.delta > 0
14026                  AND NOT EXISTS (
14027                      SELECT 1
14028                      FROM updated
14029                      WHERE updated.ready_slot = grouped.ready_slot
14030                        AND updated.queue = grouped.queue
14031                        AND updated.priority = grouped.priority
14032                        AND updated.enqueue_shard = grouped.enqueue_shard
14033                        AND updated.counter_bucket = grouped.counter_bucket
14034                  )
14035                ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket
14036                ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE
14037                SET live_terminal_count =
14038                    GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count)
14039                RETURNING 1
14040            )
14041            SELECT count(*)::bigint FROM grouped
14042            "#
14043        ))
14044        .bind(generation)
14045        .fetch_one(tx.as_mut())
14046        .await
14047        .map_err(map_sqlx_error)?;
14048
14049        let truncate_delta = sqlx::query(&format!("TRUNCATE TABLE {delta_child}"))
14050            .execute(tx.as_mut())
14051            .await;
14052
14053        match truncate_delta {
14054            Ok(_) => {
14055                tx.commit().await.map_err(map_sqlx_error)?;
14056                Ok(TerminalDeltaSlotRollup::Rolled {
14057                    delta_rows,
14058                    grouped_keys,
14059                })
14060            }
14061            Err(err) if is_lock_contention_error(&err) => {
14062                let _ = tx.rollback().await;
14063                Ok(TerminalDeltaSlotRollup::Blocked)
14064            }
14065            Err(err) => {
14066                let _ = tx.rollback().await;
14067                Err(map_sqlx_error(err))
14068            }
14069        }
14070    }
14071
14072    /// Prune the oldest sealed queue ring slot.
14073    ///
14074    /// `failed_retention` is the floor for `failed` terminal rows: rows
14075    /// with `finalized_at` inside the floor are re-homed into the live
14076    /// `done_entries` segment (still retryable) instead of being
14077    /// truncated with the slot. `Duration::ZERO` disables the floor and
14078    /// prunes every terminal row. Granularity is whole seconds,
14079    /// matching DLQ retention. The floor applies to `failed` only —
14080    /// `cancelled` is an explicit operator action and is always pruned.
14081    #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest")]
14082    pub async fn prune_oldest(
14083        &self,
14084        pool: &PgPool,
14085        failed_retention: Duration,
14086    ) -> Result<PruneOutcome, AwaError> {
14087        let schema = self.schema();
14088        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14089
14090        let state: (i32,) = sqlx::query_as(&format!(
14091            r#"
14092            SELECT current_slot
14093            FROM {schema}.queue_ring_state
14094            WHERE singleton = TRUE
14095            FOR UPDATE
14096            "#
14097        ))
14098        .fetch_one(tx.as_mut())
14099        .await
14100        .map_err(map_sqlx_error)?;
14101
14102        let target: Option<(i32, i64)> = sqlx::query_as(&format!(
14103            r#"
14104            SELECT slot, generation
14105            FROM {schema}.queue_ring_slots
14106            WHERE generation >= 0
14107              AND slot <> $1
14108            ORDER BY generation ASC, slot ASC
14109            LIMIT 1
14110            FOR UPDATE
14111            "#
14112        ))
14113        .bind(state.0)
14114        .fetch_optional(tx.as_mut())
14115        .await
14116        .map_err(map_sqlx_error)?;
14117
14118        let Some((slot, generation)) = target else {
14119            tx.commit().await.map_err(map_sqlx_error)?;
14120            return Ok(PruneOutcome::Noop);
14121        };
14122
14123        let ready_child = ready_child_name(schema, slot as usize);
14124        let claim_attempt_child = ready_claim_attempt_batch_child_name(schema, slot as usize);
14125        let done_child = done_child_name(schema, slot as usize);
14126        let tomb_child = ready_tombstone_child_name(schema, slot as usize);
14127        let segment_child = ready_segment_child_name(schema, slot as usize);
14128        let receipt_batch_child = receipt_completion_batch_child_name(schema, slot as usize);
14129        let receipt_tomb_child = receipt_completion_tombstone_child_name(schema, slot as usize);
14130        let delta_child = terminal_delta_child_name(schema, slot as usize);
14131
14132        // Queue prune has three common skip gates. Check them before
14133        // taking ACCESS EXCLUSIVE on the ready/terminal child tables so
14134        // a known-busy slot does not block hot claim reads behind a
14135        // doomed truncate attempt. These are only MVCC snapshots; the
14136        // same gates are checked again after the exclusive locks.
14137        let active_leases =
14138            queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
14139
14140        if active_leases {
14141            tx.commit().await.map_err(map_sqlx_error)?;
14142            return Ok(PruneOutcome::SkippedActive {
14143                slot,
14144                reason: SkipReason::QueueActiveLeases,
14145                count: busy_indicator(active_leases),
14146            });
14147        }
14148
14149        let pending =
14150            queue_prune_has_pending_ready_tx(&mut tx, schema, &ready_child, generation).await?;
14151
14152        if pending {
14153            tx.commit().await.map_err(map_sqlx_error)?;
14154            return Ok(PruneOutcome::SkippedActive {
14155                slot,
14156                reason: SkipReason::QueuePendingReady,
14157                count: busy_indicator(pending),
14158            });
14159        }
14160
14161        let unclosed_claim_refs =
14162            queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
14163
14164        if unclosed_claim_refs {
14165            tx.commit().await.map_err(map_sqlx_error)?;
14166            return Ok(PruneOutcome::SkippedActive {
14167                slot,
14168                reason: SkipReason::QueueUnclosedClaimRefs,
14169                count: busy_indicator(unclosed_claim_refs),
14170            });
14171        }
14172
14173        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14174
14175        let lock_tables = sqlx::query(&format!(
14176            "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"
14177        ))
14178        .execute(tx.as_mut())
14179        .await;
14180
14181        if let Err(err) = lock_tables {
14182            let _ = tx.rollback().await;
14183            if is_lock_contention_error(&err) {
14184                return Ok(PruneOutcome::Blocked { slot });
14185            }
14186            return Err(map_sqlx_error(err));
14187        }
14188
14189        // The pre-lock gates are only an MVCC snapshot. If an in-flight
14190        // producer or claimer committed while the exclusive partition
14191        // lock waited, these post-lock gates see it before TRUNCATE.
14192        let active_leases =
14193            queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?;
14194        if active_leases {
14195            tx.commit().await.map_err(map_sqlx_error)?;
14196            return Ok(PruneOutcome::SkippedActive {
14197                slot,
14198                reason: SkipReason::QueueActiveLeases,
14199                count: busy_indicator(active_leases),
14200            });
14201        }
14202
14203        let pending =
14204            queue_prune_has_pending_ready_tx(&mut tx, schema, &ready_child, generation).await?;
14205        if pending {
14206            tx.commit().await.map_err(map_sqlx_error)?;
14207            return Ok(PruneOutcome::SkippedActive {
14208                slot,
14209                reason: SkipReason::QueuePendingReady,
14210                count: busy_indicator(pending),
14211            });
14212        }
14213
14214        let unclosed_claim_refs =
14215            queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?;
14216        if unclosed_claim_refs {
14217            tx.commit().await.map_err(map_sqlx_error)?;
14218            return Ok(PruneOutcome::SkippedActive {
14219                slot,
14220                reason: SkipReason::QueueUnclosedClaimRefs,
14221                count: busy_indicator(unclosed_claim_refs),
14222            });
14223        }
14224
14225        let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX);
14226        let retention_floor = failed_retention_secs > 0;
14227
14228        // #337: failed terminal rows still inside the retention floor
14229        // are re-homed into the live done segment before the TRUNCATE
14230        // so they stay retryable. The queue_ring_state row is held FOR
14231        // UPDATE above, which serializes against rotate — the current
14232        // slot cannot move while carried rows are re-inserted into it.
14233        let carried_failed_rows = if retention_floor {
14234            let (current_generation,): (i64,) = sqlx::query_as(&format!(
14235                "SELECT generation FROM {schema}.queue_ring_slots WHERE slot = $1"
14236            ))
14237            .bind(state.0)
14238            .fetch_one(tx.as_mut())
14239            .await
14240            .map_err(map_sqlx_error)?;
14241
14242            // Carried rows are widened to the self-contained synthetic
14243            // done-row shape (the COALESCE chain mirrors
14244            // `done_row_projection`) because the ready rows they would
14245            // otherwise hydrate from are truncated with this slot.
14246            // Deliberately no ON CONFLICT: lane_seq comes from
14247            // never-resetting sequences, so a PK collision means
14248            // corrupted terminal state and must abort the prune loudly
14249            // rather than silently drop a terminal fact.
14250            let carried = sqlx::query(&format!(
14251                r#"
14252                INSERT INTO {schema}.done_entries (
14253                    ready_slot, ready_generation, job_id, kind, queue,
14254                    args, state, priority, attempt, run_lease,
14255                    max_attempts, lane_seq, enqueue_shard, run_at,
14256                    attempted_at, finalized_at, created_at, unique_key,
14257                    unique_states, payload
14258                )
14259                SELECT
14260                    $2,
14261                    $3,
14262                    done.job_id,
14263                    done.kind,
14264                    done.queue,
14265                    COALESCE(done.args, ready.args, '{{}}'::jsonb),
14266                    done.state,
14267                    done.priority,
14268                    done.attempt,
14269                    done.run_lease,
14270                    COALESCE(done.max_attempts, ready.max_attempts, 25::smallint),
14271                    done.lane_seq,
14272                    done.enqueue_shard,
14273                    COALESCE(done.run_at, ready.run_at, done.finalized_at),
14274                    COALESCE(done.attempted_at, ready.attempted_at),
14275                    done.finalized_at,
14276                    COALESCE(done.created_at, ready.created_at, done.finalized_at),
14277                    COALESCE(done.unique_key, ready.unique_key),
14278                    COALESCE(done.unique_states, ready.unique_states),
14279                    COALESCE(done.payload, ready.payload, '{{}}'::jsonb)
14280                FROM {done_child} AS done
14281                LEFT JOIN {ready_child} AS ready
14282                  ON ready.ready_slot = done.ready_slot
14283                 AND ready.ready_generation = done.ready_generation
14284                 AND ready.queue = done.queue
14285                 AND ready.priority = done.priority
14286                 AND ready.enqueue_shard = done.enqueue_shard
14287                 AND ready.lane_seq = done.lane_seq
14288                WHERE done.state = 'failed'
14289                  AND done.finalized_at >= now() - make_interval(secs => $1::bigint)
14290                "#
14291            ))
14292            .bind(failed_retention_secs)
14293            .bind(state.0)
14294            .bind(current_generation)
14295            .execute(tx.as_mut())
14296            .await
14297            .map_err(map_sqlx_error)?
14298            .rows_affected();
14299
14300            if carried > 0 {
14301                // This slot's live counter rows and delta segment are
14302                // wiped below, so the carried rows' exact-count
14303                // evidence moves with them: re-append positive deltas
14304                // under the current slot, grouped exactly like the
14305                // completion path's delta append.
14306                sqlx::query(&format!(
14307                    r#"
14308                    INSERT INTO {schema}.queue_terminal_count_deltas (
14309                        ready_slot,
14310                        ready_generation,
14311                        queue,
14312                        priority,
14313                        enqueue_shard,
14314                        counter_bucket,
14315                        terminal_delta
14316                    )
14317                    SELECT
14318                        $2,
14319                        $3,
14320                        queue,
14321                        priority,
14322                        enqueue_shard,
14323                        mod(
14324                            mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint)
14325                                + {TERMINAL_COUNTER_BUCKETS}::bigint,
14326                            {TERMINAL_COUNTER_BUCKETS}::bigint
14327                        )::smallint AS counter_bucket,
14328                        count(*)::bigint
14329                    FROM {done_child}
14330                    WHERE state = 'failed'
14331                      AND finalized_at >= now() - make_interval(secs => $1::bigint)
14332                    GROUP BY queue, priority, enqueue_shard, counter_bucket
14333                    ORDER BY queue, priority, enqueue_shard, counter_bucket
14334                    "#
14335                ))
14336                .bind(failed_retention_secs)
14337                .bind(state.0)
14338                .bind(current_generation)
14339                .execute(tx.as_mut())
14340                .await
14341                .map_err(map_sqlx_error)?;
14342            }
14343            carried
14344        } else {
14345            0
14346        };
14347
14348        // #290: scan the about-to-be-truncated partition for the rollup
14349        // fold. The rollup column is *permanent* state, so we MUST fold
14350        // from ground truth (the `{done_child}` partition itself), not
14351        // from `queue_terminal_live_counts` — the counter can drift
14352        // briefly during a rolling upgrade from a pre-counter binary
14353        // and folding drift into the rollup would bake it in forever.
14354        // The counter rows for this slot are still cleaned up after the
14355        // truncate; reads from the counter (queue_counts_exact in #305)
14356        // may transiently disagree with the rollup until the operator
14357        // runs `awa storage rebuild-terminal-counters`, but the rollup
14358        // itself stays authoritative. See PR #304 reviewer finding
14359        // "High: A1 can persist stale counter state before the
14360        // read-switch PR" for the trade-off.
14361        //
14362        // Carried failed rows are excluded from both columns: they are
14363        // still live in `done_entries`, so folding them into the
14364        // permanent rollup would double-count them.
14365        let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = sqlx::query_as(&format!(
14366            r#"
14367            WITH done_counts AS (
14368                SELECT
14369                    queue,
14370                    priority,
14371                    (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed,
14372                    (count(*) FILTER (
14373                        WHERE state = 'failed'
14374                          AND (
14375                              $2::boolean IS FALSE
14376                              OR finalized_at < now() - make_interval(secs => $1::bigint)
14377                          )
14378                    ))::bigint AS failed
14379                FROM {done_child}
14380                GROUP BY queue, priority
14381            ),
14382            -- Scan the entire partition being truncated (all generations),
14383            -- mirroring done_counts' whole-child scan above. The rotate guard
14384            -- keeps a partition to one generation at prune time, so this is the
14385            -- same set today; counting the whole child keeps the rollup fold
14386            -- conservative even if that invariant is ever weakened (a stray
14387            -- generation's batches would otherwise be truncated without being
14388            -- folded into the rollup -> terminal-count undercount).
14389            batch_rows AS (
14390                SELECT
14391                    batch.ready_generation,
14392                    batch.queue,
14393                    batch.priority,
14394                    item.job_id,
14395                    item.run_lease
14396                FROM {receipt_batch_child} AS batch
14397                CROSS JOIN LATERAL unnest(
14398                    batch.job_ids,
14399                    batch.run_leases
14400                ) AS item(job_id, run_lease)
14401            ),
14402            batch_counts AS (
14403                SELECT
14404                    batch_rows.queue,
14405                    batch_rows.priority,
14406                    count(*)::bigint AS completed
14407                FROM batch_rows
14408                WHERE NOT EXISTS (
14409                    SELECT 1
14410                    FROM {receipt_tomb_child} AS tomb
14411                    WHERE tomb.ready_generation = batch_rows.ready_generation
14412                      AND tomb.job_id = batch_rows.job_id
14413                      AND tomb.run_lease = batch_rows.run_lease
14414                )
14415                GROUP BY batch_rows.queue, batch_rows.priority
14416            ),
14417            keys AS (
14418                SELECT queue, priority FROM done_counts
14419                UNION
14420                SELECT queue, priority FROM batch_counts
14421            )
14422            SELECT
14423                keys.queue,
14424                keys.priority,
14425                (
14426                    COALESCE(done_counts.completed, 0)
14427                    + COALESCE(batch_counts.completed, 0)
14428                )::bigint AS pruned_completed_count,
14429                COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count
14430            FROM keys
14431            LEFT JOIN done_counts USING (queue, priority)
14432            LEFT JOIN batch_counts USING (queue, priority)
14433            "#
14434        ))
14435        .bind(failed_retention_secs)
14436        .bind(retention_floor)
14437        .fetch_all(tx.as_mut())
14438        .await
14439        .map_err(map_sqlx_error)?;
14440
14441        let truncate = sqlx::query(&format!(
14442            "TRUNCATE TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child}"
14443        ))
14444        .execute(tx.as_mut())
14445        .await;
14446
14447        match truncate {
14448            Ok(_) => {
14449                if !pruned_terminal_counts.is_empty() {
14450                    self.adjust_terminal_rollups_batch(&mut tx, pruned_terminal_counts.into_iter())
14451                        .await?;
14452                }
14453
14454                // #290: the live counter rows for this slot are about to
14455                // be orphans (their underlying done_entries rows just got
14456                // truncated). Delete them in the same transaction. The
14457                // rollup fold above already captured ground-truth from
14458                // the partition scan; this just cleans up the counter
14459                // index entries so a future insert into a re-rotated
14460                // slot starts from zero.
14461                sqlx::query(&format!(
14462                    "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1"
14463                ))
14464                .bind(slot)
14465                .execute(tx.as_mut())
14466                .await
14467                .map_err(map_sqlx_error)?;
14468                tx.commit().await.map_err(map_sqlx_error)?;
14469                if carried_failed_rows > 0 {
14470                    tracing::info!(
14471                        slot,
14472                        carried_failed_rows,
14473                        "Carried failed terminal rows inside the retention floor to the live done segment"
14474                    );
14475                }
14476                Ok(PruneOutcome::Pruned {
14477                    slot,
14478                    carried_failed_rows,
14479                })
14480            }
14481            Err(err) if is_lock_contention_error(&err) => {
14482                let _ = tx.rollback().await;
14483                Ok(PruneOutcome::Blocked { slot })
14484            }
14485            Err(err) => {
14486                let _ = tx.rollback().await;
14487                Err(map_sqlx_error(err))
14488            }
14489        }
14490    }
14491
14492    #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest_leases")]
14493    pub async fn prune_oldest_leases(&self, pool: &PgPool) -> Result<PruneOutcome, AwaError> {
14494        let schema = self.schema();
14495        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14496
14497        // `PruneLeasesPlan` in
14498        // `correctness/storage/AwaStorageLockOrder.tla` requires the
14499        // sequence `lease_ring_state FOR UPDATE` →
14500        // `lease_ring_slots[slot] FOR UPDATE` → `ACCESS EXCLUSIVE` on
14501        // the child. The child lock is bounded by a short
14502        // transaction-local lock_timeout because pure NOWAIT can starve
14503        // under continuous parent-partition readers, while an unbounded
14504        // wait would put maintenance at the head of the relation-lock
14505        // queue indefinitely. Without these locks a concurrent
14506        // rotator can flip the cursor under the prune's liveness check
14507        // (current_slot recheck races a CAS update) and prune what
14508        // should be the active partition.
14509        let state: (i32, i64, i32) = sqlx::query_as(&format!(
14510            r#"
14511            SELECT current_slot, generation, slot_count
14512            FROM {schema}.lease_ring_state
14513            WHERE singleton = TRUE
14514            FOR UPDATE
14515            "#
14516        ))
14517        .fetch_one(tx.as_mut())
14518        .await
14519        .map_err(map_sqlx_error)?;
14520
14521        let Some((slot, _generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2)
14522        else {
14523            tx.commit().await.map_err(map_sqlx_error)?;
14524            return Ok(PruneOutcome::Noop);
14525        };
14526
14527        let slot_locked: Option<i32> = sqlx::query_scalar(&format!(
14528            r#"
14529            SELECT slot FROM {schema}.lease_ring_slots
14530            WHERE slot = $1
14531            FOR UPDATE
14532            "#
14533        ))
14534        .bind(slot)
14535        .fetch_optional(tx.as_mut())
14536        .await
14537        .map_err(map_sqlx_error)?;
14538
14539        if slot_locked.is_none() {
14540            tx.commit().await.map_err(map_sqlx_error)?;
14541            return Ok(PruneOutcome::Noop);
14542        }
14543
14544        let lease_child = lease_child_name(schema, slot as usize);
14545
14546        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14547
14548        let lock_table = sqlx::query(&format!(
14549            "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE"
14550        ))
14551        .execute(tx.as_mut())
14552        .await;
14553
14554        if let Err(err) = lock_table {
14555            let _ = tx.rollback().await;
14556            if is_lock_contention_error(&err) {
14557                return Ok(PruneOutcome::Blocked { slot });
14558            }
14559            return Err(map_sqlx_error(err));
14560        }
14561
14562        let current_slot: i32 = sqlx::query_scalar(&format!(
14563            r#"
14564            SELECT current_slot
14565            FROM {schema}.lease_ring_state
14566            WHERE singleton = TRUE
14567            "#
14568        ))
14569        .fetch_one(tx.as_mut())
14570        .await
14571        .map_err(map_sqlx_error)?;
14572
14573        if current_slot == slot {
14574            tx.commit().await.map_err(map_sqlx_error)?;
14575            return Ok(PruneOutcome::SkippedActive {
14576                slot,
14577                reason: SkipReason::LeaseCurrent,
14578                count: 0,
14579            });
14580        }
14581
14582        let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?;
14583
14584        if active_leases {
14585            tx.commit().await.map_err(map_sqlx_error)?;
14586            return Ok(PruneOutcome::SkippedActive {
14587                slot,
14588                reason: SkipReason::LeaseActive,
14589                count: busy_indicator(active_leases),
14590            });
14591        }
14592
14593        let truncate = sqlx::query(&format!("TRUNCATE TABLE {lease_child}"))
14594            .execute(tx.as_mut())
14595            .await;
14596
14597        match truncate {
14598            Ok(_) => {
14599                tx.commit().await.map_err(map_sqlx_error)?;
14600                Ok(PruneOutcome::Pruned {
14601                    slot,
14602                    carried_failed_rows: 0,
14603                })
14604            }
14605            Err(err) if is_lock_contention_error(&err) => {
14606                let _ = tx.rollback().await;
14607                Ok(PruneOutcome::Blocked { slot })
14608            }
14609            Err(err) => {
14610                let _ = tx.rollback().await;
14611                Err(map_sqlx_error(err))
14612            }
14613        }
14614    }
14615
14616    pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> {
14617        sqlx::query(&format!("VACUUM {}", self.leases_table()))
14618            .execute(pool)
14619            .await
14620            .map_err(map_sqlx_error)?;
14621        Ok(())
14622    }
14623
14624    /// ADR-023 claim-ring rotation. Parallel of `rotate_leases`.
14625    ///
14626    /// Advances `claim_ring_state.current_slot` via compare-and-swap. Before
14627    /// flipping the cursor the target partition must be drained: the
14628    /// `lease_claims_<next>`, `lease_claim_batches_<next>`,
14629    /// `lease_claim_closures_<next>`, and
14630    /// `lease_claim_closure_batches_<next>` child tables
14631    /// must be empty. This is what the `rotate → prune → rotate` ring
14632    /// invariant requires — we only hand out a slot to new claims when a
14633    /// prior prune has truncated it.
14634    #[tracing::instrument(skip(self, pool), name = "queue_storage.rotate_claims")]
14635    pub async fn rotate_claims(&self, pool: &PgPool) -> Result<RotateOutcome, AwaError> {
14636        let schema = self.schema();
14637        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14638
14639        let state: (i32, i64, i32) = sqlx::query_as(&format!(
14640            r#"
14641            SELECT current_slot, generation, slot_count
14642            FROM {schema}.claim_ring_state
14643            WHERE singleton = TRUE
14644            FOR UPDATE
14645            "#
14646        ))
14647        .fetch_one(tx.as_mut())
14648        .await
14649        .map_err(map_sqlx_error)?;
14650
14651        let next_slot = (state.0 + 1).rem_euclid(state.2);
14652
14653        // Busy check: all children of the incoming slot must be empty.
14654        // A non-empty `lease_claims_<next>` means the previous lap's
14655        // prune hasn't run (or didn't complete); rotating anyway would
14656        // mix fresh claims with legacy rows and defeat the point of
14657        // partitioning. Non-empty closure children mean prune fell behind
14658        // on receipt-closure evidence specifically.
14659        let claim_busy =
14660            Self::relation_has_rows_tx(&mut tx, &claim_child_name(schema, next_slot as usize))
14661                .await?;
14662        let claim_batch_busy = Self::relation_has_rows_tx(
14663            &mut tx,
14664            &claim_batch_child_name(schema, next_slot as usize),
14665        )
14666        .await?;
14667        let closure_busy =
14668            Self::relation_has_rows_tx(&mut tx, &closure_child_name(schema, next_slot as usize))
14669                .await?;
14670        let closure_batch_busy = Self::relation_has_rows_tx(
14671            &mut tx,
14672            &claim_closure_batch_child_name(schema, next_slot as usize),
14673        )
14674        .await?;
14675
14676        if claim_busy || claim_batch_busy || closure_busy || closure_batch_busy {
14677            tx.commit().await.map_err(map_sqlx_error)?;
14678            return Ok(RotateOutcome::SkippedBusy {
14679                slot: next_slot,
14680                busy: BusyCounts {
14681                    claims: busy_indicator(claim_busy || claim_batch_busy),
14682                    closures: busy_indicator(closure_busy),
14683                    closure_batches: busy_indicator(closure_batch_busy),
14684                    ..Default::default()
14685                },
14686            });
14687        }
14688
14689        let next_generation = state.1 + 1;
14690
14691        let rotated = sqlx::query(&format!(
14692            r#"
14693            UPDATE {schema}.claim_ring_state
14694            SET current_slot = $1,
14695                generation = $2
14696            WHERE singleton = TRUE
14697              AND current_slot = $3
14698              AND generation = $4
14699            "#
14700        ))
14701        .bind(next_slot)
14702        .bind(next_generation)
14703        .bind(state.0)
14704        .bind(state.1)
14705        .execute(tx.as_mut())
14706        .await
14707        .map_err(map_sqlx_error)?;
14708
14709        if rotated.rows_affected() == 0 {
14710            // Lost the CAS race. Report the bounded presence evidence we
14711            // sampled before giving up.
14712            tx.commit().await.map_err(map_sqlx_error)?;
14713            return Ok(RotateOutcome::SkippedBusy {
14714                slot: next_slot,
14715                busy: BusyCounts {
14716                    claims: busy_indicator(claim_busy || claim_batch_busy),
14717                    closures: busy_indicator(closure_busy),
14718                    closure_batches: busy_indicator(closure_batch_busy),
14719                    ..Default::default()
14720                },
14721            });
14722        }
14723
14724        tx.commit().await.map_err(map_sqlx_error)?;
14725        Ok(RotateOutcome::Rotated {
14726            slot: next_slot,
14727            generation: next_generation,
14728        })
14729    }
14730
14731    /// ADR-023 claim-ring prune. Parallel of `prune_oldest_leases`.
14732    ///
14733    /// Reclaims the oldest initialized (sealed) claim-ring slot by
14734    /// `TRUNCATE`-ing its `lease_claims_<slot>`,
14735    /// `lease_claim_batches_<slot>`, `lease_claim_closures_<slot>`, and
14736    /// `lease_claim_closure_batches_<slot>` children. Takes the full ADR-023
14737    /// lock sequence:
14738    ///
14739    /// 1. `FOR UPDATE` on `claim_ring_state` (serialises with rotate).
14740    /// 2. `FOR UPDATE` on the target `claim_ring_slots` row.
14741    /// 3. Proves every claim in the sealed partition has closure evidence.
14742    /// 4. `LOCK TABLE ACCESS EXCLUSIVE` on all claim-ring children with a
14743    ///    short transaction-local `lock_timeout` (serialises with in-flight
14744    ///    claim/complete/rescue writers; gives up under sustained
14745    ///    contention).
14746    /// 5. Rechecks the slot is not the current one and that every
14747    ///    claim still has closure evidence.
14748    /// 6. `TRUNCATE` all claim-ring children in a single statement.
14749    ///
14750    /// The "every claim has closure evidence" precondition is what ADR-023
14751    /// calls `PartitionTruncateSafety`. If an open claim remains in the
14752    /// partition, prune returns `SkippedActive` and the claim has to
14753    /// drain by normal completion or be rescued by
14754    /// `rescue_stale_receipt_claims_tx` before prune will try again.
14755    #[tracing::instrument(skip(self, pool), name = "queue_storage.prune_oldest_claims")]
14756    pub async fn prune_oldest_claims(&self, pool: &PgPool) -> Result<PruneOutcome, AwaError> {
14757        let schema = self.schema();
14758        let mut tx = pool.begin().await.map_err(map_sqlx_error)?;
14759
14760        let state: (i32, i64, i32) = sqlx::query_as(&format!(
14761            r#"
14762            SELECT current_slot, generation, slot_count
14763            FROM {schema}.claim_ring_state
14764            WHERE singleton = TRUE
14765            FOR UPDATE
14766            "#
14767        ))
14768        .fetch_one(tx.as_mut())
14769        .await
14770        .map_err(map_sqlx_error)?;
14771
14772        let Some((slot, _generation)) = oldest_initialized_ring_slot(state.0, state.1, state.2)
14773        else {
14774            tx.commit().await.map_err(map_sqlx_error)?;
14775            return Ok(PruneOutcome::Noop);
14776        };
14777
14778        // Lock the slot row so concurrent rotate/prune observe the same
14779        // state machine transition.
14780        let slot_locked: Option<i32> = sqlx::query_scalar(&format!(
14781            r#"
14782            SELECT slot FROM {schema}.claim_ring_slots
14783            WHERE slot = $1
14784            FOR UPDATE
14785            "#
14786        ))
14787        .bind(slot)
14788        .fetch_optional(tx.as_mut())
14789        .await
14790        .map_err(map_sqlx_error)?;
14791
14792        if slot_locked.is_none() {
14793            tx.commit().await.map_err(map_sqlx_error)?;
14794            return Ok(PruneOutcome::Noop);
14795        }
14796
14797        let claim_child = claim_child_name(schema, slot as usize);
14798        let claim_batch_child = claim_batch_child_name(schema, slot as usize);
14799        let closure_child = closure_child_name(schema, slot as usize);
14800        let closure_batch_child = claim_closure_batch_child_name(schema, slot as usize);
14801
14802        // Before taking ACCESS EXCLUSIVE on the child partitions, prove
14803        // the sealed slot is actually reclaimable. This optimistic proof
14804        // avoids blocking claim/complete traffic behind a prune attempt
14805        // that is already known to be doomed. A claimer that read the
14806        // old current slot before rotation may still commit while the
14807        // exclusive lock waits, so the open-claim proof is repeated once
14808        // the child locks are held.
14809        let open_claims = claim_prune_has_open_claims_tx(
14810            &mut tx,
14811            schema,
14812            &claim_child,
14813            &claim_batch_child,
14814            &closure_child,
14815            &closure_batch_child,
14816        )
14817        .await?;
14818
14819        if open_claims {
14820            tx.commit().await.map_err(map_sqlx_error)?;
14821            return Ok(PruneOutcome::SkippedActive {
14822                slot,
14823                reason: SkipReason::ClaimOpen,
14824                count: busy_indicator(open_claims),
14825            });
14826        }
14827
14828        set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?;
14829
14830        let lock_tables = sqlx::query(&format!(
14831            "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE"
14832        ))
14833        .execute(tx.as_mut())
14834        .await;
14835
14836        if let Err(err) = lock_tables {
14837            let _ = tx.rollback().await;
14838            if is_lock_contention_error(&err) {
14839                return Ok(PruneOutcome::Blocked { slot });
14840            }
14841            return Err(map_sqlx_error(err));
14842        }
14843
14844        // After taking ACCESS EXCLUSIVE, recheck that the slot is still
14845        // not the current one. The ring-state lock should already make
14846        // this stable, but keeping the explicit gate documents the
14847        // truncate precondition.
14848        let current_slot: i32 = sqlx::query_scalar(&format!(
14849            r#"
14850            SELECT current_slot FROM {schema}.claim_ring_state WHERE singleton = TRUE
14851            "#
14852        ))
14853        .fetch_one(tx.as_mut())
14854        .await
14855        .map_err(map_sqlx_error)?;
14856
14857        if current_slot == slot {
14858            tx.commit().await.map_err(map_sqlx_error)?;
14859            return Ok(PruneOutcome::SkippedActive {
14860                slot,
14861                reason: SkipReason::ClaimCurrent,
14862                count: 0,
14863            });
14864        }
14865
14866        let open_claims = claim_prune_has_open_claims_tx(
14867            &mut tx,
14868            schema,
14869            &claim_child,
14870            &claim_batch_child,
14871            &closure_child,
14872            &closure_batch_child,
14873        )
14874        .await?;
14875        if open_claims {
14876            tx.commit().await.map_err(map_sqlx_error)?;
14877            return Ok(PruneOutcome::SkippedActive {
14878                slot,
14879                reason: SkipReason::ClaimOpen,
14880                count: busy_indicator(open_claims),
14881            });
14882        }
14883
14884        let truncate = sqlx::query(&format!(
14885            "TRUNCATE TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child}"
14886        ))
14887        .execute(tx.as_mut())
14888        .await;
14889
14890        match truncate {
14891            Ok(_) => {
14892                sqlx::query(&format!(
14893                    r#"
14894                    UPDATE {schema}.claim_ring_slots
14895                    SET rescue_cursor_claimed_at = '-infinity'::timestamptz,
14896                        rescue_cursor_job_id = 0,
14897                        rescue_cursor_run_lease = 0,
14898                        deadline_cursor_deadline_at = '-infinity'::timestamptz,
14899                        deadline_cursor_job_id = 0,
14900                        deadline_cursor_run_lease = 0
14901                    WHERE slot = $1
14902                    "#
14903                ))
14904                .bind(slot)
14905                .execute(tx.as_mut())
14906                .await
14907                .map_err(map_sqlx_error)?;
14908                tx.commit().await.map_err(map_sqlx_error)?;
14909                Ok(PruneOutcome::Pruned {
14910                    slot,
14911                    carried_failed_rows: 0,
14912                })
14913            }
14914            Err(err) if is_lock_contention_error(&err) => {
14915                let _ = tx.rollback().await;
14916                Ok(PruneOutcome::Blocked { slot })
14917            }
14918            Err(err) => {
14919                let _ = tx.rollback().await;
14920                Err(map_sqlx_error(err))
14921            }
14922        }
14923    }
14924
14925    fn job_id_sequence(&self) -> String {
14926        format!("{}.job_id_seq", self.schema())
14927    }
14928
14929    fn leases_table(&self) -> String {
14930        format!("{}.{}", self.schema(), self.leases_relname())
14931    }
14932
14933    fn attempt_state_table(&self) -> String {
14934        format!("{}.{}", self.schema(), self.attempt_state_relname())
14935    }
14936}