diff --git a/CHANGELOG.md b/CHANGELOG.md index 327b27c22..ad3c28fb5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -91,9 +91,12 @@ order; binaries can roll back to 0.6.13 with the patch retained. See the perform that repair, and can leave an old replacement claimable. The aging fence prevents new replacements from updated callers; it 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. + 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. ## [0.6.13] — 2026-11-06 diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 9f6d3aec8..33855063e 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -1755,6 +1755,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, @@ -10775,8 +10790,43 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, deleted: Vec, ) -> Result, AwaError> { + let (hydrated, sources) = self + .hydrate_deleted_leases_with_sources_tx(tx, deleted) + .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, + ) -> Result, BusyRetained>, AwaError> { + let (hydrated, sources) = self + .hydrate_deleted_leases_with_sources_tx(tx, deleted) + .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, + ) -> Result<(Vec, Vec), AwaError> { if deleted.is_empty() { - return Ok(Vec::new()); + return Ok((Vec::new(), Vec::new())); } let schema = self.schema(); @@ -11022,18 +11072,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 { @@ -11050,10 +11097,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)>(&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(&format!( r#" WITH refs(job_id, run_lease) AS ( @@ -11159,7 +11251,7 @@ impl QueueStorage { .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; - Ok(()) + Ok(None) } async fn close_open_receipt_claim_tx<'a>( @@ -11458,6 +11550,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; @@ -11465,7 +11559,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); @@ -11481,7 +11577,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); @@ -11500,7 +11598,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); @@ -11685,15 +11790,23 @@ 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 )"# ); // Closers take the attempt's advisory lock before their statement - // snapshot. Rescue takes the same locks in a statement of its own, so - // the statement that decides and writes the closure starts after any - // closer that held a lock first has committed, and sees its evidence. - // Only attempts locked here may be closed below. + // snapshot. Rescue takes the same locks in a statement of its own, + // so the statement that decides and writes the closure starts after + // any closer that held a lock first has committed and sees its + // evidence. Only attempts locked here may be closed below. let (locked_job_ids, locked_run_leases): (Vec, Vec) = sqlx::query_as::<_, (i64, i64)>(&format!( r#" @@ -11712,6 +11825,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)? @@ -11757,7 +11875,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 @@ -11806,7 +11924,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 @@ -11943,6 +12061,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()) @@ -11971,6 +12094,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; @@ -11978,7 +12103,9 @@ impl QueueStorage { if let Some(slot) = preferred_slot { 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); @@ -11994,7 +12121,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); @@ -12036,11 +12165,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); @@ -12148,6 +12284,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 )"# @@ -12170,6 +12314,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)? @@ -12189,7 +12338,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 @@ -12282,6 +12431,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()) @@ -15113,7 +15267,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(&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(&format!( r#" DELETE FROM {schema}.leases AS target WHERE (target.job_id, target.run_lease) IN ( @@ -15124,6 +15288,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 @@ -15148,53 +15320,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).await?; - let moved_receipts = self - .hydrate_deleted_leases_tx(&mut tx, rescued_receipts) - .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) + .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) + .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 @@ -15219,15 +15436,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(&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(&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 @@ -15251,42 +15486,92 @@ 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_claims.deadline_at` rather than on a `leases` row, so - // the receipt-plane needs its own scan; merge both 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_claims.deadline_at` rather than on a `leases` row, so + // the receipt-plane needs its own scan; merge both 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).await?; - moved.extend( - self.hydrate_deleted_leases_tx(&mut tx, receipt_deleted) - .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) + .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) + .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 @@ -15311,15 +15596,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(&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(&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 @@ -15343,27 +15646,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).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) + .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) } /// Promote up to `batch_size` due `state` rows into the ready ring. diff --git a/awa-worker/src/completion.rs b/awa-worker/src/completion.rs index 87d3646ff..50fe63e38 100644 --- a/awa-worker/src/completion.rs +++ b/awa-worker/src/completion.rs @@ -350,17 +350,66 @@ 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 { + let mut committed_pairs = HashSet::new(); + 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 _ = request @@ -1005,4 +1054,203 @@ 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.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(); + 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 carried = jobs.remove(0); + 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(&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(&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(&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::new( + pool.clone(), + CancellationToken::new(), + crate::metrics::AwaMetrics::from_global(), + RuntimeStorage::QueueStorage(runtime), + 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 + ); + 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(&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(&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 a4b09290e..e989c4a4d 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -19748,15 +19748,13 @@ async fn historical_recovery_head_contention_rolls_back(receipts: bool, deadline .await .unwrap(); let result = tokio::time::timeout( - scaled_timeout(Duration::from_secs(2)), + 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:?}" - ); + .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(&format!( "SELECT count(*) FROM {schema}.done_entries WHERE job_id=$1" )) @@ -19837,3 +19835,339 @@ async fn test_historical_tombstoned_deadline_claim_completes_and_retires_copy() async fn test_historical_tombstoned_legacy_claim_completes_and_retires_copy() { historical_tombstoned_claim_completes(false, Duration::from_secs(300)).await; } + +#[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(&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) + 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(&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(&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"] { + sqlx::query(&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(&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(&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(&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(&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_receipt_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(&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; +}