From 303b7ad9873ee9f685bec69b89b1956e8ea935e4 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 18:57:44 +1300 Subject: [PATCH 1/4] 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/4] 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"})); + } +} From e01f87f1f412cdce3c88993d72ea36d986174387 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 21:55:29 +1300 Subject: [PATCH 3/4] Keep finalizing attempts live through completion acknowledgement --- CHANGELOG.md | 6 + awa-model/src/queue_storage.rs | 246 ++++++++++++++- awa-worker/src/executor.rs | 40 ++- awa-worker/src/heartbeat.rs | 50 +++- awa-worker/src/runtime.rs | 35 ++- awa/tests/heartbeat_pass_benchmark_test.rs | 187 ++++++++++++ awa/tests/queue_storage_runtime_test.rs | 329 +++++++++++++++++++++ 7 files changed, 869 insertions(+), 24 deletions(-) create mode 100644 awa/tests/heartbeat_pass_benchmark_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 9a3dd1d7d..13a08de78 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,12 @@ Notable changes between releases. Detailed migration notes for storage transitio It no longer expands every retained claim batch for each completion flush, so bounded finalization can drain large recovery backlogs without that scan. +- Finished attempts keep heartbeating until their durable finalization finishes, + including while completion writes are stalled. Compact receipt heartbeats + address claims by batch identity and skip attempts a closer already holds, + so one stalled completion cannot block the rest of a heartbeat pass. Late + heartbeats cannot recreate state for a completed attempt. + - **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 diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 3190ebf6e..306709025 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -154,6 +154,39 @@ pub struct ClaimedRuntimeJob { pub unique_states: Option, } +/// Identity of a compact receipt claim for heartbeat purposes: addresses +/// `lease_claim_batches` by primary key and the attempt by item index, so a +/// heartbeat does not have to scan every open claim batch. Built with +/// [`ClaimHeartbeatRef::from_claim`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub struct ClaimHeartbeatRef { + pub job_id: i64, + pub run_lease: i64, + pub claim_slot: i32, + pub claim_batch_id: i64, + pub claim_batch_index: i32, + pub receipt_id: i64, +} + +impl ClaimHeartbeatRef { + /// The heartbeat identity of a compact receipt claim, or `None` for + /// legacy or row-local claims. + pub fn from_claim(job_id: i64, run_lease: i64, claim: &ClaimedEntry) -> Option { + if !claim.lease_claim_receipt { + return None; + } + Some(Self { + job_id, + run_lease, + claim_slot: claim.claim_slot, + claim_batch_id: claim.claim_batch_id?, + claim_batch_index: claim.claim_batch_index?, + receipt_id: claim.receipt_id?, + }) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)] pub struct QueueClaimerLease { pub claimer_slot: i16, @@ -9820,6 +9853,47 @@ impl QueueStorage { Ok(()) } + /// Take the closers' attempt locks without waiting, returning the attempts + /// locked. A heartbeat must not queue behind a closer: an attempt whose lock + /// is held is being finalized or rescued right now, and rescue skips it for + /// as long as the lock is held. Taking the locks in a statement of their own + /// means the statement that then decides what is still open starts after + /// every closer that held a lock first has committed (#614), so a heartbeat + /// cannot recreate state for an attempt that has just been closed. + async fn try_lock_receipt_attempt_keys_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + jobs: &[(i64, i64)], + ) -> Result, AwaError> { + if jobs.is_empty() { + return Ok(Vec::new()); + } + let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); + let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); + sqlx::query_as::<_, (i64, i64)>( + r#" + WITH attempts AS MATERIALIZED ( + SELECT DISTINCT + input.job_id, + input.run_lease, + pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', input.job_id, input.run_lease), + 0 + ) AS lock_key + FROM unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease) + ) + SELECT job_id, run_lease + FROM attempts + WHERE pg_catalog.pg_try_advisory_xact_lock(lock_key) + "#, + ) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error) + } + async fn ensure_running_leases_from_receipts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, @@ -10114,7 +10188,6 @@ impl QueueStorage { WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) - FOR UPDATE OF claim_batches ), claim_refs AS ( SELECT * FROM row_claim_refs @@ -10163,6 +10236,19 @@ impl QueueStorage { .iter() .map(|(job_id, run_lease, _)| (*job_id, *run_lease)) .collect(); + let locked: HashSet<_> = self + .try_lock_receipt_attempt_keys_tx(tx, &receipt_jobs) + .await? + .into_iter() + .collect(); + let jobs: Vec<_> = jobs + .iter() + .filter(|(job_id, run_lease, _)| locked.contains(&(*job_id, *run_lease))) + .collect(); + if jobs.is_empty() { + return Ok(0); + } + let receipt_jobs: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect(); self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?; let schema = self.schema(); @@ -10231,7 +10317,6 @@ impl QueueStorage { WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) - FOR UPDATE OF claim_batches ), claim_refs AS ( SELECT * FROM row_claim_refs @@ -13037,6 +13122,141 @@ impl QueueStorage { Ok(()) } + /// Heartbeat compact receipt claims by their batch identity. Each claim + /// batch is addressed by primary key and the attempt by item index, with + /// index-only closure probes and no batch row lock, so the cost is + /// O(in-flight attempts) rather than O(open claim batches). Attempts that + /// a closer holds right now are skipped rather than waited for. + pub async fn heartbeat_claims_batch( + &self, + pool: &PgPool, + claims: &[ClaimHeartbeatRef], + ) -> Result { + if claims.is_empty() { + return Ok(0); + } + let schema = self.schema(); + let mut tx = pool.begin().await.map_err(map_sqlx_error)?; + let pairs: Vec<(i64, i64)> = claims + .iter() + .map(|claim| (claim.job_id, claim.run_lease)) + .collect(); + let locked: HashSet<(i64, i64)> = self + .try_lock_receipt_attempt_keys_tx(&mut tx, &pairs) + .await? + .into_iter() + .collect(); + let claims: Vec<&ClaimHeartbeatRef> = claims + .iter() + .filter(|claim| locked.contains(&(claim.job_id, claim.run_lease))) + .collect(); + if claims.is_empty() { + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(0); + } + let job_ids: Vec = claims.iter().map(|claim| claim.job_id).collect(); + let run_leases: Vec = claims.iter().map(|claim| claim.run_lease).collect(); + let claim_slots: Vec = claims.iter().map(|claim| claim.claim_slot).collect(); + let batch_ids: Vec = claims.iter().map(|claim| claim.claim_batch_id).collect(); + let batch_indices: Vec = claims.iter().map(|claim| claim.claim_batch_index).collect(); + let receipt_ids: Vec = claims.iter().map(|claim| claim.receipt_id).collect(); + let updated: i64 = sqlx::query_scalar(&format!( + r#" + WITH inflight(job_id, run_lease, claim_slot, batch_id, batch_index, receipt_id) AS ( + SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::int[], $4::bigint[], $5::int[], $6::bigint[]) + ), + open_claims AS ( + SELECT inflight.job_id, inflight.run_lease + FROM inflight + JOIN {schema}.lease_claim_batches AS claim_batches + ON claim_batches.claim_slot = inflight.claim_slot + AND claim_batches.batch_id = inflight.batch_id + WHERE inflight.batch_index BETWEEN 1 AND claim_batches.claimed_count + AND claim_batches.job_ids[inflight.batch_index] = inflight.job_id + AND claim_batches.run_leases[inflight.batch_index] = inflight.run_lease + AND claim_batches.receipt_ids[inflight.batch_index] = inflight.receipt_id + AND NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closures AS closures + WHERE closures.claim_slot = inflight.claim_slot + AND closures.job_id = inflight.job_id + AND closures.run_lease = inflight.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches + WHERE closure_batches.claim_slot = inflight.claim_slot + AND closure_batches.receipt_ranges @> inflight.receipt_id + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.done_entries AS done + WHERE done.job_id = inflight.job_id AND done.run_lease = inflight.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.deferred_jobs AS deferred + WHERE deferred.job_id = inflight.job_id AND deferred.run_lease = inflight.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.dlq_entries AS dlq + WHERE dlq.job_id = inflight.job_id AND dlq.run_lease = inflight.run_lease + ) + ), + upserted AS ( + INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + SELECT open_claims.job_id, open_claims.run_lease, clock_timestamp(), clock_timestamp() + FROM open_claims + ON CONFLICT (job_id, run_lease) + DO UPDATE SET + heartbeat_at = clock_timestamp(), + updated_at = clock_timestamp() + RETURNING job_id + ) + SELECT count(*)::bigint FROM upserted + "# + )) + .bind(&job_ids) + .bind(&run_leases) + .bind(&claim_slots) + .bind(&batch_ids) + .bind(&batch_indices) + .bind(&receipt_ids) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + Ok(updated as usize) + } + + /// Lock legacy lease rows without waiting behind a finalizer. Keep this + /// separate from the updating statement: the latter must use a fresh + /// snapshot and must not mix partition pruning with row locking (PG18). + async fn lock_live_leases_for_heartbeat_tx( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + jobs: &[(i64, i64)], + ) -> Result, AwaError> { + if jobs.is_empty() { + return Ok(Vec::new()); + } + let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); + let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); + sqlx::query_as::<_, (i64, i64)>(&format!( + r#" + SELECT lease.job_id, lease.run_lease + FROM {table} AS lease + JOIN unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease) + ON lease.job_id = input.job_id AND lease.run_lease = input.run_lease + WHERE lease.state = 'running' + ORDER BY lease.job_id, lease.run_lease + FOR UPDATE OF lease SKIP LOCKED + "#, + table = self.leases_table(), + )) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error) + } + pub async fn heartbeat_batch( &self, pool: &PgPool, @@ -13057,14 +13277,18 @@ impl QueueStorage { // dead entries per write under sustained churn). Compat // reads in `LeaseJobRow` SELECTs and the `awa.jobs` view // COALESCE attempt_state.heartbeat_at first. + let locked = self.try_lock_receipt_attempt_keys_tx(&mut tx, jobs).await?; updated += self - .upsert_attempt_state_from_receipts_tx(&mut tx, jobs) + .upsert_attempt_state_from_receipts_tx(&mut tx, &locked) .await?; } else { // Legacy non-receipts mode (custom schemas with // `lease_claim_receipts=FALSE`): the `leases.heartbeat_at` // write is still the heartbeat home, since there is no // upsert_attempt_state path firing. + let jobs = self + .lock_live_leases_for_heartbeat_tx(&mut tx, jobs) + .await?; let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); let result = sqlx::query(&format!( @@ -13115,9 +13339,19 @@ impl QueueStorage { self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs) .await? } else { - // Legacy non-receipts mode keeps the old CTE shape that - // updates leases.heartbeat_at + leases.progress - // (via attempt_state upsert) in a single round-trip. + // Prelock independently so a finalizer cannot stall the entire + // heartbeat pass. Skipped progress snapshots remain unacknowledged + // in the worker and are retried on the next pass. + let keys: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect(); + let locked: HashSet<_> = self + .lock_live_leases_for_heartbeat_tx(&mut tx, &keys) + .await? + .into_iter() + .collect(); + let jobs: Vec<_> = jobs + .iter() + .filter(|(id, lease, _)| locked.contains(&(*id, *lease))) + .collect(); let schema = self.schema(); let job_ids: Vec = jobs.iter().map(|(job_id, _, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect(); diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 99b8c9c18..a392b3a31 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -365,6 +365,14 @@ impl JobExecutor { let in_flight_state = InFlightState { cancel: cancel.clone(), progress: progress_state.clone(), + claim: queue_storage_claim.as_ref().and_then(|claim| { + awa_model::queue_storage::ClaimHeartbeatRef::from_claim( + job_id, + job_run_lease, + claim, + ) + }), + finalizing: false, }; in_flight.insert((job_id, job_run_lease), in_flight_state); if let Some(counter) = queue_in_flight.get(&job_queue) { @@ -443,12 +451,17 @@ impl JobExecutor { guard.clone_latest() }; - // 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. + // The attempt stays in `in_flight`, and so keeps its heartbeat, + // until its finalize is over: a finalize that waits behind a + // completion backlog must not look like a dead worker to rescue. + // Only a bounded number of finished attempts may wait for + // finalization; while that queue is full the attempt also keeps + // its worker slot, so the dispatcher stops claiming. + in_flight.mark_finalizing((job_id, job_run_lease)); + let heartbeat = HeartbeatUntilFinalized { + in_flight: in_flight.clone(), + key: (job_id, job_run_lease), + }; let finalize_permit = match result { Some(_) => tokio::select! { permit = finalize_capacity.acquire_owned() => { @@ -458,7 +471,6 @@ impl JobExecutor { }, 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); } @@ -511,6 +523,7 @@ impl JobExecutor { outcome => break outcome, } }; + drop(heartbeat); drop(finalize_permit); match &outcome { @@ -594,6 +607,19 @@ impl JobExecutor { } } +/// Keeps a finished attempt in the heartbeat set until its finalize is over, +/// and takes it out however the finalizer ends. +struct HeartbeatUntilFinalized { + in_flight: InFlightMap, + key: (i64, i64), +} + +impl Drop for HeartbeatUntilFinalized { + fn drop(&mut self) { + self.in_flight.remove(self.key); + } +} + /// How many times a finalize transaction is re-run after Postgres aborts it /// to break a lock cycle before the failure is surfaced. const COMPLETION_DEADLOCK_RETRIES: u32 = 3; diff --git a/awa-worker/src/heartbeat.rs b/awa-worker/src/heartbeat.rs index 462b33eb7..655e02729 100644 --- a/awa-worker/src/heartbeat.rs +++ b/awa-worker/src/heartbeat.rs @@ -1,5 +1,6 @@ use crate::runtime::{InFlightMap, RunLease}; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; +use awa_model::queue_storage::ClaimHeartbeatRef; use sqlx::PgPool; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -76,8 +77,9 @@ impl HeartbeatService { fields(interval_ms = self.interval.as_millis() as u64) )] async fn heartbeat_once(&self) { - let mut all_keys: Vec<(i64, RunLease)> = self.in_flight.keys(); - all_keys.sort_unstable(); + let mut in_flight = self.in_flight.claim_refs(); + in_flight.sort_unstable_by_key(|(key, _)| *key); + let all_keys: Vec<(i64, RunLease)> = in_flight.iter().map(|(key, _)| *key).collect(); if all_keys.is_empty() { return; @@ -104,7 +106,14 @@ impl HeartbeatService { .await; } RuntimeStorage::QueueStorage(runtime) => { - self.heartbeat_once_queue_storage(runtime, &heartbeat_only, &pending_progress) + let (claimed, plain): (Vec<_>, Vec<_>) = in_flight + .iter() + .filter(|(key, _)| !progress_keys.contains(key)) + .partition(|(_, claim)| claim.is_some()); + let claimed: Vec = + claimed.into_iter().filter_map(|(_, claim)| claim).collect(); + let plain: Vec<(i64, RunLease)> = plain.into_iter().map(|(key, _)| key).collect(); + self.heartbeat_once_queue_storage(runtime, &claimed, &plain, &pending_progress) .await; } } @@ -131,7 +140,7 @@ impl HeartbeatService { AND jobs.run_lease = inflight.run_lease WHERE jobs.state = 'running' ORDER BY jobs.id - FOR UPDATE OF jobs + FOR UPDATE OF jobs SKIP LOCKED ) UPDATE awa.jobs_hot AS jobs SET heartbeat_at = now() @@ -178,7 +187,7 @@ impl HeartbeatService { AND jobs.run_lease = inflight.run_lease WHERE jobs.state = 'running' ORDER BY jobs.id - FOR UPDATE OF jobs + FOR UPDATE OF jobs SKIP LOCKED ) UPDATE awa.jobs_hot AS jobs SET heartbeat_at = now(), @@ -204,7 +213,13 @@ impl HeartbeatService { .iter() .map(|(id, lease, gen, _)| (*id, *lease, *gen)) .collect(); - self.in_flight.ack_progress(&acked); + if result.rows_affected() as usize == chunk.len() { + self.in_flight.ack_progress(&acked); + } else { + // Re-send successful snapshots too: replay is safe, + // whereas acknowledging a skipped write loses progress. + self.in_flight.clear_in_flight_progress(&acked); + } } Err(err) => { warn!(error = %err, "Failed to send heartbeat+progress batch"); @@ -222,9 +237,26 @@ impl HeartbeatService { async fn heartbeat_once_queue_storage( &self, runtime: &QueueStorageRuntime, + claimed: &[ClaimHeartbeatRef], heartbeat_only: &[(i64, RunLease)], pending_progress: &[(i64, i64, u64, serde_json::Value)], ) { + for chunk in claimed.chunks(self.batch_size) { + match runtime + .store + .heartbeat_claims_batch(&self.pool, chunk) + .await + { + Ok(updated) => { + self.metrics.heartbeat_batches.add(1, &[]); + debug!(count = chunk.len(), updated, "Heartbeat claim batch sent"); + } + Err(err) => { + warn!(error = %err, "Failed to send heartbeat claim batch"); + } + } + } + for chunk in heartbeat_only.chunks(self.batch_size) { match runtime.store.heartbeat_batch(&self.pool, chunk).await { Ok(updated) => { @@ -258,7 +290,11 @@ impl HeartbeatService { .iter() .map(|(id, lease, gen, _)| (*id, *lease, *gen)) .collect(); - self.in_flight.ack_progress(&acked); + if updated == chunk.len() { + self.in_flight.ack_progress(&acked); + } else { + self.in_flight.clear_in_flight_progress(&acked); + } } Err(err) => { warn!(error = %err, "Failed to send heartbeat+progress batch"); diff --git a/awa-worker/src/runtime.rs b/awa-worker/src/runtime.rs index 4de28684a..5a87d4428 100644 --- a/awa-worker/src/runtime.rs +++ b/awa-worker/src/runtime.rs @@ -5,6 +5,8 @@ use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::sync::RwLock; +use awa_model::queue_storage::ClaimHeartbeatRef; + const IN_FLIGHT_SHARDS: usize = 64; pub(crate) type RunLease = i64; @@ -17,6 +19,13 @@ type InFlightShard = RwLock>; pub(crate) struct InFlightState { pub cancel: CancelFlag, pub progress: Arc>, + /// Compact receipt claim identity, when the attempt was dispatched from + /// queue storage with one; lets the heartbeat address the claim batch + /// directly. + pub claim: Option, + /// The handler has returned and the attempt waits for its finalize. It is + /// still heartbeated, but its progress is settled by the finalize. + pub finalizing: bool, } /// Mutable progress buffer shared between handler and heartbeat service. @@ -170,13 +179,26 @@ impl InFlightRegistry { guard.get(&key).map(|s| s.cancel.clone()) } - pub fn keys(&self) -> Vec { - let mut keys = Vec::new(); + /// Every attempt to heartbeat, with its compact claim identity, if any. + pub fn claim_refs(&self) -> Vec<(InFlightKey, Option)> { + let mut refs = Vec::new(); for shard in &*self.shards { let guard = shard.read().expect("in_flight shard poisoned"); - keys.extend(guard.keys().copied()); + refs.extend(guard.iter().map(|(key, state)| (*key, state.claim))); + } + refs + } + + /// The handler has returned: keep heartbeating the attempt until its + /// finalize is done, and stop flushing its progress. + pub fn mark_finalizing(&self, key: InFlightKey) { + let mut guard = self + .shard_for(&key) + .write() + .expect("in_flight shard poisoned"); + if let Some(state) = guard.get_mut(&key) { + state.finalizing = true; } - keys } pub fn flags(&self) -> Vec { @@ -198,6 +220,9 @@ impl InFlightRegistry { for shard in &*self.shards { let guard = shard.read().expect("in_flight shard poisoned"); for (&(job_id, run_lease), state) in guard.iter() { + if state.finalizing { + continue; + } let mut progress = lock_progress(&state.progress); if progress.has_pending() && progress.in_flight.is_none() { if let Some(ref value) = progress.latest { @@ -283,6 +308,8 @@ mod tests { InFlightState { cancel: Arc::new(AtomicBool::new(false)), progress: progress.clone(), + claim: None, + finalizing: false, }, ); diff --git a/awa/tests/heartbeat_pass_benchmark_test.rs b/awa/tests/heartbeat_pass_benchmark_test.rs new file mode 100644 index 000000000..489f29677 --- /dev/null +++ b/awa/tests/heartbeat_pass_benchmark_test.rs @@ -0,0 +1,187 @@ +//! Heartbeat pass cost against in-flight claim count. Ignored by default: +//! +//! ```text +//! DATABASE_URL=postgres://postgres:test@localhost:15432/awa_test \ +//! cargo test --release -p awa --test heartbeat_pass_benchmark_test -- --ignored --nocapture +//! ``` +//! +//! Each size claims that many jobs through the claim path the dispatcher uses, +//! in claims of `AWA_HEARTBEAT_BENCH_CLAIM_BATCH` jobs (default 16, as a busy +//! runtime refills a few freed workers at a time), and times a heartbeat pass +//! over them in the heartbeat service's 500-attempt chunks. +//! `AWA_HEARTBEAT_BENCH_SIZES` overrides the sizes. + +use awa::model::queue_storage::ClaimHeartbeatRef; +use awa::model::{insert, migrations, QueueStorage, QueueStorageConfig}; +use awa::{InsertOpts, JobArgs}; +use serde::{Deserialize, Serialize}; +use sqlx::postgres::{PgConnection, PgPoolOptions}; +use sqlx::{Connection, PgPool}; +use std::time::{Duration, Instant}; + +const HEARTBEAT_CHUNK: usize = 500; +const PASSES: usize = 6; + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct HeartbeatBenchJob { + n: i64, +} + +fn database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) +} + +fn with_database(url: &str, database: &str) -> String { + let (base, _) = url + .rsplit_once('/') + .expect("database URL has a database name"); + format!("{base}/{database}") +} + +async fn fresh_database(name: &str) -> PgPool { + let mut admin = PgConnection::connect(&with_database(&database_url(), "postgres")) + .await + .expect("admin connection"); + sqlx::raw_sql(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&mut admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(&format!("CREATE DATABASE {name}")) + .execute(&mut admin) + .await + .expect("create database"); + let pool = PgPoolOptions::new() + .max_connections(4) + .connect(&with_database(&database_url(), name)) + .await + .expect("bench pool"); + migrations::run(&pool).await.expect("migrate"); + pool +} + +async fn drop_database(name: &str) { + let mut admin = PgConnection::connect(&with_database(&database_url(), "postgres")) + .await + .expect("admin connection"); + let _ = sqlx::raw_sql(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&mut admin) + .await; +} + +#[tokio::test] +#[ignore = "benchmark; creates and drops its own database"] +async fn heartbeat_pass_time_by_in_flight_claims() { + let sizes: Vec = std::env::var("AWA_HEARTBEAT_BENCH_SIZES") + .ok() + .map(|sizes| { + sizes + .split(',') + .filter_map(|size| size.trim().parse().ok()) + .collect() + }) + .unwrap_or_else(|| vec![1_000, 10_000, 25_000]); + let claim_batch: i64 = std::env::var("AWA_HEARTBEAT_BENCH_CLAIM_BATCH") + .ok() + .and_then(|size| size.parse().ok()) + .unwrap_or(16); + let server: String = { + let pool = fresh_database(&format!("awa_hb_bench_{}", std::process::id())).await; + let version = sqlx::query_scalar("SHOW server_version") + .fetch_one(&pool) + .await + .expect("server version"); + pool.close().await; + version + }; + for size in sizes { + let name = format!("awa_hb_bench_{}", std::process::id()); + let pool = fresh_database(&name).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + store.install(&pool).await.expect("install queue storage"); + awa_testing::setup::activate_queue_storage(&pool).await; + let queue = "heartbeat_bench"; + for first in (0..size as i64).step_by(1_000) { + let params: Vec<_> = (first..(first + 1_000).min(size as i64)) + .map(|n| { + insert::params_with( + &HeartbeatBenchJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue"); + } + let mut claimed = Vec::with_capacity(size); + while claimed.len() < size { + let batch = store + .claim_runtime_batch(&pool, queue, claim_batch, Duration::ZERO) + .await + .expect("claim"); + assert!(!batch.is_empty(), "claimed {} of {size}", claimed.len()); + claimed.extend(batch); + } + sqlx::query("ANALYZE") + .execute(&pool) + .await + .expect("analyze"); + + // The first pass creates every attempt's heartbeat row; later passes + // refresh them. + let mut passes = Vec::with_capacity(PASSES); + let mut refreshed = 0; + for _ in 0..PASSES { + let started = Instant::now(); + refreshed = heartbeat_pass(&store, &pool, &claimed).await; + passes.push(started.elapsed()); + } + let first = passes[0]; + let mut later = passes[1..].to_vec(); + later.sort_unstable(); + eprintln!( + "heartbeat pass: server={server} in_flight={size} claim_batch={claim_batch} \ + refreshed={refreshed} first_ms={} median_ms={} max_ms={}", + first.as_millis(), + later[later.len() / 2].as_millis(), + later[later.len() - 1].as_millis() + ); + assert_eq!(refreshed, size, "every in-flight attempt is refreshed"); + pool.close().await; + drop_database(&name).await; + } +} + +async fn heartbeat_pass( + store: &QueueStorage, + pool: &PgPool, + claimed: &[awa::model::ClaimedRuntimeJob], +) -> usize { + // As the heartbeat service splits them: compact receipt claims by their + // batch identity, anything else by attempt key. + let mut claims = Vec::new(); + let mut plain = Vec::new(); + for job in claimed { + match ClaimHeartbeatRef::from_claim(job.job.id, job.job.run_lease, &job.claim) { + Some(claim) => claims.push(claim), + None => plain.push((job.job.id, job.job.run_lease)), + } + } + let mut refreshed = 0; + for chunk in claims.chunks(HEARTBEAT_CHUNK) { + refreshed += store + .heartbeat_claims_batch(pool, chunk) + .await + .expect("heartbeat claims"); + } + for chunk in plain.chunks(HEARTBEAT_CHUNK) { + refreshed += store.heartbeat_batch(pool, chunk).await.expect("heartbeat"); + } + refreshed +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 4765116c1..daa942d98 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -2156,6 +2156,132 @@ async fn test_queue_storage_runtime_stops_claiming_while_finalization_is_stalled ); } +/// A finished attempt keeps heartbeating until its finalize is over, so a +/// finalize that waits behind a stalled completion batcher for several +/// staleness windows is not mistaken for a dead worker and run again. +/// Holding the first job's receipt-completion lock stalls the batcher's flush; +/// the attempts that flush has locked are skipped by rescue for as long as it +/// waits, and every other finished attempt must still be heartbeated. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_finalizing_attempts_keep_heartbeating_through_a_stalled_finalize() { + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_finalize_heartbeat"; + let store_config = QueueStorageConfig { + schema: "awa_qs_finalize_heartbeat".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"); + } + + // Only the first job's lock: a held lock also keeps the heartbeat off that + // attempt, which is right only while a closer is working on it. + let first_job: i64 = sqlx::query_scalar(&format!( + "SELECT job_id FROM {}.ready_entries WHERE queue = $1 ORDER BY lane_seq LIMIT 1", + store.schema() + )) + .bind(queue) + .fetch_one(&pool) + .await + .expect("first job"); + 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', $1::bigint, 1), 0))", + ) + .bind(first_job) + .execute(&mut *lock_holder) + .await + .expect("stall the first job's finalization"); + + let staleness = ci_timing::scaled_staleness(Duration::from_secs(1)); + 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_interval(staleness / 5) + .heartbeat_staleness(staleness) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build finalize heartbeat client"); + client + .start() + .await + .expect("start finalize heartbeat client"); + + // Hold the stall for several staleness windows once claiming has stopped. + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + while runs.lock().expect("runs").len() < WORKERS as usize * 2 { + assert!(Instant::now() < deadline, "the runtime did not claim"); + tokio::time::sleep(Duration::from_millis(25)).await; + } + tokio::time::sleep(staleness * 4).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; + + let runs = runs.lock().expect("runs"); + assert_eq!(runs.len() as i64, JOBS); + let rerun: Vec<_> = runs + .iter() + .filter(|(_, count)| **count > 1) + .take(5) + .collect(); + assert!(rerun.is_empty(), "jobs ran more than once: {rerun:?}"); +} + /// 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 @@ -15515,6 +15641,209 @@ async fn test_queue_storage_receipt_claims_are_locked_before_the_heartbeat_state .expect("heartbeat after the rewrite"); } +/// A finalizing attempt must not hold up unrelated heartbeats, and an old +/// heartbeat snapshot must not recreate state after durable completion. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_heartbeat_skips_closers_and_never_revives_closed_claims() { + let (_db_guard, pool) = setup_pool(10).await; + let config = QueueStorageConfig { + schema: "awa_qs_heartbeat_closure".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, config).await; + let queue = "heartbeat_closure"; + for id in 0..2 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + } + let jobs = store + .claim_runtime_batch(&pool, queue, 2, Duration::ZERO) + .await + .expect("claim two compact receipts"); + assert_eq!(jobs.len(), 2); + let claims: Vec<_> = jobs + .iter() + .map(|job| { + awa::model::queue_storage::ClaimHeartbeatRef::from_claim( + job.job.id, + job.job.run_lease, + &job.claim, + ) + .expect("compact claim identity") + }) + .collect(); + let mut closer = pool.begin().await.expect("begin closer"); + sqlx::query( + "SELECT pg_advisory_xact_lock(hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(jobs[0].job.id) + .bind(jobs[0].job.run_lease) + .execute(closer.as_mut()) + .await + .expect("hold one attempt's closer lock"); + let refreshed = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.heartbeat_claims_batch(&pool, &claims), + ) + .await + .expect("a closer cannot stall the heartbeat pass") + .expect("heartbeat the unlocked claim"); + assert_eq!(refreshed, 1); + closer.rollback().await.expect("release closer"); + + let completed = store + .complete_runtime_batch(&pool, &jobs) + .await + .expect("durable completion"); + assert_eq!(completed.len(), 2); + assert_eq!( + store + .heartbeat_claims_batch(&pool, &claims) + .await + .expect("a cached heartbeat snapshot after completion"), + 0, + ); + let remaining: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {}.attempt_state WHERE job_id = ANY($1)", + store.schema() + )) + .bind(jobs.iter().map(|job| job.job.id).collect::>()) + .fetch_one(&pool) + .await + .expect("read attempt state"); + assert_eq!( + remaining, 0, + "a late heartbeat must not revive a closed attempt" + ); +} + +/// A progress snapshot captured before a handler returns must not wait behind +/// its finalizer and prevent unrelated attempts from heartbeating. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_pending_progress_skips_a_closer_and_can_retry() { + progress_snapshot_skips_closer(true).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_legacy_pending_progress_skips_a_closer_and_can_retry() { + progress_snapshot_skips_closer(false).await; +} + +async fn progress_snapshot_skips_closer(receipts: bool) { + let (_db_guard, pool) = setup_pool(10).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: if receipts { + "awa_qs_pending_progress_closer".to_string() + } else { + "awa_qs_pending_progress_closer_legacy".to_string() + }, + lease_claim_receipts: receipts, + ..Default::default() + }, + ) + .await; + let queue = "pending_progress_closer"; + for id in 0..2 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + } + let jobs = store + .claim_runtime_batch(&pool, queue, 2, Duration::ZERO) + .await + .expect("claim compact receipts"); + assert_eq!(jobs.len(), 2); + let pending: Vec<_> = jobs + .iter() + .map(|job| { + ( + job.job.id, + job.job.run_lease, + serde_json::json!({"step": 2}), + ) + }) + .collect(); + let mut closer = pool.begin().await.expect("begin closer"); + if receipts { + sqlx::query( + "SELECT pg_advisory_xact_lock(hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(jobs[0].job.id) + .bind(jobs[0].job.run_lease) + .execute(closer.as_mut()) + .await + .expect("hold a closer lock after the progress snapshot"); + } else { + sqlx::query(&format!( + "SELECT 1 FROM {}.leases WHERE job_id = $1 AND run_lease = $2 FOR UPDATE", + store.schema() + )) + .bind(jobs[0].job.id) + .bind(jobs[0].job.run_lease) + .execute(closer.as_mut()) + .await + .expect("hold a legacy finalizer row lock"); + } + let updated = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.heartbeat_progress_batch(&pool, &pending), + ) + .await + .expect("a stale progress snapshot cannot block the heartbeat pass") + .expect("write the unlocked snapshot"); + assert_eq!(updated, 1); + let pairs: Vec<_> = jobs + .iter() + .map(|job| (job.job.id, job.job.run_lease)) + .collect(); + assert_eq!( + tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.heartbeat_batch(&pool, &pairs), + ) + .await + .expect("heartbeat-only passes must also skip a closer") + .expect("heartbeat unlocked claims"), + 1, + ); + closer.rollback().await.expect("release closer"); + assert_eq!( + store + .heartbeat_progress_batch(&pool, &pending) + .await + .expect("retry the unacknowledged snapshots"), + 2, + ); + let progress: Vec = sqlx::query_scalar(&format!( + "SELECT progress FROM {}.attempt_state WHERE job_id = ANY($1) ORDER BY job_id", + store.schema() + )) + .bind(jobs.iter().map(|job| job.job.id).collect::>()) + .fetch_all(&pool) + .await + .expect("read durable progress"); + assert_eq!(progress, vec![serde_json::json!({"step": 2}); 2]); +} + /// 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)] From 675188eb6c60247522f9aa0e5766a59f522beda8 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 22:45:17 +1300 Subject: [PATCH 4/4] test(worker): gate stalled completion after attempt locks --- awa/tests/queue_storage_runtime_test.rs | 75 +++++++++++++++++++------ 1 file changed, 57 insertions(+), 18 deletions(-) diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index e45b62370..e0b998538 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -2388,9 +2388,9 @@ async fn test_queue_storage_runtime_stops_claiming_while_finalization_is_stalled /// A finished attempt keeps heartbeating until its finalize is over, so a /// finalize that waits behind a stalled completion batcher for several /// staleness windows is not mistaken for a dead worker and run again. -/// Holding the first job's receipt-completion lock stalls the batcher's flush; -/// the attempts that flush has locked are skipped by rescue for as long as it -/// waits, and every other finished attempt must still be heartbeated. +/// An independent trigger gate stalls a completion after it has acquired all +/// of its attempts' locks. That batch remains protected from rescue while +/// every other finished attempt must still be heartbeated. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_finalizing_attempts_keep_heartbeating_through_a_stalled_finalize() { const WORKERS: u32 = 64; @@ -2421,27 +2421,48 @@ async fn test_queue_storage_finalizing_attempts_keep_heartbeating_through_a_stal .expect("enqueue backlog"); } - // Only the first job's lock: a held lock also keeps the heartbeat off that - // attempt, which is right only while a closer is working on it. - let first_job: i64 = sqlx::query_scalar(&format!( - "SELECT job_id FROM {}.ready_entries WHERE queue = $1 ORDER BY lane_seq LIMIT 1", - store.schema() + // The trigger runs after the completer owns every attempt lock in its + // flush. Holding an attempt lock directly would suppress its heartbeat + // before completion owns it, racing rescue when the test unlocks it. + sqlx::raw_sql(&format!( + r#" + CREATE FUNCTION {schema}.hold_completion_for_heartbeat_test() RETURNS trigger + LANGUAGE plpgsql AS $$ + BEGIN + IF NEW.outcome = 'completed' THEN + PERFORM pg_catalog.pg_advisory_xact_lock( + pg_catalog.hashtextextended('awa.test.finalize-heartbeat-gate', 0) + ); + END IF; + RETURN NEW; + END + $$; + CREATE TRIGGER hold_completion_for_heartbeat_test + BEFORE INSERT ON {schema}.lease_claim_closure_batches + FOR EACH ROW EXECUTE FUNCTION {schema}.hold_completion_for_heartbeat_test(); + "#, + schema = store.schema(), )) - .bind(queue) - .fetch_one(&pool) + .execute(&pool) .await - .expect("first job"); + .expect("install the completion gate"); let mut lock_holder = pool.acquire().await.expect("lock connection"); + let lock_holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *lock_holder) + .await + .expect("gate holder PID"); sqlx::query( - "SELECT pg_catalog.pg_advisory_lock(pg_catalog.hashtextextended( - format('awa.receipt.complete:%s:%s', $1::bigint, 1), 0))", + "SELECT pg_catalog.pg_advisory_lock( + pg_catalog.hashtextextended('awa.test.finalize-heartbeat-gate', 0))", ) - .bind(first_job) .execute(&mut *lock_holder) .await - .expect("stall the first job's finalization"); + .expect("hold the independent completion gate"); - let staleness = ci_timing::scaled_staleness(Duration::from_secs(1)); + // Sequential heartbeat chunks can need several WAL-sync round trips + // under contention. Allow margin for those writes while blocking the + // completion for four full staleness windows below. + let staleness = ci_timing::scaled_staleness(Duration::from_secs(5)); let runs: Arc>> = Arc::default(); let recorded = runs.clone(); let client = Client::builder(pool.clone()) @@ -2483,8 +2504,26 @@ async fn test_queue_storage_finalizing_attempts_keep_heartbeating_through_a_stal // Hold the stall for several staleness windows once claiming has stopped. let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); - while runs.lock().expect("runs").len() < WORKERS as usize * 2 { - assert!(Instant::now() < deadline, "the runtime did not claim"); + loop { + let completion_at_gate: bool = sqlx::query_scalar( + "SELECT EXISTS ( + SELECT 1 FROM pg_stat_activity + WHERE datname = current_database() + AND wait_event = 'advisory' + AND $1 = ANY(pg_catalog.pg_blocking_pids(pid)) + )", + ) + .bind(lock_holder_pid) + .fetch_one(&pool) + .await + .expect("observe the completion gate"); + if completion_at_gate && runs.lock().expect("runs").len() >= WORKERS as usize * 2 { + break; + } + assert!( + Instant::now() < deadline, + "the runtime did not queue attempts behind the completion gate" + ); tokio::time::sleep(Duration::from_millis(25)).await; } tokio::time::sleep(staleness * 4).await;