diff --git a/CHANGELOG.md b/CHANGELOG.md index 80fda9ffd..3dcec27d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,29 @@ 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. + +- 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 + 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. 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. - **Ring prune and terminal-count rollup report `Blocked` when any of their statements times out on a lock.** Prune sets a short `lock_timeout` (10 ms by default) before it takes `ACCESS EXCLUSIVE` on a slot's partitions, but diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 0fb6aee36..6815763e2 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, @@ -6722,14 +6755,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 @@ -6756,18 +6797,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 @@ -6804,7 +6850,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 @@ -6896,6 +6941,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)?; @@ -9994,6 +10041,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>, @@ -10288,7 +10376,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 @@ -10337,6 +10424,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(); @@ -10405,7 +10505,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 @@ -13240,6 +13339,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, @@ -13260,14 +13494,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!( @@ -13318,9 +13556,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/client.rs b/awa-worker/src/client.rs index 87ae548ed..2dacaf61a 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -1640,6 +1640,8 @@ impl Client { effective_storage.clone(), self.dlq_policy.clone(), self.finalizers.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 916edb02f..44d7d5fd7 100644 --- a/awa-worker/src/dispatcher.rs +++ b/awa-worker/src/dispatcher.rs @@ -852,6 +852,8 @@ mod tests { storage.clone(), crate::executor::DlqPolicy::default(), tokio_util::task::TaskTracker::new(), + CancellationToken::new(), + 1, )); let config = QueueConfig { max_workers: 1, diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index a4fef41c3..f9b6b67a5 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 tokio_util::task::TaskTracker; use tracing::{error, info, info_span, warn, Instrument}; @@ -269,6 +271,8 @@ pub struct JobExecutor { storage: RuntimeStorage, dlq_policy: DlqPolicy, finalizers: TaskTracker, + finalize_cancel: CancellationToken, + finalize_capacity: Arc, } impl JobExecutor { @@ -291,7 +295,12 @@ impl JobExecutor { storage: RuntimeStorage, dlq_policy: DlqPolicy, finalizers: TaskTracker, + 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, @@ -305,6 +314,8 @@ impl JobExecutor { storage, dlq_policy, finalizers, + finalize_cancel, + finalize_capacity, } } @@ -332,6 +343,8 @@ impl JobExecutor { let storage = self.storage.clone(); let dlq_policy = self.dlq_policy.clone(); let finalizers = self.finalizers.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(); @@ -357,6 +370,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) { @@ -435,23 +456,38 @@ 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. - in_flight.remove((job_id, job_run_lease)); + // 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() => { + Some(permit.expect("finalizer capacity is never closed")) + } + _ = finalize_cancel.cancelled() => None, + }, + None => None, + }; 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); - // Graceful shutdown waits for this before it stops the completion - // batcher: the attempt has already left `in_flight`, so nothing - // heartbeats it, and a finalize lost at exit leaves it to rescue. + // Graceful shutdown waits for this before stopping the completion + // batcher. The heartbeat guard keeps the attempt live until the + // durable write completes or this finalizer is abandoned. finalizers.spawn(async move { // Every finalize path is one lease-guarded transaction, and a // deadlock abort rolls the whole transaction back, so running @@ -495,6 +531,8 @@ impl JobExecutor { outcome => break outcome, } }; + drop(heartbeat); + drop(finalize_permit); match &outcome { Ok(CompletionOutcome::Applied { terminal, .. }) => { @@ -577,6 +615,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 e80c74c9e..573bb300b 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -2240,6 +2240,316 @@ async fn test_queue_storage_graceful_shutdown_finalizes_returned_attempts() { ); } +#[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" + ); +} + +/// 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. +/// 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; + 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"); + } + + // 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(), + )) + .execute(&pool) + .await + .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('awa.test.finalize-heartbeat-gate', 0))", + ) + .execute(&mut *lock_holder) + .await + .expect("hold the independent completion gate"); + + // 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()) + .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)); + 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; + + 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 @@ -16738,6 +17048,283 @@ async fn test_aged_materialized_lease_is_visible_to_admin_reads() { assert_eq!(cancelled.id, job_id); } +/// 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)] +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"})); + } +} + /// A lane move leaves immutable source rows behind. Wrap the ring so the /// destination is in an earlier partition: point reads must use live evidence, /// rather than depending on which equal-timestamp ready row the plan emits last.