From 399a52bee9f510d0daf302179e90bc9f549ae5ee Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 15:55:38 +1300 Subject: [PATCH 1/8] Fix claim aging race and retained body hydration --- CHANGELOG.md | 5 + awa-model/src/queue_storage.rs | 272 +++++++++++++-- awa/tests/queue_storage_runtime_test.rs | 445 ++++++++++++++++++++++++ 3 files changed, 686 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d52ab19ad..1b79b50b8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- Keep committed claims out of physical priority aging, even before their + cursor advances. Recovery and cancellation can hydrate an exact retained + claimed body after older aging tombstoned it, preserving original-priority + metadata and rejecting ambiguous historical bodies. + - Preserve recovered bodies and physical lane priority when manually moving historical aged terminal failures to the DLQ, so later DLQ retry retains arguments, attempt limits and priority provenance. Terminal counters are diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 4bf612398..446ddfbd8 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -2602,6 +2602,18 @@ impl ReadyJobRow { } } +#[derive(Debug, sqlx::FromRow)] +struct ReadyAgingCandidate { + ready_slot: i32, + ready_generation: i64, + queue: String, + priority: i16, + enqueue_shard: i16, + lane_seq: i64, + job_id: i64, + run_lease: i64, +} + #[derive(Debug, Clone, sqlx::FromRow)] struct ReadyTransitionRow { ready_slot: i32, @@ -3258,6 +3270,8 @@ struct ReadySnapshotRow { unique_key: Option>, unique_states: Option, payload: serde_json::Value, + run_lease: i64, + live_body: bool, } #[derive(Debug, Clone, sqlx::FromRow)] @@ -11566,11 +11580,67 @@ impl QueueStorage { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; + // Claims commit before their cursor advance. Select only a bounded + // set of identities before taking heads, without locking ready rows in + // the reverse order to a claimer. A later statement rechecks these same + // identities after every in-flight claim on their lanes has committed. + let candidates: Vec = sqlx::query_as(audited_sql(format!( + r#" + SELECT ready.ready_slot, ready.ready_generation, ready.queue, + ready.priority, ready.enqueue_shard, ready.lane_seq, + ready.job_id, ready.run_lease + FROM {schema}.ready_entries AS ready + JOIN {schema}.queue_claim_heads AS claims + ON claims.queue = ready.queue + AND claims.priority = ready.priority + AND claims.enqueue_shard = ready.enqueue_shard + WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) + AND ready.priority > 1 + AND ready.run_at <= $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + AND tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + ) + ORDER BY ready.run_at ASC, ready.lane_seq ASC + LIMIT $2 + "# + ))) + .bind(cutoff) + .bind(limit) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + if candidates.is_empty() { + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } + let job_ids: Vec = candidates.iter().map(|row| row.job_id).collect(); + self.lock_claim_heads_of_jobs_tx(&mut tx, &job_ids).await?; + let claimed = self.ready_row_claimed_sql("ready"); + let moved: Vec = sqlx::query_as(audited_sql(format!( r#" - WITH target AS ( + WITH candidates(ready_slot, ready_generation, queue, priority, + enqueue_shard, lane_seq, job_id, run_lease) AS ( + SELECT * FROM unnest($3::int[], $4::bigint[], $5::text[], $6::smallint[], + $7::smallint[], $8::bigint[], $9::bigint[], $10::bigint[]) + ), target AS ( SELECT ready.* - FROM {schema}.ready_entries AS ready + FROM candidates + JOIN {schema}.ready_entries AS ready + ON ready.ready_slot = candidates.ready_slot + AND ready.ready_generation = candidates.ready_generation + AND ready.queue = candidates.queue + AND ready.priority = candidates.priority + AND ready.enqueue_shard = candidates.enqueue_shard + AND ready.lane_seq = candidates.lane_seq + AND ready.job_id = candidates.job_id + AND ready.run_lease = candidates.run_lease JOIN {schema}.queue_claim_heads AS claims ON claims.queue = ready.queue AND claims.priority = ready.priority @@ -11578,6 +11648,7 @@ impl QueueStorage { WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) AND ready.priority > 1 AND ready.run_at <= $1 + AND NOT ({claimed}) AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.queue = ready.queue @@ -11623,6 +11694,14 @@ impl QueueStorage { ))) .bind(cutoff) .bind(limit) + .bind(candidates.iter().map(|row| row.ready_slot).collect::>()) + .bind(candidates.iter().map(|row| row.ready_generation).collect::>()) + .bind(candidates.iter().map(|row| row.queue.clone()).collect::>()) + .bind(candidates.iter().map(|row| row.priority).collect::>()) + .bind(candidates.iter().map(|row| row.enqueue_shard).collect::>()) + .bind(candidates.iter().map(|row| row.lane_seq).collect::>()) + .bind(&job_ids) + .bind(candidates.iter().map(|row| row.run_lease).collect::>()) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; @@ -12414,12 +12493,13 @@ impl QueueStorage { // Resolve moved ready bodies without opening the whole tombstone ring. // Empty input returned above; every remaining slot is validated here. - let tombstone_sources = ready_slots + let slots = ready_slots .iter() .map(|&slot| ring_slot_index(slot, self.queue_slot_count(), "queue")) - .collect::, _>>()? - .into_iter() - .map(|slot| format!("SELECT * FROM {}", ready_tombstone_child_name(schema, slot))) + .collect::, _>>()?; + let tombstone_sources = slots + .iter() + .map(|&slot| format!("SELECT * FROM {}", ready_tombstone_child_name(schema, slot))) .collect::>() .join(" UNION ALL "); @@ -12442,7 +12522,18 @@ impl QueueStorage { ready.created_at, ready.unique_key, ready.unique_states, - COALESCE(ready.payload, '{{}}'::jsonb) AS payload + COALESCE(ready.payload, '{{}}'::jsonb) AS payload, + ready.run_lease, + NOT EXISTS ( + SELECT 1 FROM ({tombstone_sources}) AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + AND tomb.job_id = ready.job_id + ) AS live_body FROM refs JOIN {schema}.ready_entries AS ready ON ready.ready_slot = refs.ready_slot @@ -12452,16 +12543,6 @@ impl QueueStorage { AND ready.lane_seq = refs.lane_seq AND ready.job_id = refs.job_id AND ready.run_lease + 1 = refs.run_lease - WHERE NOT EXISTS ( - SELECT 1 FROM ({tombstone_sources}) AS tomb - WHERE tomb.ready_slot = ready.ready_slot - AND tomb.ready_generation = ready.ready_generation - AND tomb.queue = ready.queue - AND tomb.priority = ready.priority - AND tomb.enqueue_shard = ready.enqueue_shard - AND tomb.lane_seq = ready.lane_seq - AND tomb.job_id = ready.job_id - ) "# ))) .bind(&ready_slots) @@ -12475,6 +12556,83 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; + type ReadyBodyKey = (i32, i64, String, i16, i64, i64, i64); + let mut ready_map: BTreeMap> = BTreeMap::new(); + for row in ready_rows { + ready_map + .entry(( + row.ready_slot, + row.ready_generation, + row.queue.clone(), + row.enqueue_shard, + row.lane_seq, + row.job_id, + row.run_lease + 1, + )) + .or_default() + .push(row); + } + + type ReadyBodyLaneKey = (i32, i64, String, i16, i64, i64, i64, i16); + let mut claimed_lanes: BTreeSet = BTreeSet::new(); + let ambiguous: Vec<&ReadySnapshotRow> = ready_map + .values() + .filter(|rows| rows.len() > 1 && !rows.iter().any(|row| row.live_body)) + .flatten() + .collect(); + if self.lease_claim_receipts() && !ambiguous.is_empty() { + // Normal live bodies and unique retained bodies need no extra + // relation lookup. Only ambiguous tombstoned receipt bodies need + // the emitted physical-lane evidence, scoped to their ready slots. + let attempt_sources = ambiguous + .iter() + .map(|row| row.ready_slot) + .collect::>() + .into_iter() + .map(|slot| ring_slot_index(slot, self.queue_slot_count(), "queue")) + .collect::, _>>()? + .into_iter() + .map(|slot| { + format!( + "SELECT * FROM {}", + ready_claim_attempt_batch_child_name(schema, slot) + ) + }) + .collect::>() + .join(" UNION ALL "); + let lanes: Vec = sqlx::query_as(audited_sql(format!( + r#" + WITH refs(ready_slot,ready_generation,queue,enqueue_shard,lane_seq,job_id,run_lease,priority) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], + $5::bigint[], $6::bigint[], $7::bigint[], $8::smallint[]) + ) + SELECT ready_slot,ready_generation,queue,enqueue_shard,lane_seq,job_id,run_lease,priority + FROM refs + WHERE EXISTS ( + SELECT 1 FROM ({attempt_sources}) AS attempts + WHERE attempts.ready_slot = refs.ready_slot + AND attempts.ready_generation = refs.ready_generation + AND attempts.queue = refs.queue + AND attempts.priority = refs.priority + AND attempts.enqueue_shard = refs.enqueue_shard + AND attempts.first_lane_seq <= refs.lane_seq + AND attempts.next_lane_seq > refs.lane_seq + AND attempts.lane_ranges @> int8range(refs.lane_seq,refs.lane_seq + 1,'[)') + ) + "# + ))) + .bind(ambiguous.iter().map(|row| row.ready_slot).collect::>()) + .bind(ambiguous.iter().map(|row| row.ready_generation).collect::>()) + .bind(ambiguous.iter().map(|row| row.queue.clone()).collect::>()) + .bind(ambiguous.iter().map(|row| row.enqueue_shard).collect::>()) + .bind(ambiguous.iter().map(|row| row.lane_seq).collect::>()) + .bind(ambiguous.iter().map(|row| row.job_id).collect::>()) + .bind(ambiguous.iter().map(|row| row.run_lease + 1).collect::>()) + .bind(ambiguous.iter().map(|row| row.priority).collect::>()) + .fetch_all(tx.as_mut()).await.map_err(map_sqlx_error)?; + claimed_lanes.extend(lanes); + } + let attempt_rows: Vec = sqlx::query_as(audited_sql(format!( r#" WITH refs(job_id, run_lease) AS ( @@ -12621,23 +12779,6 @@ impl QueueStorage { .map_err(map_sqlx_error)?; } - let ready_map: BTreeMap<(i32, i64, String, i16, i64, i64), ReadySnapshotRow> = ready_rows - .into_iter() - .map(|row| { - ( - ( - row.ready_slot, - row.ready_generation, - row.queue.clone(), - row.enqueue_shard, - row.lane_seq, - row.job_id, - ), - row, - ) - }) - .collect(); - let attempt_map: BTreeMap<(i64, i64), AttemptStateRow> = attempt_rows .into_iter() .map(|row| ((row.job_id, row.run_lease), row)) @@ -12645,7 +12786,7 @@ impl QueueStorage { let mut hydrated = Vec::with_capacity(deleted.len()); for deleted_row in deleted { - let ready = ready_map + let candidates = ready_map .get(&( deleted_row.ready_slot, deleted_row.ready_generation, @@ -12653,6 +12794,7 @@ impl QueueStorage { deleted_row.enqueue_shard, deleted_row.lane_seq, deleted_row.job_id, + deleted_row.run_lease, )) .ok_or_else(|| { AwaError::Validation(format!( @@ -12660,6 +12802,64 @@ impl QueueStorage { deleted_row.job_id, deleted_row.run_lease )) })?; + // Prefer a live exact-attempt body. A tombstoned body can still + // back an attempt claimed before physical aging moved the source; + // its emitted lane evidence outranks other retained moved bodies. + // Legacy ambiguous bodies have no authoritative physical priority: + // fail instead of silently choosing an unordered source payload. + let rank = |row: &ReadySnapshotRow| { + ( + row.live_body, + claimed_lanes.contains(&( + row.ready_slot, + row.ready_generation, + row.queue.clone(), + row.enqueue_shard, + row.lane_seq, + row.job_id, + row.run_lease + 1, + row.priority, + )), + ) + }; + let strongest = candidates + .iter() + .map(&rank) + .max() + .expect("ready body group is nonempty"); + let mut best = candidates.iter().filter(|row| rank(row) == strongest); + let ready = best.next().expect("ready body rank has a candidate"); + if best.next().is_some() { + return Err(AwaError::Validation(format!( + "queue storage ready body ambiguous for deleted lease job {} run_lease {}", + deleted_row.job_id, deleted_row.run_lease + ))); + } + // Claim-time aging adds this provenance to the returned JobRow, + // not the immutable ready payload. Recovery must reconstruct the + // same default when it reads that retained body after the claim. + // Mutate the cloned JSON so unrelated payload fields survive. + let mut payload = ready.payload.clone(); + if deleted_row.priority < ready.priority { + let metadata = payload + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation( + "queue storage payload must be a JSON object".to_string(), + ) + })? + .entry("metadata") + .or_insert_with(|| serde_json::json!({})) + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation( + "queue storage payload metadata must be a JSON object".to_string(), + ) + })?; + metadata + .entry("_awa_original_priority") + .or_insert_with(|| serde_json::Value::from(i64::from(ready.priority))); + } let attempt = attempt_map.get(&(deleted_row.job_id, deleted_row.run_lease)); hydrated.push(LeaseTransitionRow { @@ -12682,7 +12882,7 @@ impl QueueStorage { created_at: ready.created_at, unique_key: ready.unique_key.clone(), unique_states: ready.unique_states.clone(), - payload: ready.payload.clone(), + payload, progress: attempt.and_then(|row| row.progress.clone()), }); } diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 9679f20f5..8a1d94643 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -23050,3 +23050,448 @@ async fn test_legacy_narrow_available_cancel_retry_keeps_tombstoned_body() { .is_empty()); assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); } + +/// Commit the same installed claim function used by the runtime, deliberately +/// leaving its separate post-commit cursor update unsent. +async fn claim_before_aging_cursor_advance_in_tx( + receipts: bool, + deadline: Duration, +) -> ( + TestDbGuard, + sqlx::PgPool, + QueueStorage, + RawReceiptClaimRow, + sqlx::Transaction<'static, sqlx::Postgres>, +) { + let (guard, pool) = setup_pool(6).await; + let schema = "awa_qs_claimed_aging_body"; + let queue = "qs_claimed_aging_body"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: receipts, + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9041 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 1, + metadata: serde_json::json!({"retained":"claimed-before-aging"}), + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + // Occupy the replacement lane's first position. Old physical aging must + // append this job at seq2, not accidentally leave a live seq1 body that + // matches the already committed claim's original reference. + enqueue_job( + &pool, + &store, + &CompleteJob { id: 9042 }, + InsertOpts { + queue: queue.into(), + priority: 3, + ..Default::default() + }, + ) + .await; + let mut tx = pool.begin().await.expect("begin exact runtime claim stage"); + let claimed: Vec = sqlx::query_as(audited_sql(format!( + "SELECT ready_slot,ready_generation,job_id,priority,attempt,run_lease,lane_seq,claim_slot + FROM {schema}.claim_ready_runtime($1,1,$2,60::float8)" + ))) + .bind(queue) + .bind(deadline.as_secs_f64()) + .fetch_all(tx.as_mut()) + .await + .expect("claim old ready row"); + assert_eq!(claimed.len(), 1); + let original = claimed.into_iter().next().unwrap(); + assert_eq!(original.job_id, job_id); + assert_eq!( + original.priority, 1, + "claim-time effective aging remains enabled" + ); + assert_eq!( + claim_cursor_for(&pool, &store, queue, 4, 0).await, + original.lane_seq + ); + (guard, pool, store, original, tx) +} + +async fn claim_before_aging_cursor_advance( + receipts: bool, + deadline: Duration, +) -> (TestDbGuard, sqlx::PgPool, QueueStorage, RawReceiptClaimRow) { + let (guard, pool, store, claim, tx) = + claim_before_aging_cursor_advance_in_tx(receipts, deadline).await; + tx.commit() + .await + .expect("commit claim before advancing cursor"); + (guard, pool, store, claim) +} + +async fn assert_claimed_aging_rescue_body( + pool: &sqlx::PgPool, + store: &QueueStorage, + claim: &RawReceiptClaimRow, +) { + let rescued = store + .rescue_stale_heartbeats(pool, Duration::ZERO) + .await + .expect("hydrate retained claimed body"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].id, claim.job_id); + assert_eq!(rescued[0].state, JobState::Failed); + assert_eq!(rescued[0].args, serde_json::json!({"id":9041})); + assert_eq!(rescued[0].max_attempts, 1); + assert_eq!( + rescued[0].priority, 4, + "terminal body keeps its physical lane" + ); + assert_eq!(rescued[0].metadata["retained"], "claimed-before-aging"); + assert_eq!(rescued[0].metadata["_awa_original_priority"], 4); + let retried = store + .retry_job(pool, claim.job_id) + .await + .expect("retry rescued body") + .expect("failed attempt retries"); + assert_eq!(retried.args, rescued[0].args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(retried.priority, 4); + assert_eq!(retried.metadata["_awa_original_priority"], 4); +} + +async fn aging_skips_committed_claim_before_cursor(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; + let aged = store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("age with lagging claim cursor"); + assert!( + aged.is_empty(), + "a committed claim is not waiting, even before cursor advancement: {aged:?}" + ); + assert_eq!(ready_tombstone_count(&pool, &store).await, 0); + assert_claimed_aging_rescue_body(&pool, &store, &claim).await; +} + +async fn historical_tombstoned_claim_hydrates(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; + // Reproduce the retained source state written by pre-fix physical aging. + // Its replacement ready row has a different lane_seq and therefore cannot + // back this exact claim. Keep that unrelated body rather than deleting it. + let schema = store.schema(); + let mut tx = pool.begin().await.expect("begin historical aging shape"); + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.ready_tombstones + (ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id + FROM {schema}.ready_entries WHERE job_id=$1" + ))) + .bind(claim.job_id) + .execute(tx.as_mut()) + .await + .expect("historical source tombstone"); + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.ready_entries + (ready_slot,ready_generation,job_id,kind,queue,args,priority,attempt,run_lease,max_attempts, + lane_seq,enqueue_shard,run_at,attempted_at,created_at,unique_key,unique_states,payload) + SELECT ready_slot,ready_generation,job_id,kind,queue,args,priority-1,attempt,run_lease,max_attempts, + {schema}.reserve_enqueue_seq(queue,(priority-1)::smallint,enqueue_shard,1::bigint),enqueue_shard, + run_at,attempted_at,created_at,unique_key,unique_states, + jsonb_set(COALESCE(payload,'{{}}'::jsonb),'{{metadata}}', + COALESCE(payload->'metadata','{{}}'::jsonb) + || jsonb_build_object('_awa_original_priority',priority)) + FROM {schema}.ready_entries WHERE job_id=$1" + ))).bind(claim.job_id).execute(tx.as_mut()).await.expect("historical replacement at different position"); + tx.commit() + .await + .expect("commit historical source and replacement"); + let positions: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT lane_seq FROM {schema}.ready_entries WHERE job_id=$1 ORDER BY priority DESC" + ))) + .bind(claim.job_id) + .fetch_all(&pool) + .await + .expect("source/replacement positions"); + assert_eq!(positions, vec![1, 2]); + let provenance: Vec<(i16, Option)> = sqlx::query_as(audited_sql(format!( + "SELECT priority,(payload->'metadata'->>'_awa_original_priority')::bigint + FROM {schema}.ready_entries WHERE job_id=$1 ORDER BY priority DESC" + ))) + .bind(claim.job_id) + .fetch_all(&pool) + .await + .expect("old writer provenance placement"); + assert_eq!( + provenance, + vec![(4, None), (3, Some(4))], + "old physical aging adds provenance only to its replacement body" + ); + assert_claimed_aging_rescue_body(&pool, &store, &claim).await; +} + +#[tokio::test] +async fn test_physical_aging_skips_committed_compact_claim_before_cursor() { + aging_skips_committed_claim_before_cursor(true, Duration::ZERO).await; +} +#[tokio::test] +async fn test_physical_aging_skips_committed_deadline_claim_before_cursor() { + aging_skips_committed_claim_before_cursor(true, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_physical_aging_skips_committed_legacy_claim_before_cursor() { + aging_skips_committed_claim_before_cursor(false, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_historical_aged_tombstoned_compact_claim_rescues_body() { + historical_tombstoned_claim_hydrates(true, Duration::ZERO).await; +} +#[tokio::test] +async fn test_historical_aged_tombstoned_deadline_claim_rescues_body() { + historical_tombstoned_claim_hydrates(true, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_historical_aged_tombstoned_legacy_claim_rescues_body() { + historical_tombstoned_claim_hydrates(false, Duration::from_secs(300)).await; +} +async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool) { + let (_guard, pool) = setup_pool(6).await; + let schema = "awa_qs_moved_tomb_claim"; + let queue = "qs_moved_tomb_claim"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: receipts, + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let id = enqueue_job(&pool, &store, &CompleteJob { id: 9043 }, InsertOpts { + queue: queue.into(), priority: from, max_attempts: 1, + metadata: serde_json::json!({"retained":"moved-tomb-claim", "_awa_original_priority":from}), + ..Default::default() + }).await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + assert!(store + .set_priority(&pool, id, to) + .await + .expect("move before claim")); + let body_positions: Vec<(i16, i64)> = sqlx::query_as(audited_sql(format!( + "SELECT priority,lane_seq FROM {schema}.ready_entries WHERE job_id=$1 ORDER BY priority" + ))) + .bind(id) + .fetch_all(&pool) + .await + .expect("shared moved positions"); + assert_eq!(body_positions.len(), 2); + assert!(body_positions.iter().all(|row| row.1 == 1)); + let mut claims = Vec::new(); + for _ in 0..3 { + claims = sqlx::query_as::<_, RawReceiptClaimRow>(audited_sql(format!( + "SELECT ready_slot,ready_generation,job_id,priority,attempt,run_lease,lane_seq,claim_slot + FROM {schema}.claim_ready_runtime($1,1,300::float8,60::float8)" + ))).bind(queue).fetch_all(&pool).await.expect("claim moved body before cursor update"); + if !claims.is_empty() { + break; + } + // A higher-priority tombstoned source may first consume its spent + // cursor prefix; the actual target claim still omits cursor advance. + } + assert_eq!(claims.len(), 1); + assert_eq!(claims[0].job_id, id); + assert_eq!( + claims[0].priority, 1, + "effective priority cannot identify either physical lane" + ); + // The old source was tombstoned by the real move. Model old physical aging + // tombstoning the claimed source as well; both exact-ref bodies now remain. + tombstone_ready_job(&pool, &store, id).await; + assert_eq!(ready_tombstone_count(&pool, &store).await, 2); + if !receipts { + let err = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect_err("legacy ambiguous physical lanes must not be selected arbitrarily"); + assert!(err.to_string().contains("ready body ambiguous"), "{err}"); + assert_eq!( + lease_count(&pool, &store).await, + 1, + "ambiguous mutation rolls back" + ); + return; + } + let emitted: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT priority FROM {schema}.ready_claim_attempt_batches WHERE queue=$1" + ))) + .bind(queue) + .fetch_all(&pool) + .await + .expect("physical emitted lane proof"); + assert_eq!(emitted, vec![to]); + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("use emitted physical lane for retained tombstoned body"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].priority, to); + assert_eq!(rescued[0].args, serde_json::json!({"id":9043})); + assert_eq!(rescued[0].max_attempts, 1); + assert_eq!(rescued[0].metadata["_awa_original_priority"], from); + let retried = store + .retry_job(&pool, id) + .await + .expect("retry selected moved body") + .unwrap(); + assert_eq!(retried.priority, to); + assert_eq!(retried.args, rescued[0].args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(retried.metadata["_awa_original_priority"], from); +} + +#[tokio::test] +async fn test_historical_tombstoned_claim_uses_emitted_lane_low_to_high() { + moved_tombstoned_claim_lane(2, 4, true).await; +} +#[tokio::test] +async fn test_historical_tombstoned_claim_uses_emitted_lane_high_to_low() { + moved_tombstoned_claim_lane(4, 2, true).await; +} +#[tokio::test] +async fn test_historical_tombstoned_legacy_ambiguous_lanes_roll_back() { + moved_tombstoned_claim_lane(2, 4, false).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_claim_completes_body() { + let (_guard, pool, store, claim) = + claim_before_aging_cursor_advance(false, Duration::from_secs(300)).await; + tombstone_ready_job(&pool, &store, claim.job_id).await; + let updated = store + .complete_job_batch_by_id(&pool, &[(claim.job_id, claim.run_lease)]) + .await + .expect("complete an already tombstoned claimed body"); + assert_eq!(updated, vec![(claim.job_id, claim.run_lease)]); + let done = store + .load_job(&pool, claim.job_id) + .await + .expect("read completion") + .unwrap(); + assert_eq!(done.state, JobState::Completed); + assert_eq!(done.args, serde_json::json!({"id":9041})); + assert_eq!(done.max_attempts, 1); + assert_eq!(done.priority, 4); + assert_eq!(done.metadata["retained"], "claimed-before-aging"); + assert_eq!(done.metadata["_awa_original_priority"], 4); +} + +async fn tombstoned_claim_cancels_body(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; + tombstone_ready_job(&pool, &store, claim.job_id).await; + let done = store + .cancel_running( + &pool, + claim.job_id, + claim.run_lease, + "cancel retained body", + None, + ) + .await + .expect("cancel an already tombstoned claimed body") + .unwrap(); + assert_eq!(done.state, JobState::Cancelled); + assert_eq!(done.args, serde_json::json!({"id":9041})); + assert_eq!(done.max_attempts, 1); + assert_eq!(done.priority, 4); + assert_eq!(done.metadata["retained"], "claimed-before-aging"); + assert_eq!(done.metadata["_awa_original_priority"], 4); +} +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_cancels_body() { + tombstoned_claim_cancels_body(true, Duration::ZERO).await; +} +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_cancels_body() { + tombstoned_claim_cancels_body(true, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_historical_tombstoned_legacy_claim_cancels_body() { + tombstoned_claim_cancels_body(false, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_physical_aging_waits_for_claim_head_then_uses_fresh_snapshot() { + let (_guard, pool, store, claim, mut claim_tx) = + claim_before_aging_cursor_advance_in_tx(true, Duration::ZERO).await; + let holder: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(claim_tx.as_mut()) + .await + .expect("actual claim holder PID"); + let aging_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + (*pool.connect_options()) + .clone() + .application_name("awa_aging_head_wait_control"), + ) + .await + .expect("separate aging connection"); + let age = store.age_waiting_priorities(&aging_pool, Duration::from_secs(60), 1); + let release = async { + tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(5)), async { + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity + WHERE application_name='awa_aging_head_wait_control' + AND wait_event_type='Lock' AND $1=ANY(pg_blocking_pids(pid)))", + ) + .bind(holder) + .fetch_one(&pool) + .await + .expect("observe actual claim-head wait"); + if waiting { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("aging must wait for the actual uncommitted claim head"); + claim_tx + .commit() + .await + .expect("publish claim while aging waits"); + }; + let (aged, ()) = tokio::join!(age, release); + assert!( + aged.expect("age after committed claim").is_empty(), + "the statement after the head lock must see the committed claim" + ); + assert_eq!(ready_tombstone_count(&pool, &store).await, 0); + assert_claimed_aging_rescue_body(&pool, &store, &claim).await; + aging_pool.close().await; +} From 9e447ba3a42a3bf489a6a7e163d69281e1fc3af7 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 16:15:57 +1300 Subject: [PATCH 2/8] Preserve claimed bodies across historical aging collisions --- CHANGELOG.md | 8 +- awa-model/src/queue_storage.rs | 61 ++++++----- awa/tests/queue_storage_runtime_test.rs | 132 ++++++++++++++++++------ 3 files changed, 143 insertions(+), 58 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b79b50b8..86bca8934 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,13 @@ Notable changes between releases. Detailed migration notes for storage transitio - Keep committed claims out of physical priority aging, even before their cursor advances. Recovery and cancellation can hydrate an exact retained claimed body after older aging tombstoned it, preserving original-priority - metadata and rejecting ambiguous historical bodies. + metadata and rejecting ambiguous historical bodies. Emitted-lane evidence + takes precedence over colliding live replacements, and their recovered + terminal bodies remain self-contained through retry. Historical receipt + evidence remains usable after reopening a custom store in legacy mode. + Without that evidence, legacy or backfilled claims retain a compatibility + fallback that cannot reliably identify the original physical lane in this + historical collision. - Preserve recovered bodies and physical lane priority when manually moving historical aged terminal failures to the DLQ, so later DLQ retry retains diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 446ddfbd8..b109d3125 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -3309,6 +3309,7 @@ struct LeaseTransitionRow { unique_states: Option, payload: serde_json::Value, progress: Option, + self_contained_terminal: bool, } impl LeaseTransitionRow { @@ -3332,7 +3333,13 @@ impl LeaseTransitionRow { attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, - lane_seq: self.lane_seq, + // A competing live copy must not replace this recovered body + // when later terminal reads resolve retained ready rows. + lane_seq: if self.self_contained_terminal { + -self.job_id + } else { + self.lane_seq + }, enqueue_shard: self.enqueue_shard, run_at: self.run_at, attempted_at: self.attempted_at, @@ -12577,13 +12584,13 @@ impl QueueStorage { let mut claimed_lanes: BTreeSet = BTreeSet::new(); let ambiguous: Vec<&ReadySnapshotRow> = ready_map .values() - .filter(|rows| rows.len() > 1 && !rows.iter().any(|row| row.live_body)) + .filter(|rows| rows.len() > 1) .flatten() .collect(); - if self.lease_claim_receipts() && !ambiguous.is_empty() { - // Normal live bodies and unique retained bodies need no extra - // relation lookup. Only ambiguous tombstoned receipt bodies need - // the emitted physical-lane evidence, scoped to their ready slots. + if !ambiguous.is_empty() { + // A unique body needs no extra relation lookup. Multiple bodies + // need the emitted physical lane even if a replacement is live; + // scope that evidence to the referenced ready slots. let attempt_sources = ambiguous .iter() .map(|row| row.ready_slot) @@ -12802,25 +12809,28 @@ impl QueueStorage { deleted_row.job_id, deleted_row.run_lease )) })?; - // Prefer a live exact-attempt body. A tombstoned body can still - // back an attempt claimed before physical aging moved the source; - // its emitted lane evidence outranks other retained moved bodies. - // Legacy ambiguous bodies have no authoritative physical priority: - // fail instead of silently choosing an unordered source payload. + // Emitted physical-lane evidence wins over a live replacement. + // Without it, retain the live-body compatibility fallback; it + // cannot identify a legacy claim's original colliding lane. + let emitted = |row: &ReadySnapshotRow| { + claimed_lanes.contains(&( + row.ready_slot, + row.ready_generation, + row.queue.clone(), + row.enqueue_shard, + row.lane_seq, + row.job_id, + row.run_lease + 1, + row.priority, + )) + }; + let has_emitted = candidates.iter().any(&emitted); let rank = |row: &ReadySnapshotRow| { - ( - row.live_body, - claimed_lanes.contains(&( - row.ready_slot, - row.ready_generation, - row.queue.clone(), - row.enqueue_shard, - row.lane_seq, - row.job_id, - row.run_lease + 1, - row.priority, - )), - ) + if has_emitted { + emitted(row) + } else { + row.live_body + } }; let strongest = candidates .iter() @@ -12884,6 +12894,9 @@ impl QueueStorage { unique_states: ready.unique_states.clone(), payload, progress: attempt.and_then(|row| row.progress.clone()), + self_contained_terminal: has_emitted + && !ready.live_body + && candidates.iter().any(|row| row.live_body), }); } diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 8a1d94643..7e48878f4 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -23056,6 +23056,7 @@ async fn test_legacy_narrow_available_cancel_retry_keeps_tombstoned_body() { async fn claim_before_aging_cursor_advance_in_tx( receipts: bool, deadline: Duration, + occupy_replacement_lane: bool, ) -> ( TestDbGuard, sqlx::PgPool, @@ -23096,20 +23097,28 @@ async fn claim_before_aging_cursor_advance_in_tx( Utc::now() - chrono::Duration::minutes(5), ) .await; - // Occupy the replacement lane's first position. Old physical aging must - // append this job at seq2, not accidentally leave a live seq1 body that - // matches the already committed claim's original reference. - enqueue_job( - &pool, - &store, - &CompleteJob { id: 9042 }, - InsertOpts { - queue: queue.into(), - priority: 3, - ..Default::default() - }, - ) - .await; + if occupy_replacement_lane { + // Occupy the replacement lane's first position. Old physical aging must + // append this job at seq2, not accidentally leave a live seq1 body that + // matches the already committed claim's original reference. + enqueue_job( + &pool, + &store, + &CompleteJob { id: 9042 }, + InsertOpts { + queue: queue.into(), + priority: 3, + ..Default::default() + }, + ) + .await; + } + if !occupy_replacement_lane { + store + .prepare_queue(&pool, queue, 1) + .await + .expect("prepare empty replacement lane"); + } let mut tx = pool.begin().await.expect("begin exact runtime claim stage"); let claimed: Vec = sqlx::query_as(audited_sql(format!( "SELECT ready_slot,ready_generation,job_id,priority,attempt,run_lease,lane_seq,claim_slot @@ -23139,7 +23148,7 @@ async fn claim_before_aging_cursor_advance( deadline: Duration, ) -> (TestDbGuard, sqlx::PgPool, QueueStorage, RawReceiptClaimRow) { let (guard, pool, store, claim, tx) = - claim_before_aging_cursor_advance_in_tx(receipts, deadline).await; + claim_before_aging_cursor_advance_in_tx(receipts, deadline, true).await; tx.commit() .await .expect("commit claim before advancing cursor"); @@ -23150,6 +23159,7 @@ async fn assert_claimed_aging_rescue_body( pool: &sqlx::PgPool, store: &QueueStorage, claim: &RawReceiptClaimRow, + self_contained_terminal: bool, ) { let rescued = store .rescue_stale_heartbeats(pool, Duration::ZERO) @@ -23166,6 +23176,26 @@ async fn assert_claimed_aging_rescue_body( ); assert_eq!(rescued[0].metadata["retained"], "claimed-before-aging"); assert_eq!(rescued[0].metadata["_awa_original_priority"], 4); + let schema = store.schema(); + let terminal: (i64, bool) = sqlx::query_as(audited_sql(format!( + "SELECT lane_seq,args IS NOT NULL FROM {schema}.done_entries WHERE job_id=$1" + ))) + .bind(claim.job_id) + .fetch_one(pool) + .await + .expect("terminal backing identity"); + assert_eq!( + terminal, + ( + if self_contained_terminal { + -claim.job_id + } else { + claim.lane_seq + }, + self_contained_terminal, + ), + "only conflicting historical body recovery needs a self-contained terminal row" + ); let retried = store .retry_job(pool, claim.job_id) .await @@ -23188,14 +23218,20 @@ async fn aging_skips_committed_claim_before_cursor(receipts: bool, deadline: Dur "a committed claim is not waiting, even before cursor advancement: {aged:?}" ); assert_eq!(ready_tombstone_count(&pool, &store).await, 0); - assert_claimed_aging_rescue_body(&pool, &store, &claim).await; + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; } -async fn historical_tombstoned_claim_hydrates(receipts: bool, deadline: Duration) { - let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; - // Reproduce the retained source state written by pre-fix physical aging. - // Its replacement ready row has a different lane_seq and therefore cannot - // back this exact claim. Keep that unrelated body rather than deleting it. +async fn historical_tombstoned_claim_hydrates( + receipts: bool, + deadline: Duration, + colliding_position: bool, +) { + let (_guard, pool, store, claim, tx) = + claim_before_aging_cursor_advance_in_tx(receipts, deadline, !colliding_position).await; + tx.commit() + .await + .expect("commit claim before old physical aging"); + // Reproduce the retained source and live replacement written by old aging. let schema = store.schema(); let mut tx = pool.begin().await.expect("begin historical aging shape"); sqlx::query(audited_sql(format!( @@ -23219,7 +23255,7 @@ async fn historical_tombstoned_claim_hydrates(receipts: bool, deadline: Duration COALESCE(payload->'metadata','{{}}'::jsonb) || jsonb_build_object('_awa_original_priority',priority)) FROM {schema}.ready_entries WHERE job_id=$1" - ))).bind(claim.job_id).execute(tx.as_mut()).await.expect("historical replacement at different position"); + ))).bind(claim.job_id).execute(tx.as_mut()).await.expect("historical replacement body"); tx.commit() .await .expect("commit historical source and replacement"); @@ -23230,7 +23266,7 @@ async fn historical_tombstoned_claim_hydrates(receipts: bool, deadline: Duration .fetch_all(&pool) .await .expect("source/replacement positions"); - assert_eq!(positions, vec![1, 2]); + assert_eq!(positions, vec![1, if colliding_position { 1 } else { 2 }]); let provenance: Vec<(i16, Option)> = sqlx::query_as(audited_sql(format!( "SELECT priority,(payload->'metadata'->>'_awa_original_priority')::bigint FROM {schema}.ready_entries WHERE job_id=$1 ORDER BY priority DESC" @@ -23244,7 +23280,7 @@ async fn historical_tombstoned_claim_hydrates(receipts: bool, deadline: Duration vec![(4, None), (3, Some(4))], "old physical aging adds provenance only to its replacement body" ); - assert_claimed_aging_rescue_body(&pool, &store, &claim).await; + assert_claimed_aging_rescue_body(&pool, &store, &claim, colliding_position).await; } #[tokio::test] @@ -23261,17 +23297,26 @@ async fn test_physical_aging_skips_committed_legacy_claim_before_cursor() { } #[tokio::test] async fn test_historical_aged_tombstoned_compact_claim_rescues_body() { - historical_tombstoned_claim_hydrates(true, Duration::ZERO).await; + historical_tombstoned_claim_hydrates(true, Duration::ZERO, false).await; } #[tokio::test] async fn test_historical_aged_tombstoned_deadline_claim_rescues_body() { - historical_tombstoned_claim_hydrates(true, Duration::from_secs(300)).await; + historical_tombstoned_claim_hydrates(true, Duration::from_secs(300), false).await; } #[tokio::test] async fn test_historical_aged_tombstoned_legacy_claim_rescues_body() { - historical_tombstoned_claim_hydrates(false, Duration::from_secs(300)).await; + historical_tombstoned_claim_hydrates(false, Duration::from_secs(300), false).await; +} +#[tokio::test] +async fn test_historical_colliding_live_replacement_compact_rescues_claimed_body() { + historical_tombstoned_claim_hydrates(true, Duration::ZERO, true).await; +} +#[tokio::test] +async fn test_historical_colliding_live_replacement_deadline_rescues_claimed_body() { + historical_tombstoned_claim_hydrates(true, Duration::from_secs(300), true).await; } -async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool) { + +async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool, reopen_as_legacy: bool) { let (_guard, pool) = setup_pool(6).await; let schema = "awa_qs_moved_tomb_claim"; let queue = "qs_moved_tomb_claim"; @@ -23328,6 +23373,22 @@ async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool) { claims[0].priority, 1, "effective priority cannot identify either physical lane" ); + let store = if reopen_as_legacy { + store + .register_callback(&pool, id, claims[0].run_lease, Duration::from_secs(60)) + .await + .expect("materialize receipt before reopening"); + assert_eq!(lease_count(&pool, &store).await, 1); + QueueStorage::new(QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: false, + queue_stripe_count: 1, + ..Default::default() + }) + .expect("reopen existing custom store in legacy mode") + } else { + store + }; // The old source was tombstoned by the real move. Model old physical aging // tombstoning the claimed source as well; both exact-ref bodies now remain. tombstone_ready_job(&pool, &store, id).await; @@ -23375,15 +23436,20 @@ async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool) { #[tokio::test] async fn test_historical_tombstoned_claim_uses_emitted_lane_low_to_high() { - moved_tombstoned_claim_lane(2, 4, true).await; + moved_tombstoned_claim_lane(2, 4, true, false).await; } #[tokio::test] async fn test_historical_tombstoned_claim_uses_emitted_lane_high_to_low() { - moved_tombstoned_claim_lane(4, 2, true).await; + moved_tombstoned_claim_lane(4, 2, true, false).await; } #[tokio::test] async fn test_historical_tombstoned_legacy_ambiguous_lanes_roll_back() { - moved_tombstoned_claim_lane(2, 4, false).await; + moved_tombstoned_claim_lane(2, 4, false, false).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_receipt_materialized_then_reopened_legacy_uses_evidence() { + moved_tombstoned_claim_lane(2, 4, true, true).await; } #[tokio::test] @@ -23446,7 +23512,7 @@ async fn test_historical_tombstoned_legacy_claim_cancels_body() { #[tokio::test] async fn test_physical_aging_waits_for_claim_head_then_uses_fresh_snapshot() { let (_guard, pool, store, claim, mut claim_tx) = - claim_before_aging_cursor_advance_in_tx(true, Duration::ZERO).await; + claim_before_aging_cursor_advance_in_tx(true, Duration::ZERO, true).await; let holder: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") .fetch_one(claim_tx.as_mut()) .await @@ -23492,6 +23558,6 @@ async fn test_physical_aging_waits_for_claim_head_then_uses_fresh_snapshot() { "the statement after the head lock must see the committed claim" ); assert_eq!(ready_tombstone_count(&pool, &store).await, 0); - assert_claimed_aging_rescue_body(&pool, &store, &claim).await; + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; aging_pool.close().await; } From 19300a1227ce5d097fd8f3415264b2cfa0cfc67c Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 16:55:06 +1300 Subject: [PATCH 3/8] Recover retained active bodies and retire old aging copies --- CHANGELOG.md | 30 +- awa-model/src/queue_storage.rs | 624 +++++++++++++++--------- awa/tests/queue_storage_runtime_test.rs | 623 ++++++++++++++++++++++- 3 files changed, 1030 insertions(+), 247 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 86bca8934..955876be0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,16 +6,26 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed -- Keep committed claims out of physical priority aging, even before their - cursor advances. Recovery and cancellation can hydrate an exact retained - claimed body after older aging tombstoned it, preserving original-priority - metadata and rejecting ambiguous historical bodies. Emitted-lane evidence - takes precedence over colliding live replacements, and their recovered - terminal bodies remain self-contained through retry. Historical receipt - evidence remains usable after reopening a custom store in legacy mode. - Without that evidence, legacy or backfilled claims retain a compatibility - fallback that cannot reliably identify the original physical lane in this - historical collision. +- Keep committed claims out of physical priority aging before its batch limit, + even before their cursor advances. Active point and callback reads recover + retained tombstoned bodies left by older aging. Recovery, cancellation and + completion by ID hydrate those bodies and retire unclaimed copies of the same + attempt, preserving original-priority metadata and rejecting ambiguous + historical bodies. Emitted-lane evidence takes precedence over colliding live + replacements, and their recovered terminal bodies remain self-contained + through retry. Cancellation prefers an active attempt over its historical + available replacement. Historical receipt evidence remains usable after + reopening a custom store in legacy mode. Without that evidence, legacy or + backfilled claims retain a compatibility fallback that cannot reliably + identify the original physical lane in this historical collision. + Repair runs only when a transition hydrates a retained tombstoned body; + worker completion using its carried job body can leave an old replacement + claimable. Updated aging callers prevent new replacements, but cannot fence + an older process explicitly calling physical aging during a rolling upgrade. + Repair takes claim-head row locks without waiting and rolls back on + contention; the operation must retry. It adds weak relation locks and + tombstone writes only on this recovery path. Fresh rechecks rely on READ + COMMITTED isolation. - Preserve recovered bodies and physical lane priority when manually moving historical aged terminal failures to the DLQ, so later DLQ retry retains diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index b109d3125..05bd87a36 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -9,7 +9,7 @@ use crate::{InsertParams, JobRow, JobState}; use chrono::TimeDelta; use chrono::{DateTime, Utc}; use sqlx::Acquire; -use sqlx::{PgPool, Postgres}; +use sqlx::{FromRow, PgPool, Postgres, Row}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; @@ -2602,8 +2602,8 @@ impl ReadyJobRow { } } -#[derive(Debug, sqlx::FromRow)] -struct ReadyAgingCandidate { +#[derive(Debug, PartialEq, Eq, sqlx::FromRow)] +struct ReadyRowIdentity { ready_slot: i32, ready_generation: i64, queue: String, @@ -3274,6 +3274,53 @@ struct ReadySnapshotRow { live_body: bool, } +type ReadyBodyKey = (i32, i64, String, i16, i64, i64, i64); +type ReadyBodyLaneKey = (i32, i64, String, i16, i64, i64, i64, i16); + +fn select_claimed_ready_body<'a>( + candidates: &'a [ReadySnapshotRow], + claimed_lanes: &BTreeSet, + reference: &DeletedLeaseRow, +) -> Result<(&'a ReadySnapshotRow, bool), AwaError> { + // Emitted physical-lane evidence wins over a live replacement. + // Without it, retain the live-body compatibility fallback; it + // cannot identify a legacy claim's original colliding lane. + let emitted = |row: &ReadySnapshotRow| { + claimed_lanes.contains(&( + row.ready_slot, + row.ready_generation, + row.queue.clone(), + row.enqueue_shard, + row.lane_seq, + row.job_id, + row.run_lease + 1, + row.priority, + )) + }; + let has_emitted = candidates.iter().any(&emitted); + let rank = |row: &ReadySnapshotRow| { + if has_emitted { + emitted(row) + } else { + row.live_body + } + }; + let strongest = candidates + .iter() + .map(&rank) + .max() + .expect("ready body group is nonempty"); + let mut best = candidates.iter().filter(|row| rank(row) == strongest); + let ready = best.next().expect("ready body rank has a candidate"); + if best.next().is_some() { + return Err(AwaError::Validation(format!( + "queue storage ready body ambiguous for deleted lease job {} run_lease {}", + reference.job_id, reference.run_lease + ))); + } + Ok((ready, has_emitted)) +} + #[derive(Debug, Clone, sqlx::FromRow)] struct AttemptStateRow { job_id: i64, @@ -10170,10 +10217,10 @@ impl QueueStorage { } } } - let receipt = match locked_receipt { + match locked_receipt { Err(err) if is_lock_contention_error(&err) => return Ok(CancelStep::Busy), Err(err) => return Err(map_sqlx_error(err)), - Ok(Some(receipt)) if receipt.1 == run_lease => receipt, + Ok(Some(receipt)) if receipt.1 == run_lease => {} Ok(Some(_)) => return Ok(CancelStep::Again), Ok(None) => { // Closed by another closer, or materialized before this @@ -10195,156 +10242,17 @@ impl QueueStorage { }); } }; - let ( - claim_slot, - run_lease, - ready_slot, - ready_generation, - queue, - priority, - attempt, - max_attempts, - lane_seq, - claimed_at, - receipt_id, - compact_batch, - ) = receipt; - // Hydrate the ready row so we can synthesize the done - // row with the original args/payload. - let ready_match: Option = sqlx::query_as(audited_sql(format!( - r#" - SELECT - ready_slot, - ready_generation, - job_id, - kind, - queue, - args, - priority, - attempt, - run_lease, - max_attempts, - lane_seq, - enqueue_shard, - run_at, - attempted_at, - created_at, - unique_key, - unique_states, - COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries - WHERE job_id = $1 - AND ready_slot = $2 - AND ready_generation = $3 - AND queue = $4 - AND lane_seq = $5 - "# - ))) - .bind(job_id) - .bind(ready_slot) - .bind(ready_generation) - .bind(&queue) - .bind(lane_seq) - .fetch_optional(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - - let Some(ready) = ready_match else { - // Shouldn't happen: the claim references a ready row. - return Ok(CancelStep::Done(RunningAttemptCancel::Left)); - }; - - let done = DoneJobRow { - ready_slot, - ready_generation, - job_id, - kind: ready.kind, - queue: queue.clone(), - args: ready.args, - state: JobState::Cancelled, - priority, - attempt, - run_lease, - max_attempts, - lane_seq, - enqueue_shard: ready.enqueue_shard, - run_at: ready.run_at, - attempted_at: Some(claimed_at), - finalized_at: Utc::now(), - created_at: ready.created_at, - unique_key: ready.unique_key, - unique_states: ready.unique_states, - payload: ready.payload, - output: None, + let Some(lease) = self + .close_open_receipt_claim_tx(tx, job_id, run_lease, "cancelled", None) + .await? + else { + return Ok(CancelStep::Again); }; + let done_payload = + Self::payload_with_attempt_state(lease.payload.clone(), lease.progress.clone())?; + let done = lease.into_done_row(JobState::Cancelled, Utc::now(), done_payload); self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Running)) .await?; - // Write the closure into the same claim partition. A - // compact batch-sourced claim has no lease_claims row, so - // it closes into the batch ledger the queue prune gate - // counts via compact_count; a deadline lease claim keeps - // its explicit closure so the gate balances it against the - // lease_claims row. - if compact_batch { - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.lease_claim_closure_batches ( - claim_slot, - ready_slot, - ready_generation, - outcome, - closed_count, - receipt_ids, - receipt_ranges, - closed_at - ) - VALUES ( - $1, - $2, - $3, - 'cancelled', - 1, - ARRAY[$4::bigint], - int8multirange(int8range($4::bigint, $4::bigint + 1, '[)')), - clock_timestamp() - ) - "# - ))) - .bind(claim_slot) - .bind(ready_slot) - .bind(ready_generation) - .bind(receipt_id) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - } else { - sqlx::query(audited_sql(format!( - r#" - WITH inserted AS ( - INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at) - VALUES ($1, $2, $3, 'cancelled', clock_timestamp()) - ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING - RETURNING claim_slot, job_id, run_lease, closed_at - ), - marked AS ( - UPDATE {schema}.lease_claims AS claims - SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) - FROM inserted - WHERE claims.claim_slot = inserted.claim_slot - AND claims.job_id = inserted.job_id - AND claims.run_lease = inserted.run_lease - RETURNING claims.job_id - ) - SELECT count(*) FROM marked - "# - ))) - .bind(claim_slot) - .bind(job_id) - .bind(run_lease) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - } self.notify_cancellation_tx(tx, job_id, run_lease).await?; Ok(CancelStep::Done(RunningAttemptCancel::Cancelled(Box::new( CancelJobTxResult { @@ -10429,6 +10337,16 @@ impl QueueStorage { ) -> Result, AwaError> { let schema = self.schema(); self.lock_claim_heads_of_jobs_tx(tx, &[job_id]).await?; + if self.lease_claim_receipts() { + match self.cancel_running_attempt_tx(tx, job_id).await? { + RunningAttemptCancel::Cancelled(result) => return Ok(Some(*result)), + RunningAttemptCancel::Left => return Ok(None), + RunningAttemptCancel::NotRunning => {} + } + } else if let Some(result) = self.cancel_materialized_lease_tx(tx, job_id, None).await? { + return Ok(Some(result)); + } + let claimed = self.ready_row_claimed_sql("ready"); let ready: Option = sqlx::query_as(audited_sql(format!( r#" @@ -10519,16 +10437,6 @@ impl QueueStorage { })); } - if self.lease_claim_receipts() { - match self.cancel_running_attempt_tx(tx, job_id).await? { - RunningAttemptCancel::Cancelled(result) => return Ok(Some(*result)), - RunningAttemptCancel::Left => return Ok(None), - RunningAttemptCancel::NotRunning => {} - } - } else if let Some(result) = self.cancel_materialized_lease_tx(tx, job_id, None).await? { - return Ok(Some(result)); - } - let deferred: Option = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.deferred_jobs @@ -11032,8 +10940,19 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_ids: &[i64], + ) -> Result<(), AwaError> { + self.lock_claim_heads_of_jobs_with_wait_tx(tx, job_ids, false) + .await + } + + async fn lock_claim_heads_of_jobs_with_wait_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + job_ids: &[i64], + nowait: bool, ) -> Result<(), AwaError> { let schema = self.schema(); + let wait_clause = if nowait { "NOWAIT" } else { "" }; sqlx::query(audited_sql(format!( r#" SELECT 1 @@ -11044,7 +10963,7 @@ impl QueueStorage { WHERE ready.job_id = ANY($1) ) ORDER BY claims.queue COLLATE "C", claims.priority, claims.enqueue_shard - FOR UPDATE OF claims + FOR UPDATE OF claims {wait_clause} "# ))) .bind(job_ids) @@ -11586,12 +11505,13 @@ impl QueueStorage { .map_err(|err| AwaError::Validation(format!("invalid aging interval: {err}")))?; let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; + let claimed = self.ready_row_claimed_sql("ready"); // Claims commit before their cursor advance. Select only a bounded // set of identities before taking heads, without locking ready rows in // the reverse order to a claimer. A later statement rechecks these same // identities after every in-flight claim on their lanes has committed. - let candidates: Vec = sqlx::query_as(audited_sql(format!( + let candidates: Vec = sqlx::query_as(audited_sql(format!( r#" SELECT ready.ready_slot, ready.ready_generation, ready.queue, ready.priority, ready.enqueue_shard, ready.lane_seq, @@ -11604,6 +11524,7 @@ impl QueueStorage { WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) AND ready.priority > 1 AND ready.run_at <= $1 + AND NOT ({claimed}) AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.queue = ready.queue @@ -11628,8 +11549,6 @@ impl QueueStorage { } let job_ids: Vec = candidates.iter().map(|row| row.job_id).collect(); self.lock_claim_heads_of_jobs_tx(&mut tx, &job_ids).await?; - let claimed = self.ready_row_claimed_sql("ready"); - let moved: Vec = sqlx::query_as(audited_sql(format!( r#" WITH candidates(ready_slot, ready_generation, queue, priority, @@ -12477,18 +12396,17 @@ impl QueueStorage { Ok(targets.len()) } - /// `write_receipt_closures` is false when the caller has already closed - /// the receipt for every row in `deleted` in this transaction. - async fn hydrate_deleted_leases_tx<'a>( + async fn load_ready_snapshot_candidates_tx( &self, - tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, - deleted: Vec, - write_receipt_closures: bool, - ) -> Result, AwaError> { - if deleted.is_empty() { - return Ok(Vec::new()); - } - + tx: &mut sqlx::Transaction<'_, Postgres>, + deleted: &[DeletedLeaseRow], + ) -> Result< + ( + BTreeMap>, + BTreeSet, + ), + AwaError, + > { let schema = self.schema(); let ready_slots: Vec = deleted.iter().map(|row| row.ready_slot).collect(); let ready_generations: Vec = deleted.iter().map(|row| row.ready_generation).collect(); @@ -12563,7 +12481,6 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - type ReadyBodyKey = (i32, i64, String, i16, i64, i64, i64); let mut ready_map: BTreeMap> = BTreeMap::new(); for row in ready_rows { ready_map @@ -12580,7 +12497,6 @@ impl QueueStorage { .push(row); } - type ReadyBodyLaneKey = (i32, i64, String, i16, i64, i64, i64, i16); let mut claimed_lanes: BTreeSet = BTreeSet::new(); let ambiguous: Vec<&ReadySnapshotRow> = ready_map .values() @@ -12640,6 +12556,27 @@ impl QueueStorage { claimed_lanes.extend(lanes); } + Ok((ready_map, claimed_lanes)) + } + + /// `write_receipt_closures` is false when the caller has already closed + /// the receipt for every row in `deleted` in this transaction. + async fn hydrate_deleted_leases_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + deleted: Vec, + write_receipt_closures: bool, + ) -> Result, AwaError> { + if deleted.is_empty() { + return Ok(Vec::new()); + } + + let schema = self.schema(); + let job_ids: Vec = deleted.iter().map(|row| row.job_id).collect(); + let run_leases: Vec = deleted.iter().map(|row| row.run_lease).collect(); + let (ready_map, claimed_lanes) = + self.load_ready_snapshot_candidates_tx(tx, &deleted).await?; + let attempt_rows: Vec = sqlx::query_as(audited_sql(format!( r#" WITH refs(job_id, run_lease) AS ( @@ -12792,6 +12729,7 @@ impl QueueStorage { .collect(); let mut hydrated = Vec::with_capacity(deleted.len()); + let mut retained_sources = Vec::new(); for deleted_row in deleted { let candidates = ready_map .get(&( @@ -12809,41 +12747,19 @@ impl QueueStorage { deleted_row.job_id, deleted_row.run_lease )) })?; - // Emitted physical-lane evidence wins over a live replacement. - // Without it, retain the live-body compatibility fallback; it - // cannot identify a legacy claim's original colliding lane. - let emitted = |row: &ReadySnapshotRow| { - claimed_lanes.contains(&( - row.ready_slot, - row.ready_generation, - row.queue.clone(), - row.enqueue_shard, - row.lane_seq, - row.job_id, - row.run_lease + 1, - row.priority, - )) - }; - let has_emitted = candidates.iter().any(&emitted); - let rank = |row: &ReadySnapshotRow| { - if has_emitted { - emitted(row) - } else { - row.live_body - } - }; - let strongest = candidates - .iter() - .map(&rank) - .max() - .expect("ready body group is nonempty"); - let mut best = candidates.iter().filter(|row| rank(row) == strongest); - let ready = best.next().expect("ready body rank has a candidate"); - if best.next().is_some() { - return Err(AwaError::Validation(format!( - "queue storage ready body ambiguous for deleted lease job {} run_lease {}", - deleted_row.job_id, deleted_row.run_lease - ))); + let (ready, has_emitted) = + select_claimed_ready_body(candidates, &claimed_lanes, &deleted_row)?; + if !ready.live_body { + retained_sources.push(ReadyRowIdentity { + ready_slot: ready.ready_slot, + ready_generation: ready.ready_generation, + queue: ready.queue.clone(), + priority: ready.priority, + enqueue_shard: ready.enqueue_shard, + lane_seq: ready.lane_seq, + job_id: ready.job_id, + run_lease: ready.run_lease, + }); } // Claim-time aging adds this provenance to the returned JobRow, // not the immutable ready payload. Recovery must reconstruct the @@ -12900,9 +12816,146 @@ impl QueueStorage { }); } + if !retained_sources.is_empty() { + self.retire_retained_attempt_copies_tx(tx, &retained_sources) + .await?; + } Ok(hydrated) } + async fn retire_retained_attempt_copies_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + sources: &[ReadyRowIdentity], + ) -> Result<(), AwaError> { + let schema = self.schema(); + let mut expected = BTreeMap::new(); + for source in sources { + if expected + .insert((source.job_id, source.run_lease), source) + .is_some() + { + return Err(AwaError::Validation(format!( + "queue storage ready body ambiguous for deleted lease job {} run_lease {}", + source.job_id, + source.run_lease + 1 + ))); + } + } + let job_ids: Vec = sources.iter().map(|row| row.job_id).collect(); + let run_leases: Vec = sources.iter().map(|row| row.run_lease).collect(); + // Hydration already holds attempt locks. A contended head must abort + // instead of waiting in the reverse of the claim/mutation lock order. + self.lock_claim_heads_of_jobs_with_wait_tx(tx, &job_ids, true) + .await?; + let copies: Vec = sqlx::query_as(audited_sql(format!( + r#" + WITH refs(job_id, run_lease) AS ( + SELECT * FROM unnest($1::bigint[], $2::bigint[]) + ) + SELECT ready.ready_slot, ready.ready_generation, ready.queue, ready.priority, + ready.enqueue_shard, ready.lane_seq, ready.job_id, ready.run_lease + FROM refs + JOIN {schema}.ready_entries AS ready + ON ready.job_id = refs.job_id AND ready.run_lease = refs.run_lease + "# + ))) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let attempt_sources = copies + .iter() + .map(|row| ring_slot_index(row.ready_slot, self.queue_slot_count(), "queue")) + .collect::, _>>()? + .into_iter() + .map(|slot| { + format!( + "SELECT * FROM {}", + ready_claim_attempt_batch_child_name(schema, slot) + ) + }) + .collect::>() + .join(" UNION ALL "); + if copies.is_empty() { + return Err(AwaError::Validation( + "queue storage ready row missing for deleted lease during retirement".to_string(), + )); + } + let slots: Vec = copies.iter().map(|row| row.ready_slot).collect(); + let generations: Vec = copies.iter().map(|row| row.ready_generation).collect(); + let queues: Vec = copies.iter().map(|row| row.queue.clone()).collect(); + let priorities: Vec = copies.iter().map(|row| row.priority).collect(); + let shards: Vec = copies.iter().map(|row| row.enqueue_shard).collect(); + let sequences: Vec = copies.iter().map(|row| row.lane_seq).collect(); + let copy_job_ids: Vec = copies.iter().map(|row| row.job_id).collect(); + let copy_leases: Vec = copies.iter().map(|row| row.run_lease).collect(); + let emitted: Vec = sqlx::query_as(audited_sql(format!( + r#" + WITH refs(ready_slot, ready_generation, queue, priority, enqueue_shard, + lane_seq, job_id, run_lease) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], + $5::smallint[], $6::bigint[], $7::bigint[], $8::bigint[]) + ) + SELECT * FROM refs + WHERE EXISTS ( + SELECT 1 FROM ({attempt_sources}) AS attempts + WHERE attempts.ready_slot = refs.ready_slot + AND attempts.ready_generation = refs.ready_generation + AND attempts.queue = refs.queue AND attempts.priority = refs.priority + AND attempts.enqueue_shard = refs.enqueue_shard + AND attempts.first_lane_seq <= refs.lane_seq + AND attempts.next_lane_seq > refs.lane_seq + AND attempts.lane_ranges @> int8range(refs.lane_seq, refs.lane_seq + 1, '[)') + ) + "# + ))) + .bind(&slots) + .bind(&generations) + .bind(&queues) + .bind(&priorities) + .bind(&shards) + .bind(&sequences) + .bind(©_job_ids) + .bind(©_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let mut seen = BTreeSet::new(); + for row in &emitted { + let key = (row.job_id, row.run_lease); + if !seen.insert(key) || expected.get(&key).copied() != Some(row) { + return Err(AwaError::Validation(format!( + "queue storage ready body ambiguous for deleted lease job {} run_lease {}", + row.job_id, + row.run_lease + 1 + ))); + } + } + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.ready_tombstones ( + ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id + ) + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], + $5::smallint[], $6::bigint[], $7::bigint[]) + ON CONFLICT DO NOTHING + "# + ))) + .bind(&slots) + .bind(&generations) + .bind(&queues) + .bind(&priorities) + .bind(&shards) + .bind(&sequences) + .bind(©_job_ids) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(()) + } + /// Close one compact receipt claim by its batch identity. The claim /// batch row is addressed by primary key and the item by index, and the /// closure evidence probes are all index lookups, so the cost does not @@ -14558,6 +14611,87 @@ impl QueueStorage { Ok(rescued) } + async fn active_job_read_row_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + row: sqlx::postgres::PgRow, + ) -> Result, AwaError> { + let body_job_id: Option = row.try_get("body_job_id").map_err(map_sqlx_error)?; + if body_job_id.is_some() { + return LeaseJobRow::from_row(&row) + .map_err(map_sqlx_error)? + .into_job_row() + .map(Some); + } + let reference = DeletedLeaseRow::from_row(&row).map_err(map_sqlx_error)?; + let (ready_map, claimed_lanes) = self + .load_ready_snapshot_candidates_tx(tx, std::slice::from_ref(&reference)) + .await?; + let key = ( + reference.ready_slot, + reference.ready_generation, + reference.queue.clone(), + reference.enqueue_shard, + reference.lane_seq, + reference.job_id, + reference.run_lease, + ); + let Some(candidates) = ready_map.get(&key) else { + return Ok(None); + }; + let (ready, _) = select_claimed_ready_body(candidates, &claimed_lanes, &reference)?; + let mut payload = ready.payload.clone(); + if reference.priority < ready.priority { + let metadata = payload + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation("queue storage payload must be a JSON object".into()) + })? + .entry("metadata") + .or_insert_with(|| serde_json::json!({})) + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation( + "queue storage payload metadata must be a JSON object".into(), + ) + })?; + metadata + .entry("_awa_original_priority") + .or_insert_with(|| serde_json::Value::from(i64::from(ready.priority))); + } + LeaseJobRow { + job_id: reference.job_id, + kind: ready.kind.clone(), + queue: ready.queue.clone(), + args: ready.args.clone(), + state: reference.state, + priority: reference.priority, + attempt: reference.attempt, + run_lease: reference.run_lease, + max_attempts: reference.max_attempts, + run_at: ready.run_at, + heartbeat_at: row.try_get("heartbeat_at").map_err(map_sqlx_error)?, + deadline_at: row.try_get("deadline_at").map_err(map_sqlx_error)?, + attempted_at: reference.attempted_at, + finalized_at: None, + created_at: ready.created_at, + unique_key: ready.unique_key.clone(), + callback_id: row.try_get("callback_id").map_err(map_sqlx_error)?, + callback_timeout_at: row.try_get("callback_timeout_at").map_err(map_sqlx_error)?, + callback_filter: row.try_get("callback_filter").map_err(map_sqlx_error)?, + callback_on_complete: row + .try_get("callback_on_complete") + .map_err(map_sqlx_error)?, + callback_on_fail: row.try_get("callback_on_fail").map_err(map_sqlx_error)?, + callback_transform: row.try_get("callback_transform").map_err(map_sqlx_error)?, + payload, + progress: row.try_get("progress").map_err(map_sqlx_error)?, + callback_result: row.try_get("callback_result").map_err(map_sqlx_error)?, + } + .into_job_row() + .map(Some) + } + pub async fn load_job(&self, pool: &PgPool, job_id: i64) -> Result, AwaError> { match self.load_job_inner(pool, job_id, true).await { Err(err) if err.is_undefined_column() => { @@ -14699,14 +14833,16 @@ impl QueueStorage { candidates.push(row.into_job_row()?); } - let lease_rows: Vec = sqlx::query_as(audited_sql(format!( + let lease_rows = sqlx::query(audited_sql(format!( r#" SELECT + ready.job_id AS body_job_id, lease.ready_slot, lease.ready_generation, lease.job_id, + lease.enqueue_shard, ready.kind, - ready.queue, + lease.queue, ready.args, lease.state, lease.priority, @@ -14732,7 +14868,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {schema}.leases AS lease - JOIN {schema}.ready_entries AS ready + LEFT JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue @@ -14765,19 +14901,23 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; for row in lease_rows { - candidates.push(row.into_job_row()?); + if let Some(job) = self.active_job_read_row_tx(&mut tx, row).await? { + candidates.push(job); + } } // Report receipt-backed attempts as running by anti-joining // lease_claims against every durable closure evidence shape. - let lease_claim_rows: Vec = sqlx::query_as(audited_sql(format!( + let lease_claim_rows = sqlx::query(audited_sql(format!( r#" SELECT + ready.job_id AS body_job_id, claims.ready_slot, claims.ready_generation, claims.job_id, + claims.enqueue_shard, ready.kind, - ready.queue, + claims.queue, ready.args, 'running'::awa.job_state AS state, claims.priority, @@ -14803,7 +14943,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {schema}.lease_claims AS claims - JOIN {schema}.ready_entries AS ready + LEFT JOIN {schema}.ready_entries AS ready ON ready.ready_slot = claims.ready_slot AND ready.ready_generation = claims.ready_generation AND ready.queue = claims.queue @@ -14845,21 +14985,25 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; for row in lease_claim_rows { - candidates.push(row.into_job_row()?); + if let Some(job) = self.active_job_read_row_tx(&mut tx, row).await? { + candidates.push(job); + } } // Receipt claims are stored as compact batches. Expand // them only for this admin read, and report still-open items as // running until durable closure, terminal, or materialized-lease // evidence supersedes the claim. - let lease_claim_batch_rows: Vec = sqlx::query_as(audited_sql(format!( + let lease_claim_batch_rows = sqlx::query(audited_sql(format!( r#" SELECT + ready.job_id AS body_job_id, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, + claim_batches.enqueue_shard, ready.kind, - ready.queue, + claim_batches.queue, ready.args, 'running'::awa.job_state AS state, claim_batches.priority, @@ -14893,7 +15037,7 @@ impl QueueStorage { claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) - JOIN {schema}.ready_entries AS ready + LEFT JOIN {schema}.ready_entries AS ready ON ready.ready_slot = claim_batches.ready_slot AND ready.ready_generation = claim_batches.ready_generation AND ready.queue = claim_batches.queue @@ -14956,7 +15100,9 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; for row in lease_claim_batch_rows { - candidates.push(row.into_job_row()?); + if let Some(job) = self.active_job_read_row_tx(&mut tx, row).await? { + candidates.push(job); + } } let done_rows: Vec = sqlx::query_as(audited_sql(format!( @@ -15320,14 +15466,16 @@ impl QueueStorage { run_lease: i64, ) -> Result, AwaError> { let schema = self.schema(); - let row: Option = sqlx::query_as(audited_sql(format!( + let row = sqlx::query(audited_sql(format!( r#" SELECT + ready.job_id AS body_job_id, lease.ready_slot, lease.ready_generation, lease.job_id, + lease.enqueue_shard, ready.kind, - ready.queue, + lease.queue, ready.args, lease.state, lease.priority, @@ -15353,7 +15501,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {schema}.leases AS lease - JOIN {schema}.ready_entries AS ready + LEFT JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue @@ -15386,7 +15534,10 @@ impl QueueStorage { .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; - row.map(LeaseJobRow::into_job_row).transpose() + match row { + Some(row) => self.active_job_read_row_tx(tx, row).await, + None => Ok(None), + } } pub async fn enter_callback_wait( @@ -15530,14 +15681,16 @@ impl QueueStorage { } else { "" }; - let row: Option = sqlx::query_as(audited_sql(format!( + let row = sqlx::query(audited_sql(format!( r#" SELECT + ready.job_id AS body_job_id, lease.ready_slot, lease.ready_generation, lease.job_id, + lease.enqueue_shard, ready.kind, - ready.queue, + lease.queue, ready.args, lease.state, lease.priority, @@ -15563,7 +15716,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {} AS lease - JOIN {schema}.ready_entries AS ready + LEFT JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue @@ -15603,7 +15756,10 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - row.map(LeaseJobRow::into_job_row).transpose() + match row { + Some(row) => self.active_job_read_row_tx(tx, row).await, + None => Ok(None), + } } #[tracing::instrument(skip(self, pool, payload), name = "queue_storage.complete_external")] diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 7e48878f4..a7e5e794a 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -23196,6 +23196,13 @@ async fn assert_claimed_aging_rescue_body( ), "only conflicting historical body recovery needs a self-contained terminal row" ); + for _ in 0..8 { + let available = store + .claim_runtime_batch(pool, "qs_claimed_aging_body", 1, Duration::from_secs(300)) + .await + .expect("drain ready prefixes before explicit retry"); + assert!(available.iter().all(|row| row.job.id != claim.job_id)); + } let retried = store .retry_job(pool, claim.job_id) .await @@ -23205,6 +23212,19 @@ async fn assert_claimed_aging_rescue_body( assert_eq!(retried.max_attempts, 1); assert_eq!(retried.priority, 4); assert_eq!(retried.metadata["_awa_original_priority"], 4); + let mut claimed_retries = Vec::new(); + for _ in 0..8 { + claimed_retries.extend( + store + .claim_runtime_batch(pool, "qs_claimed_aging_body", 1, Duration::from_secs(300)) + .await + .expect("claim only the new attempt after explicit retry") + .into_iter() + .filter(|row| row.job.id == claim.job_id), + ); + } + assert_eq!(claimed_retries.len(), 1); + assert_eq!(claimed_retries[0].job.run_lease, claim.run_lease + 1); } async fn aging_skips_committed_claim_before_cursor(receipts: bool, deadline: Duration) { @@ -23221,12 +23241,12 @@ async fn aging_skips_committed_claim_before_cursor(receipts: bool, deadline: Dur assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; } -async fn historical_tombstoned_claim_hydrates( +async fn historical_tombstoned_claim_fixture( receipts: bool, deadline: Duration, colliding_position: bool, -) { - let (_guard, pool, store, claim, tx) = +) -> (TestDbGuard, sqlx::PgPool, QueueStorage, RawReceiptClaimRow) { + let (guard, pool, store, claim, tx) = claim_before_aging_cursor_advance_in_tx(receipts, deadline, !colliding_position).await; tx.commit() .await @@ -23256,6 +23276,16 @@ async fn historical_tombstoned_claim_hydrates( || jsonb_build_object('_awa_original_priority',priority)) FROM {schema}.ready_entries WHERE job_id=$1" ))).bind(claim.job_id).execute(tx.as_mut()).await.expect("historical replacement body"); + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.ready_segments + (ready_slot,ready_generation,queue,priority,enqueue_shard,first_lane_seq,next_lane_seq,first_run_at) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,lane_seq+1,run_at + FROM {schema}.ready_entries WHERE job_id=$1 AND priority=3" + ))) + .bind(claim.job_id) + .execute(tx.as_mut()) + .await + .expect("publish historical replacement segment as physical aging does"); tx.commit() .await .expect("commit historical source and replacement"); @@ -23280,6 +23310,16 @@ async fn historical_tombstoned_claim_hydrates( vec![(4, None), (3, Some(4))], "old physical aging adds provenance only to its replacement body" ); + (guard, pool, store, claim) +} + +async fn historical_tombstoned_claim_hydrates( + receipts: bool, + deadline: Duration, + colliding_position: bool, +) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, colliding_position).await; assert_claimed_aging_rescue_body(&pool, &store, &claim, colliding_position).await; } @@ -23561,3 +23601,580 @@ async fn test_physical_aging_waits_for_claim_head_then_uses_fresh_snapshot() { assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; aging_pool.close().await; } + +async fn aging_skips_claimed_candidates_before_limit(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; + let waiting = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9044 }, + InsertOpts { + queue: "qs_claimed_aging_body".into(), + priority: 4, + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + store.schema(), + waiting, + Utc::now() - chrono::Duration::minutes(4), + ) + .await; + let aged = store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("age past the oldest recorded claim"); + assert_eq!( + aged, + vec![waiting], + "claimed rows cannot consume the candidate limit" + ); + let waiting = store.load_job(&pool, waiting).await.unwrap().unwrap(); + assert_eq!(waiting.priority, 3); + assert_eq!(waiting.metadata["_awa_original_priority"], 4); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; +} + +#[tokio::test] +async fn test_physical_aging_skips_claimed_candidates_before_limit_compact() { + aging_skips_claimed_candidates_before_limit(true, Duration::ZERO).await; +} +#[tokio::test] +async fn test_physical_aging_skips_claimed_candidates_before_limit_deadline() { + aging_skips_claimed_candidates_before_limit(true, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_physical_aging_skips_claimed_candidates_before_limit_legacy() { + aging_skips_claimed_candidates_before_limit(false, Duration::from_secs(300)).await; +} + +fn assert_historical_active_body(job: &JobRow, claim: &RawReceiptClaimRow, state: JobState) { + assert_eq!(job.id, claim.job_id); + assert_eq!(job.state, state); + assert_eq!(job.run_lease, claim.run_lease); + assert_eq!(job.priority, claim.priority); + assert_eq!(job.attempt, claim.attempt); + assert_eq!(job.args, serde_json::json!({"id":9041})); + assert_eq!(job.max_attempts, 1); + assert_eq!(job.metadata["retained"], "claimed-before-aging"); + assert_eq!(job.metadata["_awa_original_priority"], 4); +} + +async fn historical_tombstoned_claim_point_read(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("read historically tombstoned active claim") + .expect("claimed job remains visible"); + assert_historical_active_body(&loaded, &claim, JobState::Running); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_point_read_is_running() { + historical_tombstoned_claim_point_read(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_point_read_is_running() { + historical_tombstoned_claim_point_read(true, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_lease_point_read_is_running() { + historical_tombstoned_claim_point_read(false, Duration::from_secs(300)).await; +} + +async fn historical_tombstoned_claim_callback_reads(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + let callback_id = store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("register callback for retained claim"); + let schema = store.schema(); + let progress = serde_json::json!({"phase":"retained-body"}); + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.attempt_state + (job_id,run_lease,progress,callback_filter,callback_on_complete,callback_on_fail,callback_transform,updated_at) + VALUES ($1,$2,$3,'payload.ok','payload.done','payload.fail','payload.value',clock_timestamp()) + ON CONFLICT (job_id,run_lease) DO UPDATE SET + progress=EXCLUDED.progress,callback_filter=EXCLUDED.callback_filter, + callback_on_complete=EXCLUDED.callback_on_complete,callback_on_fail=EXCLUDED.callback_on_fail, + callback_transform=EXCLUDED.callback_transform" + ))).bind(claim.job_id).bind(claim.run_lease).bind(&progress) + .execute(&pool).await.expect("set mutable callback state"); + let mut tx = pool.begin().await.unwrap(); + let running = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect("raw active lease lookup") + .expect("running lease still has retained body"); + assert_historical_active_body(&running, &claim, JobState::Running); + assert_eq!(running.callback_id, Some(callback_id)); + assert_eq!(running.progress, Some(progress.clone())); + assert!(store + .enter_callback_wait_in_tx(&mut tx, claim.job_id, claim.run_lease, callback_id) + .await + .expect("park retained claim")); + let parked = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect("read post-park row in caller transaction") + .expect("parked body"); + assert_historical_active_body(&parked, &claim, JobState::WaitingExternal); + assert_eq!(parked.callback_id, Some(callback_id)); + assert_eq!(parked.callback_timeout_at, running.callback_timeout_at); + assert_eq!(parked.heartbeat_at, None); + assert_eq!(parked.deadline_at, None); + for for_update in [false, true] { + let callback = store + .callback_job_in_tx(&mut tx, callback_id, Some(claim.run_lease), for_update) + .await + .expect("lookup retained callback body") + .expect("callback exists"); + assert_historical_active_body(&callback, &claim, JobState::WaitingExternal); + assert_eq!(callback.callback_id, Some(callback_id)); + assert_eq!(callback.callback_timeout_at, parked.callback_timeout_at); + assert_eq!(callback.progress, Some(progress.clone())); + assert_eq!(callback.callback_filter.as_deref(), Some("payload.ok")); + assert_eq!( + callback.callback_on_complete.as_deref(), + Some("payload.done") + ); + assert_eq!(callback.callback_on_fail.as_deref(), Some("payload.fail")); + assert_eq!( + callback.callback_transform.as_deref(), + Some("payload.value") + ); + } + tx.commit().await.expect("commit parked lease"); + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("point read parked lease") + .unwrap(); + assert_historical_active_body(&loaded, &claim, JobState::WaitingExternal); + let mut locked = pool.begin().await.unwrap(); + store + .callback_job_in_tx(&mut locked, callback_id, None, true) + .await + .expect("lock callback lease despite missing live body") + .unwrap(); + let mut contender = pool.begin().await.unwrap(); + let err = sqlx::query(audited_sql(format!( + "SELECT job_id FROM {schema}.leases WHERE job_id=$1 AND run_lease=$2 FOR UPDATE NOWAIT" + ))) + .bind(claim.job_id) + .bind(claim.run_lease) + .fetch_one(contender.as_mut()) + .await + .expect_err("callback lookup must retain its lease row lock"); + assert_eq!( + err.as_database_error() + .and_then(|err| err.code()) + .as_deref(), + Some("55P03") + ); + contender.rollback().await.unwrap(); + locked.rollback().await.unwrap(); + let resumed = store + .complete_external( + &pool, + callback_id, + Some(serde_json::json!({"ok":true})), + Some(claim.run_lease), + true, + ) + .await + .expect("resume callback and read retained body inside update transaction"); + assert_historical_active_body(&resumed, &claim, JobState::Running); + assert_eq!(resumed.callback_id, None); + assert_eq!(resumed.progress, Some(progress)); + assert_eq!( + resumed.metadata["_awa_callback_result"], + serde_json::json!({"ok":true}) + ); +} + +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_callback_reads_and_resume() { + historical_tombstoned_claim_callback_reads(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_callback_reads_and_resume() { + historical_tombstoned_claim_callback_reads(true, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_lease_callback_reads_and_resume() { + historical_tombstoned_claim_callback_reads(false, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_colliding_receipt_point_reads_use_emitted_lane() { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(true, Duration::ZERO, true).await; + tombstone_ready_job(&pool, &store, claim.job_id).await; + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("read emitted claimed lane") + .unwrap(); + assert_historical_active_body(&loaded, &claim, JobState::Running); + let callback_id = store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("materialize retained emitted claim"); + let mut tx = pool.begin().await.unwrap(); + let active = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect("load retained emitted lease") + .unwrap(); + assert_historical_active_body(&active, &claim, JobState::Running); + let callback = store + .callback_job_in_tx(&mut tx, callback_id, None, true) + .await + .expect("callback resolves emitted lane") + .unwrap(); + assert_historical_active_body(&callback, &claim, JobState::Running); + tx.rollback().await.unwrap(); +} + +#[tokio::test] +async fn test_historical_tombstoned_colliding_legacy_point_reads_fail_ambiguous() { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(false, Duration::from_secs(300), true).await; + let callback_id = store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("register legacy callback"); + tombstone_ready_job(&pool, &store, claim.job_id).await; + let expected = format!( + "queue storage ready body ambiguous for deleted lease job {} run_lease {}", + claim.job_id, claim.run_lease + ); + let err = store + .load_job(&pool, claim.job_id) + .await + .expect_err("point read cannot choose legacy lane"); + assert!( + matches!(err, AwaError::Validation(ref message) if message == &expected), + "{err}" + ); + let mut tx = pool.begin().await.unwrap(); + let err = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect_err("active lookup cannot choose legacy lane"); + assert!( + matches!(err, AwaError::Validation(ref message) if message == &expected), + "{err}" + ); + let err = store + .callback_job_in_tx(&mut tx, callback_id, None, true) + .await + .expect_err("callback lookup cannot choose legacy lane"); + assert!( + matches!(err, AwaError::Validation(ref message) if message == &expected), + "{err}" + ); + tx.rollback().await.unwrap(); + assert_eq!(lease_count(&pool, &store).await, 1); +} + +async fn historical_generic_cancel_live_replacement(deadline: Duration, colliding_position: bool) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(true, deadline, colliding_position).await; + let schema = store.schema(); + let live_bodies_sql = format!( + "SELECT count(*)::bigint FROM {schema}.ready_entries AS ready + WHERE ready.job_id=$1 AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot=ready.ready_slot AND tomb.ready_generation=ready.ready_generation + AND tomb.queue=ready.queue AND tomb.priority=ready.priority + AND tomb.enqueue_shard=ready.enqueue_shard AND tomb.lane_seq=ready.lane_seq + AND tomb.job_id=ready.job_id)" + ); + let live_before: i64 = sqlx::query_scalar(audited_sql(live_bodies_sql.clone())) + .bind(claim.job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(live_before, 1); + let cancelled = store + .cancel_job(&pool, claim.job_id) + .await + .expect("cancel running receipt ahead of its historical live replacement") + .expect("claimed attempt is cancelled"); + assert_eq!(cancelled.id, claim.job_id); + assert_eq!(cancelled.state, JobState::Cancelled); + assert_eq!(cancelled.run_lease, claim.run_lease); + assert_eq!(cancelled.args, serde_json::json!({"id":9041})); + assert_eq!(cancelled.max_attempts, 1); + assert_eq!(cancelled.priority, 4); + assert_eq!(cancelled.metadata["retained"], "claimed-before-aging"); + assert_eq!(cancelled.metadata["_awa_original_priority"], 4); + let live_after: i64 = sqlx::query_scalar(audited_sql(live_bodies_sql.clone())) + .bind(claim.job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(live_after, 0); + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("read cancelled receipt") + .expect("cancelled receipt remains visible"); + assert_eq!(loaded.state, JobState::Cancelled); + assert_eq!(loaded.run_lease, claim.run_lease); + assert_eq!(loaded.args, cancelled.args); + assert_eq!(loaded.priority, 4); + assert_eq!(loaded.metadata["_awa_original_priority"], 4); + assert!(store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("cancelled receipt is closed") + .is_empty()); + for _ in 0..8 { + let before_retry = store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 1, deadline) + .await + .expect("claim any unrelated ready job after cancel"); + assert!(before_retry + .iter() + .all(|entry| entry.job.id != claim.job_id)); + } + let retried = store + .retry_job(&pool, claim.job_id) + .await + .expect("retry generic receipt cancellation") + .expect("cancelled receipt retries"); + assert_eq!(retried.args, cancelled.args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(retried.priority, 4); + assert_eq!(retried.metadata["_awa_original_priority"], 4); + let mut claimed_retry = Vec::new(); + for _ in 0..8 { + claimed_retry.extend( + store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 1, deadline) + .await + .expect("claim the explicit retry"), + ); + } + assert_eq!(claimed_retry.len(), 1); + assert_eq!(claimed_retry[0].job.id, claim.job_id); + assert_eq!(claimed_retry[0].job.run_lease, claim.run_lease + 1); + assert_eq!(claimed_retry[0].job.priority, 4); + assert!(store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 10, deadline) + .await + .expect("old historical copy cannot emit another claim") + .is_empty()); +} + +#[tokio::test] +async fn test_generic_cancel_historical_colliding_live_compact_replacement_keeps_claimed_lane() { + historical_generic_cancel_live_replacement(Duration::ZERO, true).await; +} + +#[tokio::test] +async fn test_generic_cancel_historical_live_deadline_replacement_keeps_claimed_lane() { + historical_generic_cancel_live_replacement(Duration::from_secs(300), false).await; +} + +async fn historical_recovery_head_contention_rolls_back(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + let schema = store.schema(); + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(audited_sql(format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue='qs_claimed_aging_body' AND priority=3 AND enqueue_shard=0 FOR UPDATE" + ))) + .execute(blocker.as_mut()) + .await + .unwrap(); + let result = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.rescue_stale_heartbeats(&pool, Duration::ZERO), + ) + .await + .expect("contended recovery head fails without waiting"); + assert!( + matches!(result, Err(AwaError::Database(ref error)) if error.as_database_error().is_some_and(|error| error.code().as_deref() == Some("55P03"))), + "recovery must fail with lock_not_available: {result:?}" + ); + let done: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.done_entries WHERE job_id=$1" + ))) + .bind(claim.job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(done, 0); + assert_eq!(ready_tombstone_count(&pool, &store).await, 1); + let active = store.load_job(&pool, claim.job_id).await.unwrap().unwrap(); + assert_historical_active_body(&active, &claim, JobState::Running); + blocker.rollback().await.unwrap(); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; +} + +#[tokio::test] +async fn test_historical_compact_recovery_head_contention_rolls_back() { + historical_recovery_head_contention_rolls_back(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_legacy_recovery_head_contention_rolls_back() { + historical_recovery_head_contention_rolls_back(false, Duration::from_secs(300)).await; +} + +async fn historical_tombstoned_claim_completes(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + if receipts { + store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("materialize retained claim before completion"); + } + let completed = store + .complete_job_batch_by_id(&pool, &[(claim.job_id, claim.run_lease)]) + .await + .expect("complete retained body after old aging race"); + assert_eq!(completed, vec![(claim.job_id, claim.run_lease)]); + let loaded = store.load_job(&pool, claim.job_id).await.unwrap().unwrap(); + assert_eq!(loaded.state, JobState::Completed); + assert_eq!(loaded.run_lease, claim.run_lease); + assert_eq!(loaded.args, serde_json::json!({"id":9041})); + assert_eq!(loaded.max_attempts, 1); + assert_eq!(loaded.priority, 4); + assert_eq!(loaded.metadata["_awa_original_priority"], 4); + for _ in 0..8 { + let available = store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 10, deadline) + .await + .expect("old aging replacement cannot claim after hydrated completion"); + assert!(available.iter().all(|row| row.job.id != claim.job_id)); + } + assert!(store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .unwrap() + .iter() + .all(|row| row.id != claim.job_id)); +} + +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_completes_and_retires_copy() { + historical_tombstoned_claim_completes(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_completes_and_retires_copy() { + historical_tombstoned_claim_completes(true, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_claim_completes_and_retires_copy() { + historical_tombstoned_claim_completes(false, Duration::from_secs(300)).await; +} + +/// A fast completion closes an attempt under its completion lock without +/// locking claim rows. A cancel of the same running attempt must not commit +/// alongside it: each attempt ends either completed or cancelled. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_cancel_racing_a_fast_completion_leaves_one_outcome() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_cancel_fast_race"; + let schema = "awa_qs_cancel_fast_race"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await; + + let mut both = 0; + let mut neither = 0; + for id in 0..60 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim"); + assert_eq!(claimed.len(), 1); + let job_id = claimed[0].job.id; + let (completed, cancelled) = tokio::join!( + store.complete_runtime_batch(&pool, &claimed), + store.cancel_job(&pool, job_id) + ); + completed.expect("complete"); + cancelled.expect("cancel"); + let cancelled_rows: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*)::bigint FROM {schema}.done_entries WHERE job_id = $1 AND state = 'cancelled'" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + let completion_batches: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*)::bigint FROM {schema}.receipt_completion_batches WHERE job_ids @> ARRAY[$1]::bigint[]" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + if cancelled_rows > 0 && completion_batches > 0 { + both += 1; + } + if cancelled_rows == 0 && completion_batches == 0 { + neither += 1; + } + } + assert_eq!(both, 0, "attempts both completed and cancelled"); + assert_eq!(neither, 0, "attempts neither completed nor cancelled"); +} From 753cec3781b3702ccc51d014c1fc972f56576d8f Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 17:43:33 +1300 Subject: [PATCH 4/8] Use installed ready slots for retained body recovery --- CHANGELOG.md | 1 + awa-model/src/queue_storage.rs | 17 ++++--- awa/tests/queue_storage_runtime_test.rs | 63 +++++++++++++++++++++++-- 3 files changed, 69 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 955876be0..86e85e4ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- Schema-only admin clients use installed ready-slot numbers for point reads, cancellation and retained-body recovery, including rings larger than the default. - Keep committed claims out of physical priority aging before its batch limit, even before their cursor advances. Active point and callback reads recover retained tombstoned bodies left by older aging. Recovery, cancellation and diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 05bd87a36..aa3cf7c44 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -941,12 +941,16 @@ fn claim_closure_batch_child_name(schema: &str, slot: usize) -> String { format!("{schema}.lease_claim_closure_batches_{slot}") } -fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result { - let index = usize::try_from(slot).map_err(|_| { +fn ring_slot_number(slot: i32, ring: &str) -> Result { + usize::try_from(slot).map_err(|_| { AwaError::Validation(format!( "invalid {ring} ring slot {slot}: slot must be non-negative" )) - })?; + }) +} + +fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result { + let index = ring_slot_number(slot, ring)?; if index >= slot_count { return Err(AwaError::Validation(format!( "invalid {ring} ring slot {slot}: configured slot count is {slot_count}" @@ -12417,10 +12421,9 @@ impl QueueStorage { let run_leases: Vec = deleted.iter().map(|row| row.run_lease).collect(); // Resolve moved ready bodies without opening the whole tombstone ring. - // Empty input returned above; every remaining slot is validated here. let slots = ready_slots .iter() - .map(|&slot| ring_slot_index(slot, self.queue_slot_count(), "queue")) + .map(|&slot| ring_slot_number(slot, "queue")) .collect::, _>>()?; let tombstone_sources = slots .iter() @@ -12512,7 +12515,7 @@ impl QueueStorage { .map(|row| row.ready_slot) .collect::>() .into_iter() - .map(|slot| ring_slot_index(slot, self.queue_slot_count(), "queue")) + .map(|slot| ring_slot_number(slot, "queue")) .collect::, _>>()? .into_iter() .map(|slot| { @@ -12867,7 +12870,7 @@ impl QueueStorage { .map_err(map_sqlx_error)?; let attempt_sources = copies .iter() - .map(|row| ring_slot_index(row.ready_slot, self.queue_slot_count(), "queue")) + .map(|row| ring_slot_number(row.ready_slot, "queue")) .collect::, _>>()? .into_iter() .map(|slot| { diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index a7e5e794a..d9864f216 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -23356,7 +23356,13 @@ async fn test_historical_colliding_live_replacement_deadline_rescues_claimed_bod historical_tombstoned_claim_hydrates(true, Duration::from_secs(300), true).await; } -async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool, reopen_as_legacy: bool) { +async fn moved_tombstoned_claim_lane( + from: i16, + to: i16, + receipts: bool, + reopen_as_legacy: bool, + discovered_wide: bool, +) { let (_guard, pool) = setup_pool(6).await; let schema = "awa_qs_moved_tomb_claim"; let queue = "qs_moved_tomb_claim"; @@ -23365,11 +23371,34 @@ async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool, reopen_ QueueStorageConfig { schema: schema.into(), lease_claim_receipts: receipts, + queue_slot_count: if discovered_wide { + 20 + } else { + QueueStorageConfig::default().queue_slot_count + }, queue_stripe_count: 1, ..Default::default() }, ) .await; + if discovered_wide { + enqueue_job( + &pool, + &store, + &CompleteJob { id: 9044 }, + InsertOpts { + queue: "qs_wide_anchor".into(), + ..Default::default() + }, + ) + .await; + for _ in 0..16 { + assert!(matches!( + store.rotate(&pool).await.unwrap(), + RotateOutcome::Rotated { .. } + )); + } + } let id = enqueue_job(&pool, &store, &CompleteJob { id: 9043 }, InsertOpts { queue: queue.into(), priority: from, max_attempts: 1, metadata: serde_json::json!({"retained":"moved-tomb-claim", "_awa_original_priority":from}), @@ -23409,6 +23438,9 @@ async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool, reopen_ } assert_eq!(claims.len(), 1); assert_eq!(claims[0].job_id, id); + if discovered_wide { + assert_eq!(claims[0].ready_slot, 16); + } assert_eq!( claims[0].priority, 1, "effective priority cannot identify either physical lane" @@ -23426,6 +23458,12 @@ async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool, reopen_ ..Default::default() }) .expect("reopen existing custom store in legacy mode") + } else if discovered_wide { + QueueStorage::new(QueueStorageConfig { + schema: schema.into(), + ..Default::default() + }) + .expect("open installed schema without its writer configuration") } else { store }; @@ -23433,6 +23471,16 @@ async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool, reopen_ // tombstoning the claimed source as well; both exact-ref bodies now remain. tombstone_ready_job(&pool, &store, id).await; assert_eq!(ready_tombstone_count(&pool, &store).await, 2); + if discovered_wide { + let loaded = store + .load_job(&pool, id) + .await + .expect("read retained body in installed slot") + .unwrap(); + assert_eq!(loaded.state, JobState::Running); + assert_eq!(loaded.args, serde_json::json!({"id":9043})); + assert_eq!(loaded.max_attempts, 1); + } if !receipts { let err = store .rescue_stale_heartbeats(&pool, Duration::ZERO) @@ -23476,20 +23524,25 @@ async fn moved_tombstoned_claim_lane(from: i16, to: i16, receipts: bool, reopen_ #[tokio::test] async fn test_historical_tombstoned_claim_uses_emitted_lane_low_to_high() { - moved_tombstoned_claim_lane(2, 4, true, false).await; + moved_tombstoned_claim_lane(2, 4, true, false, false).await; } #[tokio::test] async fn test_historical_tombstoned_claim_uses_emitted_lane_high_to_low() { - moved_tombstoned_claim_lane(4, 2, true, false).await; + moved_tombstoned_claim_lane(4, 2, true, false, false).await; } #[tokio::test] async fn test_historical_tombstoned_legacy_ambiguous_lanes_roll_back() { - moved_tombstoned_claim_lane(2, 4, false, false).await; + moved_tombstoned_claim_lane(2, 4, false, false, false).await; } #[tokio::test] async fn test_historical_tombstoned_receipt_materialized_then_reopened_legacy_uses_evidence() { - moved_tombstoned_claim_lane(2, 4, true, true).await; + moved_tombstoned_claim_lane(2, 4, true, true, false).await; +} + +#[tokio::test] +async fn test_schema_only_store_recovers_retained_body_in_installed_slot() { + moved_tombstoned_claim_lane(2, 4, true, false, true).await; } #[tokio::test] From 18b66cd7884c8f5f3a3e4c3ca6cfc63e15b777b4 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 18:43:53 +1300 Subject: [PATCH 5/8] Isolate retained recovery contention in rescue and completion batches --- CHANGELOG.md | 9 +- awa-model/src/queue_storage.rs | 584 +++++++++++++++++++----- awa-worker/src/completion.rs | 271 ++++++++++- awa/tests/queue_storage_runtime_test.rs | 345 +++++++++++++- 4 files changed, 1080 insertions(+), 129 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 86e85e4ba..7c2e29fad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,9 +23,12 @@ Notable changes between releases. Detailed migration notes for storage transitio worker completion using its carried job body can leave an old replacement claimable. Updated aging callers prevent new replacements, but cannot fence an older process explicitly calling physical aging during a rolling upgrade. - Repair takes claim-head row locks without waiting and rolls back on - contention; the operation must retry. It adds weak relation locks and - tombstone writes only on this recovery path. Fresh rechecks rely on READ + Repair takes claim-head row locks without waiting. Rescue batches defer + busy historical attempts while progressing other attempts; receipt rescue + cursors keep deferred attempts eligible. Worker completion batches isolate + lock contention to the affected requests. Direct completion and + single-attempt mutations still fail on contention and must retry. It adds + weak relation locks and tombstone writes only on this recovery path. Fresh rechecks rely on READ COMMITTED isolation. - Preserve recovered bodies and physical lane priority when manually moving diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index aa3cf7c44..e1d203fb1 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -2618,6 +2618,21 @@ struct ReadyRowIdentity { run_lease: i64, } +struct BusyRetained(Vec<(i64, i64)>); + +impl BusyRetained { + fn exclude(self, excluded: &mut BTreeSet<(i64, i64)>) -> Result<(), AwaError> { + let previous = excluded.len(); + excluded.extend(self.0); + if excluded.len() == previous { + return Err(AwaError::Validation( + "queue storage retained attempt retry made no progress".to_string(), + )); + } + Ok(()) + } +} + #[derive(Debug, Clone, sqlx::FromRow)] struct ReadyTransitionRow { ready_slot: i32, @@ -12570,8 +12585,45 @@ impl QueueStorage { deleted: Vec, write_receipt_closures: bool, ) -> Result, AwaError> { + let (hydrated, sources) = self + .hydrate_deleted_leases_with_sources_tx(tx, deleted, write_receipt_closures) + .await?; + if !sources.is_empty() { + self.retire_retained_attempt_copies_tx(tx, &sources, false) + .await?; + } + Ok(hydrated) + } + + async fn try_hydrate_deleted_leases_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + deleted: Vec, + write_receipt_closures: bool, + ) -> Result, BusyRetained>, AwaError> { + let (hydrated, sources) = self + .hydrate_deleted_leases_with_sources_tx(tx, deleted, write_receipt_closures) + .await?; + if sources.is_empty() { + return Ok(Ok(hydrated)); + } + if let Some(busy) = self + .retire_retained_attempt_copies_tx(tx, &sources, true) + .await? + { + return Ok(Err(busy)); + } + Ok(Ok(hydrated)) + } + + async fn hydrate_deleted_leases_with_sources_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + deleted: Vec, + write_receipt_closures: bool, + ) -> Result<(Vec, Vec), AwaError> { if deleted.is_empty() { - return Ok(Vec::new()); + return Ok((Vec::new(), Vec::new())); } let schema = self.schema(); @@ -12819,18 +12871,15 @@ impl QueueStorage { }); } - if !retained_sources.is_empty() { - self.retire_retained_attempt_copies_tx(tx, &retained_sources) - .await?; - } - Ok(hydrated) + Ok((hydrated, retained_sources)) } async fn retire_retained_attempt_copies_tx( &self, tx: &mut sqlx::Transaction<'_, Postgres>, sources: &[ReadyRowIdentity], - ) -> Result<(), AwaError> { + skip_busy: bool, + ) -> Result, AwaError> { let schema = self.schema(); let mut expected = BTreeMap::new(); for source in sources { @@ -12847,10 +12896,55 @@ impl QueueStorage { } let job_ids: Vec = sources.iter().map(|row| row.job_id).collect(); let run_leases: Vec = sources.iter().map(|row| row.run_lease).collect(); - // Hydration already holds attempt locks. A contended head must abort - // instead of waiting in the reverse of the claim/mutation lock order. - self.lock_claim_heads_of_jobs_with_wait_tx(tx, &job_ids, true) - .await?; + if skip_busy { + let busy = sqlx::query_as::<_, (i64, i64)>(audited_sql(format!( + r#" + WITH refs(job_id, run_lease) AS ( + SELECT * FROM unnest($1::bigint[], $2::bigint[]) + ), locked_heads AS MATERIALIZED ( + SELECT heads.queue, heads.priority, heads.enqueue_shard + FROM {schema}.queue_claim_heads AS heads + WHERE EXISTS ( + SELECT 1 + FROM refs + JOIN {schema}.ready_entries AS ready + ON ready.job_id = refs.job_id AND ready.run_lease = refs.run_lease + WHERE ready.queue = heads.queue AND ready.priority = heads.priority + AND ready.enqueue_shard = heads.enqueue_shard + ) + ORDER BY heads.queue COLLATE "C", heads.priority, heads.enqueue_shard + FOR UPDATE OF heads SKIP LOCKED + ) + SELECT refs.job_id, refs.run_lease + 1 + FROM refs + WHERE EXISTS ( + SELECT 1 + FROM {schema}.ready_entries AS ready + JOIN {schema}.queue_claim_heads AS heads + ON heads.queue = ready.queue AND heads.priority = ready.priority + AND heads.enqueue_shard = ready.enqueue_shard + WHERE ready.job_id = refs.job_id AND ready.run_lease = refs.run_lease + AND NOT EXISTS ( + SELECT 1 FROM locked_heads + WHERE locked_heads.queue = heads.queue + AND locked_heads.priority = heads.priority + AND locked_heads.enqueue_shard = heads.enqueue_shard + ) + ) + "# + ))) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + if !busy.is_empty() { + return Ok(Some(BusyRetained(busy))); + } + } else { + self.lock_claim_heads_of_jobs_with_wait_tx(tx, &job_ids, true) + .await?; + } let copies: Vec = sqlx::query_as(audited_sql(format!( r#" WITH refs(job_id, run_lease) AS ( @@ -12956,7 +13050,7 @@ impl QueueStorage { .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; - Ok(()) + Ok(None) } /// Close one compact receipt claim by its batch identity. The claim @@ -13411,6 +13505,8 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, cutoff: DateTime, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { let mut rescued = Vec::new(); let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; @@ -13418,7 +13514,9 @@ impl QueueStorage { if let Some(slot) = preferred_slot { let mut slot_rescued = self - .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) + .rescue_stale_receipt_claims_for_slot_tx( + tx, slot, cutoff, remaining, excluded, selected, + ) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -13434,7 +13532,9 @@ impl QueueStorage { } let mut slot_rescued = self - .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) + .rescue_stale_receipt_claims_for_slot_tx( + tx, slot, cutoff, remaining, excluded, selected, + ) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -13453,7 +13553,14 @@ impl QueueStorage { slot: i32, cutoff: DateTime, rescue_limit: i64, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); let schema = self.schema(); let claim_child = claim_child_name(schema, slot as usize); let claim_batch_child = claim_batch_child_name(schema, slot as usize); @@ -13638,6 +13745,14 @@ impl QueueStorage { WHERE NOT is_closed AND NOT is_lease_managed AND is_stale + AND NOT EXISTS ( + SELECT 1 FROM unnest($5::bigint[], $6::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = candidates.job_id AND excluded.run_lease = candidates.run_lease + ) + AND ($7::boolean OR EXISTS ( + SELECT 1 FROM unnest($8::bigint[], $9::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = candidates.job_id AND selected.run_lease = candidates.run_lease + )) ORDER BY rn LIMIT $4 )"# @@ -13665,6 +13780,11 @@ impl QueueStorage { .bind(cutoff) .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? @@ -13710,7 +13830,7 @@ impl QueueStorage { ) AND (claims.job_id, claims.run_lease) IN ( SELECT locked.job_id, locked.run_lease - FROM unnest($5::bigint[], $6::bigint[]) AS locked(job_id, run_lease) + FROM unnest($10::bigint[], $11::bigint[]) AS locked(job_id, run_lease) ) -- A claim that already materialized into `leases` is -- on the lease-side heartbeat-rescue path (see @@ -13759,7 +13879,7 @@ impl QueueStorage { ) AND (stale_candidates.job_id, stale_candidates.run_lease) IN ( SELECT locked.job_id, locked.run_lease - FROM unnest($5::bigint[], $6::bigint[]) AS locked(job_id, run_lease) + FROM unnest($10::bigint[], $11::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease @@ -13896,6 +14016,11 @@ impl QueueStorage { .bind(cutoff) .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) @@ -13933,6 +14058,8 @@ impl QueueStorage { async fn rescue_expired_receipt_deadlines_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { let mut rescued = Vec::new(); let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; @@ -13951,7 +14078,7 @@ impl QueueStorage { for slot in slots { let mut slot_rescued = self - .rescue_expired_batch_deadlines_for_slot_tx(tx, slot, remaining) + .rescue_expired_batch_deadlines_for_slot_tx(tx, slot, remaining, excluded, selected) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -13960,7 +14087,9 @@ impl QueueStorage { } let mut slot_rescued = self - .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining) + .rescue_expired_receipt_deadlines_for_slot_tx( + tx, slot, remaining, excluded, selected, + ) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -14007,11 +14136,18 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, rescue_limit: i64, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { if rescue_limit <= 0 { return Ok(Vec::new()); } + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); let schema = self.schema(); let claim_child = claim_child_name(schema, slot as usize); let closure_child = closure_child_name(schema, slot as usize); @@ -14119,6 +14255,14 @@ impl QueueStorage { WHERE NOT is_closed AND NOT is_lease_managed AND is_expired + AND NOT EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = candidates.job_id AND excluded.run_lease = candidates.run_lease + ) + AND ($6::boolean OR EXISTS ( + SELECT 1 FROM unnest($7::bigint[], $8::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = candidates.job_id AND selected.run_lease = candidates.run_lease + )) ORDER BY rn LIMIT $3 )"# @@ -14141,6 +14285,11 @@ impl QueueStorage { .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? @@ -14160,7 +14309,7 @@ impl QueueStorage { AND claims.deadline_at < clock_timestamp() AND (claims.job_id, claims.run_lease) IN ( SELECT locked.job_id, locked.run_lease - FROM unnest($4::bigint[], $5::bigint[]) AS locked(job_id, run_lease) + FROM unnest($9::bigint[], $10::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease @@ -14253,6 +14402,11 @@ impl QueueStorage { .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) @@ -14286,11 +14440,18 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, rescue_limit: i64, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { if rescue_limit <= 0 { return Ok(Vec::new()); } + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); let schema = self.schema(); let claim_batch_child = claim_batch_child_name(schema, slot as usize); let closure_child = closure_child_name(schema, slot as usize); @@ -14430,6 +14591,14 @@ impl QueueStorage { FROM members WHERE NOT is_closed AND NOT is_lease_managed + AND NOT EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = members.job_id AND excluded.run_lease = members.run_lease + ) + AND ($6::boolean OR EXISTS ( + SELECT 1 FROM unnest($7::bigint[], $8::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = members.job_id AND selected.run_lease = members.run_lease + )) ORDER BY rn, receipt_id LIMIT $3 )"# @@ -14452,6 +14621,11 @@ impl QueueStorage { .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? @@ -14495,7 +14669,7 @@ impl QueueStorage { ) AND (expired_members.job_id, expired_members.run_lease) IN ( SELECT locked.job_id, locked.run_lease - FROM unnest($4::bigint[], $5::bigint[]) AS locked(job_id, run_lease) + FROM unnest($9::bigint[], $10::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease @@ -14606,6 +14780,11 @@ impl QueueStorage { .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) @@ -17769,7 +17948,17 @@ impl QueueStorage { // `(state, callback_timeout_at)` indexes, followed by a heap // recheck of the COALESCE. Bounded by running-lease count and // called at 30s cadence — cheap. - let deleted: Vec = sqlx::query_as(audited_sql(format!( + let mut excluded = BTreeSet::new(); + let mut selected: Option> = None; + loop { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .as_ref() + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); + let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; + let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases AS target WHERE (target.job_id, target.run_lease) IN ( @@ -17780,6 +17969,14 @@ impl QueueStorage { AND attempt.run_lease = lease.run_lease WHERE lease.state = 'running' AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1 + AND NOT EXISTS ( + SELECT 1 FROM unnest($2::bigint[], $3::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = lease.job_id AND excluded.run_lease = lease.run_lease + ) + AND ($4::boolean OR EXISTS ( + SELECT 1 FROM unnest($5::bigint[], $6::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = lease.job_id AND selected.run_lease = lease.run_lease + )) ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC LIMIT 500 FOR UPDATE OF lease SKIP LOCKED @@ -17804,55 +18001,98 @@ impl QueueStorage { "# ))) .bind(cutoff) - .fetch_all(tx.as_mut()) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .fetch_all(attempt.as_mut()) .await .map_err(map_sqlx_error)?; - let rescued_receipts = if self.lease_claim_receipts() { - self.rescue_stale_receipt_claims_tx(&mut tx, cutoff).await? - } else { - Vec::new() - }; + let receipt_excluded: BTreeSet<_> = excluded + .iter() + .copied() + .chain(deleted.iter().map(|row| (row.job_id, row.run_lease))) + .collect(); + let rescued_receipts = if self.lease_claim_receipts() { + self.rescue_stale_receipt_claims_tx( + &mut attempt, + cutoff, + &receipt_excluded, + selected.as_ref(), + ) + .await? + } else { + Vec::new() + }; - if deleted.is_empty() && rescued_receipts.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); - } + selected.get_or_insert_with(|| { + deleted + .iter() + .chain(rescued_receipts.iter()) + .map(|row| (row.job_id, row.run_lease)) + .collect() + }); - let moved_leases = self - .hydrate_deleted_leases_tx(&mut tx, deleted, true) - .await?; - let moved_receipts = self - .hydrate_deleted_leases_tx(&mut tx, rescued_receipts, true) - .await?; + if deleted.is_empty() && rescued_receipts.is_empty() { + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } - let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len()); - for row in moved_leases { - let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut tx, - row, - "heartbeat stale: worker presumed dead", - dlq_reason, - ) - .await?, - ); - } - for row in moved_receipts { - let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut tx, - row, - "receipt claim stale: worker presumed dead", - dlq_reason, - ) - .await?, - ); + let moved_leases = match self + .try_hydrate_deleted_leases_tx(&mut attempt, deleted, true) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + let moved_receipts = match self + .try_hydrate_deleted_leases_tx(&mut attempt, rescued_receipts, true) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + + let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len()); + for row in moved_leases { + let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "heartbeat stale: worker presumed dead", + dlq_reason, + ) + .await?, + ); + } + for row in moved_receipts { + let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "receipt claim stale: worker presumed dead", + dlq_reason, + ) + .await?, + ); + } + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(rescued); } - tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) } /// Rescue attempts past their hard deadline. Exhausted attempts fail @@ -17877,15 +18117,33 @@ impl QueueStorage { ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let deleted: Vec = sqlx::query_as(audited_sql(format!( + let mut excluded = BTreeSet::new(); + let mut selected: Option> = None; + loop { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .as_ref() + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); + let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; + let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases - WHERE job_id IN ( - SELECT job_id + WHERE (job_id, run_lease) IN ( + SELECT job_id, run_lease FROM {schema}.leases WHERE state = 'running' AND deadline_at IS NOT NULL AND deadline_at < clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease + ) + AND ($3::boolean OR EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease + )) ORDER BY deadline_at ASC LIMIT 500 FOR UPDATE SKIP LOCKED @@ -17909,45 +18167,93 @@ impl QueueStorage { callback_timeout_at "# ))) - .fetch_all(tx.as_mut()) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .fetch_all(attempt.as_mut()) .await .map_err(map_sqlx_error)?; - // Receipts-mode short-path claims hold their deadline on - // `lease_claim_batches.deadline_at` (compact batches, #246) or on - // legacy row-local `lease_claims.deadline_at` rather than on a - // `leases` row, so the receipt-plane needs its own scan; merge - // all populations into one `moved` set so the maintenance caller - // observes a single rescue batch per tick. - let receipt_deleted = if self.lease_claim_receipts() { - self.rescue_expired_receipt_deadlines_tx(&mut tx).await? - } else { - Vec::new() - }; + // Receipts-mode short-path claims hold their deadline on + // `lease_claim_batches.deadline_at` (compact batches, #246) or on + // legacy row-local `lease_claims.deadline_at` rather than on a + // `leases` row, so the receipt-plane needs its own scan; merge + // all populations into one `moved` set so the maintenance caller + // observes a single rescue batch per tick. + let receipt_excluded: BTreeSet<_> = excluded + .iter() + .copied() + .chain(deleted.iter().map(|row| (row.job_id, row.run_lease))) + .collect(); + let receipt_deleted = if self.lease_claim_receipts() { + self.rescue_expired_receipt_deadlines_tx( + &mut attempt, + &receipt_excluded, + selected.as_ref(), + ) + .await? + } else { + Vec::new() + }; - if deleted.is_empty() && receipt_deleted.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); - } + selected.get_or_insert_with(|| { + deleted + .iter() + .chain(receipt_deleted.iter()) + .map(|row| (row.job_id, row.run_lease)) + .collect() + }); - let mut moved = self - .hydrate_deleted_leases_tx(&mut tx, deleted, true) - .await?; - moved.extend( - self.hydrate_deleted_leases_tx(&mut tx, receipt_deleted, true) - .await?, - ); + if deleted.is_empty() && receipt_deleted.is_empty() { + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } - let mut rescued = Vec::with_capacity(moved.len()); - for row in moved { - let dlq_reason = dlq_enabled(&row.queue).then_some("deadline_exceeded"); - rescued.push( - self.settle_rescued_attempt_tx(&mut tx, row, "hard deadline exceeded", dlq_reason) - .await?, + let mut moved = match self + .try_hydrate_deleted_leases_tx(&mut attempt, deleted, true) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + moved.extend( + match self + .try_hydrate_deleted_leases_tx(&mut attempt, receipt_deleted, true) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }, ); + + let mut rescued = Vec::with_capacity(moved.len()); + for row in moved { + let dlq_reason = dlq_enabled(&row.queue).then_some("deadline_exceeded"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "hard deadline exceeded", + dlq_reason, + ) + .await?, + ); + } + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(rescued); } - tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) } /// Rescue attempts whose callback timed out. Exhausted attempts fail @@ -17972,15 +18278,33 @@ impl QueueStorage { ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let deleted: Vec = sqlx::query_as(audited_sql(format!( + let mut excluded = BTreeSet::new(); + let mut selected: Option> = None; + loop { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .as_ref() + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); + let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; + let deleted: Vec = sqlx::query_as(audited_sql(format!( r#" DELETE FROM {schema}.leases - WHERE job_id IN ( - SELECT job_id + WHERE (job_id, run_lease) IN ( + SELECT job_id, run_lease FROM {schema}.leases WHERE state = 'waiting_external' AND callback_timeout_at IS NOT NULL AND callback_timeout_at < clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease + ) + AND ($3::boolean OR EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease + )) ORDER BY callback_timeout_at ASC LIMIT 500 FOR UPDATE SKIP LOCKED @@ -18004,29 +18328,57 @@ impl QueueStorage { callback_timeout_at "# ))) - .fetch_all(tx.as_mut()) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .fetch_all(attempt.as_mut()) .await .map_err(map_sqlx_error)?; - if deleted.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); - } + selected.get_or_insert_with(|| { + deleted + .iter() + .map(|row| (row.job_id, row.run_lease)) + .collect() + }); - let moved = self - .hydrate_deleted_leases_tx(&mut tx, deleted, true) - .await?; + if deleted.is_empty() { + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } - let mut rescued = Vec::with_capacity(moved.len()); - for row in moved { - let dlq_reason = dlq_enabled(&row.queue).then_some("callback_timeout"); - rescued.push( - self.settle_rescued_attempt_tx(&mut tx, row, "callback timed out", dlq_reason) + let moved = match self + .try_hydrate_deleted_leases_tx(&mut attempt, deleted, true) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + + let mut rescued = Vec::with_capacity(moved.len()); + for row in moved { + let dlq_reason = dlq_enabled(&row.queue).then_some("callback_timeout"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "callback timed out", + dlq_reason, + ) .await?, - ); + ); + } + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(rescued); } - tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) } pub async fn promote_due( diff --git a/awa-worker/src/completion.rs b/awa-worker/src/completion.rs index ddb3f7916..5fc0fa531 100644 --- a/awa-worker/src/completion.rs +++ b/awa-worker/src/completion.rs @@ -412,17 +412,65 @@ impl CompletionWorker { ); } Err(err) => { - self.metrics.record_completion_flush( - self.shard_id, - job_ids.len() as u64, - flush_start.elapsed(), - ); warn!( error = %err, shard = self.shard_id, batch_size = job_ids.len(), "Failed to flush completed job batch" ); + if matches!( + &err, + AwaError::Database(sqlx::Error::Database(error)) + if error.code().as_deref() == Some("55P03") + ) { + if let RuntimeStorage::QueueStorage(runtime) = &self.storage { + for request in batch { + let pair = (request.job_id, request.run_lease); + let response = if committed_pairs.contains(&pair) { + Ok(true) + } else { + let updated = if let Some(runtime_job) = &request.runtime_job { + runtime + .store + .complete_runtime_batch( + &self.pool, + std::slice::from_ref(runtime_job), + ) + .await + } else if let Some(claim) = &request.claim { + runtime + .store + .complete_claimed_batch( + &self.pool, + std::slice::from_ref(claim), + ) + .await + } else { + runtime + .store + .complete_job_batch_by_id(&self.pool, &[pair]) + .await + }; + updated.map(|updated| updated.contains(&pair)) + }; + if matches!(&response, Ok(true)) { + committed_pairs.insert(pair); + } + let _ = request.response.send(response); + } + self.metrics.record_completion_flush( + self.shard_id, + job_ids.len() as u64, + flush_start.elapsed(), + ); + return; + } + } + self.metrics.record_completion_flush( + self.shard_id, + job_ids.len() as u64, + flush_start.elapsed(), + ); let message = format!("completion batch flush failed: {err}"); for request in batch { let response = if committed_job_ids.contains(&request.job_id) { @@ -1073,4 +1121,217 @@ mod tests { jobs.len() as f64 / elapsed.as_secs_f64() ); } + async fn queue_storage_busy_completion_responses(with_claim: bool) { + let pool = setup(6).await; + let schema = format!("awa_complete_busy_{}", uuid::Uuid::new_v4().simple()); + let runtime = crate::storage::QueueStorageRuntime::new( + awa_model::QueueStorageConfig { + schema: schema.clone(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: false, + }, + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(); + let store = runtime.store.clone(); + store.install(&pool).await.unwrap(); + let mut jobs = Vec::new(); + for queue in ["busy", "ordinary", "stale", "carried"] { + store.enqueue_batch(&pool, queue, 4, 1).await.unwrap(); + let mut claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(300)) + .await + .unwrap(); + assert_eq!(claimed.len(), 1); + jobs.push(claimed.remove(0)); + } + let historical = jobs.remove(0); + let ordinary = jobs.remove(0); + let stale = jobs.remove(0); + let mut carried = jobs.remove(0); + carried.job.output = Some(serde_json::json!({"result":"preserved"})); + assert_eq!( + store + .complete_runtime_batch(&pool, std::slice::from_ref(&stale)) + .await + .unwrap(), + vec![(stale.job.id, stale.job.run_lease)] + ); + store.enqueue_batch(&pool, "busy", 3, 1).await.unwrap(); + sqlx::query(awa_model::audited_sql(format!( + "INSERT INTO {schema}.ready_tombstones + (ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id + FROM {schema}.ready_entries WHERE job_id=$1;" + ))) + .bind(historical.job.id) + .execute(&pool) + .await + .unwrap(); + sqlx::query(awa_model::audited_sql(format!( + "INSERT INTO {schema}.ready_entries + (ready_slot,ready_generation,job_id,kind,queue,args,priority,attempt,run_lease,max_attempts, + lane_seq,enqueue_shard,run_at,attempted_at,created_at,unique_key,unique_states,payload) + SELECT ready_slot,ready_generation,job_id,kind,queue,args,3::smallint,attempt,run_lease,max_attempts, + {schema}.reserve_enqueue_seq(queue,3::smallint,enqueue_shard,1::bigint),enqueue_shard, + run_at,attempted_at,created_at,unique_key,unique_states,payload + FROM {schema}.ready_entries WHERE job_id=$1" + ))) + .bind(historical.job.id) + .execute(&pool) + .await + .unwrap(); + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(awa_model::audited_sql(format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue='busy' AND priority=3 AND enqueue_shard=0 FOR UPDATE" + ))) + .fetch_one(blocker.as_mut()) + .await + .unwrap(); + let pairs = [ + (historical.job.id, historical.job.run_lease), + (ordinary.job.id, ordinary.job.run_lease), + ]; + let strict = if with_claim { + store + .complete_claimed_batch(&pool, &[historical.claim.clone(), ordinary.claim.clone()]) + .await + } else { + store.complete_job_batch_by_id(&pool, &pairs).await + }; + assert!(matches!(strict, Err(AwaError::Database(ref error)) + if error.as_database_error().is_some_and(|error| error.code().as_deref()==Some("55P03")))); + let mut caller_tx = pool.begin().await.unwrap(); + let strict = store + .complete_job_batch_by_id_in_tx(&mut caller_tx, &pairs) + .await; + assert!(matches!(strict, Err(AwaError::Database(ref error)) + if error.as_database_error().is_some_and(|error| error.code().as_deref()==Some("55P03")))); + caller_tx.rollback().await.unwrap(); + assert_eq!( + store + .load_job(&pool, ordinary.job.id) + .await + .unwrap() + .unwrap() + .state, + awa_model::JobState::Running + ); + + let (mut batcher, _) = CompletionBatcher::with_batching( + pool.clone(), + CancellationToken::new(), + crate::metrics::AwaMetrics::from_global(), + RuntimeStorage::QueueStorage(runtime), + 1, + 4, + Duration::from_millis(1), + ); + let worker = batcher.workers.pop().unwrap(); + let mut batch = Vec::new(); + let mut responses = Vec::new(); + for job in [&historical, &ordinary, &stale, &carried] { + let (response, rx) = oneshot::channel(); + batch.push(CompletionRequest { + job_id: job.job.id, + run_lease: job.job.run_lease, + claim: with_claim.then(|| job.claim.clone()), + runtime_job: (job.job.id == carried.job.id).then(|| job.clone()), + response, + }); + responses.push(rx); + } + tokio::time::timeout(Duration::from_secs(5), worker.flush(&mut batch)) + .await + .expect("contended batch isolates singleton errors without waiting"); + let busy = responses.remove(0).await.unwrap(); + assert!( + matches!(busy, Err(AwaError::Database(ref error)) + if error.as_database_error().is_some_and(|error| error.code().as_deref()==Some("55P03"))), + "busy attempt must report its real error, never stale: {busy:?}" + ); + assert!(responses.remove(0).await.unwrap().unwrap()); + assert!(!responses.remove(0).await.unwrap().unwrap()); + assert!(responses.remove(0).await.unwrap().unwrap()); + assert_eq!( + store + .load_job(&pool, historical.job.id) + .await + .unwrap() + .unwrap() + .state, + awa_model::JobState::Running + ); + assert_eq!( + store + .load_job(&pool, ordinary.job.id) + .await + .unwrap() + .unwrap() + .state, + awa_model::JobState::Completed + ); + assert_eq!( + store + .load_job(&pool, carried.job.id) + .await + .unwrap() + .unwrap() + .output, + carried.job.output + ); + blocker.rollback().await.unwrap(); + let (response, rx) = oneshot::channel(); + worker + .flush(&mut vec![CompletionRequest { + job_id: historical.job.id, + run_lease: historical.job.run_lease, + claim: with_claim.then(|| historical.claim.clone()), + runtime_job: None, + response, + }]) + .await; + assert!(rx.await.unwrap().unwrap()); + let terminals: Vec<(i64, i64, i64)> = sqlx::query_as(awa_model::audited_sql(format!( + "SELECT job_id,run_lease,count(*) FROM {schema}.done_entries + WHERE job_id=ANY($1) GROUP BY job_id,run_lease ORDER BY job_id" + ))) + .bind(vec![ + historical.job.id, + ordinary.job.id, + stale.job.id, + carried.job.id, + ]) + .fetch_all(&pool) + .await + .unwrap(); + let mut expected: Vec<_> = [&historical, &ordinary, &stale, &carried] + .into_iter() + .map(|job| (job.job.id, job.job.run_lease, 1_i64)) + .collect(); + expected.sort_unstable(); + assert_eq!(terminals, expected); + sqlx::query(awa_model::audited_sql(format!( + "DROP SCHEMA {schema} CASCADE" + ))) + .execute(&pool) + .await + .unwrap(); + pool.close().await; + } + + #[tokio::test] + async fn queue_storage_id_completion_contention_isolates_responses() { + queue_storage_busy_completion_responses(false).await; + } + + #[tokio::test] + async fn queue_storage_claim_completion_contention_isolates_responses() { + queue_storage_busy_completion_responses(true).await; + } } diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index d9864f216..b66ad8335 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -24075,11 +24075,9 @@ async fn historical_recovery_head_contention_rolls_back(receipts: bool, deadline store.rescue_stale_heartbeats(&pool, Duration::ZERO), ) .await - .expect("contended recovery head fails without waiting"); - assert!( - matches!(result, Err(AwaError::Database(ref error)) if error.as_database_error().is_some_and(|error| error.code().as_deref() == Some("55P03"))), - "recovery must fail with lock_not_available: {result:?}" - ); + .expect("contended recovery head is skipped without waiting") + .expect("busy historical attempt is deferred"); + assert!(result.is_empty(), "busy attempt remains open"); let done: i64 = sqlx::query_scalar(audited_sql(format!( "SELECT count(*) FROM {schema}.done_entries WHERE job_id=$1" ))) @@ -24231,3 +24229,340 @@ async fn test_queue_storage_cancel_racing_a_fast_completion_leaves_one_outcome() assert_eq!(both, 0, "attempts both completed and cancelled"); assert_eq!(neither, 0, "attempts neither completed nor cancelled"); } + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RetainedBatchTransition { + Heartbeat, + Deadline, + Callback, +} + +async fn retained_batch_transition( + pool: &sqlx::PgPool, + store: &QueueStorage, + transition: RetainedBatchTransition, +) -> Vec<(i64, i64)> { + tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(2)), async { + match transition { + RetainedBatchTransition::Heartbeat => store + .rescue_stale_heartbeats(pool, Duration::ZERO) + .await + .expect("rescue unblocked stale attempts") + .into_iter() + .map(|job| (job.id, job.run_lease)) + .collect(), + RetainedBatchTransition::Deadline => store + .rescue_expired_deadlines(pool) + .await + .expect("rescue unblocked expired attempts") + .into_iter() + .map(|job| (job.id, job.run_lease)) + .collect(), + RetainedBatchTransition::Callback => store + .rescue_expired_callbacks(pool) + .await + .expect("rescue unblocked expired callbacks") + .into_iter() + .map(|job| (job.id, job.run_lease)) + .collect(), + } + }) + .await + .expect("busy historical head cannot delay ordinary transitions") +} + +async fn retained_receipt_cursor( + pool: &sqlx::PgPool, + store: &QueueStorage, + slot: i32, +) -> serde_json::Value { + sqlx::query_scalar(audited_sql(format!( + "SELECT jsonb_build_array(rescue_cursor_claimed_at, rescue_cursor_job_id, + rescue_cursor_run_lease, deadline_cursor_deadline_at, deadline_cursor_job_id, + deadline_cursor_run_lease, batch_deadline_cursor_deadline_at, + batch_deadline_cursor_batch_id) + FROM {}.claim_ring_slots WHERE slot=$1", + store.schema() + ))) + .bind(slot) + .fetch_one(pool) + .await + .expect("read receipt rescue cursors") +} + +async fn retained_closed_receipts( + pool: &sqlx::PgPool, + store: &QueueStorage, + pair: (i64, i64), +) -> i64 { + let schema = store.schema(); + sqlx::query_scalar(audited_sql(format!( + "SELECT + (SELECT count(*) FROM {schema}.lease_claim_closures WHERE job_id=$1 AND run_lease=$2) + + (SELECT count(*) FROM {schema}.lease_claim_closure_batches AS closures + WHERE EXISTS ( + SELECT 1 FROM {schema}.lease_claim_batches AS claims + CROSS JOIN LATERAL unnest(claims.job_ids,claims.run_leases,claims.receipt_ids) + AS members(job_id,run_lease,receipt_id) + WHERE members.job_id=$1 AND members.run_lease=$2 + AND closures.receipt_ranges @> members.receipt_id)) + + (SELECT count(*) FROM {schema}.lease_claims + WHERE job_id=$1 AND run_lease=$2 AND closed_at IS NOT NULL)" + ))) + .bind(pair.0) + .bind(pair.1) + .fetch_one(pool) + .await + .expect("count exact receipt closure evidence") +} + +async fn historical_busy_attempt_does_not_block_batch( + transition: RetainedBatchTransition, + materialized: bool, +) { + let deadline = if transition == RetainedBatchTransition::Heartbeat { + Duration::ZERO + } else { + Duration::from_secs(300) + }; + let (_guard, pool, store, historical) = + historical_tombstoned_claim_fixture(true, deadline, false).await; + let queue = "qs_unblocked_batch"; + let ordinary_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9042 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 1, + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, deadline) + .await + .expect("claim ordinary batch member"); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].job.id, ordinary_id); + let historical_pair = (historical.job_id, historical.run_lease); + let ordinary_pair = (ordinary_id, claimed[0].job.run_lease); + let pairs = [historical_pair, ordinary_pair]; + if materialized || transition == RetainedBatchTransition::Callback { + for &(job_id, run_lease) in &pairs { + let callback = store + .register_callback(&pool, job_id, run_lease, Duration::from_secs(300)) + .await + .expect("materialize attempt"); + if transition == RetainedBatchTransition::Callback { + assert!(store + .enter_callback_wait(&pool, job_id, run_lease, callback) + .await + .expect("park callback attempt")); + } + } + } + let schema = store.schema(); + let progress = serde_json::json!({"phase":"before-retained-head"}); + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.attempt_state (job_id,run_lease,progress,heartbeat_at,updated_at) + VALUES ($1,$2,$3,clock_timestamp()-interval '1 hour',clock_timestamp()) + ON CONFLICT (job_id,run_lease) DO UPDATE SET progress=EXCLUDED.progress, + heartbeat_at=EXCLUDED.heartbeat_at" + ))) + .bind(historical_pair.0) + .bind(historical_pair.1) + .bind(&progress) + .execute(&pool) + .await + .expect("seed mutable attempt state"); + if transition == RetainedBatchTransition::Deadline { + for table in ["leases", "lease_claims", "lease_claim_batches"] { + sqlx::query(audited_sql(format!( + "UPDATE {schema}.{table} SET deadline_at=clock_timestamp()-interval '1 hour' + WHERE deadline_at IS NOT NULL" + ))) + .execute(&pool) + .await + .expect("expire test deadlines"); + } + } + if transition == RetainedBatchTransition::Callback { + sqlx::query(audited_sql(format!( + "UPDATE {schema}.leases SET callback_timeout_at=clock_timestamp()-interval '1 hour' + WHERE state='waiting_external'" + ))) + .execute(&pool) + .await + .expect("expire test callbacks"); + } + let cursor_before = retained_receipt_cursor(&pool, &store, historical.claim_slot).await; + assert_eq!( + retained_closed_receipts(&pool, &store, historical_pair).await, + 0 + ); + let mut blocker = pool.begin().await.expect("begin retained lane blocker"); + sqlx::query(audited_sql(format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue='qs_claimed_aging_body' AND priority=3 AND enqueue_shard=0 FOR UPDATE" + ))) + .fetch_one(blocker.as_mut()) + .await + .expect("hold one required historical head"); + + assert_eq!( + retained_batch_transition(&pool, &store, transition).await, + vec![ordinary_pair] + ); + let active = store + .load_job(&pool, historical.job_id) + .await + .unwrap() + .unwrap(); + let expected_state = if transition == RetainedBatchTransition::Callback { + JobState::WaitingExternal + } else { + JobState::Running + }; + assert_historical_active_body(&active, &historical, expected_state); + assert_eq!(active.progress, Some(progress.clone())); + let stored_progress: serde_json::Value = sqlx::query_scalar(audited_sql(format!( + "SELECT progress FROM {schema}.attempt_state WHERE job_id=$1 AND run_lease=$2" + ))) + .bind(historical_pair.0) + .bind(historical_pair.1) + .fetch_one(&pool) + .await + .expect("blocked attempt state survives rollback"); + assert_eq!(stored_progress, progress); + assert_eq!( + retained_closed_receipts(&pool, &store, historical_pair).await, + 0 + ); + if !materialized { + assert_eq!( + retained_receipt_cursor(&pool, &store, historical.claim_slot).await, + cursor_before, + "an open first receipt remains eligible at its original cursor" + ); + } + assert_eq!(ready_tombstone_count(&pool, &store).await, 1); + assert!( + terminal_counts_by_priority(&pool, schema, "qs_claimed_aging_body") + .await + .is_empty() + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(4, 1)] + ); + blocker.rollback().await.expect("release retained head"); + assert_eq!( + retained_batch_transition(&pool, &store, transition).await, + vec![historical_pair] + ); + assert!(retained_batch_transition(&pool, &store, transition) + .await + .is_empty()); + let terminals: Vec<(i64, i64, i64)> = sqlx::query_as(audited_sql(format!( + "SELECT job_id,run_lease,count(*) FROM {schema}.done_entries + WHERE job_id=ANY($1) GROUP BY job_id,run_lease ORDER BY job_id" + ))) + .bind(vec![historical.job_id, ordinary_id]) + .fetch_all(&pool) + .await + .expect("read terminal attempt uniqueness"); + let mut expected = vec![ + (historical_pair.0, historical_pair.1, 1), + (ordinary_pair.0, ordinary_pair.1, 1), + ]; + expected.sort_unstable(); + assert_eq!(terminals, expected); + assert_eq!( + terminal_counts_by_priority(&pool, schema, "qs_claimed_aging_body").await, + vec![(4, 1)] + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(4, 1)] + ); + for _ in 0..8 { + assert!(store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 10, deadline) + .await + .expect("retired old replacement stays unclaimable") + .iter() + .all(|row| row.job.id != historical.job_id)); + } +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_compact_heartbeat_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Heartbeat, false).await; +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_compact_deadline_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Deadline, false).await; +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_materialized_deadline_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Deadline, true).await; +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_callback_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Callback, true).await; +} + +#[tokio::test] +async fn test_busy_retained_head_cancel_preserves_lock_not_available() { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(true, Duration::ZERO, false).await; + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(audited_sql(format!( + "SELECT 1 FROM {}.queue_claim_heads + WHERE queue='qs_claimed_aging_body' AND priority=3 AND enqueue_shard=0 FOR UPDATE", + store.schema() + ))) + .fetch_one(blocker.as_mut()) + .await + .unwrap(); + let result = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.cancel_running( + &pool, + claim.job_id, + claim.run_lease, + "test cancellation", + None, + ), + ) + .await + .expect("single attempt cancellation fails without waiting"); + assert!( + matches!(result, Err(AwaError::Database(ref error)) if error.as_database_error().is_some_and(|error| error.code().as_deref() == Some("55P03"))), + "single attempt preserves lock_not_available: {result:?}" + ); + assert_eq!( + retained_closed_receipts(&pool, &store, (claim.job_id, claim.run_lease)).await, + 0 + ); + assert_eq!( + store + .load_job(&pool, claim.job_id) + .await + .unwrap() + .unwrap() + .state, + JobState::Running + ); + blocker.rollback().await.unwrap(); +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_materialized_heartbeat_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Heartbeat, true).await; +} From b3034df43fcba47e2f2d19eb2e8a8cb79f55f249 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 19:11:02 +1300 Subject: [PATCH 6/8] Keep completion contention fixtures from changing shared storage routing --- awa-worker/src/completion.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/awa-worker/src/completion.rs b/awa-worker/src/completion.rs index 5fc0fa531..4bc8c359c 100644 --- a/awa-worker/src/completion.rs +++ b/awa-worker/src/completion.rs @@ -1138,7 +1138,7 @@ mod tests { ) .unwrap(); let store = runtime.store.clone(); - store.install(&pool).await.unwrap(); + store.prepare_schema(&pool).await.unwrap(); let mut jobs = Vec::new(); for queue in ["busy", "ordinary", "stale", "carried"] { store.enqueue_batch(&pool, queue, 4, 1).await.unwrap(); From f5cfb4610988a15c0737018dc58a9604a6bcc69a Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 20:00:10 +1300 Subject: [PATCH 7/8] Retry queue-storage finalization on head contention --- CHANGELOG.md | 9 +- awa-worker/src/executor.rs | 562 ++++++++++++++++++++++++++++++++++++- 2 files changed, 558 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7c2e29fad..c090dbb43 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- Retry queue-storage worker finalization after transient lock contention within + its existing 120-second deadline, preserving the saved handler outcome. Keep + attempt heartbeats active until finalization settles while releasing handler + capacity. Direct store APIs still return contention errors for caller retries. + - Schema-only admin clients use installed ready-slot numbers for point reads, cancellation and retained-body recovery, including rings larger than the default. - Keep committed claims out of physical priority aging before its batch limit, even before their cursor advances. Active point and callback reads recover @@ -26,8 +31,8 @@ Notable changes between releases. Detailed migration notes for storage transitio Repair takes claim-head row locks without waiting. Rescue batches defer busy historical attempts while progressing other attempts; receipt rescue cursors keep deferred attempts eligible. Worker completion batches isolate - lock contention to the affected requests. Direct completion and - single-attempt mutations still fail on contention and must retry. It adds + lock contention to the affected requests. Direct store completion and + single-attempt store mutations still fail on contention and must retry. It adds weak relation locks and tombstone writes only on this recovery path. Fresh rechecks rely on READ COMMITTED isolation. diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 894881c2e..c99a07501 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -411,6 +411,10 @@ impl JobExecutor { }), }; in_flight.insert((job_id, job_run_lease), in_flight_state); + let heartbeat_guard = storage.queue_storage().map(|_| FinalizationHeartbeatGuard { + in_flight: in_flight.clone(), + key: (job_id, job_run_lease), + }); if let Some(counter) = queue_in_flight.get(&job_queue) { counter.fetch_add(1, Ordering::SeqCst); } @@ -523,7 +527,9 @@ impl JobExecutor { None => None, }; - in_flight.remove((job_id, job_run_lease)); + if heartbeat_guard.is_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); } @@ -632,7 +638,9 @@ impl JobExecutor { &job_kind, &job_queue, COMPLETION_RECONNECT_WINDOW, - ) + +storage.queue_storage().is_some(), +) .await }; if !output_kept.into_inner() { @@ -758,6 +766,7 @@ impl JobExecutor { None }, }; + drop(heartbeat_guard); drop(finalize_permit); if let Some(event) = event { tokio::select! { @@ -773,14 +782,22 @@ impl JobExecutor { } } +struct FinalizationHeartbeatGuard { + in_flight: InFlightMap, + key: (i64, i64), +} + +impl Drop for FinalizationHeartbeatGuard { + 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; -/// Short exponential backoff between deadlock retries (10ms, 20ms, 40ms, ...). -/// Postgres aborts only one party of a cycle, so the survivor has usually -/// committed by the time the retry runs. -fn completion_deadlock_backoff(retry: u32) -> Duration { +fn completion_contention_backoff(retry: u32) -> Duration { Duration::from_millis(10u64 << retry.saturating_sub(1).min(4)) } @@ -805,6 +822,7 @@ struct FinalizeAck { elapsed: Duration, } +#[allow(clippy::too_many_arguments)] async fn finalize_with_retry( mut finalize: F, metrics: &crate::metrics::AwaMetrics, @@ -812,6 +830,7 @@ async fn finalize_with_retry( kind: &str, queue: &str, window: Duration, + retry_lock_contention: bool, ) -> Result, FinalizeError> where F: FnMut(CompletionPath) -> Fut, @@ -821,6 +840,8 @@ where let deadline = started + window; let mut deadlock_retries = 0; let mut reconnect_retries = 0; + let mut lock_retries = 0; + let mut warned = false; let mut schema_retried = false; let mut completion_path = CompletionPath::Batched; loop { @@ -839,7 +860,10 @@ where }; let delay = if error.is_deadlock() && deadlock_retries < COMPLETION_DEADLOCK_RETRIES { deadlock_retries += 1; - completion_deadlock_backoff(deadlock_retries) + completion_contention_backoff(deadlock_retries) + } else if retry_lock_contention && is_completion_lock_contention(&error) { + lock_retries += 1; + completion_contention_backoff(lock_retries) } else if error.is_connection_lost() { reconnect_retries += 1; completion_reconnect_backoff(reconnect_retries) @@ -851,7 +875,8 @@ where } else { return Err(error.into()); }; - if deadlock_retries + reconnect_retries == 1 { + if !warned { + warned = true; warn!(job_id, error = %error, delay_ms = delay.as_millis() as u64, "Durable finalization failed; retrying"); } else { debug!(job_id, error = %error, delay_ms = delay.as_millis() as u64, "Durable finalization failed; retrying"); @@ -861,14 +886,19 @@ where .map_err(|_| FinalizeError::Deadline)?; if error.is_connection_lost() { metrics.record_completion_reconnect_retry(kind, queue); - } else { + } else if error.is_deadlock() { metrics.record_completion_deadlock_retry(kind, queue); } completion_path = CompletionPath::Direct; } } -/// Reconnect retry budget measured from the first finalize attempt. Rescue +fn is_completion_lock_contention(error: &AwaError) -> bool { + matches!(error, AwaError::Database(sqlx::Error::Database(database)) + if database.code().as_deref() == Some("55P03")) +} + +/// Finalization retry budget measured from the first finalize attempt. Rescue /// can win before this expires; the lease guard then rejects the old result. const COMPLETION_RECONNECT_WINDOW: Duration = Duration::from_secs(120); @@ -3451,6 +3481,7 @@ mod tests { "test", "test", Duration::from_secs(2), + false, ) .await; assert!(matches!(result, Err(FinalizeError::Deadline))); @@ -3471,6 +3502,7 @@ mod tests { "test", "test", Duration::from_secs(2), + false, ) .await; assert!(matches!(result, Err(FinalizeError::Deadline))); @@ -3491,6 +3523,509 @@ mod tests { "test", "test", Duration::from_secs(120), + false, + ) + .await; + assert!(matches!(result, Err(FinalizeError::Operation(_)))); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[derive(Debug)] + struct FinalizeDatabaseError(&'static str); + + impl std::fmt::Display for FinalizeDatabaseError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "sqlstate {}", self.0) + } + } + + impl std::error::Error for FinalizeDatabaseError {} + + impl sqlx::error::DatabaseError for FinalizeDatabaseError { + fn message(&self) -> &str { + "finalize transaction rejected" + } + fn code(&self) -> Option> { + Some(std::borrow::Cow::Borrowed(self.0)) + } + fn as_error(&self) -> &(dyn std::error::Error + Send + Sync + 'static) { + self + } + fn as_error_mut(&mut self) -> &mut (dyn std::error::Error + Send + Sync + 'static) { + self + } + fn into_error(self: Box) -> Box { + self + } + fn kind(&self) -> sqlx::error::ErrorKind { + sqlx::error::ErrorKind::Other + } + } + + fn finalize_sqlstate(code: &'static str) -> AwaError { + AwaError::Database(sqlx::Error::Database(Box::new(FinalizeDatabaseError(code)))) + } + + #[tokio::test(start_paused = true)] + async fn finalize_lock_contention_preserves_the_sealed_outcome() { + let calls = AtomicU32::new(0); + let applied = AtomicU32::new(0); + let paths = std::sync::Mutex::new(Vec::new()); + let sealed = serde_json::json!({ + "outcome":"retry_after", "delay_secs":60, + "progress":{"checkpoint":7}, "followups":[{"id":42}], "output":{"value":9} + }); + let acknowledgement = finalize_with_retry( + |path| { + paths.lock().unwrap().push(path); + std::future::ready(if calls.fetch_add(1, Ordering::SeqCst) < 4 { + Err(finalize_sqlstate("55P03")) + } else { + applied.fetch_add(1, Ordering::SeqCst); + Ok(sealed.clone()) + }) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_secs(2), + true, + ) + .await + .expect("contention must retry the saved outcome"); + assert_eq!(acknowledgement.outcome, sealed); + assert_eq!(acknowledgement.reconnect_retries, 0); + assert_eq!(calls.load(Ordering::SeqCst), 5); + assert_eq!(applied.load(Ordering::SeqCst), 1); + assert_eq!( + *paths.lock().unwrap(), + [ + CompletionPath::Batched, + CompletionPath::Direct, + CompletionPath::Direct, + CompletionPath::Direct, + CompletionPath::Direct, + ] + ); + } + + #[tokio::test(start_paused = true)] + async fn finalize_lock_contention_keeps_the_existing_deadline() { + let calls = AtomicU32::new(0); + let started = tokio::time::Instant::now(); + let result = finalize_with_retry( + |_| { + calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err::<(), _>(finalize_sqlstate("55P03"))) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_millis(100), + true, + ) + .await; + assert!(matches!(result, Err(FinalizeError::Deadline)), "{result:?}"); + assert_eq!(started.elapsed(), Duration::from_millis(100)); + assert!((3..=5).contains(&calls.load(Ordering::SeqCst))); + } + + #[tokio::test(start_paused = true)] + async fn finalize_lock_contention_does_not_retry_unrelated_statement_errors() { + for code in ["57014", "23505", "55P04"] { + let calls = AtomicU32::new(0); + let result = finalize_with_retry( + |_| { + calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err::<(), _>(finalize_sqlstate(code))) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_secs(2), + true, + ) + .await; + assert!( + matches!(result, Err(FinalizeError::Operation(_))), + "{code}: {result:?}" + ); + assert_eq!(calls.load(Ordering::SeqCst), 1, "{code}"); + } + } + + #[tokio::test] + async fn queue_finalization_retains_heartbeat_without_holding_worker_capacity() { + let metrics = test_metrics(); + let now = chrono::Utc::now(); + let job: JobRow = serde_json::from_value(serde_json::json!({ + "id":1,"kind":"reconnect_parent","queue":"test","args":{}, + "state":"running","priority":2,"attempt":1,"run_lease":1,"max_attempts":5, + "run_at":now,"attempted_at":now,"created_at":now,"metadata":{},"tags":[], + })) + .unwrap(); + let key = (job.id, job.run_lease); + let runtime = QueueStorageRuntime::new( + awa_model::QueueStorageConfig::default(), + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(); + for storage in [ + RuntimeStorage::Canonical, + RuntimeStorage::QueueStorage(runtime), + ] { + let queue_storage = matches!(storage, RuntimeStorage::QueueStorage(_)); + let pool = sqlx::postgres::PgPoolOptions::new() + .connect_lazy("postgres://postgres:test@localhost/awa_test") + .unwrap(); + let finalizers = TaskTracker::new(); + let cancel = CancellationToken::new(); + let in_flight = Arc::new(crate::runtime::InFlightRegistry::default()); + let count = Arc::new(AtomicU32::new(0)); + let (_unstarted, handle) = crate::completion::CompletionBatcher::new( + pool.clone(), + cancel.clone(), + metrics.clone(), + storage.clone(), + 1, + ); + let executor = JobExecutor::new( + pool, + Arc::new(HashMap::from([( + "reconnect_parent".into(), + Box::new(ReconnectWorker(Arc::new(AtomicU32::new(0)))) as BoxedWorker, + )])), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + in_flight.clone(), + Arc::new(HashMap::from([("test".into(), count.clone())])), + Arc::new(HashMap::new()), + metrics.clone(), + handle, + storage, + DlqPolicy::default(), + finalizers.clone(), + cancel.clone(), + 1, + ); + executor + .execute_task( + DispatchedJob { + job: job.clone(), + queue_storage_claim: queue_storage.then(|| ClaimedEntry { + queue: "test".into(), + priority: 2, + lane_seq: 1, + ready_slot: 0, + ready_generation: 0, + lease_slot: 0, + lease_generation: 0, + claim_slot: 0, + receipt_id: Some(1), + claim_batch_id: Some(1), + claim_batch_index: Some(1), + lease_claim_receipt: true, + enqueue_shard: 0, + }), + queue_storage_unique_states: None, + abandon_after: None, + }, + Arc::new(AtomicBool::new(false)), + ) + .await; + assert_eq!( + count.load(Ordering::SeqCst), + 0, + "handler capacity must be released" + ); + assert_eq!(finalizers.len(), 1); + assert_eq!( + in_flight.get_cancel(key).is_some(), + queue_storage, + "queue-store heartbeats remain registered until finalization settles" + ); + cancel.cancel(); + finalizers.close(); + tokio::time::timeout(Duration::from_secs(2), finalizers.wait()) + .await + .unwrap(); + assert!( + in_flight.get_cancel(key).is_none(), + "cancellation must unregister the exact lease" + ); + } + } + + struct RetainedFinalizeFollowup; + + impl crate::enqueue_specs::EnqueueFollowUp for RetainedFinalizeFollowup { + fn run<'a>( + &'a self, + connection: &'a mut sqlx::PgConnection, + job: &'a JobRow, + _context: Option<&'a crate::enqueue_specs::OutcomeContext>, + ) -> std::pin::Pin> + Send + 'a>> + { + Box::pin(async move { + sqlx::query("INSERT INTO public.finalize_followups(parent) VALUES ($1)") + .bind(job.id) + .execute(connection) + .await?; + Ok(()) + }) + } + } + + #[sqlx::test(migrations = false)] + async fn retained_head_contention_replays_the_worker_outcome_once(pool: PgPool) { + let _database_guard = RECONNECT_DB_LOCK.lock().await; + awa_model::migrations::run(&pool).await.unwrap(); + sqlx::query("CREATE TABLE public.finalize_followups(parent bigint PRIMARY KEY)") + .execute(&pool) + .await + .unwrap(); + let runtime = QueueStorageRuntime::new( + awa_model::QueueStorageConfig { + queue_stripe_count: 1, + ..Default::default() + }, + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(); + let queue = "retained_finalize"; + runtime.store.prepare_queue(&pool, queue, 1).await.unwrap(); + runtime + .store + .enqueue_params_batch( + &pool, + &[awa_model::job::InsertParams { + kind: "retained_finalize_parent".into(), + args: serde_json::json!({"body":"retained"}), + opts: awa_model::InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 3, + metadata: serde_json::json!({"preserved":true}), + ..Default::default() + }, + }], + ) + .await + .unwrap(); + let claimed = runtime + .store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .unwrap() + .pop() + .unwrap(); + let pair = (claimed.job.id, claimed.job.run_lease); + let schema = runtime.store.schema(); + let mut historical = pool.begin().await.unwrap(); + sqlx::query(awa_model::audited_sql(format!( + "INSERT INTO {schema}.ready_tombstones + (ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id + FROM {schema}.ready_entries WHERE job_id=$1" + ))) + .bind(pair.0) + .execute(historical.as_mut()) + .await + .unwrap(); + sqlx::query(awa_model::audited_sql(format!( + "INSERT INTO {schema}.ready_entries + (ready_slot,ready_generation,job_id,kind,queue,args,priority,attempt,run_lease,max_attempts, + lane_seq,enqueue_shard,run_at,attempted_at,created_at,unique_key,unique_states,payload) + SELECT ready_slot,ready_generation,job_id,kind,queue,args,3,attempt,run_lease,max_attempts, + {schema}.reserve_enqueue_seq(queue,3::smallint,enqueue_shard,1::bigint),enqueue_shard, + run_at,attempted_at,created_at,unique_key,unique_states,payload + FROM {schema}.ready_entries WHERE job_id=$1" + ))).bind(pair.0).execute(historical.as_mut()).await.unwrap(); + sqlx::query(awa_model::audited_sql(format!( + "INSERT INTO {schema}.ready_segments + (ready_slot,ready_generation,queue,priority,enqueue_shard,first_lane_seq,next_lane_seq,first_run_at,created_at) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,lane_seq+1,run_at,created_at + FROM {schema}.ready_entries WHERE job_id=$1 AND priority=3 ON CONFLICT DO NOTHING" + ))).bind(pair.0).execute(historical.as_mut()).await.unwrap(); + historical.commit().await.unwrap(); + let progress = serde_json::json!({"checkpoint":7}); + runtime + .store + .flush_progress(&pool, pair.0, pair.1, progress.clone()) + .await + .unwrap(); + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(awa_model::audited_sql(format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue=$1 AND priority=3 AND enqueue_shard=0 FOR UPDATE" + ))) + .bind(queue) + .fetch_one(blocker.as_mut()) + .await + .unwrap(); + let metrics = test_metrics(); + let storage = RuntimeStorage::QueueStorage(runtime.clone()); + let (_unstarted, handle) = crate::completion::CompletionBatcher::new( + pool.clone(), + CancellationToken::new(), + metrics.clone(), + storage.clone(), + 1, + ); + let specs: Arc = Arc::new(HashMap::from([( + crate::enqueue_specs::Outcome::Retried, + HashMap::from([( + "retained_finalize_parent".into(), + vec![Arc::new(RetainedFinalizeFollowup) as crate::enqueue_specs::BoxedEnqueueSpec], + )]), + )])); + let intended = Ok(JobResult::RetryAfter(Duration::from_secs(60))); + let calls = AtomicU32::new(0); + let busy = tokio::sync::Notify::new(); + let finalize = finalize_with_retry( + |path| { + calls.fetch_add(1, Ordering::SeqCst); + let (pool, job, claim, intended, progress, handle, specs, storage, metrics, busy) = ( + &pool, + &claimed.job, + &claimed.claim, + &intended, + &progress, + &handle, + &specs, + &storage, + &metrics, + &busy, + ); + async move { + let result = complete_job( + pool, + job, + Some(claim), + None, + intended, + &[], + handle, + path, + Some(progress.clone()), + Duration::ZERO, + true, + specs, + storage, + false, + metrics, + ) + .await; + if let Err(error) = &result { + assert!( + is_completion_lock_contention(error), + "expected actual held-head55P03: {error}" + ); + busy.notify_one(); + } + result + } + }, + &metrics, + pair.0, + &claimed.job.kind, + queue, + Duration::from_secs(5), + true, + ); + let release = async { + tokio::time::timeout(Duration::from_secs(3), busy.notified()) + .await + .expect("worker transition must encounter the held historical head"); + let active = runtime + .store + .load_job(&pool, pair.0) + .await + .unwrap() + .unwrap(); + assert_eq!(active.state, JobState::Running); + assert_eq!(active.run_lease, pair.1); + assert_eq!(active.progress, Some(progress.clone())); + let effects: i64 = sqlx::query_scalar("SELECT count(*) FROM public.finalize_followups") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 0, "failed transaction cannot publish its followup"); + blocker.rollback().await.unwrap(); + }; + let (ack, ()) = tokio::join!(finalize, release); + let ack = ack.expect("saved worker outcome must survive transient head contention"); + assert_eq!(ack.reconnect_retries, 0); + assert!(calls.load(Ordering::SeqCst) >= 2); + let CompletionOutcome::Applied { + event: Some(UntypedJobEvent::Retried { job, .. }), + terminal: false, + } = ack.outcome + else { + panic!("worker must acknowledge its intended retry outcome") + }; + assert_eq!(job.id, pair.0); + assert_eq!(job.run_lease, pair.1); + assert_eq!(job.state, JobState::Retryable); + assert_eq!(job.args, claimed.job.args); + assert_eq!(job.max_attempts, claimed.job.max_attempts); + assert_eq!(job.progress, Some(progress.clone())); + assert_eq!(job.metadata["preserved"], true); + let replay = complete_job( + &pool, + &claimed.job, + Some(&claimed.claim), + None, + &intended, + &[], + &handle, + CompletionPath::Direct, + Some(progress), + Duration::ZERO, + true, + &specs, + &storage, + false, + &metrics, + ) + .await + .unwrap(); + assert!(matches!(replay, CompletionOutcome::IgnoredStale)); + let effects: i64 = + sqlx::query_scalar("SELECT count(*) FROM public.finalize_followups WHERE parent=$1") + .bind(pair.0) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + effects, 1, + "the acknowledged outcome publishes exactly one followup" + ); + let deferred:i64 = sqlx::query_scalar(awa_model::audited_sql(format!( + "SELECT count(*) FROM {schema}.deferred_jobs WHERE job_id=$1 AND run_lease=$2 AND state='retryable'" + ))).bind(pair.0).bind(pair.1).fetch_one(&pool).await.unwrap(); + assert_eq!(deferred, 1); + } + + #[tokio::test(start_paused = true)] + async fn canonical_finalization_does_not_retry_lock_contention() { + let calls = AtomicU32::new(0); + let result = finalize_with_retry( + |_| { + calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err::<(), _>(finalize_sqlstate("55P03"))) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_secs(2), + false, ) .await; assert!(matches!(result, Err(FinalizeError::Operation(_)))); @@ -3542,6 +4077,7 @@ mod tests { "test", "test", Duration::from_secs(120), + false, ) .with_subscriber(subscriber) .await @@ -3768,6 +4304,7 @@ mod tests { &job.kind, &job.queue, Duration::from_secs(5), + false, ) .await .unwrap(); @@ -3822,7 +4359,9 @@ mod tests { } }, &metrics, job.id, &job.kind, &job.queue, Duration::from_secs(5), - ).await.unwrap(); + +false, +).await.unwrap(); assert!(matches!(outcome.outcome, CompletionOutcome::IgnoredStale)); let current: JobRow = sqlx::query_as("SELECT * FROM awa.jobs_hot WHERE id = $1") .bind(job.id) @@ -3902,6 +4441,7 @@ mod tests { &job.kind, &job.queue, Duration::from_secs(5), + false, ) .await .unwrap(); From 753cbaa39532f9ac10a652473c511b6ec7bb2570 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sat, 7 Nov 2026 20:44:02 +1300 Subject: [PATCH 8/8] Retain accepted completions across finalization deadlines --- CHANGELOG.md | 9 +- awa-worker/src/completion.rs | 372 +++++++++++++++++++++++++++++++++-- awa-worker/src/executor.rs | 57 ++++-- 3 files changed, 412 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c090dbb43..f7be9872f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,12 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed - Retry queue-storage worker finalization after transient lock contention within - its existing 120-second deadline, preserving the saved handler outcome. Keep - attempt heartbeats active until finalization settles while releasing handler - capacity. Direct store APIs still return contention errors for caller retries. + its 120-second running-work budget, preserving the saved handler outcome. + Waiting for an accepted completion batch pauses that budget; admission, + direct writes, post-acknowledgement reads and retries use the remaining time. + Keep attempt heartbeats active until finalization settles while releasing + handler capacity. Runtime shutdown retains its existing drain limit. Direct + store APIs still return contention errors for caller retries. - Schema-only admin clients use installed ready-slot numbers for point reads, cancellation and retained-body recovery, including rings larger than the default. - Keep committed claims out of physical priority aging before its batch limit, diff --git a/awa-worker/src/completion.rs b/awa-worker/src/completion.rs index 4bc8c359c..8b289b18e 100644 --- a/awa-worker/src/completion.rs +++ b/awa-worker/src/completion.rs @@ -4,8 +4,10 @@ use awa_model::{AwaError, ClaimedEntry, ClaimedRuntimeJob}; use sqlx::PgPool; use std::collections::HashSet; use std::env; +use std::sync::{Arc, Mutex}; use std::time::Duration; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{mpsc, oneshot, Notify}; +use tokio::time::Instant; use tokio_util::sync::CancellationToken; use tracing::{debug, warn}; @@ -92,6 +94,144 @@ struct CompletionRequest { response: oneshot::Sender>, } +#[derive(Clone)] +pub(crate) struct FinalizationBudget { + inner: Arc, +} + +struct BudgetClock { + state: Mutex, + changed: Notify, +} + +struct BudgetState { + deadline: Instant, + paused_at: Option, + acknowledgement_waits: usize, +} + +#[derive(Debug)] +pub(crate) struct FinalizationBudgetElapsed; + +impl From for FinalizationBudget { + fn from(window: Duration) -> Self { + Self { + inner: Arc::new(BudgetClock { + state: Mutex::new(BudgetState { + deadline: Instant::now() + window, + paused_at: None, + acknowledgement_waits: 0, + }), + changed: Notify::new(), + }), + } + } +} + +impl FinalizationBudget { + pub(crate) fn is_elapsed(&self) -> bool { + self.deadline() + .is_some_and(|deadline| Instant::now() >= deadline) + } + + fn deadline(&self) -> Option { + let state = self.inner.state.lock().unwrap(); + state.paused_at.is_none().then_some(state.deadline) + } + + fn wait_for_acknowledgement(&self) -> AcknowledgementWait { + let mut state = self.inner.state.lock().unwrap(); + if state.acknowledgement_waits == 0 { + state.paused_at = Some(Instant::now()); + } + state.acknowledgement_waits += 1; + drop(state); + self.inner.changed.notify_waiters(); + AcknowledgementWait { + budget: self.clone(), + } + } + + pub(crate) async fn timeout(&self, future: F) -> Result + where + F: std::future::Future, + { + tokio::pin!(future); + loop { + let changed = self.inner.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + match self.deadline() { + Some(deadline) => { + tokio::select! { + biased; + _ = tokio::time::sleep_until(deadline) => { + if self.deadline().is_some_and(|deadline| Instant::now() >= deadline) { + return Err(FinalizationBudgetElapsed); + } + } + outcome = &mut future => return Ok(outcome), + _ = &mut changed => {}, + } + } + None => { + tokio::select! { + biased; + outcome = &mut future => return Ok(outcome), + _ = &mut changed => {}, + } + } + } + } + } + + pub(crate) async fn timeout_unpaused( + &self, + future: F, + ) -> Result + where + F: std::future::Future, + { + let deadline = self.deadline().expect("unpaused finalization budget"); + tokio::time::timeout_at(deadline, future) + .await + .map_err(|_| FinalizationBudgetElapsed) + } +} + +struct AcknowledgementWait { + budget: FinalizationBudget, +} + +impl Drop for AcknowledgementWait { + fn drop(&mut self) { + let mut state = self.budget.inner.state.lock().unwrap(); + state.acknowledgement_waits -= 1; + if state.acknowledgement_waits == 0 { + let paused_at = state.paused_at.take().unwrap(); + state.deadline += Instant::now() - paused_at; + } + drop(state); + self.budget.inner.changed.notify_waiters(); + } +} + +struct CompletionReceipt { + response: oneshot::Receiver>, + acknowledgement_wait: Option, +} + +impl CompletionReceipt { + async fn acknowledge(self) -> Result { + let result = self + .response + .await + .map_err(|_| AwaError::Validation("completion batcher dropped response".into()))?; + drop(self.acknowledgement_wait); + result + } +} + fn completion_sort_key(request: &CompletionRequest) -> (i32, i64, RunLease) { let claim_slot = request .runtime_job @@ -106,9 +246,25 @@ fn completion_sort_key(request: &CompletionRequest) -> (i32, i64, RunLease) { pub(crate) struct CompletionBatcherHandle { shards: Vec>, pipeline_capacity: usize, + acknowledgement_budget: Option, } impl CompletionBatcherHandle { + pub(crate) fn with_acknowledgement_budget(&self, budget: FinalizationBudget) -> Self { + Self { + acknowledgement_budget: Some(budget), + ..self.clone() + } + } + + #[cfg(test)] + pub(crate) fn pending_request_count(&self) -> usize { + self.shards + .iter() + .map(|shard| shard.max_capacity() - shard.capacity()) + .sum() + } + pub fn finalizer_capacity(&self, runtime_worker_capacity: u32) -> usize { self.pipeline_capacity .max(runtime_worker_capacity.max(1) as usize) @@ -139,22 +295,41 @@ impl CompletionBatcherHandle { claim: Option, runtime_job: Option, ) -> Result { + self.submit(job_id, run_lease, claim, runtime_job) + .await? + .acknowledge() + .await + } + + async fn submit( + &self, + job_id: i64, + run_lease: RunLease, + claim: Option, + runtime_job: Option, + ) -> Result { let shard = (job_id.rem_euclid(self.shards.len() as i64)) as usize; let (response_tx, response_rx) = oneshot::channel(); - self.shards[shard] - .send(CompletionRequest { - job_id, - run_lease, - claim, - runtime_job, - response: response_tx, - }) + let permit = self.shards[shard] + .reserve() .await .map_err(|_| AwaError::Validation("completion batcher stopped".into()))?; + let acknowledgement_wait = self + .acknowledgement_budget + .as_ref() + .map(FinalizationBudget::wait_for_acknowledgement); + permit.send(CompletionRequest { + job_id, + run_lease, + claim, + runtime_job, + response: response_tx, + }); - response_rx - .await - .map_err(|_| AwaError::Validation("completion batcher dropped response".into()))? + Ok(CompletionReceipt { + response: response_rx, + acknowledgement_wait, + }) } } @@ -218,6 +393,7 @@ impl CompletionBatcher { shards, // Each shard needs an in-flight batch and a queued batch. pipeline_capacity: batch_size.saturating_mul(shard_count).saturating_mul(2), + acknowledgement_budget: None, }, ) } @@ -493,6 +669,178 @@ mod tests { use std::sync::Arc; use std::time::{Duration, Instant}; + fn controlled_completions( + capacity: usize, + ) -> (CompletionBatcherHandle, mpsc::Receiver) { + let (sender, receiver) = mpsc::channel(capacity); + ( + CompletionBatcherHandle { + shards: vec![sender], + pipeline_capacity: capacity, + acknowledgement_budget: None, + }, + receiver, + ) + } + + #[tokio::test(start_paused = true)] + async fn accepted_completion_keeps_its_acknowledgement_after_the_budget() { + use crate::events::{BoxedUntypedEventHandler, UntypedJobEvent}; + use std::collections::HashMap; + use std::sync::atomic::{AtomicU32, Ordering}; + + let (handle, mut requests) = controlled_completions(1); + let budget = FinalizationBudget::from(Duration::from_millis(100)); + let handle = handle.with_acknowledgement_budget(budget.clone()); + let events = Arc::new(AtomicU32::new(0)); + let observed = events.clone(); + let handler: BoxedUntypedEventHandler = Arc::new(move |event| { + assert!(matches!(event, UntypedJobEvent::Completed { job, .. } + if job.id == 1 && job.run_lease == 2)); + observed.fetch_add(1, Ordering::SeqCst); + Box::pin(async {}) + }); + let handlers = HashMap::from([("accepted_parent".into(), vec![handler])]); + let task = tokio::spawn(async move { + let updated = budget + .timeout(handle.complete(1, 2)) + .await + .unwrap() + .unwrap(); + assert!(updated); + let now = chrono::Utc::now(); + let job = serde_json::from_value(serde_json::json!({ + "id":1,"kind":"accepted_parent","queue":"test","args":{}, + "state":"completed","priority":2,"attempt":1,"run_lease":2, + "max_attempts":5,"run_at":now,"attempted_at":now,"created_at":now, + "finalized_at":now,"metadata":{},"tags":[], + })) + .unwrap(); + crate::executor::dispatch_lifecycle_event( + &handlers, + "accepted_parent", + UntypedJobEvent::Completed { + job, + duration: Duration::ZERO, + }, + ) + .await; + }); + let request = requests.recv().await.unwrap(); + assert_eq!((request.job_id, request.run_lease), (1, 2)); + tokio::time::advance(Duration::from_millis(200)).await; + tokio::task::yield_now().await; + assert!(!task.is_finished(), "accepted work must keep its waiter"); + assert_eq!(events.load(Ordering::SeqCst), 0); + assert!(request.response.send(Ok(true)).is_ok()); + task.await.unwrap(); + assert_eq!(events.load(Ordering::SeqCst), 1); + assert!(requests.try_recv().is_err(), "one accepted request only"); + } + + #[tokio::test(start_paused = true)] + async fn completion_admission_still_expires_when_the_channel_is_full() { + let (handle, mut requests) = controlled_completions(1); + assert!(handle.shards[0] + .send(completion_request_with_claim_slot(1, 1, 0)) + .await + .is_ok()); + let budget = FinalizationBudget::from(Duration::from_millis(100)); + let handle = handle.with_acknowledgement_budget(budget.clone()); + let started = tokio::time::Instant::now(); + assert!(budget.timeout(handle.complete(2, 1)).await.is_err()); + assert_eq!(started.elapsed(), Duration::from_millis(100)); + assert_eq!(requests.recv().await.unwrap().job_id, 1); + assert!( + requests.try_recv().is_err(), + "expired admission cannot send" + ); + } + + #[tokio::test(start_paused = true)] + async fn expired_budget_cannot_admit_work_to_a_ready_channel() { + let (handle, mut requests) = controlled_completions(1); + let budget = FinalizationBudget::from(Duration::from_millis(100)); + let handle = handle.with_acknowledgement_budget(budget.clone()); + tokio::time::advance(Duration::from_millis(100)).await; + assert!(budget.timeout(handle.complete(1, 2)).await.is_err()); + assert!(requests.try_recv().is_err()); + } + + async fn assert_remaining_budget_after_acknowledgement(response: Result) { + let (handle, mut requests) = controlled_completions(1); + let budget = FinalizationBudget::from(Duration::from_millis(100)); + let handle = handle.with_acknowledgement_budget(budget.clone()); + let started = tokio::time::Instant::now(); + let task = tokio::spawn(async move { + budget + .timeout(async move { + tokio::time::sleep(Duration::from_millis(30)).await; + let outcome = handle.complete(1, 2).await; + tokio::time::sleep(Duration::from_millis(80)).await; + outcome + }) + .await + }); + let request = requests.recv().await.unwrap(); + assert_eq!(started.elapsed(), Duration::from_millis(30)); + tokio::time::advance(Duration::from_millis(200)).await; + tokio::task::yield_now().await; + assert!(!task.is_finished()); + assert!(request.response.send(response).is_ok()); + assert!(task.await.unwrap().is_err()); + assert_eq!( + started.elapsed(), + Duration::from_millis(300), + "30ms before admission leaves 70ms after the 200ms acknowledgement wait" + ); + } + + #[tokio::test(start_paused = true)] + async fn post_acknowledgement_work_uses_only_the_remaining_budget() { + assert_remaining_budget_after_acknowledgement(Ok(true)).await; + } + + #[tokio::test(start_paused = true)] + async fn failed_acknowledgement_does_not_exempt_direct_fallback() { + assert_remaining_budget_after_acknowledgement(Err(AwaError::Validation( + "completion flush failed".into(), + ))) + .await; + } + + #[tokio::test(start_paused = true)] + async fn dropping_an_accepted_receipt_resumes_its_budget() { + let (handle, mut requests) = controlled_completions(1); + let budget = FinalizationBudget::from(Duration::from_millis(100)); + let handle = handle.with_acknowledgement_budget(budget.clone()); + let started = tokio::time::Instant::now(); + tokio::time::sleep(Duration::from_millis(30)).await; + let receipt = handle.submit(1, 2, None, None).await.unwrap(); + let request = requests.recv().await.unwrap(); + tokio::time::advance(Duration::from_millis(200)).await; + drop(receipt); + assert!(request.response.send(Ok(true)).is_err()); + assert!(budget + .timeout(tokio::time::sleep(Duration::from_millis(80))) + .await + .is_err()); + assert_eq!(started.elapsed(), Duration::from_millis(300)); + } + + #[tokio::test(start_paused = true)] + async fn unscoped_batch_acknowledgement_retains_the_original_timeout() { + let (handle, mut requests) = controlled_completions(1); + let budget = FinalizationBudget::from(Duration::from_millis(100)); + let started = tokio::time::Instant::now(); + let task = + tokio::spawn(async move { budget.timeout_unpaused(handle.complete(1, 2)).await }); + let request = requests.recv().await.unwrap(); + assert!(task.await.unwrap().is_err()); + assert_eq!(started.elapsed(), Duration::from_millis(100)); + assert!(request.response.send(Ok(true)).is_err()); + } + fn base_database_url() -> String { std::env::var("DATABASE_URL") .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index c99a07501..37719ff1e 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -1,4 +1,4 @@ -use crate::completion::CompletionBatcherHandle; +use crate::completion::{CompletionBatcherHandle, FinalizationBudget}; use crate::context::{CallbackGuard, JobContext}; use crate::events::{BoxedUntypedEventHandler, UntypedJobEvent}; use crate::runtime::{InFlightMap, InFlightState, ProgressState}; @@ -581,6 +581,12 @@ impl JobExecutor { }); let output_kept = AtomicBool::new(true); let outcome = { + let budget = FinalizationBudget::from(COMPLETION_RECONNECT_WINDOW); + let completion_batcher = if storage.queue_storage().is_some() { + completion_batcher.with_acknowledgement_budget(budget.clone()) + } else { + completion_batcher.clone() + }; let (job, stripped_job, output_kept) = (&job, &stripped_job, &output_kept); let (pool, result, storage) = (&pool, &result, &storage); let (staged_follow_ups, completion_batcher) = @@ -637,7 +643,7 @@ impl JobExecutor { job_id, &job_kind, &job_queue, - COMPLETION_RECONNECT_WINDOW, + budget, storage.queue_storage().is_some(), ) @@ -829,7 +835,7 @@ async fn finalize_with_retry( job_id: i64, kind: &str, queue: &str, - window: Duration, + window: impl Into, retry_lock_contention: bool, ) -> Result, FinalizeError> where @@ -837,7 +843,7 @@ where Fut: std::future::Future>, { let started = tokio::time::Instant::now(); - let deadline = started + window; + let budget = window.into(); let mut deadlock_retries = 0; let mut reconnect_retries = 0; let mut lock_retries = 0; @@ -845,9 +851,16 @@ where let mut schema_retried = false; let mut completion_path = CompletionPath::Batched; loop { - let outcome = tokio::time::timeout_at(deadline, finalize(completion_path)) - .await - .map_err(|_| FinalizeError::Deadline)?; + if retry_lock_contention && budget.is_elapsed() { + return Err(FinalizeError::Deadline); + } + let attempt = finalize(completion_path); + let outcome = if retry_lock_contention { + budget.timeout(attempt).await + } else { + budget.timeout_unpaused(attempt).await + } + .map_err(|_| FinalizeError::Deadline)?; let error = match outcome { Ok(outcome) => { return Ok(FinalizeAck { @@ -881,9 +894,13 @@ where } else { debug!(job_id, error = %error, delay_ms = delay.as_millis() as u64, "Durable finalization failed; retrying"); } - tokio::time::timeout_at(deadline, tokio::time::sleep(delay)) - .await - .map_err(|_| FinalizeError::Deadline)?; + let backoff = tokio::time::sleep(delay); + if retry_lock_contention { + budget.timeout(backoff).await + } else { + budget.timeout_unpaused(backoff).await + } + .map_err(|_| FinalizeError::Deadline)?; if error.is_connection_lost() { metrics.record_completion_reconnect_retry(kind, queue); } else if error.is_deadlock() { @@ -3657,7 +3674,7 @@ mod tests { } } - #[tokio::test] + #[tokio::test(start_paused = true)] async fn queue_finalization_retains_heartbeat_without_holding_worker_capacity() { let metrics = test_metrics(); let now = chrono::Utc::now(); @@ -3693,6 +3710,7 @@ mod tests { storage.clone(), 1, ); + let pending_completions = handle.clone(); let executor = JobExecutor::new( pool, Arc::new(HashMap::from([( @@ -3748,6 +3766,23 @@ mod tests { queue_storage, "queue-store heartbeats remain registered until finalization settles" ); + if queue_storage { + tokio::time::timeout(Duration::from_secs(1), async { + while pending_completions.pending_request_count() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the completion must be accepted before advancing the clock"); + tokio::time::advance(COMPLETION_RECONNECT_WINDOW + Duration::from_secs(1)).await; + tokio::task::yield_now().await; + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(finalizers.len(), 1); + assert!( + in_flight.get_cancel(key).is_some(), + "an accepted completion must retain its lease past the running-work budget" + ); + } cancel.cancel(); finalizers.close(); tokio::time::timeout(Duration::from_secs(2), finalizers.wait())