diff --git a/CHANGELOG.md b/CHANGELOG.md index 80fda9ffd..1db9c1170 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,23 @@ 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 + 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..b253338d3 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -6722,14 +6722,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 +6764,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 +6817,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 +6908,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-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..aa6c5da9b 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(); @@ -435,17 +448,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); @@ -495,6 +519,7 @@ impl JobExecutor { outcome => break outcome, } }; + drop(finalize_permit); match &outcome { Ok(CompletionOutcome::Applied { terminal, .. }) => { diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index e80c74c9e..241ba2ee7 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -2240,6 +2240,151 @@ 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" + ); +} + /// 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 @@ -16848,3 +16993,77 @@ async fn test_point_read_uses_live_ready_queue_after_ring_wrap() { "point read must ignore the tombstoned source" ); } + +/// 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"})); + } +}