diff --git a/CHANGELOG.md b/CHANGELOG.md index 065d50db1..d52ab19ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,75 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- Preserve recovered bodies and physical lane priority when manually moving + historical aged terminal failures to the DLQ, so later DLQ retry retains + arguments, attempt limits and priority provenance. Terminal counters are + decremented at their original stored priority; automatic rescue-to-DLQ + policy keeps the effective priority. + +- Recover retained terminal bodies from the exact prior attempt when older + workers stored an aged effective priority in a narrow terminal key. Live + moved-lane metadata supplies defaults without replacing stored terminal + metadata, errors, tags or progress. Manual terminal retry restarts aging from + the physical lane; counter removal keeps the stored terminal key and failed + carry derives replacement counters from its inserted rows. The installed + `terminal_jobs` SQL view is unchanged for older mismatched-priority rows; + bodies whose ready backing was already pruned cannot be reconstructed. +- Exclude tombstoned source lanes and require the prior run lease when + hydrating rescued attempts, retaining physical priority only for terminal + identity. Deferred retry and DLQ policy keep effective priority. The lookup + reads only validated referenced tombstone children; this adds read locks. + +- Preserve the original job body when manually retrying a priority-aged attempt + that exhausted its retries during rescue. Retained failed reads and failed-row + carry through pruning now use the physical ready-lane identity. + +- Close read connections if their transaction start is cancelled or fails before + SQLx acknowledges it, so later reads cannot inherit an untracked transaction. + +- **Point reads show the live queue and priority after a lane move.** Retained + ready rows from before a priority or queue change could win a read's tie + after the ring wrapped, returning the old fields even though the batch + operation succeeded. Point reads now exclude tombstoned source rows. + +- **Callback completion, `get_job` and lifecycle events find running jobs + that priority aging promoted** (the Rust side of + [#569](https://github.com/hardbyte/awa/issues/569)). Claims and leases + record a job's aged effective priority while its ready row keeps the lane + priority, and the job and admin reads joined them on priority. With aging on (the + default, every 60 s), a job that waited past a minute before it was claimed + could not be completed by callback with an expected kind, failed the + mixed-transition callback routing check, and read back from `get_job` in + the wrong state. Admin job lists, kind counts and cancellation by unique + key also missed those running jobs. They now match the attempt (`job_id` and `run_lease`), + skipping the tombstoned source row a lane move leaves behind. +- **`get_job` no longer misses a job that moves while it reads.** It read + ready, deferred, leased, claimed, completed and DLQ rows in separate + statements, so a job promoted, claimed or completed between two of them + could read as missing; the weekly mixed-version rehearsal hit this on + 2026-11-01. It now reads all of them from one read-only snapshot, after + locking the ready, lease and claim rings so prune cannot truncate a partition the + snapshot still needs. +- **Ring prune and terminal-count rollup report `Blocked` when any of their + statements times out on a lock.** Prune sets a short `lock_timeout` (10 ms + by default) before it takes `ACCESS EXCLUSIVE` on a slot's partitions, but + only those `LOCK TABLE` statements and the final `TRUNCATE` turned a + timeout into `PruneOutcome::Blocked`. The timeout covers every later + statement in the transaction too: the post-lock liveness checks read the + lease and claim partitions, the failed-row carry inserts into the current + slot's `done_entries` and terminal-count deltas, and the fold appends + `queue_terminal_rollup_deltas` and deletes `queue_terminal_live_counts` + rows. When one of those waited out the timeout, for example behind + autovacuum truncating a relation or a hosted Postgres's own background + workers, prune failed with a `lock_not_available` (`55P03`) database error, + was logged as an error and skipped the prune backoff. A lock timeout in any + statement after a queue, lease or claim prune, or a terminal-count rollup + slot, sets its `lock_timeout` now rolls the whole transaction back and + reports `Blocked` + (`blocked_slots` for rollup), so maintenance backs off and retries; nothing + the attempt wrote survives the rollback. A rollup pass stops at its first + blocked slot rather than waiting out the timeout again on each remaining + one. - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 71cd067ea..7be7913e0 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -376,9 +376,22 @@ pub fn queue_storage_current_jobs_cte(schema: &str) -> String { ON ready.ready_slot = leases.ready_slot AND ready.ready_generation = leases.ready_generation AND ready.queue = leases.queue - AND ready.priority = leases.priority + -- Leases keep the aged effective priority; match the ready attempt. + AND ready.job_id = leases.job_id + AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) UNION ALL -- Receipt claims that have not materialised into `leases` are -- running too (#246 / #416). They live on legacy row-local @@ -404,9 +417,22 @@ pub fn queue_storage_current_jobs_cte(schema: &str) -> String { ON ready.ready_slot = claims.ready_slot AND ready.ready_generation = claims.ready_generation AND ready.queue = claims.queue - AND ready.priority = claims.priority + -- Claims keep the aged effective priority; match the ready attempt. + AND ready.job_id = claims.job_id + AND ready.run_lease + 1 = claims.run_lease AND ready.enqueue_shard = claims.enqueue_shard AND ready.lane_seq = claims.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) WHERE claims.closed_at IS NULL AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS cx @@ -466,6 +492,16 @@ pub fn queue_storage_current_jobs_cte(schema: &str) -> String { AND ready.enqueue_shard = batches.enqueue_shard AND ready.lane_seq = items.lane_seq AND ready.job_id = items.job_id + AND ready.run_lease + 1 = items.run_lease + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) WHERE NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS cx WHERE cx.claim_slot = batches.claim_slot @@ -564,9 +600,22 @@ fn queue_storage_recent_jobs_cte(schema: &str, since: &str) -> String { ON ready.ready_slot = leases.ready_slot AND ready.ready_generation = leases.ready_generation AND ready.queue = leases.queue - AND ready.priority = leases.priority + -- Leases keep the aged effective priority; match the ready attempt. + AND ready.job_id = leases.job_id + AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) WHERE ready.created_at >= {since} UNION ALL -- A queue or priority move copies the row with its times, so @@ -685,9 +734,22 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { ON ready.ready_slot = leases.ready_slot AND ready.ready_generation = leases.ready_generation AND ready.queue = leases.queue - AND ready.priority = leases.priority + -- Leases keep the aged effective priority; match the ready attempt. + AND ready.job_id = leases.job_id + AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) GROUP BY ready.kind, leases.queue ), deferred_by_kind AS MATERIALIZED ( @@ -931,9 +993,22 @@ fn queue_storage_job_id_candidates_sql( ON ready.ready_slot = leases.ready_slot AND ready.ready_generation = leases.ready_generation AND ready.queue = leases.queue - AND ready.priority = leases.priority + -- Leases keep the aged effective priority; match the ready attempt. + AND ready.job_id = leases.job_id + AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) WHERE TRUE{predicates} ORDER BY leases.job_id DESC LIMIT {limit} @@ -1091,12 +1166,15 @@ async fn list_queue_storage_jobs( }; } // As in `queue_overviews`: compiling costs more than running. - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; let ids: Vec = query.fetch_all(tx.as_mut()).await?; tx.commit().await?; + drop(begin_guard); if ids.is_empty() { break; } @@ -1408,9 +1486,22 @@ fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { WHERE ready.ready_slot = leases.ready_slot AND ready.ready_generation = leases.ready_generation AND ready.queue = leases.queue - AND ready.priority = leases.priority + -- Leases keep the aged effective priority; match the ready attempt. + AND ready.job_id = leases.job_id + AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) AND ready.unique_key = $1 OFFSET 0 ) AS ready ON TRUE @@ -3141,7 +3232,9 @@ pub async fn queue_overviews(pool: &PgPool) -> Result, AwaErr // JIT compiles this statement for longer than it runs (about 1.2 s // against 0.35 s with 100k retained completions). - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; @@ -3149,6 +3242,7 @@ pub async fn queue_overviews(pool: &PgPool) -> Result, AwaErr .fetch_all(tx.as_mut()) .await?; tx.commit().await?; + drop(begin_guard); return Ok(rows); } @@ -3637,7 +3731,9 @@ pub async fn state_counts(pool: &PgPool) -> Result, AwaEr waiting_external, ): (i64, i64, i64, i64, i64, i64, i64, i64) = { // As in `queue_overviews`: compiling costs more than running. - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; @@ -3645,6 +3741,7 @@ pub async fn state_counts(pool: &PgPool) -> Result, AwaEr .fetch_one(tx.as_mut()) .await?; tx.commit().await?; + drop(begin_guard); counts }; @@ -3750,7 +3847,9 @@ pub async fn job_kind_overviews(pool: &PgPool) -> Result, A ); // As in `queue_overviews`: compiling costs more than running. - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; @@ -3758,6 +3857,7 @@ pub async fn job_kind_overviews(pool: &PgPool) -> Result, A .fetch_all(tx.as_mut()) .await?; tx.commit().await?; + drop(begin_guard); return Ok(rows); } @@ -4037,7 +4137,9 @@ pub async fn state_timeseries( ORDER BY bucket", queue_storage_recent_jobs_cte(store.schema(), "$1") ); - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; @@ -4055,6 +4157,7 @@ pub async fn state_timeseries( .fetch_all(tx.as_mut()) .await?; tx.commit().await?; + drop(begin_guard); return Ok(rows .into_iter() diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 66a5dd4ee..8edc4f099 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -15,6 +15,7 @@ pub mod migrations; pub mod output; pub mod partitioned_queue; pub mod queue_storage; +mod read_transaction; pub mod reschedule; pub mod storage; mod storage_footprint_sql; diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 339001705..4bf612398 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -713,6 +713,35 @@ async fn set_prune_lock_timeout_tx( Ok(()) } +/// Commit or roll back a maintenance transaction that has set the prune +/// `lock_timeout`. +/// +/// The timeout applies to every later statement in the transaction, not only +/// the explicit `LOCK TABLE`: an insert, update or read of another relation can +/// also queue behind a conflicting lock (autovacuum's heap truncation takes +/// `ACCESS EXCLUSIVE`, for example). A lock timeout anywhere therefore rolls the +/// whole transaction back and reports `blocked`, so the caller backs off and +/// retries instead of surfacing an error. +async fn settle_prune_tx( + tx: sqlx::Transaction<'_, Postgres>, + result: Result, + blocked: T, +) -> Result { + match result { + Ok(outcome) => { + tx.commit().await.map_err(map_sqlx_error)?; + Ok(outcome) + } + Err(err) => { + let _ = tx.rollback().await; + match err { + AwaError::Database(ref db_err) if is_lock_contention_error(db_err) => Ok(blocked), + err => Err(err), + } + } + } +} + fn validate_ident(ident: &str) -> Result<(), AwaError> { let mut chars = ident.chars(); match chars.next() { @@ -778,20 +807,84 @@ fn terminal_delta_child_name(schema: &str, slot: usize) -> String { } fn done_ready_join(schema: &str, done_alias: &str, ready_alias: &str) -> String { + done_ready_join_from( + done_alias, + ready_alias, + &format!("{schema}.ready_entries"), + &format!("{schema}.ready_tombstones"), + ) +} + +// Old writers could store effective priority in a narrow terminal row. Prefer +// the live body of that exact claimed attempt, while retaining same-lease, +// exact-lane bodies tombstoned by an available-job cancellation. Choose one +// body so moved source rows cannot multiply terminal mutations or counters. +fn done_ready_join_from( + done_alias: &str, + ready_alias: &str, + ready_relation: &str, + tombstone_relation: &str, +) -> String { format!( r#" - LEFT JOIN {schema}.ready_entries AS {ready_alias} - ON {ready_alias}.ready_slot = {done_alias}.ready_slot - AND {ready_alias}.ready_generation = {done_alias}.ready_generation - AND {ready_alias}.queue = {done_alias}.queue - AND {ready_alias}.priority = {done_alias}.priority - AND {ready_alias}.enqueue_shard = {done_alias}.enqueue_shard - AND {ready_alias}.lane_seq = {done_alias}.lane_seq + LEFT JOIN LATERAL ( + SELECT candidate.* + FROM ( + SELECT body.*, + ( + body.run_lease + 1 = {done_alias}.run_lease + AND NOT EXISTS ( + SELECT 1 FROM {tombstone_relation} AS tomb + WHERE tomb.ready_slot = body.ready_slot + AND tomb.ready_generation = body.ready_generation + AND tomb.queue = body.queue + AND tomb.priority = body.priority + AND tomb.enqueue_shard = body.enqueue_shard + AND tomb.lane_seq = body.lane_seq + AND tomb.job_id = body.job_id + ) + ) AS live_attempt_body + FROM {ready_relation} AS body + WHERE body.ready_slot = {done_alias}.ready_slot + AND body.ready_generation = {done_alias}.ready_generation + AND body.queue = {done_alias}.queue + AND body.enqueue_shard = {done_alias}.enqueue_shard + AND body.lane_seq = {done_alias}.lane_seq + AND body.job_id = {done_alias}.job_id + AND (body.run_lease = {done_alias}.run_lease + OR body.run_lease + 1 = {done_alias}.run_lease) + ) AS candidate + WHERE candidate.priority = {done_alias}.priority + OR candidate.live_attempt_body + ORDER BY candidate.live_attempt_body DESC, + (candidate.priority = {done_alias}.priority) DESC, + candidate.priority + LIMIT 1 + ) AS {ready_alias} ON TRUE "# ) } +fn done_payload_projection(done_alias: &str, ready_alias: &str) -> String { + // A narrow old terminal payload can lack metadata from its live moved + // body. Supply only metadata defaults; terminal errors/tags/progress and + // explicit metadata overrides remain authoritative. Keep default {} elision. + format!( + r#"CASE + WHEN {done_alias}.args IS NULL + AND {ready_alias}.live_attempt_body + AND COALESCE({ready_alias}.payload -> 'metadata', '{{}}'::jsonb) <> '{{}}'::jsonb + THEN COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) + || jsonb_build_object('metadata', + COALESCE({ready_alias}.payload -> 'metadata', '{{}}'::jsonb) + || COALESCE({done_alias}.payload -> 'metadata', '{{}}'::jsonb)) + ELSE COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) + END"# + ) +} + fn done_row_projection(done_alias: &str, ready_alias: &str) -> String { + let payload = done_payload_projection(done_alias, ready_alias); format!( r#" {done_alias}.ready_slot, @@ -813,7 +906,7 @@ fn done_row_projection(done_alias: &str, ready_alias: &str) -> String { COALESCE({done_alias}.created_at, {ready_alias}.created_at, {done_alias}.finalized_at) AS created_at, COALESCE({done_alias}.unique_key, {ready_alias}.unique_key) AS unique_key, COALESCE({done_alias}.unique_states, {ready_alias}.unique_states) AS unique_states, - COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) AS payload + {payload} AS payload "# ) } @@ -2935,6 +3028,15 @@ struct DoneJobRow { output: Option, } +// Keep the stored terminal priority for decrementing its old counter key. +// Manual retry and terminal-to-DLQ moves recover the retained physical lane. +#[derive(Debug, sqlx::FromRow)] +struct HydratedTerminalRow { + #[sqlx(flatten)] + terminal: DoneJobRow, + body_priority: i16, +} + impl DoneJobRow { fn into_job_row(self) -> Result { let payload = RuntimePayload::from_json(self.payload)?; @@ -3148,6 +3250,7 @@ struct ReadySnapshotRow { kind: String, queue: String, args: serde_json::Value, + priority: i16, lane_seq: i64, enqueue_shard: i16, run_at: DateTime, @@ -3179,6 +3282,7 @@ struct LeaseTransitionRow { args: serde_json::Value, state: JobState, priority: i16, + ready_priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, @@ -3208,7 +3312,9 @@ impl LeaseTransitionRow { queue: self.queue, args: self.args, state, - priority: self.priority, + // Terminal rows retain their ready-lane identity for body hydration. + // The effective attempt priority is still used by deferred/DLQ rows. + priority: self.ready_priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, @@ -9331,7 +9437,7 @@ impl QueueStorage { let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); - let terminal: Option = sqlx::query_as(audited_sql(format!( + let terminal: Option = sqlx::query_as(audited_sql(format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries @@ -9346,7 +9452,7 @@ impl QueueStorage { ) RETURNING * ) - SELECT {done_projection} + SELECT {done_projection}, COALESCE(ready.priority, done.priority) AS body_priority FROM deleted AS done {ready_join} "# @@ -9356,7 +9462,9 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - if let Some(terminal) = terminal { + if let Some(hydrated) = terminal { + let body_priority = hydrated.body_priority; + let terminal = hydrated.terminal; self.ensure_terminal_removed_receipt_closures_tx(tx, std::slice::from_ref(&terminal)) .await?; // The DELETE FROM done_entries above removes one terminal row; @@ -9371,7 +9479,7 @@ impl QueueStorage { kind: terminal.kind, queue: terminal.queue.clone(), args: terminal.args, - priority: terminal.priority, + priority: body_priority, attempt: 0, run_lease: terminal.run_lease, max_attempts: terminal.max_attempts, @@ -12304,10 +12412,21 @@ impl QueueStorage { let job_ids: Vec = deleted.iter().map(|row| row.job_id).collect(); let run_leases: Vec = deleted.iter().map(|row| row.run_lease).collect(); + // Resolve moved ready bodies without opening the whole tombstone ring. + // Empty input returned above; every remaining slot is validated here. + let tombstone_sources = ready_slots + .iter() + .map(|&slot| ring_slot_index(slot, self.queue_slot_count(), "queue")) + .collect::, _>>()? + .into_iter() + .map(|slot| format!("SELECT * FROM {}", ready_tombstone_child_name(schema, slot))) + .collect::>() + .join(" UNION ALL "); + let ready_rows: Vec = sqlx::query_as(audited_sql(format!( r#" - WITH refs(ready_slot, ready_generation, queue, enqueue_shard, lane_seq, job_id) AS ( - SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::bigint[], $6::bigint[]) + WITH refs(ready_slot, ready_generation, queue, enqueue_shard, lane_seq, job_id, run_lease) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::bigint[], $6::bigint[], $7::bigint[]) ) SELECT ready.ready_slot, @@ -12316,6 +12435,7 @@ impl QueueStorage { ready.kind, ready.queue, ready.args, + ready.priority, ready.lane_seq, ready.enqueue_shard, ready.run_at, @@ -12331,6 +12451,17 @@ impl QueueStorage { AND ready.enqueue_shard = refs.enqueue_shard AND ready.lane_seq = refs.lane_seq AND ready.job_id = refs.job_id + AND ready.run_lease + 1 = refs.run_lease + WHERE NOT EXISTS ( + SELECT 1 FROM ({tombstone_sources}) AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + AND tomb.job_id = ready.job_id + ) "# ))) .bind(&ready_slots) @@ -12339,6 +12470,7 @@ impl QueueStorage { .bind(&enqueue_shards) .bind(&lane_seqs) .bind(&job_ids) + .bind(&run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; @@ -12539,6 +12671,7 @@ impl QueueStorage { args: ready.args.clone(), state: deleted_row.state, priority: deleted_row.priority, + ready_priority: ready.priority, attempt: deleted_row.attempt, run_lease: deleted_row.run_lease, max_attempts: deleted_row.max_attempts, @@ -14251,6 +14384,37 @@ impl QueueStorage { let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let mut candidates = Vec::new(); + // A job moves between these relations as it is promoted, claimed, + // completed or dead-lettered. Read them all from one snapshot so a + // move between two reads cannot hide the job or surface a stale row. + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool) + .await + .map_err(map_sqlx_error)?; + let mut tx = sqlx::Connection::begin_with( + &mut *begin_guard.connection, + "BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY", + ) + .await + .map_err(map_sqlx_error)?; + begin_guard.begin_pending = false; + // TRUNCATE is not MVCC-safe: a partition that prune empties after the + // snapshot reads as empty, while the completion that let prune run is + // newer than the snapshot, so the job could read as available or + // missing. A query takes its snapshot before it locks the relations it + // reads, so lock every ring prune can truncate first; LOCK TABLE takes + // no snapshot. Queue prune takes the ready child with NOWAIT before its + // done and receipt-completion children, so holding `ready_entries` + // covers them and that prune reports `Blocked`. Lease and claim prune + // wait or time out to `Blocked`; the claim tables are listed in claim + // prune's order, so the two cannot deadlock. + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.ready_entries, {schema}.leases, {schema}.lease_claims, \ + {schema}.lease_claim_batches, {schema}.lease_claim_closures, \ + {schema}.lease_claim_closure_batches IN ACCESS SHARE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; let ready_rows: Vec = sqlx::query_as(audited_sql(format!( r#" @@ -14269,13 +14433,22 @@ impl QueueStorage { unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries - WHERE job_id = $1 + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC "#, ))) .bind(job_id) - .fetch_all(pool) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; for row in ready_rows { @@ -14306,7 +14479,7 @@ impl QueueStorage { "#, ))) .bind(job_id) - .fetch_all(pool) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; for row in deferred_rows { @@ -14350,9 +14523,23 @@ impl QueueStorage { ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue - AND ready.priority = lease.priority + -- Claims and leases record the aged effective priority while the + -- ready row keeps its lane priority, so match the attempt instead. + AND ready.job_id = lease.job_id + AND ready.run_lease + 1 = lease.run_lease AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease @@ -14361,7 +14548,7 @@ impl QueueStorage { "#, ))) .bind(job_id) - .fetch_all(pool) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; for row in lease_rows { @@ -14407,9 +14594,23 @@ impl QueueStorage { ON ready.ready_slot = claims.ready_slot AND ready.ready_generation = claims.ready_generation AND ready.queue = claims.queue - AND ready.priority = claims.priority + -- Claims and leases record the aged effective priority while the + -- ready row keeps its lane priority, so match the attempt instead. + AND ready.job_id = claims.job_id + AND ready.run_lease + 1 = claims.run_lease AND ready.enqueue_shard = claims.enqueue_shard AND ready.lane_seq = claims.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = claims.job_id AND attempt.run_lease = claims.run_lease @@ -14427,7 +14628,7 @@ impl QueueStorage { "#, ))) .bind(job_id) - .fetch_all(pool) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; for row in lease_claim_rows { @@ -14486,6 +14687,16 @@ impl QueueStorage { AND ready.enqueue_shard = claim_batches.enqueue_shard AND ready.lane_seq = items.lane_seq AND ready.job_id = items.job_id + AND ready.run_lease + 1 = items.run_lease + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = items.job_id AND attempt.run_lease = items.run_lease @@ -14528,7 +14739,7 @@ impl QueueStorage { "#, ))) .bind(job_id) - .fetch_all(pool) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; for row in lease_claim_batch_rows { @@ -14564,7 +14775,7 @@ impl QueueStorage { "#, ))) .bind(job_id) - .fetch_all(pool) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; for row in done_rows { @@ -14599,13 +14810,15 @@ impl QueueStorage { "#, ))) .bind(job_id) - .fetch_all(pool) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; for row in dlq_rows { candidates.push(row.into_job_row()?); } + tx.commit().await.map_err(map_sqlx_error)?; + drop(begin_guard); Ok(candidates.into_iter().max_by_key(|job| { ( job.run_lease, @@ -14931,9 +15144,23 @@ impl QueueStorage { ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue - AND ready.priority = lease.priority + -- Claims and leases record the aged effective priority while the + -- ready row keeps its lane priority, so match the attempt instead. + AND ready.job_id = lease.job_id + AND ready.run_lease + 1 = lease.run_lease AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease @@ -15127,9 +15354,23 @@ impl QueueStorage { ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue - AND ready.priority = lease.priority + -- Claims and leases record the aged effective priority while the + -- ready row keeps its lane priority, so match the attempt instead. + AND ready.job_id = lease.job_id + AND ready.run_lease + 1 = lease.run_lease AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease @@ -16564,7 +16805,7 @@ impl QueueStorage { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); - let moved: Option = sqlx::query_as(audited_sql(format!( + let moved: Option = sqlx::query_as(audited_sql(format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries @@ -16579,7 +16820,7 @@ impl QueueStorage { ) RETURNING * ) - SELECT {done_projection} + SELECT {done_projection}, COALESCE(ready.priority, done.priority) AS body_priority FROM deleted AS done {ready_join} "# @@ -16589,10 +16830,11 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - let Some(moved) = moved else { + let Some(hydrated) = moved else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(None); }; + let mut moved = hydrated.terminal; self.ensure_terminal_removed_receipt_closures_tx(&mut tx, std::slice::from_ref(&moved)) .await?; @@ -16603,6 +16845,7 @@ impl QueueStorage { &Self::done_rows_to_counter_keys(std::slice::from_ref(&moved)), ) .await?; + moved.priority = hydrated.body_priority; let dlq_row = moved .clone() .into_dlq_row(dlq_reason.to_string(), Utc::now()); @@ -16628,7 +16871,7 @@ impl QueueStorage { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); - let moved: Vec = sqlx::query_as(audited_sql(format!( + let moved: Vec = sqlx::query_as(audited_sql(format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries @@ -16637,7 +16880,7 @@ impl QueueStorage { AND ($2::text IS NULL OR queue = $2) RETURNING * ) - SELECT {done_projection} + SELECT {done_projection}, COALESCE(ready.priority, done.priority) AS body_priority FROM deleted AS done {ready_join} "# @@ -16653,6 +16896,10 @@ impl QueueStorage { return Ok(0); } + let (moved, body_priorities): (Vec<_>, Vec<_>) = moved + .into_iter() + .map(|row| (row.terminal, row.body_priority)) + .unzip(); self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &moved) .await?; // Bulk DLQ move deletes from done_entries. Append negative deltas by @@ -16662,7 +16909,11 @@ impl QueueStorage { let dlq_at = Utc::now(); let rows: Vec = moved .into_iter() - .map(|row| row.into_dlq_row(dlq_reason.to_string(), dlq_at)) + .zip(body_priorities) + .map(|(mut row, priority)| { + row.priority = priority; + row.into_dlq_row(dlq_reason.to_string(), dlq_at) + }) .collect(); self.insert_dlq_rows_tx(&mut tx, &rows, Some(JobState::Failed)) .await?; @@ -18163,7 +18414,10 @@ impl QueueStorage { break; } TerminalDeltaSlotRollup::Blocked => { + // The next slot would most likely wait out the same lock + // timeout, and the leader loop is blocked while it does. outcome.blocked_slots += 1; + break; } } } @@ -18350,237 +18604,215 @@ impl QueueStorage { return Ok(TerminalDeltaSlotRollup::SkippedActive); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - - let lock_delta = sqlx::query(audited_sql(format!( - "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - - match lock_delta { - Ok(_) => {} - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - return Ok(TerminalDeltaSlotRollup::Blocked); - } - Err(err) => { - let _ = tx.rollback().await; - return Err(map_sqlx_error(err)); - } - } - - let active_leases = - queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::SkippedActive); - } + sqlx::query(audited_sql(format!( + "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - let unclosed_claim_refs = - queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; + let active_leases = + queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; - if unclosed_claim_refs { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::SkippedActive); - } + if active_leases { + return Ok(TerminalDeltaSlotRollup::SkippedActive); + } - let delta_rows: i64 = sqlx::query_scalar(audited_sql(format!( - r#" - SELECT count(*)::bigint - FROM {delta_child} - WHERE ready_generation = $1 - "# - ))) - .bind(generation) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + let unclosed_claim_refs = + queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; - if delta_rows == 0 { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::Empty); - } + if unclosed_claim_refs { + return Ok(TerminalDeltaSlotRollup::SkippedActive); + } - // In ledger authority each touched lane collapses to one bucket-0 row - // holding its old total (summed over every bucket) plus the signed - // delta, which keeps a negative delta exact when the matching positive - // count sits in another bucket. Before the flip, N-1 maintenance can - // still win the election and folds per job-id bucket, so the fold - // keeps that shape until the flip fences it (ADR-041). - let grouped_keys: i64 = match authority { - RingAuthority::Ledger => sqlx::query_scalar(audited_sql(format!( - r#" - WITH grouped AS MATERIALIZED ( - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - SUM(terminal_delta)::bigint AS delta - FROM {delta_child} - WHERE ready_generation = $1 - GROUP BY ready_slot, queue, priority, enqueue_shard - HAVING SUM(terminal_delta) <> 0 - ), - collapsed AS ( - DELETE FROM {schema}.queue_terminal_live_counts AS counts - USING grouped - WHERE counts.ready_slot = grouped.ready_slot - AND counts.queue = grouped.queue - AND counts.priority = grouped.priority - AND counts.enqueue_shard = grouped.enqueue_shard - RETURNING - counts.ready_slot, - counts.queue, - counts.priority, - counts.enqueue_shard, - counts.live_terminal_count - ), - folded AS ( - SELECT - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - (grouped.delta + COALESCE(SUM(collapsed.live_terminal_count), 0))::bigint - AS live_terminal_count - FROM grouped - LEFT JOIN collapsed - ON collapsed.ready_slot = grouped.ready_slot - AND collapsed.queue = grouped.queue - AND collapsed.priority = grouped.priority - AND collapsed.enqueue_shard = grouped.enqueue_shard - GROUP BY - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - grouped.delta - ), - inserted AS ( - INSERT INTO {schema}.queue_terminal_live_counts ( - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - live_terminal_count - ) - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - 0::smallint, - live_terminal_count - FROM folded - WHERE live_terminal_count <> 0 - ORDER BY ready_slot, queue, priority, enqueue_shard - RETURNING 1 - ) - SELECT count(*)::bigint FROM grouped - "# - ))) - .bind(generation) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?, - RingAuthority::Columns => sqlx::query_scalar(audited_sql(format!( + let delta_rows: i64 = sqlx::query_scalar(audited_sql(format!( r#" - WITH grouped AS MATERIALIZED ( - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - SUM(terminal_delta)::bigint AS delta + SELECT count(*)::bigint FROM {delta_child} WHERE ready_generation = $1 - GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket - HAVING SUM(terminal_delta) <> 0 - ), - updated AS ( - UPDATE {schema}.queue_terminal_live_counts AS counts - SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta) - FROM grouped - WHERE counts.ready_slot = grouped.ready_slot - AND counts.queue = grouped.queue - AND counts.priority = grouped.priority - AND counts.enqueue_shard = grouped.enqueue_shard - AND counts.counter_bucket = grouped.counter_bucket - RETURNING - counts.ready_slot, - counts.queue, - counts.priority, - counts.enqueue_shard, - counts.counter_bucket - ), - inserted AS ( - INSERT INTO {schema}.queue_terminal_live_counts AS counts ( - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - live_terminal_count - ) - SELECT - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - grouped.counter_bucket, - grouped.delta - FROM grouped - WHERE grouped.delta > 0 - AND NOT EXISTS ( - SELECT 1 - FROM updated - WHERE updated.ready_slot = grouped.ready_slot - AND updated.queue = grouped.queue - AND updated.priority = grouped.priority - AND updated.enqueue_shard = grouped.enqueue_shard - AND updated.counter_bucket = grouped.counter_bucket - ) - ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket - ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE - SET live_terminal_count = - GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count) - RETURNING 1 - ) - SELECT count(*)::bigint FROM grouped - "# + "# ))) .bind(generation) .fetch_one(tx.as_mut()) .await - .map_err(map_sqlx_error)?, - }; - - let truncate_delta = sqlx::query(audited_sql(format!("TRUNCATE TABLE {delta_child}"))) - .execute(tx.as_mut()) - .await; + .map_err(map_sqlx_error)?; - match truncate_delta { - Ok(_) => { - tx.commit().await.map_err(map_sqlx_error)?; - Ok(TerminalDeltaSlotRollup::Rolled { - delta_rows, - grouped_keys, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(TerminalDeltaSlotRollup::Blocked) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) + if delta_rows == 0 { + return Ok(TerminalDeltaSlotRollup::Empty); } + + // In ledger authority each touched lane collapses to one bucket-0 row + // holding its old total (summed over every bucket) plus the signed + // delta, which keeps a negative delta exact when the matching positive + // count sits in another bucket. Before the flip, N-1 maintenance can + // still win the election and folds per job-id bucket, so the fold + // keeps that shape until the flip fences it (ADR-041). + let grouped_keys: i64 = match authority { + RingAuthority::Ledger => sqlx::query_scalar(audited_sql(format!( + r#" + WITH grouped AS MATERIALIZED ( + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + SUM(terminal_delta)::bigint AS delta + FROM {delta_child} + WHERE ready_generation = $1 + GROUP BY ready_slot, queue, priority, enqueue_shard + HAVING SUM(terminal_delta) <> 0 + ), + collapsed AS ( + DELETE FROM {schema}.queue_terminal_live_counts AS counts + USING grouped + WHERE counts.ready_slot = grouped.ready_slot + AND counts.queue = grouped.queue + AND counts.priority = grouped.priority + AND counts.enqueue_shard = grouped.enqueue_shard + RETURNING + counts.ready_slot, + counts.queue, + counts.priority, + counts.enqueue_shard, + counts.live_terminal_count + ), + folded AS ( + SELECT + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + (grouped.delta + COALESCE(SUM(collapsed.live_terminal_count), 0))::bigint + AS live_terminal_count + FROM grouped + LEFT JOIN collapsed + ON collapsed.ready_slot = grouped.ready_slot + AND collapsed.queue = grouped.queue + AND collapsed.priority = grouped.priority + AND collapsed.enqueue_shard = grouped.enqueue_shard + GROUP BY + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + grouped.delta + ), + inserted AS ( + INSERT INTO {schema}.queue_terminal_live_counts ( + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + live_terminal_count + ) + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + 0::smallint, + live_terminal_count + FROM folded + WHERE live_terminal_count <> 0 + ORDER BY ready_slot, queue, priority, enqueue_shard + RETURNING 1 + ) + SELECT count(*)::bigint FROM grouped + "# + ))) + .bind(generation) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?, + RingAuthority::Columns => sqlx::query_scalar(audited_sql(format!( + r#" + WITH grouped AS MATERIALIZED ( + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + SUM(terminal_delta)::bigint AS delta + FROM {delta_child} + WHERE ready_generation = $1 + GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket + HAVING SUM(terminal_delta) <> 0 + ), + updated AS ( + UPDATE {schema}.queue_terminal_live_counts AS counts + SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta) + FROM grouped + WHERE counts.ready_slot = grouped.ready_slot + AND counts.queue = grouped.queue + AND counts.priority = grouped.priority + AND counts.enqueue_shard = grouped.enqueue_shard + AND counts.counter_bucket = grouped.counter_bucket + RETURNING + counts.ready_slot, + counts.queue, + counts.priority, + counts.enqueue_shard, + counts.counter_bucket + ), + inserted AS ( + INSERT INTO {schema}.queue_terminal_live_counts AS counts ( + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + live_terminal_count + ) + SELECT + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + grouped.counter_bucket, + grouped.delta + FROM grouped + WHERE grouped.delta > 0 + AND NOT EXISTS ( + SELECT 1 + FROM updated + WHERE updated.ready_slot = grouped.ready_slot + AND updated.queue = grouped.queue + AND updated.priority = grouped.priority + AND updated.enqueue_shard = grouped.enqueue_shard + AND updated.counter_bucket = grouped.counter_bucket + ) + ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket + ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE + SET live_terminal_count = + GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count) + RETURNING 1 + ) + SELECT count(*)::bigint FROM grouped + "# + ))) + .bind(generation) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?, + }; + + sqlx::query(audited_sql(format!("TRUNCATE TABLE {delta_child}"))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(TerminalDeltaSlotRollup::Rolled { + delta_rows, + grouped_keys, + }) } + .await; + + settle_prune_tx(tx, locked, TerminalDeltaSlotRollup::Blocked).await } /// Fold pending `queue_terminal_rollup_deltas` rows into the @@ -18830,158 +19062,120 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - // Claims read segments before ready rows, while producers write in - // the opposite order. Queue behind segment readers, then abandon - // the transaction rather than wait with an exclusive child held. - let lock_tables = async { - sqlx::query(audited_sql(format!( - "LOCK TABLE {segment_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await?; - sqlx::query(audited_sql(format!( - "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE NOWAIT" - ))) - .execute(tx.as_mut()) - .await - } - .await; - let lock_duration = lock_started.elapsed(); - - if let Err(err) = lock_tables { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); + let lock_started = Instant::now(); + // Claims read segments before ready rows, while producers write in + // the opposite order. Queue behind segment readers, then abandon + // the transaction rather than wait with an exclusive child held. + async { + sqlx::query(audited_sql(format!( + "LOCK TABLE {segment_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await?; + sqlx::query(audited_sql(format!( + "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE NOWAIT" + ))) + .execute(tx.as_mut()) + .await + } + .await + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); + + // The pre-lock gates are only an MVCC snapshot. If an in-flight + // producer or claimer committed while the exclusive partition + // lock waited, these post-lock gates see it before TRUNCATE. + let active_leases = + queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; + if active_leases { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueueActiveLeases, + count: busy_indicator(active_leases), + }); } - return Err(map_sqlx_error(err)); - } - - // The pre-lock gates are only an MVCC snapshot. If an in-flight - // producer or claimer committed while the exclusive partition - // lock waited, these post-lock gates see it before TRUNCATE. - let active_leases = - queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueueActiveLeases, - count: busy_indicator(active_leases), - }); - } - let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; - if pending { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueuePendingReady, - count: busy_indicator(pending), - }); - } + let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; + if pending { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueuePendingReady, + count: busy_indicator(pending), + }); + } - let unclosed_claim_refs = - queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; - if unclosed_claim_refs { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueueUnclosedClaimRefs, - count: busy_indicator(unclosed_claim_refs), - }); - } + let unclosed_claim_refs = + queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; + if unclosed_claim_refs { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueueUnclosedClaimRefs, + count: busy_indicator(unclosed_claim_refs), + }); + } - let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); - let retention_floor = failed_retention_secs > 0; - - // #337: failed terminal rows still inside the retention floor - // are re-homed into the live done segment before the TRUNCATE - // so they stay retryable. The ring rotation lock is held above, - // which serializes against rotate — the current slot (and its - // generation, read from the ledger cursor) cannot move while - // carried rows are re-inserted into it. - let carried_failed_rows = if retention_floor { - // Carried rows are widened to the self-contained synthetic - // done-row shape (the COALESCE chain mirrors - // `done_row_projection`) because the ready rows they would - // otherwise hydrate from are truncated with this slot. - // Deliberately no ON CONFLICT: lane_seq comes from - // never-resetting sequences, so a PK collision means - // corrupted terminal state and must abort the prune loudly - // rather than silently drop a terminal fact. - let carried = sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.done_entries ( - ready_slot, ready_generation, job_id, kind, queue, - args, state, priority, attempt, run_lease, - max_attempts, lane_seq, enqueue_shard, run_at, - attempted_at, finalized_at, created_at, unique_key, - unique_states, payload - ) - SELECT - $2, - $3, - done.job_id, - done.kind, - done.queue, - COALESCE(done.args, ready.args, '{{}}'::jsonb), - done.state, - done.priority, - done.attempt, - done.run_lease, - COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), - done.lane_seq, - done.enqueue_shard, - COALESCE(done.run_at, ready.run_at, done.finalized_at), - COALESCE(done.attempted_at, ready.attempted_at), - done.finalized_at, - COALESCE(done.created_at, ready.created_at, done.finalized_at), - COALESCE(done.unique_key, ready.unique_key), - COALESCE(done.unique_states, ready.unique_states), - COALESCE(done.payload, ready.payload, '{{}}'::jsonb) - FROM {done_child} AS done - LEFT JOIN {ready_child} AS ready - ON ready.ready_slot = done.ready_slot - AND ready.ready_generation = done.ready_generation - AND ready.queue = done.queue - AND ready.priority = done.priority - AND ready.enqueue_shard = done.enqueue_shard - AND ready.lane_seq = done.lane_seq - WHERE done.state = 'failed' - AND done.finalized_at >= now() - make_interval(secs => $1::bigint) - "# - ))) - .bind(failed_retention_secs) - .bind(state.0) - .bind(current_generation) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)? - .rows_affected(); - - if carried > 0 { - // This slot's live counter rows and delta segment are - // wiped below, so the carried rows' exact-count - // evidence moves with them: re-append positive deltas - // under the current slot, grouped exactly like the - // completion path's delta append. - sqlx::query(audited_sql(format!( + let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); + let retention_floor = failed_retention_secs > 0; + + // #337: failed terminal rows still inside the retention floor + // are re-homed into the live done segment before the TRUNCATE + // so they stay retryable. The ring rotation lock is held above, + // which serializes against rotate — the current slot (and its + // generation, read from the ledger cursor) cannot move while + // carried rows are re-inserted into it. + let carried_failed_rows = if retention_floor { + // Carried rows are widened to the self-contained synthetic + // done-row shape (the COALESCE chain mirrors + // `done_row_projection`) because the ready rows they would + // otherwise hydrate from are truncated with this slot. + // Deliberately no ON CONFLICT: lane_seq comes from + // never-resetting sequences, so a PK collision means + // corrupted terminal state and must abort the prune loudly + // rather than silently drop a terminal fact. + let carry_ready_join = + done_ready_join_from("done", "ready", &ready_child, &tomb_child); + let carry_payload = done_payload_projection("done", "ready"); + let groups: Vec<(String, i16, i16, i16, i64)> = sqlx::query_as(audited_sql(format!( r#" - INSERT INTO {schema}.queue_terminal_count_deltas ( - ready_slot, - ready_generation, - queue, - priority, - enqueue_shard, - counter_bucket, - terminal_delta + WITH carried AS ( + INSERT INTO {schema}.done_entries ( + ready_slot, ready_generation, job_id, kind, queue, + args, state, priority, attempt, run_lease, + max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, finalized_at, created_at, unique_key, + unique_states, payload ) SELECT $2, $3, + done.job_id, + done.kind, + done.queue, + COALESCE(done.args, ready.args, '{{}}'::jsonb), + done.state, + COALESCE(ready.priority, done.priority), + done.attempt, + done.run_lease, + COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), + done.lane_seq, + done.enqueue_shard, + COALESCE(done.run_at, ready.run_at, done.finalized_at), + COALESCE(done.attempted_at, ready.attempted_at), + done.finalized_at, + COALESCE(done.created_at, ready.created_at, done.finalized_at), + COALESCE(done.unique_key, ready.unique_key), + COALESCE(done.unique_states, ready.unique_states), + {carry_payload} + FROM {done_child} AS done + {carry_ready_join} + WHERE done.state = 'failed' + AND done.finalized_at >= now() - make_interval(secs => $1::bigint) + RETURNING queue, priority, enqueue_shard, job_id + ) + SELECT queue, priority, enqueue_shard, @@ -18991,9 +19185,7 @@ impl QueueStorage { {TERMINAL_COUNTER_BUCKETS}::bigint )::smallint AS counter_bucket, count(*)::bigint - FROM {done_child} - WHERE state = 'failed' - AND finalized_at >= now() - make_interval(secs => $1::bigint) + FROM carried GROUP BY queue, priority, enqueue_shard, counter_bucket ORDER BY queue, priority, enqueue_shard, counter_bucket "# @@ -19001,189 +19193,200 @@ impl QueueStorage { .bind(failed_retention_secs) .bind(state.0) .bind(current_generation) - .execute(tx.as_mut()) + .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; - } - carried - } else { - 0 - }; - // #290: scan the about-to-be-truncated partition for the rollup - // fold. The rollup column is *permanent* state, so we MUST fold - // from ground truth (the `{done_child}` partition itself), not - // from `queue_terminal_live_counts` — the counter can drift - // briefly during a rolling upgrade from a pre-counter binary - // and folding drift into the rollup would bake it in forever. - // The counter rows for this slot are still cleaned up after the - // truncate; reads from the counter (queue_counts_exact in #305) - // may transiently disagree with the rollup until the operator - // runs `awa storage rebuild-terminal-counters`, but the rollup - // itself stays authoritative. See PR #304 reviewer finding - // "High: A1 can persist stale counter state before the - // read-switch PR" for the trade-off. - // - // Carried failed rows are excluded from both columns: they are - // still live in `done_entries`, so folding them into the - // permanent rollup would double-count them. - let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = - sqlx::query_as(audited_sql(format!( - r#" - WITH done_counts AS ( - SELECT - queue, - priority, - (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed, - (count(*) FILTER ( - WHERE state = 'failed' - AND ( - $2::boolean IS FALSE - OR finalized_at < now() - make_interval(secs => $1::bigint) - ) - ))::bigint AS failed - FROM {done_child} - GROUP BY queue, priority - ), - -- Scan the entire partition being truncated (all generations), - -- mirroring done_counts' whole-child scan above. The rotate guard - -- keeps a partition to one generation at prune time, so this is the - -- same set today; counting the whole child keeps the rollup fold - -- conservative even if that invariant is ever weakened (a stray - -- generation's batches would otherwise be truncated without being - -- folded into the rollup -> terminal-count undercount). - batch_rows AS ( - SELECT - batch.ready_generation, - batch.queue, - batch.priority, - item.job_id, - item.run_lease - FROM {receipt_batch_child} AS batch - CROSS JOIN LATERAL unnest( - batch.job_ids, - batch.run_leases - ) AS item(job_id, run_lease) - ), - batch_counts AS ( + // The replacement delta uses each inserted row's recovered + // priority. The old slot's counter evidence is truncated below. + let carried: u64 = groups.iter().map(|(_, _, _, _, count)| *count as u64).sum(); + let deltas = groups + .into_iter() + .map(|(queue, priority, shard, bucket, count)| { + ((state.0, current_generation, queue, priority, shard, bucket), count) + }) + .collect(); + self.append_terminal_count_deltas_tx(&mut tx, deltas).await?; + carried + } else { + 0 + }; + + // #290: scan the about-to-be-truncated partition for the rollup + // fold. The rollup column is *permanent* state, so we MUST fold + // from ground truth (the `{done_child}` partition itself), not + // from `queue_terminal_live_counts` — the counter can drift + // briefly during a rolling upgrade from a pre-counter binary + // and folding drift into the rollup would bake it in forever. + // The counter rows for this slot are still cleaned up after the + // truncate; reads from the counter (queue_counts_exact in #305) + // may transiently disagree with the rollup until the operator + // runs `awa storage rebuild-terminal-counters`, but the rollup + // itself stays authoritative. See PR #304 reviewer finding + // "High: A1 can persist stale counter state before the + // read-switch PR" for the trade-off. + // + // Carried failed rows are excluded from both columns: they are + // still live in `done_entries`, so folding them into the + // permanent rollup would double-count them. + let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = + sqlx::query_as(audited_sql(format!( + r#" + WITH done_counts AS ( + SELECT + queue, + priority, + (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed, + (count(*) FILTER ( + WHERE state = 'failed' + AND ( + $2::boolean IS FALSE + OR finalized_at < now() - make_interval(secs => $1::bigint) + ) + ))::bigint AS failed + FROM {done_child} + GROUP BY queue, priority + ), + -- Scan the entire partition being truncated (all generations), + -- mirroring done_counts' whole-child scan above. The rotate guard + -- keeps a partition to one generation at prune time, so this is the + -- same set today; counting the whole child keeps the rollup fold + -- conservative even if that invariant is ever weakened (a stray + -- generation's batches would otherwise be truncated without being + -- folded into the rollup -> terminal-count undercount). + batch_rows AS ( + SELECT + batch.ready_generation, + batch.queue, + batch.priority, + item.job_id, + item.run_lease + FROM {receipt_batch_child} AS batch + CROSS JOIN LATERAL unnest( + batch.job_ids, + batch.run_leases + ) AS item(job_id, run_lease) + ), + batch_counts AS ( + SELECT + batch_rows.queue, + batch_rows.priority, + count(*)::bigint AS completed + FROM batch_rows + WHERE NOT EXISTS ( + SELECT 1 + FROM {receipt_tomb_child} AS tomb + WHERE tomb.ready_generation = batch_rows.ready_generation + AND tomb.job_id = batch_rows.job_id + AND tomb.run_lease = batch_rows.run_lease + ) + GROUP BY batch_rows.queue, batch_rows.priority + ), + keys AS ( + SELECT queue, priority FROM done_counts + UNION + SELECT queue, priority FROM batch_counts + ) SELECT - batch_rows.queue, - batch_rows.priority, - count(*)::bigint AS completed - FROM batch_rows - WHERE NOT EXISTS ( - SELECT 1 - FROM {receipt_tomb_child} AS tomb - WHERE tomb.ready_generation = batch_rows.ready_generation - AND tomb.job_id = batch_rows.job_id - AND tomb.run_lease = batch_rows.run_lease + keys.queue, + keys.priority, + ( + COALESCE(done_counts.completed, 0) + + COALESCE(batch_counts.completed, 0) + )::bigint AS pruned_completed_count, + COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count + FROM keys + LEFT JOIN done_counts USING (queue, priority) + LEFT JOIN batch_counts USING (queue, priority) + "# + ))) + .bind(failed_retention_secs) + .bind(retention_floor) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + let truncate_started = Instant::now(); + self + .truncate_used_children_tx( + &mut tx, + &[ + &ready_child, + &claim_attempt_child, + &done_child, + &tomb_child, + &segment_child, + &receipt_batch_child, + &receipt_tomb_child, + &delta_child, + ], ) - GROUP BY batch_rows.queue, batch_rows.priority - ), - keys AS ( - SELECT queue, priority FROM done_counts - UNION - SELECT queue, priority FROM batch_counts - ) - SELECT - keys.queue, - keys.priority, - ( - COALESCE(done_counts.completed, 0) - + COALESCE(batch_counts.completed, 0) - )::bigint AS pruned_completed_count, - COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count - FROM keys - LEFT JOIN done_counts USING (queue, priority) - LEFT JOIN batch_counts USING (queue, priority) - "# + .await + .map_err(map_sqlx_error)?; + let truncate_duration = truncate_started.elapsed(); + if !pruned_terminal_counts.is_empty() { + // #371: append the pruned counts as delta rows instead + // of upserting `queue_terminal_rollups` here — the + // upsert under a pinned MVCC horizon left one + // unreclaimable dead rollup version per prune. The + // maintenance leader folds the deltas + // (`fold_terminal_rollup_deltas`) when the horizon is + // clear; exact readers add unfolded delta sums to + // their rollup reads. + self.append_terminal_rollup_deltas_tx( + &mut tx, + pruned_terminal_counts.into_iter(), + ) + .await?; + } + + // #290: the live counter rows for this slot are about to + // be orphans (their underlying done_entries rows just got + // truncated). Delete them in the same transaction. The + // rollup fold above already captured ground-truth from + // the partition scan; this just cleans up the counter + // index entries so a future insert into a re-rotated + // slot starts from zero. + sqlx::query(audited_sql(format!( + "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" ))) - .bind(failed_retention_secs) - .bind(retention_floor) - .fetch_all(tx.as_mut()) + .bind(slot) + .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { + slot, + generation, + carried_failed_rows, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) + } + .await; - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx( - &mut tx, - &[ - &ready_child, - &claim_attempt_child, - &done_child, - &tomb_child, - &segment_child, - &receipt_batch_child, - &receipt_tomb_child, - &delta_child, - ], - ) - .await; - let truncate_duration = truncate_started.elapsed(); - - match truncate { - Ok(_) => { - if !pruned_terminal_counts.is_empty() { - // #371: append the pruned counts as delta rows instead - // of upserting `queue_terminal_rollups` here — the - // upsert under a pinned MVCC horizon left one - // unreclaimable dead rollup version per prune. The - // maintenance leader folds the deltas - // (`fold_terminal_rollup_deltas`) when the horizon is - // clear; exact readers add unfolded delta sums to - // their rollup reads. - self.append_terminal_rollup_deltas_tx( - &mut tx, - pruned_terminal_counts.into_iter(), - ) - .await?; - } - - // #290: the live counter rows for this slot are about to - // be orphans (their underlying done_entries rows just got - // truncated). Delete them in the same transaction. The - // rollup fold above already captured ground-truth from - // the partition scan; this just cleans up the counter - // index entries so a future insert into a re-rotated - // slot starts from zero. - sqlx::query(audited_sql(format!( - "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" - ))) - .bind(slot) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - if carried_failed_rows > 0 { - tracing::info!( - slot, - carried_failed_rows, - "Carried failed terminal rows inside the retention floor to the live done segment" - ); - } - Ok(PruneOutcome::Pruned { + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { + carried_failed_rows, + durations, + .. + } = &mut outcome + { + durations.commit = commit_started.elapsed(); + let carried_failed_rows = *carried_failed_rows; + if carried_failed_rows > 0 { + tracing::info!( slot, - generation, carried_failed_rows, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) + "Carried failed terminal rows inside the retention floor to the live done segment" + ); } } + Ok(outcome) } #[tracing::instrument( @@ -19259,82 +19462,68 @@ impl QueueStorage { let lease_child = lease_child_name(schema, slot as usize); - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - let lock_table = sqlx::query(audited_sql(format!( - "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - let lock_duration = lock_started.elapsed(); + let lock_started = Instant::now(); + sqlx::query(audited_sql(format!( + "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); - if let Err(err) = lock_table { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); - } - return Err(map_sqlx_error(err)); - } + // The ring rotation serialization held above keeps the cursor + // stable, but keep the explicit recheck to document the truncate + // precondition. + let (current_slot, _) = self + .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) + .await?; - // The ring rotation serialization held above keeps the cursor - // stable, but keep the explicit recheck to document the truncate - // precondition. - let (current_slot, _) = self - .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) - .await?; + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::LeaseCurrent, + count: 0, + }); + } - if current_slot == slot { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::LeaseCurrent, - count: 0, - }); - } + let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?; - let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?; + if active_leases { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::LeaseActive, + count: busy_indicator(active_leases), + }); + } - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { + let truncate_started = Instant::now(); + self.truncate_used_children_tx(&mut tx, &[&lease_child]) + .await + .map_err(map_sqlx_error)?; + let truncate_duration = truncate_started.elapsed(); + Ok(PruneOutcome::Pruned { slot, - reason: SkipReason::LeaseActive, - count: busy_indicator(active_leases), - }); + generation, + carried_failed_rows: 0, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) } + .await; - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx(&mut tx, &[&lease_child]) - .await; - let truncate_duration = truncate_started.elapsed(); - - match truncate { - Ok(_) => { - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - Ok(PruneOutcome::Pruned { - slot, - generation, - carried_failed_rows: 0, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) - } + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { durations, .. } = &mut outcome { + durations.commit = commit_started.elapsed(); } + Ok(outcome) } pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> { @@ -19638,116 +19827,103 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - let lock_tables = sqlx::query(audited_sql(format!( - "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - let lock_duration = lock_started.elapsed(); - - if let Err(err) = lock_tables { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); - } - return Err(map_sqlx_error(err)); - } - - // After taking ACCESS EXCLUSIVE, recheck that the slot is still - // not the current one. The ring rotation serialization should - // already make this stable, but keeping the explicit gate documents - // the truncate precondition. - let (current_slot, _) = self - .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) - .await?; + let lock_started = Instant::now(); + sqlx::query(audited_sql(format!( + "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); - if current_slot == slot { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::ClaimCurrent, - count: 0, - }); - } + // After taking ACCESS EXCLUSIVE, recheck that the slot is still + // not the current one. The ring rotation serialization should + // already make this stable, but keeping the explicit gate documents + // the truncate precondition. + let (current_slot, _) = self + .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) + .await?; - let open_claims = claim_prune_has_open_claims_tx( - &mut tx, - schema, - &claim_child, - &claim_batch_child, - &closure_child, - &closure_batch_child, - ) - .await?; - if open_claims { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::ClaimOpen, - count: busy_indicator(open_claims), - }); - } + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimCurrent, + count: 0, + }); + } - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx( + let open_claims = claim_prune_has_open_claims_tx( &mut tx, - &[ - &claim_child, - &claim_batch_child, - &closure_child, - &closure_batch_child, - ], + schema, + &claim_child, + &claim_batch_child, + &closure_child, + &closure_batch_child, ) - .await; - let truncate_duration = truncate_started.elapsed(); + .await?; + if open_claims { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimOpen, + count: busy_indicator(open_claims), + }); + } - match truncate { - Ok(_) => { - sqlx::query(audited_sql(format!( - r#" - UPDATE {schema}.claim_ring_slots - SET rescue_cursor_claimed_at = '-infinity'::timestamptz, - rescue_cursor_job_id = 0, - rescue_cursor_run_lease = 0, - deadline_cursor_deadline_at = '-infinity'::timestamptz, - deadline_cursor_job_id = 0, - deadline_cursor_run_lease = 0, - batch_deadline_cursor_deadline_at = '-infinity'::timestamptz, - batch_deadline_cursor_batch_id = 0 - WHERE slot = $1 - "# - ))) - .bind(slot) - .execute(tx.as_mut()) + let truncate_started = Instant::now(); + self + .truncate_used_children_tx( + &mut tx, + &[ + &claim_child, + &claim_batch_child, + &closure_child, + &closure_batch_child, + ], + ) .await .map_err(map_sqlx_error)?; - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - Ok(PruneOutcome::Pruned { - slot, - generation, - carried_failed_rows: 0, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) - } + let truncate_duration = truncate_started.elapsed(); + sqlx::query(audited_sql(format!( + r#" + UPDATE {schema}.claim_ring_slots + SET rescue_cursor_claimed_at = '-infinity'::timestamptz, + rescue_cursor_job_id = 0, + rescue_cursor_run_lease = 0, + deadline_cursor_deadline_at = '-infinity'::timestamptz, + deadline_cursor_job_id = 0, + deadline_cursor_run_lease = 0, + batch_deadline_cursor_deadline_at = '-infinity'::timestamptz, + batch_deadline_cursor_batch_id = 0 + WHERE slot = $1 + "# + ))) + .bind(slot) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { + slot, + generation, + carried_failed_rows: 0, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) } + .await; + + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { durations, .. } = &mut outcome { + durations.commit = commit_started.elapsed(); + } + Ok(outcome) } fn job_id_sequence(&self) -> String { diff --git a/awa-model/src/read_transaction.rs b/awa-model/src/read_transaction.rs new file mode 100644 index 000000000..88155e7fb --- /dev/null +++ b/awa-model/src/read_transaction.rs @@ -0,0 +1,31 @@ +//! Protect owned read connections while SQLx is still acknowledging BEGIN. +//! +//! SQLx records transaction depth after BEGIN completes. Dropping that future +//! earlier can leave a transaction that its normal rollback guard cannot see. + +use sqlx::pool::PoolConnection; +use sqlx::{PgPool, Postgres}; + +pub(crate) struct BeginGuard { + // Callers borrow this field directly so the separate flag can be cleared + // while the resulting transaction retains its connection borrow. + pub(crate) connection: PoolConnection, + pub(crate) begin_pending: bool, +} + +impl BeginGuard { + pub(crate) async fn acquire(pool: &PgPool) -> Result { + Ok(Self { + connection: pool.acquire().await?, + begin_pending: true, + }) + } +} + +impl Drop for BeginGuard { + fn drop(&mut self) { + if self.begin_pending { + self.connection.close_on_drop(); + } + } +} diff --git a/awa/tests/admin_read_cancel_test.rs b/awa/tests/admin_read_cancel_test.rs new file mode 100644 index 000000000..88b017003 --- /dev/null +++ b/awa/tests/admin_read_cancel_test.rs @@ -0,0 +1,452 @@ +//! Cancel a real public read after PostgreSQL has entered its transaction, +//! before SQLx receives ReadyForQuery. No sleeps or custom BEGIN statements. + +mod ci_timing; + +use awa::model::admin::{self, ListJobsFilter}; +use awa::model::{insert_many, migrations, InsertOpts, InsertParams, QueueStorage}; +use awa::{AwaError, JobState, QueueStorageConfig}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode}; +use sqlx::PgPool; +use std::io; +use std::str::FromStr; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{watch, Notify}; +use tokio::task::{JoinHandle, JoinSet}; + +type TestResult = Result>; + +fn bound() -> Duration { + ci_timing::scaled_timeout(Duration::from_secs(5)) +} + +#[derive(Default)] +struct BeginGate { + armed: AtomicBool, + held: watch::Sender>, + release: Notify, +} + +impl BeginGate { + async fn held_pid(&self) -> TestResult { + let mut observed = self.held.subscribe(); + loop { + if let Some(pid) = *observed.borrow_and_update() { + return Ok(pid); + } + observed.changed().await?; + } + } +} + +struct Proxy { + port: u16, + gate: Arc, + task: JoinHandle<()>, +} + +impl Proxy { + async fn start(options: &PgConnectOptions) -> TestResult { + // The proxy is local and test-only; credentials and backend secrets are + // forwarded unchanged and never logged. SSL is disabled only on it. + assert!( + options.get_socket().is_none(), + "this TCP fixture needs a TCP DATABASE_URL" + ); + let upstream = (options.get_host().to_string(), options.get_port()); + let listener = TcpListener::bind(("127.0.0.1", 0)).await?; + let port = listener.local_addr()?.port(); + let gate = Arc::new(BeginGate::default()); + let connections_gate = gate.clone(); + let task = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let Ok((client, _)) = accepted else { break }; + let upstream = upstream.clone(); + let gate = connections_gate.clone(); + connections.spawn(async move { relay(client, upstream, gate).await }); + } + Some(_finished) = connections.join_next(), if !connections.is_empty() => {} + } + } + }); + Ok(Self { port, gate, task }) + } + + async fn stop(self) { + self.gate.release.notify_one(); + self.task.abort(); + let _ = self.task.await; + // Dropping the listener's JoinSet aborts every relay and closes both + // socket halves, including a relay still holding ReadyForQuery. + } +} + +async fn frame(stream: &mut R) -> io::Result)>> { + let tag = match stream.read_u8().await { + Ok(tag) => tag, + Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), + Err(error) => return Err(error), + }; + let length = stream.read_u32().await?; + if !(4..=1024 * 1024).contains(&length) { + return Err(io::Error::other("invalid test proxy frame length")); + } + let mut payload = vec![0; (length - 4) as usize]; + stream.read_exact(&mut payload).await?; + Ok(Some((tag, payload))) +} + +async fn forward(stream: &mut W, tag: u8, payload: &[u8]) -> io::Result<()> { + stream.write_u8(tag).await?; + stream.write_u32((payload.len() + 4) as u32).await?; + stream.write_all(payload).await +} + +async fn relay( + mut client: TcpStream, + upstream: (String, u16), + gate: Arc, +) -> io::Result<()> { + let mut server = TcpStream::connect((upstream.0.as_str(), upstream.1)).await?; + // Startup packets precede the normal tagged PostgreSQL messages. + let length = client.read_u32().await?; + if !(8..=64 * 1024).contains(&length) { + return Err(io::Error::other("invalid startup packet")); + } + let mut startup = vec![0; (length - 4) as usize]; + client.read_exact(&mut startup).await?; + server.write_u32(length).await?; + server.write_all(&startup).await?; + + let (mut client_read, mut client_write) = client.into_split(); + let (mut server_read, mut server_write) = server.into_split(); + let awaiting_begin = Arc::new(AtomicBool::new(false)); + let frontend_gate = gate.clone(); + let frontend_begin = awaiting_begin.clone(); + let frontend = async move { + while let Some((tag, payload)) = frame(&mut client_read).await? { + if tag == b'Q' { + let sql = String::from_utf8_lossy(&payload); + let is_begin = sql + .trim_end_matches('\0') + .split_whitespace() + .next() + .is_some_and(|first| first.eq_ignore_ascii_case("BEGIN")); + if is_begin && frontend_gate.armed.swap(false, Ordering::SeqCst) { + frontend_begin.store(true, Ordering::SeqCst); + } + } + forward(&mut server_write, tag, &payload).await?; + } + Ok::<_, io::Error>(()) + }; + let backend = async move { + let mut pid = 0; + let mut begin_completed = false; + while let Some((tag, payload)) = frame(&mut server_read).await? { + if tag == b'K' && payload.len() == 8 { + pid = i32::from_be_bytes(payload[..4].try_into().expect("four PID bytes")); + } + if tag == b'C' + && payload.as_slice() == b"BEGIN\0" + && awaiting_begin.load(Ordering::SeqCst) + { + begin_completed = true; + } + if tag == b'Z' && begin_completed && awaiting_begin.swap(false, Ordering::SeqCst) { + if payload.as_slice() != b"T" || pid <= 0 { + return Err(io::Error::other("BEGIN did not enter a server transaction")); + } + // CommandComplete was forwarded, but the real transaction + // status frame stays here until the read future is dropped. + gate.held.send_replace(Some(pid)); + gate.release.notified().await; + begin_completed = false; + } + forward(&mut client_write, tag, &payload).await?; + } + Ok::<_, io::Error>(()) + }; + tokio::select! { result = frontend => result, result = backend => result } +} + +struct Fixture { + admin: PgPool, + observer: PgPool, + reader: PgPool, + store: Arc, + database: String, + job_id: i64, + proxy: Proxy, +} + +impl Fixture { + async fn new() -> TestResult { + let url = awa_testing::setup::database_url(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&url) + .await?; + let database = format!("awa_read_cancel_{}", uuid::Uuid::new_v4().simple()); + sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE {database}"))) + .execute(&admin) + .await?; + let options = PgConnectOptions::from_str(&url)?.database(&database); + let observer = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.clone()) + .await?; + migrations::run(&observer).await?; + awa_testing::setup::activate_queue_storage(&observer).await; + let store = Arc::new(QueueStorage::new(QueueStorageConfig::default())?); + let rows = insert_many( + &observer, + &[InsertParams { + kind: "admin_read_cancel".to_string(), + args: serde_json::json!({}), + opts: InsertOpts { + queue: "admin_read_cancel".to_string(), + ..Default::default() + }, + }], + ) + .await?; + let proxy = Proxy::start(&options).await?; + let reader = PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(bound()) + .connect_with( + options + .host("127.0.0.1") + .port(proxy.port) + .ssl_mode(PgSslMode::Disable), + ) + .await?; + Ok(Self { + admin, + observer, + reader, + store, + database, + job_id: rows[0].id, + proxy, + }) + } + + async fn close(self) -> TestResult { + self.proxy.gate.release.notify_one(); + // Stop relays even if a broken connection cannot be returned normally. + let closed = tokio::time::timeout(bound(), self.reader.close()).await; + self.proxy.stop().await; + if closed.is_err() { + tokio::time::timeout(bound(), self.reader.close()).await?; + } + self.observer.close().await; + tokio::time::timeout( + bound(), + sqlx::query(sqlx::AssertSqlSafe(format!( + "DROP DATABASE {} WITH (FORCE)", + self.database + ))) + .execute(&self.admin), + ) + .await??; + self.admin.close().await; + Ok(()) + } +} + +#[derive(Clone, Copy, Debug)] +enum Read { + Job, + Counts, + Queue, + Kinds, + Timeseries, + Jobs, +} + +async fn public_read(which: Read, store: &QueueStorage, pool: &PgPool, job_id: i64) -> TestResult { + match which { + Read::Job => { + let job = store.load_job(pool, job_id).await?.expect("job exists"); + assert_eq!((job.id, job.state), (job_id, JobState::Available)); + } + Read::Counts => assert_eq!( + admin::state_counts(pool).await?.get(&JobState::Available), + Some(&1) + ), + Read::Queue => assert_eq!( + admin::queue_overview(pool, "admin_read_cancel") + .await? + .unwrap() + .available, + 1 + ), + Read::Kinds => assert_eq!( + admin::job_kind_overviews(pool) + .await? + .iter() + .find(|kind| kind.kind == "admin_read_cancel") + .unwrap() + .job_count, + 1 + ), + Read::Timeseries => assert_eq!( + admin::state_timeseries(pool, 60) + .await? + .iter() + .filter(|bucket| bucket.state == JobState::Available) + .map(|bucket| bucket.count) + .sum::(), + 1 + ), + Read::Jobs => assert_eq!( + admin::list_jobs(pool, &ListJobsFilter::default()) + .await? + .iter() + .map(|job| job.id) + .collect::>(), + vec![job_id] + ), + } + Ok(()) +} + +async fn backend_is_clean(observer: &PgPool, pid: i32) -> TestResult { + tokio::time::timeout(bound(), async { + loop { + let state: Option<(String, bool)> = sqlx::query_as( + "SELECT state, xact_start IS NULL FROM pg_stat_activity WHERE pid = $1", + ) + .bind(pid) + .fetch_optional(observer) + .await?; + if state + .as_ref() + .is_none_or(|(state, no_tx)| state == "idle" && *no_tx) + { + return Ok::<_, sqlx::Error>(()); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .map_err(|_| { + io::Error::other(format!( + "backend {pid} retained a transaction after read cancellation" + )) + })??; + Ok(()) +} + +async fn cancelled_begin(which: Read) -> TestResult { + let fixture = Fixture::new().await?; + let reader = fixture.reader.clone(); + let observer = fixture.observer.clone(); + let store = fixture.store.clone(); + let gate = fixture.proxy.gate.clone(); + let job_id = fixture.job_id; + let tested = tokio::spawn(async move { + gate.armed.store(true, Ordering::SeqCst); + let mut read = Box::pin(public_read(which, &store, &reader, job_id)); + let pid = tokio::select! { + held = tokio::time::timeout(bound(), gate.held_pid()) => held??, + result = &mut read => return Err(io::Error::other(format!("{which:?} completed before real BEGIN gate: {result:?}")).into()), + }; + let state: (String, bool) = sqlx::query_as( + "SELECT state, xact_start IS NOT NULL FROM pg_stat_activity WHERE pid = $1") + .bind(pid).fetch_one(&observer).await?; + assert_eq!(state, ("idle in transaction".to_string(), true), "real BEGIN reached server"); + drop(read); + gate.release.notify_one(); + backend_is_clean(&observer, pid).await?; + tokio::time::timeout(bound(), public_read(which, &store, &reader, job_id)).await??; + Ok::<_, Box>(()) + }).await; + let cleanup = fixture.close().await; + tested??; + cleanup +} + +#[tokio::test] +async fn cancelled_load_job_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Job).await +} + +#[tokio::test] +async fn cancelled_state_counts_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Counts).await +} + +#[tokio::test] +async fn cancelled_queue_overview_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Queue).await +} + +#[tokio::test] +async fn cancelled_kind_overviews_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Kinds).await +} + +#[tokio::test] +async fn cancelled_timeseries_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Timeseries).await +} + +#[tokio::test] +async fn cancelled_job_list_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Jobs).await +} + +#[tokio::test] +async fn completed_begin_and_query_error_preserve_healthy_connection() -> TestResult { + let fixture = Fixture::new().await?; + let reader = fixture.reader.clone(); + let observer = fixture.observer.clone(); + let store = fixture.store.clone(); + let job_id = fixture.job_id; + let tested = tokio::spawn(async move { + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&reader) + .await?; + public_read(Read::Job, &store, &reader, job_id).await?; + public_read(Read::Counts, &store, &reader, job_id).await?; + sqlx::query("ALTER TABLE awa.deferred_jobs RENAME TO cancel_test_deferred") + .execute(&observer) + .await?; + let error = store + .load_job(&reader, job_id) + .await + .expect_err("real post-BEGIN query error"); + assert!( + matches!(error, AwaError::Database(sqlx::Error::Database(ref db)) + if db.code().as_deref() == Some("42P01")), + "expected undefined table, got {error:?}" + ); + backend_is_clean(&observer, pid).await?; + sqlx::query("ALTER TABLE awa.cancel_test_deferred RENAME TO deferred_jobs") + .execute(&observer) + .await?; + public_read(Read::Job, &store, &reader, job_id).await?; + let reused: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&reader) + .await?; + assert_eq!( + reused, pid, + "completed BEGIN and ordinary query errors retain healthy reuse" + ); + Ok::<_, Box>(()) + }) + .await; + let cleanup = fixture.close().await; + tested??; + cleanup +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index bf51546fe..9679f20f5 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -3,6 +3,7 @@ //! These tests exercise the full dispatcher/worker/maintenance wiring with the //! queue_storage backend enabled. +mod ci_timing; #[path = "../../awa-model/src/storage_footprint_sql.rs"] mod storage_footprint_sql; @@ -13503,13 +13504,43 @@ async fn test_queue_storage_prune_past_floor_folds_failed_into_rollup() { RotateOutcome::Rotated { slot: 1, .. } => {} other => panic!("expected Rotated {{ slot: 1 }}, got {other:?}"), } - // Duration::ZERO disables the floor: every failed row is treated as - // past-floor and pruned for good — no carry. - match store - .prune_oldest(&pool, Duration::ZERO) - .await - .expect("prune_oldest") - { + let mut reader = pool.begin().await.expect("begin transient terminal reader"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.done_entries_0 IN ACCESS SHARE MODE" + ))) + .execute(reader.as_mut()) + .await + .expect("hold terminal partition reader"); + + // A valid transient Blocked must leave both the rows and rollup untouched. + // This test checks retention accounting, not immediate exclusive-lock access. + assert!(matches!( + store + .prune_oldest(&pool, Duration::ZERO) + .await + .expect("blocked prune"), + PruneOutcome::Blocked { slot: 0 } + )); + assert_eq!(done_entries_count(&pool, schema, queue).await, 3); + assert_eq!(pruned_rollup_sums(&pool, schema, queue).await, (0, 0)); + reader.rollback().await.expect("release terminal reader"); + + // Duration::ZERO disables the floor. A background reader may still briefly + // hold a partition, so require the fold to finish within an assertion bound. + // Only Blocked may defer it; every error or other outcome still fails below. + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + let outcome = loop { + let outcome = store + .prune_oldest(&pool, Duration::ZERO) + .await + .expect("prune_oldest"); + if matches!(outcome, PruneOutcome::Blocked { slot: 0 }) && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(10)).await; + } else { + break outcome; + } + }; + match outcome { PruneOutcome::Pruned { slot: 0, carried_failed_rows: 0, @@ -15632,6 +15663,204 @@ async fn test_queue_storage_prune_oldest_blocks_on_reader_lock() { ); } +/// The prune `lock_timeout` covers every statement after it is set, not +/// only the `LOCK TABLE` on the slot's children. A statement that waits on +/// another table's lock — here the failed-row carry, its count deltas, the +/// rollup delta append and the live-counter cleanup — must roll the whole prune back +/// and report `Blocked`, leaving the slot to the next attempt. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_prune_lock_timeout_after_child_locks_reports_blocked() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_prune_late_lock_timeout"; + let schema = "awa_qs_prune_late_lock_timeout"; + let store = create_store(&pool, schema).await; + + store + .enqueue_batch(&pool, queue, 1, 1) + .await + .expect("enqueue completed job"); + let claimed = store + .claim_batch(&pool, queue, 1) + .await + .expect("claim completed job"); + store + .complete_batch(&pool, &claimed) + .await + .expect("complete job"); + + let failed_job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim failing job") + .into_iter() + .next() + .expect("missing claimed job"); + store + .fail_terminal(&pool, failed_job_id, claimed.job.run_lease, "boom", None) + .await + .expect("fail job") + .expect("running job should fail"); + + let rotated = store.rotate(&pool).await.expect("rotate"); + assert!( + matches!(rotated, RotateOutcome::Rotated { slot: 1, .. }), + "unexpected rotate outcome: {rotated:?}" + ); + + let retention = Duration::from_secs(3600); + for held in [ + format!("{schema}.done_entries_1"), + format!("{schema}.queue_terminal_count_deltas_1"), + format!("{schema}.queue_terminal_rollup_deltas"), + format!("{schema}.queue_terminal_live_counts"), + ] { + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(audited_sql(format!("LOCK TABLE {held} IN SHARE MODE"))) + .execute(holder.as_mut()) + .await + .expect("hold share lock"); + + let outcome = store.prune_oldest(&pool, retention).await; + assert!( + matches!(outcome, Ok(PruneOutcome::Blocked { slot: 0 })), + "prune waiting on {held} should report Blocked, got {outcome:?}" + ); + + holder.rollback().await.expect("release share lock"); + + assert_eq!( + done_entries_count(&pool, schema, queue).await, + 2, + "blocked prune on {held} must not truncate or carry rows" + ); + let rollup_rows: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*)::bigint FROM {schema}.queue_terminal_rollup_deltas" + ))) + .fetch_one(&pool) + .await + .expect("count rollups"); + assert_eq!( + rollup_rows, 0, + "blocked prune on {held} must not fold rollups" + ); + assert_invariant_holds( + &pool, + schema, + queue, + &format!("after prune blocked on {held}"), + ) + .await; + } + + // Other activity can briefly hold a partition after the test's own + // locks are gone. Only Blocked may defer the prune; every error or other + // outcome still fails below. + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + let pruned = loop { + let pruned = store + .prune_oldest(&pool, retention) + .await + .expect("prune after locks released"); + if matches!(pruned, PruneOutcome::Blocked { slot: 0 }) && Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(10)).await; + } else { + break pruned; + } + }; + assert!( + matches!( + pruned, + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } + ), + "unexpected prune outcome after locks released: {pruned:?}" + ); + assert_invariant_holds(&pool, schema, queue, "after prune").await; + let counts = store.queue_counts(&pool, queue).await.expect("counts"); + assert_eq!(counts.terminal, 2, "rollup (1) + carried live row (1)"); +} + +/// Terminal-delta rollup shares prune's lock timeout. A timeout on the +/// live-counter fold, after the delta child is locked, must report a blocked +/// slot rather than fail the rollup pass, and the pass stops there instead of +/// waiting out the same timeout on every remaining slot. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_terminal_count_delta_rollup_lock_timeout_after_child_lock_reports_blocked() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_terminal_delta_rollup_late_lock"; + let schema = "awa_qs_terminal_delta_rollup_late_lock"; + let store = create_store(&pool, schema).await; + + for _ in 0..2 { + store + .enqueue_batch(&pool, queue, 1, 3) + .await + .expect("enqueue"); + let claimed = store.claim_batch(&pool, queue, 3).await.expect("claim"); + store + .complete_batch(&pool, &claimed) + .await + .expect("complete"); + match store.rotate(&pool).await.expect("rotate") { + RotateOutcome::Rotated { .. } => {} + other => panic!("expected Rotated, got {other:?}"), + } + } + + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.queue_terminal_live_counts IN SHARE MODE" + ))) + .execute(holder.as_mut()) + .await + .expect("hold share lock"); + + let outcome = store.rollup_terminal_count_deltas(&pool, 4).await; + holder.rollback().await.expect("release share lock"); + let outcome = outcome.expect("rollup waiting on live counts should not error"); + assert_eq!(outcome.rolled_slots, 0); + assert_eq!( + outcome.blocked_slots, 1, + "the pass stops at the first blocked slot" + ); + assert_eq!(live_count_sum(&pool, schema, queue).await, 0); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 6); + + // As above, a transient lock may briefly block a slot after ours is + // released. Slots rolled by a partly blocked pass stay rolled, so count + // them across attempts. + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + let mut rolled_slots = 0; + loop { + let outcome = store + .rollup_terminal_count_deltas(&pool, 4) + .await + .expect("rollup after lock released"); + rolled_slots += outcome.rolled_slots; + if outcome.blocked_slots == 0 || Instant::now() >= deadline { + assert_eq!(outcome.blocked_slots, 0); + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!(rolled_slots, 2); + assert_eq!(live_count_sum(&pool, schema, queue).await, 6); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 0); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_prune_allows_segment_reader_to_finish_claim_lock_order() { let (_db_guard, pool) = setup_pool(8).await; @@ -21093,3 +21322,1731 @@ async fn test_queue_storage_legacy_row_claims_are_locked_before_the_heartbeat_st .expect("join heartbeat") .expect("heartbeat after the rewrite"); } + +/// Claims and leases record a job's aged effective priority while its ready +/// row keeps the lane priority. Reads that join the two must still find the +/// running attempt: `load_job` (behind `get_job` and lifecycle events) and +/// the callback lookup behind external completion. +async fn aged_claim_is_visible_to_job_and_callback_reads(receipts: bool) { + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_aged_claim_reads"; + let schema = if receipts { + "awa_qs_aged_claim_reads_receipts" + } else { + "awa_qs_aged_claim_reads_leases" + }; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + lease_claim_receipts: receipts, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 902 }, + InsertOpts { + queue: queue.to_string(), + priority: 4, + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + let claimed = store + .claim_runtime_batch_with_aging_for_instance( + &pool, + queue, + 1, + Duration::ZERO, + Duration::from_secs(60), + Uuid::new_v4(), + 4, + Duration::from_secs(30), + Duration::from_secs(1), + ) + .await + .expect("claim with aging"); + assert_eq!(claimed.len(), 1); + let run_lease = claimed[0].job.run_lease; + assert!( + claimed[0].job.priority < 4, + "the claim must have aged the job for this test to mean anything" + ); + + let loaded = store + .load_job(&pool, job_id) + .await + .expect("load aged job") + .expect("aged job exists"); + assert_eq!(loaded.state, JobState::Running); + assert_eq!(loaded.run_lease, run_lease); + + let callback_id = store + .register_callback(&pool, job_id, run_lease, Duration::from_secs(300)) + .await + .expect("register callback on aged job"); + let mut tx = pool.begin().await.unwrap(); + let by_callback = store + .callback_job_in_tx(&mut tx, callback_id, Some(run_lease), false) + .await + .expect("callback lookup") + .expect("callback lookup finds the aged job"); + tx.rollback().await.unwrap(); + assert_eq!(by_callback.id, job_id); + assert_eq!(by_callback.run_lease, run_lease); +} + +#[tokio::test] +async fn test_aged_receipt_claim_is_visible_to_job_and_callback_reads() { + aged_claim_is_visible_to_job_and_callback_reads(true).await; +} + +#[tokio::test] +async fn test_aged_materialized_lease_is_visible_to_job_and_callback_reads() { + aged_claim_is_visible_to_job_and_callback_reads(false).await; +} + +/// The admin reads join a materialized lease to its ready row. A lease records +/// the aged effective priority, so joining on priority hid a job that aged +/// before its claim from the job list, the kinds page and unique-key cancel. +#[tokio::test] +async fn test_aged_materialized_lease_is_visible_to_admin_reads() { + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_aged_lease_admin_reads"; + let schema = "awa_qs_aged_lease_admin_reads"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + lease_claim_receipts: false, + ..unique_cancel_store_config(schema) + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &unique_cancel_job("aged"), + InsertOpts { + priority: 4, + ..unique_cancel_opts(queue, None) + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + let claimed = store + .claim_runtime_batch_with_aging_for_instance( + &pool, + queue, + 1, + Duration::ZERO, + Duration::from_secs(60), + Uuid::new_v4(), + 4, + Duration::from_secs(30), + Duration::from_secs(1), + ) + .await + .expect("claim with aging"); + assert_eq!(claimed.len(), 1); + assert!( + claimed[0].job.priority < 4, + "the claim must have aged the job for this test to mean anything" + ); + + let running = admin::list_jobs( + &pool, + &admin::ListJobsFilter { + state: Some(JobState::Running), + queue: Some(queue.to_string()), + ..Default::default() + }, + ) + .await + .expect("list running jobs"); + assert_eq!( + running.iter().map(|job| job.id).collect::>(), + vec![job_id], + "the aged running job must be listed" + ); + + let overview = admin::job_kind_overviews(&pool) + .await + .expect("job kind overviews"); + let kind = overview + .iter() + .find(|row| row.kind == UniqueCancelJob::kind()) + .expect("the aged job's kind is listed"); + assert_eq!(kind.job_count, 1, "the aged running job must be counted"); + + let cancelled = cancel_unique(&pool, Some(queue), "aged") + .await + .expect("cancel by unique key finds the aged running job"); + assert_eq!(cancelled.id, job_id); +} + +/// Aging moves a waiting job by tombstoning its ready row and re-inserting the +/// same attempt on the next lane. In a quiet queue the new lane hands out the +/// same `lane_seq`, so a lease joined to its ready row by attempt alone matched +/// both rows and the admin reads counted the running job twice. +async fn moved_ready_row_is_joined_once_by_attempt_reads(receipts: bool) { + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_moved_ready_lease_reads"; + let schema = if receipts { + "awa_qs_moved_ready_compact_reads" + } else { + "awa_qs_moved_ready_lease_reads" + }; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + lease_claim_receipts: receipts, + ..unique_cancel_store_config(schema) + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &unique_cancel_job("moved"), + InsertOpts { + priority: 2, + ..unique_cancel_opts(queue, None) + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + let aged = store + .age_waiting_priorities(&pool, Duration::from_secs(60), 10) + .await + .expect("age waiting priorities"); + assert_eq!(aged, vec![job_id]); + + let same_position: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(DISTINCT priority)::bigint + FROM {schema}.ready_entries + WHERE job_id = $1 + GROUP BY ready_slot, ready_generation, enqueue_shard, lane_seq" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("ready rows for the moved job"); + assert_eq!( + same_position, 2, + "the source and destination rows must share a position for this test to mean anything" + ); + + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim moved job"); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].job.id, job_id); + assert_eq!(claimed[0].claim.claim_batch_id.is_some(), receipts); + + // This is the production CTE used by current-job projections, not only + // the point-read candidate list. A compact item must hydrate exactly once. + let projected: Vec = sqlx::query_scalar(audited_sql(format!( + "{} SELECT job_id FROM current_jobs WHERE queue = $1 AND state = 'running'", + admin::queue_storage_current_jobs_cte(schema) + ))) + .bind(queue) + .fetch_all(&pool) + .await + .expect("project current running jobs"); + assert_eq!(projected, vec![job_id], "hydrate only the live ready row"); + + let running = admin::list_jobs( + &pool, + &admin::ListJobsFilter { + state: Some(JobState::Running), + queue: Some(queue.to_string()), + ..Default::default() + }, + ) + .await + .expect("list running jobs"); + assert_eq!( + running.iter().map(|job| job.id).collect::>(), + vec![job_id], + "the moved running job is listed once" + ); + + let overview = admin::job_kind_overviews(&pool) + .await + .expect("job kind overviews"); + let kind = overview + .iter() + .find(|row| row.kind == UniqueCancelJob::kind()) + .expect("the moved job's kind is listed"); + assert_eq!(kind.job_count, 1, "the moved running job is counted once"); + + let loaded = store + .load_job(&pool, job_id) + .await + .expect("load moved job") + .expect("moved job exists"); + assert_eq!(loaded.state, JobState::Running); + assert_eq!(loaded.priority, 1, "reads the live destination row"); + assert_eq!( + loaded.metadata.get("_awa_original_priority"), + Some(&serde_json::json!(2)), + "payload comes from the aged destination, not its tombstoned source" + ); +} + +#[tokio::test] +async fn test_moved_ready_row_is_joined_once_by_lease_reads() { + moved_ready_row_is_joined_once_by_attempt_reads(false).await; +} + +#[tokio::test] +async fn test_moved_ready_row_is_joined_once_by_compact_receipt_reads() { + moved_ready_row_is_joined_once_by_attempt_reads(true).await; +} + +/// A point read's snapshot can outlive completion. TRUNCATE is not MVCC-safe: +/// pruning its claim/lease partition before the reader visits it would erase +/// the old running evidence, while the new completion is outside that snapshot. +async fn point_read_snapshot_survives_ring_prune(receipts: bool) { + let (_db_guard, pool) = setup_pool(6).await; + let schema = if receipts { + "awa_qs_snapshot_claim_prune" + } else { + "awa_qs_snapshot_lease_prune" + }; + let store = Arc::new( + create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + lease_claim_receipts: receipts, + ..Default::default() + }, + ) + .await, + ); + let queue = "snapshot_prune"; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim the snapshot probe"); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].claim.claim_batch_id.is_some(), receipts); + assert_eq!( + store.load_job(&pool, job_id).await.unwrap().unwrap().state, + JobState::Running + ); + + // Gate the reader's second query without blocking completion's own + // deferred-state checks. Only this reader's application name takes the + // advisory lock; the view's scalar initplan runs even with no deferred rows. + sqlx::raw_sql(audited_sql(format!( + r#"CREATE FUNCTION {schema}.snapshot_read_gate() RETURNS boolean + LANGUAGE plpgsql VOLATILE AS $$ + BEGIN + IF current_setting('application_name') = 'awa_snapshot_reader' THEN + PERFORM pg_advisory_xact_lock(60006, 1); + END IF; + RETURN TRUE; + END; + $$; + ALTER TABLE {schema}.deferred_jobs RENAME TO deferred_rows; + CREATE VIEW {schema}.deferred_jobs AS + SELECT * FROM {schema}.deferred_rows + WHERE (SELECT {schema}.snapshot_read_gate())"# + ))) + .execute(&pool) + .await + .expect("install reader-only query gate"); + let mut gate = pool.begin().await.expect("begin read gate"); + sqlx::query("SELECT pg_advisory_xact_lock(60006, 1)") + .execute(gate.as_mut()) + .await + .expect("hold the second read"); + let reader_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + (*pool.connect_options()) + .clone() + .application_name("awa_snapshot_reader"), + ) + .await + .expect("connect isolated snapshot reader"); + let reader_store = store.clone(); + let reader_connection = reader_pool.clone(); + let reader = + tokio::spawn(async move { reader_store.load_job(&reader_connection, job_id).await }); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(10)); + loop { + let at_gate: bool = sqlx::query_scalar( + "SELECT EXISTS ( + SELECT 1 FROM pg_locks + WHERE locktype = 'advisory' + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) + AND classid = 60006 AND objid = 1 AND NOT granted + )", + ) + .fetch_one(&pool) + .await + .expect("observe the reader-only gate"); + if at_gate { + break; + } + assert!( + Instant::now() < deadline, + "point read did not reach its second query" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + + let during_read = + tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(10)), async { + assert_eq!( + store + .complete_runtime_batch(&pool, &claimed) + .await + .expect("complete outside the snapshot") + .len(), + 1 + ); + let rotation = if receipts { + store.rotate_claims(&pool).await + } else { + store.rotate_leases(&pool).await + } + .expect("seal the completed attempt's slot"); + assert!(matches!(rotation, RotateOutcome::Rotated { slot: 1, .. })); + if receipts { + store.prune_oldest_claims(&pool).await + } else { + store.prune_oldest_leases(&pool).await + } + .expect("prune while the snapshot is live") + }) + .await; + + // Always release the gate and finish the reader before asserting the + // outcome, so the negative control also reports the inconsistent state. + gate.rollback().await.expect("release the read gate"); + let snapshot = tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(10)), reader) + .await + .expect("point read finished") + .expect("join point read") + .expect("point read succeeds") + .expect("the job remains in its snapshot"); + let during_read = + during_read.expect("completion/prune must not wait on the deferred read gate"); + assert_eq!( + ( + matches!(during_read, PruneOutcome::Blocked { slot: 0 }), + snapshot.state + ), + (true, JobState::Running), + "prune must retain the snapshot's running evidence: {during_read:?}" + ); + assert_eq!(snapshot.run_lease, claimed[0].job.run_lease); + reader_pool.close().await; + let after_read = if receipts { + store.prune_oldest_claims(&pool).await + } else { + store.prune_oldest_leases(&pool).await + } + .expect("prune after the reader releases its snapshot"); + assert!(matches!(after_read, PruneOutcome::Pruned { slot: 0, .. })); + assert_eq!( + store.load_job(&pool, job_id).await.unwrap().unwrap().state, + JobState::Completed + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_point_read_snapshot_blocks_compact_claim_prune() { + point_read_snapshot_survives_ring_prune(true).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_point_read_snapshot_blocks_lease_prune() { + point_read_snapshot_survives_ring_prune(false).await; +} + +/// A lane move leaves immutable source rows behind. Wrap the ring so the +/// destination is in an earlier partition: point reads must use live evidence, +/// rather than depending on which equal-timestamp ready row the plan emits last. +async fn ready_point_read_wrap_fixture() -> (TestDbGuard, sqlx::PgPool, QueueStorage, i64) { + let (guard, pool) = setup_pool(4).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_point_read_wrap".to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + ..Default::default() + }, + ) + .await; + // A fresh empty ring can be positioned just before wrap without retaining + // unrelated jobs. This is the state reached after three rotations/prunes. + let schema = store.schema(); + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.queue_ring_rotations (generation, slot) VALUES (3, 3)" + ))) + .execute(&pool) + .await + .expect("position empty ring before wrap"); + sqlx::query(audited_sql(format!( + "UPDATE {schema}.queue_ring_slots SET generation = 3 WHERE slot = 3" + ))) + .execute(&pool) + .await + .expect("set current slot generation"); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: "point_read_source".to_string(), + priority: 4, + ..Default::default() + }, + ) + .await; + assert!(matches!( + store.rotate(&pool).await.expect("wrap to slot zero"), + RotateOutcome::Rotated { slot: 0, .. } + )); + (guard, pool, store, job_id) +} + +async fn assert_live_ready_row( + pool: &sqlx::PgPool, + store: &QueueStorage, + job_id: i64, + expected_queue: &str, + expected_priority: i16, +) { + let schema = store.schema(); + let live: Vec<(String, i16, i32)> = sqlx::query_as(audited_sql(format!( + "SELECT ready.queue, ready.priority, ready.ready_slot + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + )" + ))) + .bind(job_id) + .fetch_all(pool) + .await + .expect("read live ready evidence"); + assert_eq!( + live, + vec![(expected_queue.to_string(), expected_priority, 0)] + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_priority_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .set_priority(&pool, job_id, 1) + .await + .expect("set priority")); + assert_live_ready_row(&pool, &store, job_id, "point_read_source", 1).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.priority, 1, + "point read must ignore the tombstoned source" + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_queue_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .move_queue(&pool, job_id, "point_read_destination", None) + .await + .expect("move queue")); + assert_live_ready_row(&pool, &store, job_id, "point_read_destination", 4).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.queue, "point_read_destination", + "point read must ignore the tombstoned source" + ); +} + +/// A reader's first query can establish its snapshot before waiting on its +/// ready relation lock. Queue prune must finish its failed-row carry before +/// that snapshot, or TRUNCATE hides the source and the new carried version. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_point_read_first_snapshot_follows_ready_prune_carry() { + let (_db_guard, pool) = setup_pool(10).await; + let schema = "awa_qs_first_snapshot_ready_carry"; + let queue = "first_snapshot_ready_carry"; + let bound = ci_timing::scaled_timeout(Duration::from_secs(10)); + let store = Arc::new( + create_store(&pool, schema) + .await + .with_prune_lock_timeout(ci_timing::scaled_timeout(Duration::from_secs(20))) + .expect("bound the controlled prune lock wait"), + ); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 315 }, + InsertOpts { + queue: queue.into(), + metadata: serde_json::json!({"ready_prune_guard": "retained payload"}), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(300)) + .await + .expect("claim the retained failed job"); + assert_eq!(claimed.len(), 1); + let run_lease = claimed[0].job.run_lease; + store + .fail_terminal(&pool, job_id, run_lease, "retained failure", None) + .await + .expect("fail the retained job") + .expect("the attempt was running"); + let before = store + .load_job(&pool, job_id) + .await + .expect("load the original failed job") + .expect("the failed job exists"); + assert_eq!(before.state, JobState::Failed); + assert!(matches!( + store.rotate(&pool).await.expect("seal slot zero"), + RotateOutcome::Rotated { slot: 1, .. } + )); + + // Distinct one-connection pools keep the observed backend identities + // stable across the two public operations and their relation waits. + let pruner_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + (*pool.connect_options()) + .clone() + .application_name("awa_ready_carry_pruner"), + ) + .await + .expect("connect the controlled pruner"); + let reader_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + (*pool.connect_options()) + .clone() + .application_name("awa_ready_carry_reader"), + ) + .await + .expect("connect the controlled reader"); + let pruner_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&pruner_pool) + .await + .expect("identify the pruner"); + let reader_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&reader_pool) + .await + .expect("identify the reader"); + + let mut carry_gate = pool.begin().await.expect("begin the carry gate"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.done_entries_1 IN SHARE MODE" + ))) + .execute(carry_gate.as_mut()) + .await + .expect("hold the live carry destination"); + let prune_store = store.clone(); + let prune_connection = pruner_pool.clone(); + let mut pruner = tokio::spawn(async move { + prune_store + .prune_oldest(&prune_connection, Duration::from_secs(3600)) + .await + }); + let mut reader = None; + let wait_sql = r#" + SELECT activity.query + FROM pg_stat_activity AS activity + JOIN pg_locks AS locks ON locks.pid = activity.pid + WHERE activity.pid = $1 + AND activity.datid = (SELECT oid FROM pg_database + WHERE datname = current_database()) + AND activity.state = 'active' + AND activity.wait_event_type = 'Lock' + AND activity.query LIKE $3 + AND locks.locktype = 'relation' + AND locks.relation = to_regclass($2) + AND NOT locks.granted + LIMIT 1 + "#; + let observed = tokio::time::timeout(bound, async { + loop { + let waiting: Option = sqlx::query_scalar(wait_sql) + .bind(pruner_pid) + .bind(format!("{schema}.done_entries_1")) + .bind(format!("%INSERT INTO {schema}.done_entries%")) + .fetch_optional(&pool) + .await?; + if waiting.is_some() { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + + let read_store = store.clone(); + let read_connection = reader_pool.clone(); + reader = Some(tokio::spawn(async move { + read_store.load_job(&read_connection, job_id).await + })); + loop { + let waiting: Option = sqlx::query_scalar(wait_sql) + .bind(reader_pid) + .bind(format!("{schema}.ready_entries_0")) + .bind(format!("%{schema}.ready_entries%")) + .fetch_optional(&pool) + .await?; + if let Some(query) = waiting { + return Ok::<_, sqlx::Error>(query); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await; + + // Release every gate and join both tasks before outcome assertions, so + // the old-source negative control cannot strand a read or lock holder. + let released = tokio::time::timeout(bound, carry_gate.rollback()).await; + let prune_result = match tokio::time::timeout(bound, &mut pruner).await { + Ok(result) => Some(result), + Err(_) => { + pruner.abort(); + let _ = pruner.await; + None + } + }; + let read_result = match reader { + Some(mut task) => match tokio::time::timeout(bound, &mut task).await { + Ok(result) => Some(result), + Err(_) => { + task.abort(); + let _ = task.await; + None + } + }, + None => None, + }; + let pruner_closed = tokio::time::timeout(bound, pruner_pool.close()).await; + let reader_closed = tokio::time::timeout(bound, reader_pool.close()).await; + + released + .expect("release the carry gate within the bound") + .expect("roll back the carry gate"); + pruner_closed.expect("close the pruner pool"); + reader_closed.expect("close the reader pool"); + let waiting_query = observed + .expect("observe both exact relation waits within the bound") + .expect("query the relation waits"); + let pruned = prune_result + .expect("prune finishes within the bound") + .expect("join the pruner") + .expect("prune succeeds after the carry gate releases"); + assert!( + matches!( + pruned, + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } + ), + "the real prune must carry the failed row: {pruned:?}" + ); + let loaded = read_result + .expect("reader finishes within the bound") + .expect("join the reader") + .expect("the public read succeeds") + .unwrap_or_else(|| { + panic!("retained failed job disappeared; reader waited in {waiting_query}") + }); + let fresh = store + .load_job(&pool, job_id) + .await + .expect("load after prune commits") + .expect("the retained failed job still exists"); + for row in [&loaded, &fresh] { + assert_eq!(row.id, job_id); + assert_eq!(row.state, JobState::Failed); + assert_eq!(row.run_lease, run_lease); + assert_eq!(row.args, before.args); + assert_eq!(row.metadata, before.metadata); + assert_eq!(row.finalized_at, before.finalized_at); + } +} + +async fn claim_aged_rescue_body_job( + schema: &str, + receipts: bool, + deadline: Duration, + max_attempts: i16, +) -> ( + TestDbGuard, + sqlx::PgPool, + QueueStorage, + awa::model::ClaimedRuntimeJob, +) { + let (guard, pool) = setup_pool(6).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: receipts, + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let queue = "qs_aged_rescue_body"; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9027 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts, + metadata: serde_json::json!({"retained": "aged-rescue-body"}), + ..Default::default() + }, + ) + .await; + // Control elapsed queue age without waiting five minutes; all claim, + // rescue, terminal and retry transitions below use their real APIs. + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + let mut claimed = store + .claim_runtime_batch_with_aging(&pool, queue, 1, deadline, Duration::from_secs(60)) + .await + .expect("claim aged job"); + assert_eq!(claimed.len(), 1); + let original = claimed.remove(0); + assert_eq!(original.job.id, job_id); + assert_eq!(original.claim.priority, 4, "physical ready lane"); + assert_eq!(original.job.priority, 1, "aged effective priority"); + (guard, pool, store, original) +} + +#[derive(Clone, Copy)] +enum AgedFailedConsumer { + Retry, + Carry, + View, +} + +async fn aged_rescued_failure_keeps_body( + schema: &str, + receipts: bool, + deadline: Duration, + consumer: AgedFailedConsumer, +) { + let (_guard, pool, store, original) = + claim_aged_rescue_body_job(schema, receipts, deadline, 1).await; + let job_id = original.job.id; + // Zero staleness makes this committed attempt eligible without sleeps + // or rewriting receipt timestamps. No worker is running in this test. + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue exhausted aged attempt"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].state, JobState::Failed); + assert_eq!(rescued[0].args, original.job.args); + assert_eq!(rescued[0].max_attempts, 1); + let raw_body = done_body_columns(&pool, &store, job_id).await; + assert!( + raw_body.0.is_none() && raw_body.1.is_none(), + "ready-backed failure stays narrow" + ); + + if matches!(consumer, AgedFailedConsumer::Carry) { + assert!(matches!( + store.rotate(&pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + let prune_deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("carry failed row") + { + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } => break, + PruneOutcome::Blocked { slot: 0 } if Instant::now() < prune_deadline => { + tokio::time::sleep(Duration::from_millis(20)).await; + } + other => panic!("unexpected carry prune outcome: {other:?}"), + } + } + let carried = done_body_columns(&pool, &store, job_id).await; + assert_eq!( + carried.0, + Some(original.job.args.clone()), + "carry must hydrate before truncating ready backing" + ); + assert_eq!(carried.1, Some(original.job.max_attempts)); + assert_eq!(carried.3, Some(original.job.created_at)); + let ready_count: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.ready_entries WHERE job_id = $1" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("count pruned ready body"); + assert_eq!( + ready_count, 0, + "carried body must survive without ready backing" + ); + } + + // This is the installed SQL view, independent of the Rust hydration helper. + let view = terminal_view_body_columns(&pool, &store, job_id).await; + if matches!(consumer, AgedFailedConsumer::View) { + assert_eq!( + view.0, original.job.args, + "terminal_jobs must hydrate the aged failure" + ); + assert_eq!(view.1, original.job.max_attempts); + } + let retained = store + .load_job(&pool, job_id) + .await + .expect("load retained failure") + .expect("retained failure"); + let retry = store + .retry_job(&pool, job_id) + .await + .expect("retry aged failure") + .expect("failed job retries"); + assert_eq!( + retry.args, original.job.args, + "manual retry must retain original arguments" + ); + assert_eq!( + retry.max_attempts, original.job.max_attempts, + "manual retry must not default to25" + ); + assert_eq!(retry.created_at, original.job.created_at); + assert_eq!(retry.metadata["retained"], "aged-rescue-body"); + assert_eq!(retry.state, JobState::Available); + assert_eq!(retry.attempt, 0); + assert_eq!(view.0, original.job.args); + assert_eq!(view.1, original.job.max_attempts); + assert_eq!(view.3, original.job.created_at); + assert_eq!(retained.args, original.job.args); + assert_eq!(retained.max_attempts, original.job.max_attempts); + let rerun = store + .claim_runtime_batch(&pool, &original.job.queue, 1, Duration::from_secs(300)) + .await + .expect("claim manual retry"); + assert_eq!(rerun.len(), 1); + assert_eq!(rerun[0].job.args, original.job.args); + assert_eq!(rerun[0].job.max_attempts, original.job.max_attempts); + assert_eq!(rerun[0].job.run_lease, original.job.run_lease + 1); +} + +#[tokio::test] +async fn test_aged_rescued_failure_compact_retry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_compact", + true, + Duration::ZERO, + AgedFailedConsumer::Retry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_deadline_retry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_deadline", + true, + Duration::from_secs(300), + AgedFailedConsumer::Retry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_legacy_retry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_legacy", + false, + Duration::from_secs(300), + AgedFailedConsumer::Retry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_carry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_carry", + true, + Duration::ZERO, + AgedFailedConsumer::Carry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_terminal_view_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_view", + true, + Duration::ZERO, + AgedFailedConsumer::View, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_expired_deadline_retry_preserves_body() { + let (_guard, pool, store, original) = claim_aged_rescue_body_job( + "awa_qs_aged_rescue_expired_deadline", + true, + Duration::from_micros(1), + 1, + ) + .await; + // The real positive deadline has expired by the time its claim commits. + // This exercises deadline rescue itself, separately from heartbeat rescue. + let rescued = store + .rescue_expired_deadlines(&pool) + .await + .expect("rescue expired aged attempt"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].id, original.job.id); + assert_eq!(rescued[0].state, JobState::Failed); + let view = terminal_view_body_columns(&pool, &store, original.job.id).await; + let retry = store + .retry_job(&pool, original.job.id) + .await + .expect("retry expired aged failure") + .expect("failed job retries"); + assert_eq!(retry.args, original.job.args); + assert_eq!(retry.max_attempts, original.job.max_attempts); + assert_eq!(retry.created_at, original.job.created_at); + assert_eq!(view.0, original.job.args); + assert_eq!(view.1, original.job.max_attempts); +} + +async fn aged_rescue_keeps_effective_policy_priority(dlq: bool) { + let schema = if dlq { + "awa_qs_aged_rescue_dlq" + } else { + "awa_qs_aged_rescue_retryable" + }; + let max_attempts = if dlq { 1 } else { 3 }; + let (_guard, pool, store, original) = + claim_aged_rescue_body_job(schema, true, Duration::ZERO, max_attempts).await; + let rescued = store + .rescue_stale_heartbeats_with_dlq(&pool, Duration::ZERO, &|_| dlq) + .await + .expect("rescue aged policy job"); + assert_eq!(rescued.len(), 1); + assert_eq!( + rescued[0].priority, 1, + "deferred/DLQ policy keeps effective priority" + ); + assert_eq!(rescued[0].args, original.job.args); + assert_eq!(rescued[0].max_attempts, max_attempts); + let table = if dlq { "dlq_entries" } else { "deferred_jobs" }; + let stored: (i16, serde_json::Value, i16) = sqlx::query_as(audited_sql(format!( + "SELECT priority, args, max_attempts FROM {schema}.{table} WHERE job_id = $1" + ))) + .bind(original.job.id) + .fetch_one(&pool) + .await + .expect("load self-contained policy row"); + assert_eq!(stored, (1, original.job.args, max_attempts)); +} + +#[tokio::test] +async fn test_aged_rescued_failure_dlq_keeps_effective_priority() { + aged_rescue_keeps_effective_policy_priority(true).await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_retryable_keeps_effective_priority() { + aged_rescue_keeps_effective_policy_priority(false).await; +} + +async fn moved_rescued_failure_keeps_live_lane( + initial_priority: i16, + runtime_aging: Duration, + consumer: AgedFailedConsumer, +) { + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_moved_rescue_identity"; + let queue = "qs_moved_rescue_identity"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9028 }, + InsertOpts { + queue: queue.into(), + priority: initial_priority, + max_attempts: 1, + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + assert_eq!( + store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("move queued job"), + vec![job_id] + ); + let claimed = store + .claim_runtime_batch_with_aging(&pool, queue, 1, Duration::ZERO, runtime_aging) + .await + .expect("claim moved job"); + assert_eq!(claimed.len(), 1); + let physical_priority = initial_priority - 1; + assert_eq!(claimed[0].claim.priority, physical_priority); + assert_eq!(claimed[0].job.priority, 1); + // A physical reorder of the disposable fixture preserves all logical + // tuples while exposing both legal fetch orders of the retained body. + sqlx::raw_sql(audited_sql(format!( + "CLUSTER {schema}.ready_entries_0 USING ready_entries_0_pkey" + ))) + .execute(&pool) + .await + .expect("order retained fixture rows by physical lane"); + let physical_order: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT priority FROM {schema}.ready_entries_0 WHERE job_id = $1 ORDER BY ctid" + ))) + .bind(job_id) + .fetch_all(&pool) + .await + .expect("verify fixture row order"); + assert_eq!(physical_order, vec![physical_priority, initial_priority]); + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue moved job"); + assert_eq!(rescued.len(), 1); + if matches!(consumer, AgedFailedConsumer::Carry) { + // Queued aging leaves its tombstoned source lane's cursor behind. + // A real empty claim consumes it before that queue slot can prune. + let remaining = store + .claim_runtime_batch(&pool, queue, 10, Duration::ZERO) + .await + .expect("advance tombstoned source lane"); + assert!( + remaining.is_empty(), + "closed moved attempt must not run again" + ); + assert!(matches!( + store.rotate(&pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("carry moved failure") + { + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } => break, + PruneOutcome::Blocked { slot: 0 } if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(20)).await + } + other => panic!("unexpected moved-failure prune outcome: {other:?}"), + } + } + let body = done_body_columns(&pool, &store, job_id).await; + assert_eq!(body.0, Some(claimed[0].job.args.clone())); + assert_eq!(body.1, Some(1)); + } + let terminal: (i16, serde_json::Value, i16) = sqlx::query_as(audited_sql(format!( + "SELECT priority,args,max_attempts FROM {schema}.terminal_jobs WHERE job_id=$1" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("read moved terminal view"); + if matches!(consumer, AgedFailedConsumer::View) { + assert_eq!( + terminal.0, physical_priority, + "terminal view must retain the live lane" + ); + assert_eq!(terminal.1, claimed[0].job.args); + assert_eq!(terminal.2, 1); + } + let retried = store + .retry_job(&pool, job_id) + .await + .expect("retry moved failure") + .expect("failed job retries"); + assert_eq!( + retried.priority, claimed[0].claim.priority, + "retry must not select the tombstoned source lane" + ); + assert_eq!(retried.args, claimed[0].job.args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(terminal.0, physical_priority); + assert_eq!(terminal.1, claimed[0].job.args); + assert_eq!( + retried.metadata.get("_awa_original_priority"), + Some(&serde_json::json!(initial_priority)) + ); +} + +#[tokio::test] +async fn test_moved_rescued_failure_retry_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(2, Duration::ZERO, AgedFailedConsumer::Retry).await; +} + +#[tokio::test] +async fn test_moved_rescued_failure_aged_retry_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(4, Duration::from_secs(60), AgedFailedConsumer::Retry) + .await; +} + +#[tokio::test] +async fn test_moved_rescued_failure_aged_carry_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(4, Duration::from_secs(60), AgedFailedConsumer::Carry) + .await; +} + +#[tokio::test] +async fn test_moved_rescued_failure_aged_view_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(4, Duration::from_secs(60), AgedFailedConsumer::View) + .await; +} + +async fn terminal_counts_by_priority( + pool: &sqlx::PgPool, + schema: &str, + queue: &str, +) -> Vec<(i16, i64)> { + sqlx::query_as(audited_sql(format!( + r#" + SELECT priority, sum(n)::bigint + FROM ( + SELECT priority, live_terminal_count AS n + FROM {schema}.queue_terminal_live_counts WHERE queue = $1 + UNION ALL + SELECT priority, terminal_delta AS n + FROM {schema}.queue_terminal_count_deltas WHERE queue = $1 + ) AS counters + GROUP BY priority HAVING sum(n) <> 0 ORDER BY priority + "# + ))) + .bind(queue) + .fetch_all(pool) + .await + .expect("read per-priority terminal accounting") +} + +#[derive(Clone, Copy, Debug)] +enum LegacyFailedConsumer { + Retry, + Carry, + SingleDlq, + BulkDlq, + Discard, +} + +async fn legacy_aged_failure_recovers_body(moved: bool, consumer: LegacyFailedConsumer) { + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_legacy_aged_body"; + let queue = "qs_legacy_aged_body"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9029 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 1, + metadata: serde_json::json!({"retained": "legacy-aged-body", "ready_only": true, "shared": "ready"}), + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + if moved { + assert_eq!( + store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("move queued job"), + vec![job_id] + ); + } + let claimed = store + .claim_runtime_batch_with_aging(&pool, queue, 1, Duration::ZERO, Duration::from_secs(60)) + .await + .expect("claim aged legacy-shape fixture"); + assert_eq!(claimed.len(), 1); + let original = &claimed[0].job; + let physical_priority = if moved { 3 } else { 4 }; + assert_eq!(claimed[0].claim.priority, physical_priority); + assert_eq!(original.priority, 1); + assert_eq!( + store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue aged fixture") + .len(), + 1 + ); + + // Reproduce the persisted shape of an older writer: its narrow failed + // row and positive terminal delta used the effective attempt priority. + // The ready body and all real claim/closure evidence remain intact. + assert_eq!( + sqlx::query(audited_sql(format!( + "UPDATE {schema}.done_entries SET priority = $2 WHERE job_id = $1" + ))) + .bind(job_id) + .bind(original.priority) + .execute(&pool) + .await + .expect("seed old terminal priority") + .rows_affected(), + 1 + ); + assert_eq!( + sqlx::query(audited_sql(format!( + "UPDATE {schema}.queue_terminal_count_deltas SET priority = $2 WHERE queue = $1" + ))) + .bind(queue) + .bind(original.priority) + .execute(&pool) + .await + .expect("seed matching old terminal delta") + .rows_affected(), + 1 + ); + // An old moved-body lookup could retain the source payload, losing + // live-lane provenance. Keep terminal errors/custom overrides authoritative. + let terminal_payload: serde_json::Value = sqlx::query_scalar(audited_sql(format!( + r#"UPDATE {schema}.done_entries + SET payload = jsonb_set( + COALESCE(payload, '{{}}'::jsonb), '{{metadata}}', + (COALESCE(payload -> 'metadata', '{{}}'::jsonb) - '_awa_original_priority' - 'ready_only') + || '{{"terminal_only":true,"shared":"terminal"}}'::jsonb + ) || '{{"tags":["terminal-tag"],"progress":{{"terminal":true}}}}'::jsonb + WHERE job_id = $1 RETURNING payload"# + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("seed old terminal payload without live provenance"); + assert!(terminal_payload["metadata"] + .get("_awa_original_priority") + .is_none()); + assert!(!terminal_payload["errors"].as_array().unwrap().is_empty()); + let body = done_body_columns(&pool, &store, job_id).await; + assert!( + body.0.is_none() && body.1.is_none(), + "old row requires its ready body" + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(1, 1)] + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 1); + + if matches!(consumer, LegacyFailedConsumer::Carry) { + let remaining = store + .claim_runtime_batch(&pool, queue, 10, Duration::ZERO) + .await + .expect("advance any tombstoned source cursor"); + assert!(remaining.is_empty()); + assert!(matches!( + store.rotate(&pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("carry old failed row") + { + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } => break, + PruneOutcome::Blocked { slot: 0 } if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(20)).await; + } + other => panic!("unexpected legacy carry outcome: {other:?}"), + } + } + let body = done_body_columns(&pool, &store, job_id).await; + assert_eq!(body.0, Some(original.args.clone())); + assert_eq!(body.1, Some(original.max_attempts)); + let ready_count: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.ready_entries WHERE job_id = $1" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + ready_count, 0, + "carry must preserve the body before its backing is pruned" + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(physical_priority, 1)], + "carried delta must match its replacement lane" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 1); + } + let retry = match consumer { + LegacyFailedConsumer::Retry | LegacyFailedConsumer::Carry => store + .retry_job(&pool, job_id) + .await + .expect("retry legacy failure") + .expect("legacy failed row retries"), + LegacyFailedConsumer::SingleDlq | LegacyFailedConsumer::BulkDlq => { + if matches!(consumer, LegacyFailedConsumer::SingleDlq) { + let moved = store + .move_failed_to_dlq(&pool, job_id, "manual-legacy-test") + .await + .expect("move old failed row to DLQ") + .expect("one row moves"); + assert_eq!(moved.id, job_id); + assert_eq!(moved.args, original.args); + assert_eq!(moved.max_attempts, original.max_attempts); + } else { + assert_eq!( + store + .bulk_move_failed_to_dlq( + &pool, + Some(CompleteJob::kind()), + Some(queue), + "manual-legacy-test", + ) + .await + .expect("bulk move old failed row to DLQ"), + 1, + "one terminal row must not multiply through moved ready bodies" + ); + } + let dlq_rows: Vec<(i16, serde_json::Value, i16, serde_json::Value)> = + sqlx::query_as(audited_sql(format!( + "SELECT priority,args,max_attempts,COALESCE(payload,'{{}}'::jsonb) FROM {schema}.dlq_entries WHERE job_id=$1" + ))).bind(job_id).fetch_all(&pool).await.expect("read actual stored DLQ row"); + assert_eq!(dlq_rows.len(), 1); + assert_eq!(dlq_rows[0].0, physical_priority); + assert_eq!(dlq_rows[0].1, original.args); + assert_eq!(dlq_rows[0].2, original.max_attempts); + assert_eq!(dlq_rows[0].3["metadata"]["_awa_original_priority"], 4); + assert_eq!(dlq_rows[0].3["metadata"]["shared"], "terminal"); + assert_eq!(dlq_rows[0].3["errors"], terminal_payload["errors"]); + assert_eq!(done_entries_count(&pool, schema, queue).await, 0); + assert!( + terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty(), + "DLQ move must decrement the stored priority-1 terminal key" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); + store + .retry_from_dlq(&pool, job_id, &awa::model::RetryFromDlqOpts::default()) + .await + .expect("retry recovered DLQ row") + .expect("DLQ row retries") + } + LegacyFailedConsumer::Discard => { + assert_eq!( + store + .discard_failed_by_kind(&pool, CompleteJob::kind()) + .await + .expect("discard old failed row"), + 1, + "one terminal row must not multiply through moved ready bodies" + ); + assert_eq!(done_entries_count(&pool, schema, queue).await, 0); + let dlq_count: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.dlq_entries WHERE job_id=$1" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(dlq_count, 0); + assert!( + terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty(), + "discard must decrement exactly its stored terminal lane" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); + return; + } + }; + assert_eq!(retry.args, original.args); + assert_eq!(retry.max_attempts, original.max_attempts); + assert_eq!(retry.created_at, original.created_at); + assert_eq!( + retry.priority, physical_priority, + "manual retry restarts aging from its live lane" + ); + assert_eq!(retry.metadata["retained"], "legacy-aged-body"); + assert_eq!( + retry.metadata["ready_only"], true, + "recover ready metadata defaults" + ); + assert_eq!(retry.metadata["shared"], "terminal", "stored metadata wins"); + assert_eq!(retry.metadata["terminal_only"], true); + if moved { + assert_eq!(retry.metadata["_awa_original_priority"], 4); + } + assert_eq!(retry.tags, vec!["terminal-tag"]); + assert_eq!( + serde_json::to_value(&retry.errors).unwrap(), + terminal_payload["errors"] + ); + let expected_progress = if matches!( + consumer, + LegacyFailedConsumer::SingleDlq | LegacyFailedConsumer::BulkDlq + ) { + None // Existing DLQ retry clears progress before making the job ready. + } else { + Some(serde_json::json!({"terminal": true})) + }; + assert_eq!(retry.progress, expected_progress); + assert_eq!(retry.attempt, 0); + assert_eq!(retry.state, JobState::Available); + assert!( + terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty(), + "remove the stored terminal lane, not the recovered retry lane" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); +} + +#[tokio::test] +async fn test_legacy_aged_failure_retry_recovers_body() { + legacy_aged_failure_recovers_body(false, LegacyFailedConsumer::Retry).await; +} + +#[tokio::test] +async fn test_legacy_aged_moved_failure_retry_recovers_live_body() { + legacy_aged_failure_recovers_body(true, LegacyFailedConsumer::Retry).await; +} + +#[tokio::test] +async fn test_legacy_aged_moved_failure_carry_recovers_body_and_counters() { + legacy_aged_failure_recovers_body(true, LegacyFailedConsumer::Carry).await; +} + +#[tokio::test] +async fn test_legacy_aged_moved_terminal_dlq_and_discard_consumers() { + tokio::task::LocalSet::new() + .run_until(async { + let mut failures = Vec::new(); + for consumer in [ + LegacyFailedConsumer::SingleDlq, + LegacyFailedConsumer::BulkDlq, + LegacyFailedConsumer::Discard, + ] { + // Exercise every parameter even after another panics. These + // database futures do not require Send on a local task. + if let Err(error) = + tokio::task::spawn_local(legacy_aged_failure_recovers_body(true, consumer)) + .await + { + failures.push(format!("{consumer:?}: {error}")); + } + } + assert!( + failures.is_empty(), + "terminal consumer failures: {failures:?}" + ); + }) + .await; +} + +#[tokio::test] +async fn test_legacy_narrow_available_cancel_retry_keeps_tombstoned_body() { + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_legacy_cancel_body"; + let queue = "qs_legacy_cancel_body"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9030 }, + InsertOpts { + queue: queue.into(), + priority: 3, + max_attempts: 7, + metadata: serde_json::json!({"retained": "legacy-cancel-body"}), + ..Default::default() + }, + ) + .await; + let cancelled = store + .cancel_job(&pool, job_id) + .await + .expect("cancel available job") + .expect("available job cancels"); + assert_eq!(cancelled.state, JobState::Cancelled); + let backing: (i64, i64, bool) = sqlx::query_as(audited_sql(format!( + r#" + SELECT done.run_lease, ready.run_lease, EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq AND tomb.job_id = ready.job_id + ) + FROM {schema}.done_entries done JOIN {schema}.ready_entries ready + ON ready.job_id = done.job_id AND ready.priority = done.priority + WHERE done.job_id = $1 + "# + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("inspect cancelled body identity"); + assert_eq!( + backing.0, backing.1, + "unclaimed cancellation keeps the same run lease" + ); + assert!(backing.2, "cancellation tombstones its retained body"); + // Older narrow cancellations legitimately hydrate a same-lease body + // even though cancellation has already tombstoned that exact lane. + assert_eq!(sqlx::query(audited_sql(format!( + "UPDATE {schema}.done_entries SET args=NULL, max_attempts=NULL, run_at=NULL, created_at=NULL WHERE job_id=$1" + ))).bind(job_id).execute(&pool).await.unwrap().rows_affected(), 1); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(3, 1)] + ); + let retry = store + .retry_job(&pool, job_id) + .await + .expect("retry narrow cancellation") + .expect("cancelled job retries"); + assert_eq!(retry.args, cancelled.args); + assert_eq!(retry.max_attempts, 7); + assert_eq!(retry.priority, 3); + assert_eq!(retry.created_at, cancelled.created_at); + assert_eq!(retry.metadata["retained"], "legacy-cancel-body"); + assert!(terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty()); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); +}