From 303b7ad9873ee9f685bec69b89b1956e8ea935e4 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 18:57:44 +1300 Subject: [PATCH 1/2] fix(worker): bound finished attempts waiting for finalization (0.6) The executor removed an attempt from the heartbeat set and released its worker slot as soon as the handler returned, then queued its finalize on the completion batcher. When finalization was slower than claiming, the dispatcher kept claiming and the queue of finished, unheartbeated attempts grew without bound. Those that waited past heartbeat_staleness were rescued and run again. Admit at most the completion pipeline's depth (2 x batch size x shards, and at least one per worker) of attempts to finalization. Past that, a finished attempt keeps its worker slot and heartbeat until there is room, which stops the dispatcher claiming. This is the bound main already has. --- CHANGELOG.md | 12 ++ awa-worker/src/client.rs | 2 + awa-worker/src/completion.rs | 19 ++- awa-worker/src/dispatcher.rs | 2 + awa-worker/src/executor.rs | 35 ++++- awa/tests/ci_timing.rs | 186 ++++++++++++++++++++++++ awa/tests/queue_storage_runtime_test.rs | 147 +++++++++++++++++++ 7 files changed, 397 insertions(+), 6 deletions(-) create mode 100644 awa/tests/ci_timing.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..9527d89a4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,18 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **A completion backlog no longer gets finished jobs rescued and run + again.** An attempt stops heartbeating when its handler returns, and its + finalize then waits for the completion batcher. The runtime released the + worker slot at the same moment and kept claiming, so when finalization was + slower than claiming (for example, retried jobs, whose completions take + the slower per-row path) the queue of finished, unheartbeated attempts grew + without bound. Attempts that waited longer than `heartbeat_staleness` were + rescued and executed again, some of them twice. At most `2 × batch size × + completion shards` attempts (1,024 by default), and at least one per + worker, now wait for finalization. Past that, a finished attempt keeps its + worker slot and its heartbeat until there is room, so the dispatcher stops + claiming until finalization catches up. - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index e5a498dbb..42ff2aee2 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -1566,6 +1566,8 @@ impl Client { completion_handle, effective_storage.clone(), self.dlq_policy.clone(), + self.service_cancel.clone(), + runtime_worker_capacity, )); // Admin cancellation listener: fires the in-flight cancel flag diff --git a/awa-worker/src/completion.rs b/awa-worker/src/completion.rs index 87d3646ff..ac54468ca 100644 --- a/awa-worker/src/completion.rs +++ b/awa-worker/src/completion.rs @@ -105,9 +105,19 @@ fn completion_sort_key(request: &CompletionRequest) -> (i32, i64, RunLease) { #[derive(Clone)] pub(crate) struct CompletionBatcherHandle { shards: Vec>, + pipeline_capacity: usize, } impl CompletionBatcherHandle { + /// How many finished attempts may wait for finalization at once: enough + /// to keep every shard's in-flight and next batch full, and at least one + /// per worker. + pub fn finalizer_capacity(&self, runtime_worker_capacity: u32) -> usize { + self.pipeline_capacity + .max(runtime_worker_capacity.max(1) as usize) + .min(tokio::sync::Semaphore::MAX_PERMITS) + } + pub async fn complete(&self, job_id: i64, run_lease: RunLease) -> Result { self.complete_inner(job_id, run_lease, None, None).await } @@ -185,7 +195,14 @@ impl CompletionBatcher { }); } - (Self { workers }, CompletionBatcherHandle { shards }) + ( + Self { workers }, + CompletionBatcherHandle { + shards, + // Each shard needs an in-flight batch and a queued batch. + pipeline_capacity: batch_size.saturating_mul(shard_count).saturating_mul(2), + }, + ) } pub fn spawn(self) -> Vec> { diff --git a/awa-worker/src/dispatcher.rs b/awa-worker/src/dispatcher.rs index 11e3e9dbd..08d55c02b 100644 --- a/awa-worker/src/dispatcher.rs +++ b/awa-worker/src/dispatcher.rs @@ -851,6 +851,8 @@ mod tests { completion, storage.clone(), crate::executor::DlqPolicy::default(), + CancellationToken::new(), + 1, )); let config = QueueConfig { max_workers: 1, diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 3d99ec8df..99b8c9c18 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -11,6 +11,8 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::sync::Arc; use std::time::Duration; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; use tracing::{error, info, info_span, warn, Instrument}; /// Result of executing a job handler. @@ -267,6 +269,8 @@ pub struct JobExecutor { completion_batcher: CompletionBatcherHandle, storage: RuntimeStorage, dlq_policy: DlqPolicy, + finalize_cancel: CancellationToken, + finalize_capacity: Arc, } impl JobExecutor { @@ -288,7 +292,12 @@ impl JobExecutor { completion_batcher: CompletionBatcherHandle, storage: RuntimeStorage, dlq_policy: DlqPolicy, + finalize_cancel: CancellationToken, + runtime_worker_capacity: u32, ) -> Self { + let finalize_capacity = Arc::new(Semaphore::new( + completion_batcher.finalizer_capacity(runtime_worker_capacity), + )); Self { pool, workers, @@ -301,6 +310,8 @@ impl JobExecutor { completion_batcher, storage, dlq_policy, + finalize_cancel, + finalize_capacity, } } @@ -327,6 +338,8 @@ impl JobExecutor { let completion_batcher = self.completion_batcher.clone(); let storage = self.storage.clone(); let dlq_policy = self.dlq_policy.clone(); + let finalize_cancel = self.finalize_cancel.clone(); + let finalize_capacity = self.finalize_capacity.clone(); let job_id = job.id; let job_run_lease = job.run_lease; let job_kind = job.kind.clone(); @@ -430,17 +443,28 @@ impl JobExecutor { guard.clone_latest() }; - // Remove from in-flight immediately after the handler returns and - // the progress snapshot is captured. This keeps local worker - // capacity tied to active handler execution, not to the tail - // latency of durable completion bookkeeping. + // Leaving `in_flight` ends the attempt's heartbeat, so only a + // bounded number of finished attempts may wait for finalization. + // While that queue is full the attempt keeps its worker slot and + // its heartbeat, and the dispatcher stops claiming: a completion + // backlog cannot outgrow `heartbeat_staleness` and get finished + // attempts rescued and run again. + let finalize_permit = match result { + Some(_) => tokio::select! { + permit = finalize_capacity.acquire_owned() => { + Some(permit.expect("finalizer capacity is never closed")) + } + _ = finalize_cancel.cancelled() => None, + }, + None => None, + }; in_flight.remove((job_id, job_run_lease)); if let Some(counter) = queue_in_flight.get(&job_queue) { counter.fetch_sub(1, Ordering::SeqCst); } metrics.record_in_flight_change(&job_queue, -1); - let Some(result) = result else { + let (Some(result), Some(finalize_permit)) = (result, finalize_permit) else { return; }; let dlq_enabled = dlq_policy.enabled_for(&job_queue); @@ -487,6 +511,7 @@ impl JobExecutor { outcome => break outcome, } }; + drop(finalize_permit); match &outcome { Ok(CompletionOutcome::Applied { terminal, .. }) => { diff --git a/awa/tests/ci_timing.rs b/awa/tests/ci_timing.rs new file mode 100644 index 000000000..3704a0ba7 --- /dev/null +++ b/awa/tests/ci_timing.rs @@ -0,0 +1,186 @@ +//! Shared CI contention scaling for the nightly chaos and benchmark suites. +//! +//! The nightly suites run on shared GitHub runners whose CPU allocation +//! varies run to run. Three assertion shapes are sensitive to that and +//! flaked repeatedly through 2026-07 (#399, #434): +//! +//! 1. **Wall-clock waits.** A test waits N seconds for a state the +//! runtime reaches in milliseconds when it has a core to itself. +//! Scale with [`scaled_timeout`]. +//! +//! 2. **Aggressive rescue cadences.** A chaos client sets +//! `heartbeat_staleness` a few multiples above `heartbeat_interval` +//! so its own backdating triggers rescue promptly. Under contention +//! a *live* worker's heartbeat can miss that window, so the runtime +//! correctly rescues a healthy attempt — and the test sees a genuine +//! duplicate completion, which reads as a correctness failure. Scale +//! with [`scaled_staleness`]: the interval stays fast so rescue is +//! still exercised, but the staleness window gains the same margin +//! as the waits around it. +//! +//! 3. **Minimum-progress floors.** A gate asserts "at least N of these +//! happened" to catch a stalled subsystem. The floor must sit far +//! below the *observed* operating point, not just below the nominal +//! one. Scale with [`contention_floor`]. +//! +//! Scaling only ever loosens a bound, and only on CI. A local run keeps +//! the strict values, so a real regression still fails fast on a +//! developer machine. Override with `AWA_CHAOS_TIMEOUT_MULTIPLIER`. +#![allow(dead_code)] + +use std::time::Duration; + +/// Upper bound on the override. `Duration::mul_f64` panics on a non-finite +/// or overflowing result, so an unbounded override would turn a typo into a +/// failed nightly rather than a loose bound. +const MAX_MULTIPLIER: f64 = 100.0; + +/// How much slack to give contention-sensitive bounds. `1.0` locally, +/// `3.0` on CI, or the `AWA_CHAOS_TIMEOUT_MULTIPLIER` override (clamped to +/// `1.0 ..= MAX_MULTIPLIER` so the override can only ever loosen, and can +/// never produce a duration `mul_f64` refuses). +pub fn chaos_timeout_multiplier() -> f64 { + let override_var = std::env::var("AWA_CHAOS_TIMEOUT_MULTIPLIER").ok(); + multiplier_from(override_var.as_deref(), std::env::var_os("CI").is_some()) +} + +/// The multiplier decision, with the environment passed in so it can be +/// tested without mutating process-global state. +fn multiplier_from(override_var: Option<&str>, is_ci: bool) -> f64 { + // `"inf"` and `"nan"` parse successfully as f64, and both reach + // `Duration::mul_f64`, which panics rather than saturating. Treat a + // non-finite override as no override at all. + if let Some(parsed) = override_var + .and_then(|raw| raw.parse::().ok()) + .filter(|parsed| parsed.is_finite()) + { + return parsed.clamp(1.0, MAX_MULTIPLIER); + } + + if is_ci { + 3.0 + } else { + 1.0 + } +} + +/// Grow a wait deadline by the contention multiplier. +pub fn scaled_timeout(timeout: Duration) -> Duration { + timeout.mul_f64(chaos_timeout_multiplier()) +} + +/// Grow a heartbeat-staleness window by the contention multiplier. +/// +/// Distinct from [`scaled_timeout`] only in intent: this one is passed to +/// `ClientBuilder::heartbeat_staleness`, where the cost of being too tight +/// is a spurious rescue of a live attempt rather than a timeout. Paired +/// heartbeat/rescue *intervals* are deliberately left unscaled so the +/// rescue path still runs at chaos cadence. +pub fn scaled_staleness(staleness: Duration) -> Duration { + staleness.mul_f64(chaos_timeout_multiplier()) +} + +/// Shrink a minimum-progress floor by the contention multiplier. +/// +/// Use for "this subsystem must have advanced at least N times" gates. +/// The returned floor is at least 1: the regression these gates exist to +/// catch is a fully stalled subsystem, so zero progress must still fail. +pub fn contention_floor(nominal: i64) -> i64 { + let scaled = (nominal as f64 / chaos_timeout_multiplier()).floor() as i64; + scaled.max(1) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn local_runs_keep_strict_bounds() { + assert_eq!(multiplier_from(None, false), 1.0); + } + + #[test] + fn ci_runs_get_margin() { + assert_eq!(multiplier_from(None, true), 3.0); + } + + #[test] + fn override_can_only_loosen() { + assert_eq!(multiplier_from(Some("0.1"), true), 1.0); + assert_eq!(multiplier_from(Some("-5"), true), 1.0); + assert_eq!(multiplier_from(Some("6"), false), 6.0); + } + + #[test] + fn unparseable_override_falls_back_to_the_environment() { + assert_eq!(multiplier_from(Some("banana"), true), 3.0); + assert_eq!(multiplier_from(Some(""), false), 1.0); + } + + #[test] + fn non_finite_override_falls_back_to_the_environment() { + // These all parse as f64 and would otherwise reach `mul_f64`, which + // panics on a non-finite result instead of saturating. + for raw in ["inf", "-inf", "infinity", "NaN", "nan"] { + assert_eq!(multiplier_from(Some(raw), true), 3.0, "raw={raw}"); + assert_eq!(multiplier_from(Some(raw), false), 1.0, "raw={raw}"); + } + } + + #[test] + fn oversized_override_is_clamped() { + assert_eq!(multiplier_from(Some("1e308"), true), MAX_MULTIPLIER); + assert_eq!(multiplier_from(Some("10000"), true), MAX_MULTIPLIER); + } + + #[test] + fn every_multiplier_produces_a_usable_duration() { + // `Duration::mul_f64` panics rather than saturating, so the clamp is + // what keeps a bad override from failing the nightly it was set to + // rescue. + for raw in [ + "inf", "-inf", "NaN", "1e308", "10000", "0", "-1", "banana", "", + ] { + for is_ci in [true, false] { + let m = multiplier_from(Some(raw), is_ci); + assert!(m.is_finite() && (1.0..=MAX_MULTIPLIER).contains(&m)); + // Would panic if the multiplier were unbounded. + let _ = Duration::from_secs(60).mul_f64(m); + } + } + } + + #[test] + fn scaling_a_timeout_grows_it() { + // 250ms staleness against a 50ms heartbeat interval is a 5x margin + // locally; on CI it becomes 15x, which is what stops a contended + // runner's live worker from being rescued as if it had died. + assert_eq!( + Duration::from_millis(250).mul_f64(multiplier_from(None, true)), + Duration::from_millis(750) + ); + } + + #[test] + fn progress_floor_keeps_margin_under_the_observed_rate() { + // The receipt gate's nominal floor is duration_secs / 4 = 45 for a + // 180s run. #399 saw a healthy run produce 41, so the CI floor has + // to sit well under that. + let ci_floor = (45.0_f64 / multiplier_from(None, true)).floor() as i64; + assert_eq!(ci_floor, 15); + assert!( + ci_floor < 41, + "floor must clear the observed operating point" + ); + } + + #[test] + fn progress_floor_never_reaches_zero() { + // Whatever the multiplier, a fully pinned ring (zero rotations) has + // to keep failing — that is the regression the gate exists for. + assert_eq!(contention_floor(45), contention_floor(45).max(1)); + let absurd = (45.0_f64 / multiplier_from(Some("10000"), true)).floor() as i64; + assert_eq!(absurd.max(1), 1); + assert_eq!(contention_floor(0), 1); + } +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index a060369f8..8174c7d66 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -3,6 +3,8 @@ //! These tests exercise the full dispatcher/worker/maintenance wiring with the //! queue_storage backend enabled. +mod ci_timing; + use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, BatchOperationSpec, PruneOutcome, QueueStorage, QueueStorageConfig, RotateOutcome, SkipReason, @@ -2009,6 +2011,151 @@ impl Worker for MultiClientTrackingWorker { } } +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct FinalizeBacklogJob { + n: i64, +} + +/// A runtime's finished attempts stop heartbeating once they are handed to +/// finalization, so it must stop claiming when finalization falls behind. +/// Otherwise the backlog of finished, unheartbeated attempts grows with the +/// claim rate, outlives `heartbeat_staleness`, and rescue runs those jobs +/// again. Holding the first jobs' receipt-completion locks stalls the +/// completion batcher, standing in for a batcher that cannot keep up. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_stops_claiming_while_finalization_is_stalled() { + // The default completion pipeline: one shard of 512-job batches, one + // batch in flight and one queued. + const FINALIZER_CAPACITY: usize = 1_024; + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_finalize_backlog"; + let store_config = QueueStorageConfig { + schema: "awa_qs_finalize_backlog".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + for first in (0..JOBS).step_by(500) { + let params: Vec<_> = (first..(first + 500).min(JOBS)) + .map(|n| { + insert::params_with( + &FinalizeBacklogJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue backlog"); + } + + let first_jobs: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {}.ready_entries WHERE queue = $1 ORDER BY lane_seq LIMIT $2", + store.schema() + )) + .bind(queue) + .bind(i64::from(WORKERS) * 2) + .fetch_all(&pool) + .await + .expect("first jobs"); + let mut lock_holder = pool.acquire().await.expect("lock connection"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_lock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', job_id, 1), 0)) + FROM unnest($1::bigint[]) AS job_id", + ) + .bind(&first_jobs) + .execute(&mut *lock_holder) + .await + .expect("stall the first jobs' finalization"); + + let runs: Arc>> = Arc::default(); + let recorded = runs.clone(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: WORKERS, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: WORKERS as usize, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register::(move |_, ctx| { + *recorded + .lock() + .expect("runs") + .entry(ctx.job.id) + .or_insert(0) += 1; + async { Ok(JobResult::Completed) } + }) + .heartbeat_staleness(Duration::from_secs(300)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build finalize backlog client"); + client.start().await.expect("start finalize backlog client"); + + // Wait for claiming to stop, or to run past the bound. A fixed runtime + // can never pass it, so a pause in claiming only weakens the check. + let limit = FINALIZER_CAPACITY + WORKERS as usize; + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + let mut last = (0, Instant::now()); + let started_while_stalled = loop { + let started = runs.lock().expect("runs").len(); + if started > limit { + break started; + } + if started != last.0 { + last = (started, Instant::now()); + } else if started > 0 && last.1.elapsed() >= Duration::from_secs(5) { + break started; + } + assert!( + Instant::now() < deadline, + "claiming neither stopped nor passed {limit}" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + }; + + sqlx::query("SELECT pg_catalog.pg_advisory_unlock_all()") + .execute(&mut *lock_holder) + .await + .expect("release the stalled finalization"); + drop(lock_holder); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(60)); + while completed_terminal_count(&pool, &store, queue).await < JOBS { + assert!(Instant::now() < deadline, "the backlog did not drain"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + client.shutdown(Duration::from_secs(10)).await; + + assert!( + started_while_stalled <= limit, + "{started_while_stalled} jobs started while finalization was stalled; \ + at most {limit} may wait without a heartbeat or hold a worker" + ); + let runs = runs.lock().expect("runs"); + assert_eq!(runs.len() as i64, JOBS); + assert!( + runs.values().all(|count| *count == 1), + "a job ran more than once" + ); +} + /// ADR-023 claim-ring control-plane smoke test. Exercises /// `rotate_claims`, the busy-check, and `prune_oldest_claims` on an /// empty schema: rotate cycles through every slot, prune is a noop From 9361839829ce9528ee52cfce46bcfc76ab3089f1 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 21:52:37 +1300 Subject: [PATCH 2/2] perf(queue-storage): complete metadata jobs by compact claim identity Address compact claims by their batch ID and item index, validating the job, attempt and receipt before closing. Avoid expanding every retained compact batch for each metadata or retry completion flush, and avoid locking immutable shared batch rows. Preserve per-attempt advisory serialization, row-local prelocking and the fresh deciding snapshot. A 50,000-claim recovery rehearsal with heartbeat protection timed out waiting 120 seconds for retries on PostgreSQL 17 and 18 before this change. With keyed completion, the same full rehearsal passes in 116.35s and 123.30s total. A new regression rejects a wrong item index, completes while the immutable batch row is locked elsewhere, and checks retained tags and metadata. Validation: cargo fmt --all; SQLX_OFFLINE=true cargo clippy --all-targets --all-features -- -D warnings; SQLX_OFFLINE=true cargo build --workspace; eight focused metadata, stale-completion, cancellation, idempotence and finalization-bound tests plus the shutdown-waits test on PostgreSQL 17 and 18. --- CHANGELOG.md | 11 +++- awa-model/src/queue_storage.rs | 44 ++++++++++----- awa/tests/queue_storage_runtime_test.rs | 74 +++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9527d89a4..9a3dd1d7d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- Queue-storage completion of retried jobs or jobs carrying metadata now + addresses compact claims by their batch identity and item index, as on main. + It no longer expands every retained claim batch for each completion flush, + so bounded finalization can drain large recovery backlogs without that scan. + - **A completion backlog no longer gets finished jobs rescued and run again.** An attempt stops heartbeating when its handler returns, and its finalize then waits for the completion batcher. The runtime released the @@ -13,9 +18,9 @@ Notable changes between releases. Detailed migration notes for storage transitio slower than claiming (for example, retried jobs, whose completions take the slower per-row path) the queue of finished, unheartbeated attempts grew without bound. Attempts that waited longer than `heartbeat_staleness` were - rescued and executed again, some of them twice. At most `2 × batch size × - completion shards` attempts (1,024 by default), and at least one per - worker, now wait for finalization. Past that, a finished attempt keeps its + rescued and executed again, some of them twice. The bound on attempts waiting + for finalization is `max(worker count, 2 × batch size × completion shards)`. + Past that, a finished attempt keeps its worker slot and its heartbeat until there is room, so the dispatcher stops claiming until finalization catches up. - **`admin::job_kind_overviews` on queue storage no longer depends on the diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 930aab715..3190ebf6e 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -6534,14 +6534,22 @@ impl QueueStorage { .expect("receipt-backed slow completion requires receipt_id") }) .collect(); + let receipt_claim_batch_ids: Vec> = receipt_claimed + .iter() + .map(|entry| entry.claim.claim_batch_id) + .collect(); + let receipt_claim_batch_indices: Vec> = receipt_claimed + .iter() + .map(|entry| entry.claim.claim_batch_index) + .collect(); let closure_rel = format!("{schema}.lease_claim_closures"); let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); let closed_evidence = receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); let updated: Vec<(i64, i64)> = sqlx::query_as(&format!( r#" - WITH completed(claim_slot, job_id, run_lease, receipt_id) AS ( - SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[]) + WITH completed(claim_slot, job_id, run_lease, receipt_id, claim_batch_id, claim_batch_index) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::bigint[], $4::bigint[], $5::bigint[], $6::int[]) ), locked_row_claims AS ( SELECT claims.claim_slot, claims.job_id, claims.run_lease @@ -6568,18 +6576,23 @@ impl QueueStorage { items.job_id, items.run_lease, items.receipt_id - FROM {schema}.lease_claim_batches AS claim_batches - CROSS JOIN LATERAL unnest( - claim_batches.job_ids, - claim_batches.run_leases, - claim_batches.receipt_ids - ) AS items(job_id, run_lease, receipt_id) - JOIN completed - ON completed.claim_slot = claim_batches.claim_slot - AND completed.job_id = items.job_id - AND completed.run_lease = items.run_lease - AND completed.receipt_id = items.receipt_id - WHERE NOT EXISTS ( + FROM completed + JOIN {schema}.lease_claim_batches AS claim_batches + ON claim_batches.claim_slot = completed.claim_slot + AND claim_batches.batch_id = completed.claim_batch_id + CROSS JOIN LATERAL ( + SELECT + claim_batches.job_ids[completed.claim_batch_index] AS job_id, + claim_batches.run_leases[completed.claim_batch_index] AS run_lease, + claim_batches.receipt_ids[completed.claim_batch_index] AS receipt_id + ) AS items + WHERE completed.claim_batch_id IS NOT NULL + AND completed.claim_batch_index IS NOT NULL + AND completed.claim_batch_index BETWEEN 1 AND claim_batches.claimed_count + AND completed.job_id = items.job_id + AND completed.run_lease = items.run_lease + AND completed.receipt_id = items.receipt_id + AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot @@ -6616,7 +6629,6 @@ impl QueueStorage { WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) - FOR UPDATE OF claim_batches ), locked_claims AS ( SELECT claim_slot, job_id, run_lease FROM locked_row_claims @@ -6708,6 +6720,8 @@ impl QueueStorage { .bind(&receipt_job_ids) .bind(&receipt_run_leases) .bind(&receipt_receipt_ids) + .bind(&receipt_claim_batch_ids) + .bind(&receipt_claim_batch_indices) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 8174c7d66..4765116c1 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -15514,3 +15514,77 @@ async fn test_queue_storage_receipt_claims_are_locked_before_the_heartbeat_state .expect("join heartbeat") .expect("heartbeat after the rewrite"); } + +/// Metadata completions validate the exact compact item and do not lock the +/// immutable batch shared with other attempts. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_metadata_completion_validates_and_does_not_lock_compact_batch() { + let (_db_guard, pool) = setup_pool(10).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_metadata_completion_identity".into(), + ..Default::default() + }, + ) + .await; + let queue = "metadata_completion_identity"; + for id in 0..2 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + tags: vec!["retained".into()], + metadata: serde_json::json!({"source": "recovery"}), + ..Default::default() + }, + ) + .await; + } + let jobs = store + .claim_runtime_batch(&pool, queue, 2, Duration::ZERO) + .await + .expect("claim compact jobs with metadata"); + assert_eq!(jobs.len(), 2); + assert_eq!(jobs[0].claim.claim_batch_id, jobs[1].claim.claim_batch_id); + let mut invalid = jobs[0].clone(); + invalid.claim.claim_batch_index = jobs[1].claim.claim_batch_index; + assert!(store + .complete_runtime_batch(&pool, &[invalid]) + .await + .expect("reject a different batch item") + .is_empty()); + + let mut holder = pool.begin().await.expect("begin batch holder"); + sqlx::query(&format!( + "SELECT batch_id FROM {}.lease_claim_batches \ + WHERE claim_slot = $1 AND batch_id = $2 FOR UPDATE", + store.schema() + )) + .bind(jobs[0].claim.claim_slot) + .bind(jobs[0].claim.claim_batch_id.expect("compact batch")) + .execute(holder.as_mut()) + .await + .expect("hold the immutable batch row"); + let completed = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.complete_runtime_batch(&pool, &jobs), + ) + .await + .expect("completion must not wait on an unrelated batch-row lock") + .expect("complete metadata jobs"); + assert_eq!(completed.len(), 2); + holder.rollback().await.expect("release batch row"); + for job in jobs { + let durable = store + .load_job(&pool, job.job.id) + .await + .expect("read completed job") + .expect("completed job"); + assert_eq!(durable.state, JobState::Completed); + assert_eq!(durable.tags, vec!["retained".to_string()]); + assert_eq!(durable.metadata, serde_json::json!({"source": "recovery"})); + } +}