diff --git a/CHANGELOG.md b/CHANGELOG.md index 065d50db1..bc3fbb6b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,29 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- Ordinary rescue ticks retain their cadence under scheduler jitter, while + backlog re-arms keep the heartbeat owner-evidence gate. + +- Keep confirmed-owner bulk heartbeat rescue's first cursor window small so + short passes can reach follow-up batches. Unconfirmed owners and + configurations that disable acceleration retain the wider initial scan. + +- Queue-storage rescue drains large abandoned-claim backlogs in bounded + transactions, controlled by `rescue_pass_budget` and a soft + `rescue_pass_max_rows` cap. Cursor sweeps continue past closed or freshly + heartbeated history and stop after a full traversal. Accelerated heartbeat + rescue requires settled runtime evidence; ordinary first batches keep their + configured cadence. Retries retain their backoffs, and due times can cluster: + the row cap bounds rescue writes, not promotion throughput. Zero for either + pass control keeps one batch per ordinary interval. See the + [backlog-drain configuration](docs/configuration.md#rescue-backlog). + +- Finished attempts keep heartbeating until their durable finalization finishes, + including while completion writes are stalled. Compact receipt heartbeats + address claims by batch identity and skip attempts a closer already holds, + so one stalled completion cannot block the rest of a heartbeat pass. Late + heartbeats cannot recreate state for a completed attempt. + - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 1a1c9097b..1b831c341 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -67,6 +67,8 @@ pub mod names { pub const COMPLETION_FLUSH_DURATION: &str = "awa.completion.flush_duration"; pub const HEARTBEAT_BATCHES: &str = "awa.heartbeat.batches"; pub const MAINTENANCE_RESCUES: &str = "awa.maintenance.rescues"; + pub const MAINTENANCE_RESCUE_PASS_SIZE: &str = "awa.maintenance.rescue.pass_size"; + pub const MAINTENANCE_RESCUE_OLDEST_AGE: &str = "awa.maintenance.rescue.oldest_attempt_age"; pub const MAINTENANCE_PROMOTE_BATCHES: &str = "awa.maintenance.promote_batches"; pub const MAINTENANCE_PROMOTE_BATCH_SIZE: &str = "awa.maintenance.promote_batch_size"; pub const MAINTENANCE_PROMOTE_DURATION: &str = "awa.maintenance.promote_duration"; @@ -198,6 +200,14 @@ pub struct AwaMetrics { pub heartbeat_batches: Counter, /// Total maintenance rescue operations. pub maintenance_rescues: Counter, + /// Attempts rescued per queue-storage rescue pass, attributed by + /// `awa.rescue.kind` and why the pass stopped (`awa.rescue.pass_end`). + pub maintenance_rescue_pass_size: Histogram, + /// Age of the oldest attempt the latest queue-storage rescue pass + /// rescued, per `awa.rescue.kind`; 0 when it rescued nothing. Rescue + /// works oldest claim first, so a value well above the heartbeat + /// staleness or deadline means a backlog of abandoned claims. + pub maintenance_rescue_oldest_attempt_age_seconds: Gauge, /// Total jobs parked for external callback. pub jobs_waiting_external: Counter, /// Current queue depth per state — how many jobs are in each state per queue. @@ -494,6 +504,19 @@ impl AwaMetrics { .with_description("Number of jobs rescued by maintenance") .with_unit("{job}") .build(), + maintenance_rescue_pass_size: meter + .u64_histogram(names::MAINTENANCE_RESCUE_PASS_SIZE) + .with_description("Attempts rescued per queue-storage rescue pass") + .with_unit("{job}") + .with_boundaries(COUNT_BUCKETS.to_vec()) + .build(), + maintenance_rescue_oldest_attempt_age_seconds: meter + .f64_gauge(names::MAINTENANCE_RESCUE_OLDEST_AGE) + .with_description( + "Age of the oldest attempt rescued by the latest queue-storage rescue pass", + ) + .with_unit("s") + .build(), jobs_waiting_external: meter .u64_counter(names::JOB_WAITING_EXTERNAL) .with_description("Number of jobs parked for external callback") @@ -894,6 +917,28 @@ impl AwaMetrics { .record(duration.as_secs_f64(), &attrs); } + /// Record one queue-storage rescue pass: how many attempts it rescued, + /// why it stopped, and the age of the oldest attempt it rescued. + pub fn record_rescue_pass( + &self, + kind: &'static str, + pass_end: &'static str, + rescued: u64, + oldest_attempt_age: Duration, + ) { + self.maintenance_rescue_pass_size.record( + rescued, + &[ + opentelemetry::KeyValue::new("awa.rescue.kind", kind), + opentelemetry::KeyValue::new("awa.rescue.pass_end", pass_end), + ], + ); + self.maintenance_rescue_oldest_attempt_age_seconds.record( + oldest_attempt_age.as_secs_f64(), + &[opentelemetry::KeyValue::new("awa.rescue.kind", kind)], + ); + } + /// Record a scheduled/retryable promotion batch. pub fn record_promotion_batch(&self, state: &str, batch_size: u64, duration: Duration) { let attrs = [opentelemetry::KeyValue::new( @@ -1504,7 +1549,7 @@ mod tests { &JOB_DURATION_BUCKETS_SECONDS, ), ]; - let counts: [(&str, &Histogram); 7] = [ + let counts: [(&str, &Histogram); 8] = [ (names::ENQUEUE_BATCH_SIZE, &metrics.enqueue_batch_size), ( names::DISPATCH_CAPACITY_AVAILABLE, @@ -1519,6 +1564,10 @@ mod tests { names::MAINTENANCE_PROMOTE_BATCH_SIZE, &metrics.promotion_batch_size, ), + ( + names::MAINTENANCE_RESCUE_PASS_SIZE, + &metrics.maintenance_rescue_pass_size, + ), ( names::MAINTENANCE_ROTATE_SKIPPED_ROWS, &metrics.maintenance_rotate_skipped_rows, diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 66a5dd4ee..e31b2153e 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -54,7 +54,8 @@ pub use partitioned_queue::{ }; pub use queue_storage::{ ClaimedEntry, ClaimedRuntimeJob, PruneDurations, PruneOutcome, QueueCounts, QueueStorage, - QueueStorageConfig, RingLedgerFoldOutcome, RingRotationCadence, RotateOutcome, SkipReason, + QueueStorageConfig, RescueBatch, RescueBatchOptions, RescueScanCursor, RescueScanProgress, + RescueSlotOrder, RingLedgerFoldOutcome, RingRotationCadence, RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, TerminalRollupFoldOutcome, }; pub use storage::StorageStatus; diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 339001705..102c4e61a 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -39,6 +39,7 @@ pub(crate) const COPY_MIN_ROWS: usize = 2048; const CLAIM_TRANSACTION_BEGIN: &str = "BEGIN; SET LOCAL synchronous_commit = off"; const TERMINAL_COUNTER_BUCKETS: i16 = 256; const RECEIPT_RESCUE_BATCH_LIMIT: i64 = 500; +const RESCUED_DUPLICATE_ERROR: &str = "rescued as duplicate: unique claim held by a newer job"; const DRAIN_BATCH_SIZE: i64 = 1_000; // The planner overestimates the drain statements enough to JIT-compile them, // which takes seconds against milliseconds of execution. @@ -599,6 +600,114 @@ pub struct PruneDurations { pub commit: Duration, } +/// Which claim slot a rescue batch visits first. A batch stops at its row +/// limit, so whichever slot it visits first gets most of a busy batch. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RescueSlotOrder { + /// The oldest initialized claim slot, then the rest. The claim ring + /// cannot rotate back into that slot until its claims are closed. + #[default] + OldestFirst, + /// Every other slot, starting at this offset (modulo their count), then + /// the oldest, so a long drain of the oldest slot does not starve the rest. + OthersFirst(usize), +} + +/// Options for one rescue batch ([`QueueStorage::rescue_stale_heartbeats_batch`] +/// and its deadline and callback counterparts). +#[derive(Debug, Clone, Default)] +pub struct RescueBatchOptions { + pub slot_order: RescueSlotOrder, + /// Claim slots to leave out, typically ones an earlier batch of the same + /// pass found nothing in ([`RescueBatch::idle_slots`]). + pub skip_slots: Vec, + /// Earliest `run_at` for a rescued attempt that goes back to `retryable`. + /// Its backoff still applies when that is later. + pub retry_not_before: Option>, + /// Maximum cursor-window rows (default 10,000): claims for the heartbeat + /// and row-local deadline walks, batches for compact deadline claims. + /// Follow-up windows can be smaller to reduce closure-evidence checks. + pub cursor_scan_limit: Option, +} + +/// The outcome of one rescue batch, which runs in one transaction. +#[derive(Debug, Default)] +pub struct RescueBatch { + pub rescued: Vec, + /// A scan window or rescue quota was exhausted, so another bounded + /// batch may find eligible attempts. + pub saturated: bool, + /// Claim slots whose bounded scans completed without eligible attempts. + pub idle_slots: Vec, + /// Last cursor reached by each bounded receipt scan. A caller continuing + /// batches can stop a sweep once its cursor returns to its first window. + pub scan_progress: Vec, +} + +/// Cursor progress from a bounded receipt rescue scan. +#[derive(Debug, Clone)] +pub struct RescueScanProgress { + pub slot: i32, + /// One cursor per representation, in a stable order. A slot is drained + /// only after every representation has completed its sweep. + pub cursors: Vec, +} + +/// Progress of one independently ordered receipt representation. +#[derive(Debug, Clone)] +pub struct RescueScanCursor { + pub position: Option<(DateTime, i64, i64)>, + pub saturated: bool, +} + +/// Receipt claims one rescue batch took, and what it found on the way. +#[derive(Default)] +struct ReceiptRescueBatch { + rescued: Vec, + saturated: bool, + idle_slots: Vec, + scan_progress: Vec, +} + +type ReceiptRescueScan = ( + Vec, + bool, + Option<(DateTime, i64, i64)>, +); +type ReceiptRescueScanRow = ( + sqlx::types::Json>, + bool, + Option>, + Option, + Option, +); + +/// Where a rescue sends an exhausted attempt on a queue whose DLQ is enabled. +struct RescueDlq<'a> { + reason: &'static str, + enabled: &'a (dyn Fn(&str) -> bool + Sync), +} + +/// Split rescued attempts by the state they left, keeping first-seen order, +/// because each terminal or deferred insert releases unique claims for one +/// previous state. +fn group_by_previous_state( + rows: Vec<(LeaseTransitionRow, T)>, +) -> Vec<(JobState, Vec<(LeaseTransitionRow, T)>)> { + let mut groups: Vec<(JobState, Vec<(LeaseTransitionRow, T)>)> = Vec::new(); + for (row, settled) in rows { + let state = row.state; + match groups + .iter_mut() + .find(|(group_state, _)| *group_state == state) + { + Some((_, group)) => group.push((row, settled)), + None => groups.push((state, vec![(row, settled)])), + } + } + groups +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub struct TerminalDeltaRollupOutcome { pub rolled_slots: usize, @@ -3124,7 +3233,7 @@ struct ExistingReadyRow { payload: serde_json::Value, } -#[derive(Debug, Clone, sqlx::FromRow)] +#[derive(Debug, Clone, serde::Deserialize, sqlx::FromRow)] struct DeletedLeaseRow { ready_slot: i32, ready_generation: i64, @@ -11783,6 +11892,47 @@ impl QueueStorage { Ok(()) } + /// Take the closers' attempt locks without waiting, returning the attempts + /// locked. A heartbeat must not queue behind a closer: an attempt whose lock + /// is held is being finalized or rescued right now, and rescue skips it for + /// as long as the lock is held. Taking the locks in a statement of their own + /// means the statement that then decides what is still open starts after + /// every closer that held a lock first has committed (#614), so a heartbeat + /// cannot recreate state for an attempt that has just been closed. + async fn try_lock_receipt_attempt_keys_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + jobs: &[(i64, i64)], + ) -> Result, AwaError> { + if jobs.is_empty() { + return Ok(Vec::new()); + } + let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); + let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); + sqlx::query_as::<_, (i64, i64)>( + r#" + WITH attempts AS MATERIALIZED ( + SELECT DISTINCT + input.job_id, + input.run_lease, + pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', input.job_id, input.run_lease), + 0 + ) AS lock_key + FROM unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease) + ) + SELECT job_id, run_lease + FROM attempts + WHERE pg_catalog.pg_try_advisory_xact_lock(lock_key) + "#, + ) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error) + } + async fn ensure_running_leases_from_receipts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, @@ -12086,7 +12236,6 @@ impl QueueStorage { WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) - FOR UPDATE OF claim_batches ), claim_refs AS ( SELECT * FROM row_claim_refs @@ -12135,6 +12284,19 @@ impl QueueStorage { .iter() .map(|(job_id, run_lease, _)| (*job_id, *run_lease)) .collect(); + let locked: HashSet<_> = self + .try_lock_receipt_attempt_keys_tx(tx, &receipt_jobs) + .await? + .into_iter() + .collect(); + let jobs: Vec<_> = jobs + .iter() + .filter(|(job_id, run_lease, _)| locked.contains(&(*job_id, *run_lease))) + .collect(); + if jobs.is_empty() { + return Ok(0); + } + let receipt_jobs: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect(); self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?; let schema = self.schema(); @@ -12203,7 +12365,6 @@ impl QueueStorage { WHERE dlq.job_id = items.job_id AND dlq.run_lease = items.run_lease ) - FOR UPDATE OF claim_batches ), claim_refs AS ( SELECT * FROM row_claim_refs @@ -13009,40 +13170,73 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, cutoff: DateTime, - ) -> Result, AwaError> { - let mut rescued = Vec::new(); + options: &RescueBatchOptions, + batch: &mut ReceiptRescueBatch, + ) -> Result<(), AwaError> { let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; - let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?; - - if let Some(slot) = preferred_slot { - let mut slot_rescued = self - .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) + for slot in self.rescue_slot_order_tx(tx, options).await? { + let (slot_rescued, scan_full, cursor) = self + .rescue_stale_receipt_claims_for_slot_tx( + tx, + slot, + cutoff, + options + .cursor_scan_limit + .map_or(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT, |limit| { + limit.max(1) as i64 + }), + remaining, + ) .await?; + batch.saturated |= scan_full; + if !scan_full && slot_rescued.is_empty() { + batch.idle_slots.push(slot); + } + batch.scan_progress.push(RescueScanProgress { + slot, + cursors: vec![RescueScanCursor { + position: cursor, + saturated: scan_full, + }], + }); remaining = remaining.saturating_sub(slot_rescued.len() as i64); - rescued.append(&mut slot_rescued); + batch.rescued.extend(slot_rescued); if remaining == 0 { - return Ok(rescued); + batch.saturated = true; + break; } } + Ok(()) + } - for slot in 0..self.claim_slot_count() { - let slot = slot as i32; - if Some(slot) == preferred_slot { - continue; + /// The claim slots one rescue batch visits, in order. Slots go from the + /// oldest initialized one to the current one, so older claims come first. + async fn rescue_slot_order_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + options: &RescueBatchOptions, + ) -> Result, AwaError> { + let (oldest, by_age) = self.claim_slots_by_age_tx(tx).await?; + let others: Vec = by_age + .into_iter() + .filter(|slot| Some(*slot) != oldest) + .collect(); + let mut slots = Vec::with_capacity(self.claim_slot_count()); + match options.slot_order { + RescueSlotOrder::OldestFirst => { + slots.extend(oldest); + slots.extend(others); } - - let mut slot_rescued = self - .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) - .await?; - remaining = remaining.saturating_sub(slot_rescued.len() as i64); - rescued.append(&mut slot_rescued); - - if remaining == 0 { - break; + RescueSlotOrder::OthersFirst(offset) => { + if !others.is_empty() { + let start = offset % others.len(); + slots.extend(others[start..].iter().chain(&others[..start])); + } + slots.extend(oldest); } } - - Ok(rescued) + slots.retain(|slot| !options.skip_slots.contains(slot)); + Ok(slots) } async fn rescue_stale_receipt_claims_for_slot_tx<'a>( @@ -13050,8 +13244,9 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, cutoff: DateTime, + scan_limit: i64, rescue_limit: i64, - ) -> Result, AwaError> { + ) -> Result { 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); @@ -13068,8 +13263,33 @@ impl QueueStorage { WHERE slot = $1 FOR UPDATE ), - claim_source AS MATERIALIZED ( - SELECT + -- The walk takes the first $3 claims after the cursor in + -- (claimed_at, job_id, run_lease) order, then wraps to the + -- first claims at or before it. Each side reads only the claims + -- it can return: row-local claims through the cursor index, and + -- compact batches (whose members share claimed_at) only up to + -- the batch at which $3 members are certainly covered. + after_batch_bound AS ( + SELECT min(running.claimed_at) AS claimed_at + FROM ( + SELECT + claim_batches.claimed_at, + sum(claim_batches.claimed_count) FILTER ( + WHERE claim_batches.claimed_at > cursor_row.rescue_cursor_claimed_at + ) OVER ( + ORDER BY claim_batches.claimed_at, claim_batches.batch_id + ) AS covered + FROM {claim_batch_child} AS claim_batches + CROSS JOIN cursor_row + WHERE claim_batches.claim_slot = $1 + AND claim_batches.claimed_at < $2 + AND claim_batches.claimed_at >= cursor_row.rescue_cursor_claimed_at + ) AS running + WHERE running.covered >= $3 + ), + after_source AS ( + ( + SELECT claims.claim_slot, NULL::bigint AS batch_id, claims.ready_slot, @@ -13086,8 +13306,19 @@ impl QueueStorage { claims.claimed_at, claims.closed_at, false AS compact_batch - FROM {claim_child} AS claims - WHERE claims.claim_slot = $1 + FROM {claim_child} AS claims + CROSS JOIN cursor_row + WHERE claims.claim_slot = $1 + AND claims.claimed_at < $2 + AND (claims.claimed_at, claims.job_id, claims.run_lease) + > ( + cursor_row.rescue_cursor_claimed_at, + cursor_row.rescue_cursor_job_id, + cursor_row.rescue_cursor_run_lease + ) + ORDER BY claims.claimed_at, claims.job_id, claims.run_lease + LIMIT $3 + ) UNION ALL SELECT claim_batches.claim_slot, @@ -13107,6 +13338,8 @@ impl QueueStorage { NULL::timestamptz AS closed_at, true AS compact_batch FROM {claim_batch_child} AS claim_batches + CROSS JOIN cursor_row + CROSS JOIN after_batch_bound CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, @@ -13116,119 +13349,197 @@ impl QueueStorage { claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) WHERE claim_batches.claim_slot = $1 - ), - after_cursor AS MATERIALIZED ( - SELECT - claim_source.claim_slot, - claim_source.job_id, - claim_source.run_lease, - row_number() OVER ( - ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease - ) AS rn - FROM claim_source - CROSS JOIN cursor_row - WHERE claim_source.claimed_at < $2 - AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease) + AND claim_batches.claimed_at < $2 + AND claim_batches.claimed_at >= cursor_row.rescue_cursor_claimed_at + AND claim_batches.claimed_at + <= COALESCE(after_batch_bound.claimed_at, 'infinity'::timestamptz) + AND (claim_batches.claimed_at, items.job_id, items.run_lease) > ( cursor_row.rescue_cursor_claimed_at, cursor_row.rescue_cursor_job_id, cursor_row.rescue_cursor_run_lease ) - ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease + ), + after_cursor AS MATERIALIZED ( + SELECT + after_source.*, + row_number() OVER ( + ORDER BY after_source.claimed_at, after_source.job_id, after_source.run_lease + ) AS rn + FROM after_source + ORDER BY after_source.claimed_at, after_source.job_id, after_source.run_lease LIMIT $3 ), - after_count AS ( + after_count AS MATERIALIZED ( SELECT count(*)::bigint AS count FROM after_cursor ), - before_cursor AS MATERIALIZED ( + before_batch_bound AS ( + SELECT min(running.claimed_at) AS claimed_at + FROM ( + SELECT + claim_batches.claimed_at, + sum(claim_batches.claimed_count) FILTER ( + WHERE claim_batches.claimed_at < cursor_row.rescue_cursor_claimed_at + ) OVER ( + ORDER BY claim_batches.claimed_at, claim_batches.batch_id + ) AS covered + FROM {claim_batch_child} AS claim_batches + CROSS JOIN cursor_row + WHERE claim_batches.claim_slot = $1 + AND claim_batches.claimed_at < $2 + AND claim_batches.claimed_at <= cursor_row.rescue_cursor_claimed_at + ) AS running + WHERE running.covered >= $3 + ), + before_source AS ( + ( + SELECT + claims.claim_slot, + NULL::bigint AS batch_id, + claims.ready_slot, + claims.ready_generation, + claims.job_id, + claims.queue, + claims.priority, + claims.attempt, + claims.run_lease, + claims.max_attempts, + claims.lane_seq, + claims.enqueue_shard, + claims.receipt_id, + claims.claimed_at, + claims.closed_at, + false AS compact_batch + FROM {claim_child} AS claims + CROSS JOIN cursor_row + WHERE (SELECT count FROM after_count) < $3 + AND claims.claim_slot = $1 + AND claims.claimed_at < $2 + AND (claims.claimed_at, claims.job_id, claims.run_lease) + <= ( + cursor_row.rescue_cursor_claimed_at, + cursor_row.rescue_cursor_job_id, + cursor_row.rescue_cursor_run_lease + ) + ORDER BY claims.claimed_at, claims.job_id, claims.run_lease + LIMIT $3 + ) + UNION ALL SELECT - claim_source.claim_slot, - claim_source.job_id, - claim_source.run_lease, - after_count.count + row_number() OVER ( - ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease - ) AS rn - FROM claim_source + claim_batches.claim_slot, + claim_batches.batch_id, + claim_batches.ready_slot, + claim_batches.ready_generation, + items.job_id, + claim_batches.queue, + claim_batches.priority, + items.attempt, + items.run_lease, + items.max_attempts, + items.lane_seq, + claim_batches.enqueue_shard, + items.receipt_id, + claim_batches.claimed_at, + NULL::timestamptz AS closed_at, + true AS compact_batch + FROM {claim_batch_child} AS claim_batches CROSS JOIN cursor_row - CROSS JOIN after_count - WHERE after_count.count < $3 - AND claim_source.claimed_at < $2 - AND (claim_source.claimed_at, claim_source.job_id, claim_source.run_lease) + CROSS JOIN before_batch_bound + CROSS JOIN LATERAL unnest( + claim_batches.job_ids, + claim_batches.run_leases, + claim_batches.receipt_ids, + claim_batches.lane_seqs, + claim_batches.attempts, + claim_batches.max_attempts + ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) + WHERE (SELECT count FROM after_count) < $3 + AND claim_batches.claim_slot = $1 + AND claim_batches.claimed_at < $2 + AND claim_batches.claimed_at <= cursor_row.rescue_cursor_claimed_at + AND claim_batches.claimed_at + <= COALESCE(before_batch_bound.claimed_at, 'infinity'::timestamptz) + AND (claim_batches.claimed_at, items.job_id, items.run_lease) <= ( cursor_row.rescue_cursor_claimed_at, cursor_row.rescue_cursor_job_id, cursor_row.rescue_cursor_run_lease ) - ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease + ), + before_cursor AS MATERIALIZED ( + SELECT + before_source.*, + after_count.count + row_number() OVER ( + ORDER BY before_source.claimed_at, before_source.job_id, before_source.run_lease + ) AS rn + FROM before_source + CROSS JOIN after_count + ORDER BY before_source.claimed_at, before_source.job_id, before_source.run_lease LIMIT (SELECT GREATEST($3 - count, 0) FROM after_count) ), - candidate_keys AS MATERIALIZED ( - SELECT claim_slot, job_id, run_lease, rn FROM after_cursor + candidate_rows AS MATERIALIZED ( + SELECT * FROM after_cursor UNION ALL - SELECT claim_slot, job_id, run_lease, rn FROM before_cursor + SELECT * FROM before_cursor ), candidates AS MATERIALIZED ( SELECT - claim_source.claim_slot, - claim_source.batch_id, - claim_source.ready_slot, - claim_source.ready_generation, - claim_source.job_id, - claim_source.queue, - claim_source.priority, - claim_source.attempt, - claim_source.run_lease, - claim_source.max_attempts, - claim_source.lane_seq, - claim_source.enqueue_shard, - claim_source.receipt_id, - claim_source.claimed_at, - claim_source.closed_at, - claim_source.compact_batch, - COALESCE(attempt.heartbeat_at, claim_source.claimed_at) < $2 AS is_stale, + candidate_rows.claim_slot, + candidate_rows.batch_id, + candidate_rows.ready_slot, + candidate_rows.ready_generation, + candidate_rows.job_id, + candidate_rows.queue, + candidate_rows.priority, + candidate_rows.attempt, + candidate_rows.run_lease, + candidate_rows.max_attempts, + candidate_rows.lane_seq, + candidate_rows.enqueue_shard, + candidate_rows.receipt_id, + candidate_rows.claimed_at, + candidate_rows.closed_at, + candidate_rows.compact_batch, + COALESCE(attempt.heartbeat_at, candidate_rows.claimed_at) < $2 AS is_stale, ( - claim_source.closed_at IS NOT NULL + candidate_rows.closed_at IS NOT NULL OR EXISTS ( SELECT 1 FROM {closure_child} AS closures - WHERE closures.claim_slot = claim_source.claim_slot - AND closures.job_id = claim_source.job_id - AND closures.run_lease = claim_source.run_lease + WHERE closures.claim_slot = candidate_rows.claim_slot + AND closures.job_id = candidate_rows.job_id + AND closures.run_lease = candidate_rows.run_lease ) OR EXISTS ( SELECT 1 FROM {closure_batch_child} AS closure_batches - WHERE closure_batches.receipt_ranges @> claim_source.receipt_id + WHERE closure_batches.receipt_ranges @> candidate_rows.receipt_id ) OR EXISTS ( SELECT 1 FROM {schema}.done_entries AS done - WHERE done.job_id = claim_source.job_id - AND done.run_lease = claim_source.run_lease + WHERE done.job_id = candidate_rows.job_id + AND done.run_lease = candidate_rows.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.deferred_jobs AS deferred - WHERE deferred.job_id = claim_source.job_id - AND deferred.run_lease = claim_source.run_lease + WHERE deferred.job_id = candidate_rows.job_id + AND deferred.run_lease = candidate_rows.run_lease ) OR EXISTS ( SELECT 1 FROM {schema}.dlq_entries AS dlq - WHERE dlq.job_id = claim_source.job_id - AND dlq.run_lease = claim_source.run_lease + WHERE dlq.job_id = candidate_rows.job_id + AND dlq.run_lease = candidate_rows.run_lease ) ) AS is_closed, EXISTS ( SELECT 1 FROM {schema}.leases AS lease - WHERE lease.job_id = claim_source.job_id - AND lease.run_lease = claim_source.run_lease + WHERE lease.job_id = candidate_rows.job_id + AND lease.run_lease = candidate_rows.run_lease ) AS is_lease_managed, - candidate_keys.rn - FROM candidate_keys - JOIN claim_source - ON claim_source.claim_slot = candidate_keys.claim_slot - AND claim_source.job_id = candidate_keys.job_id - AND claim_source.run_lease = candidate_keys.run_lease + candidate_rows.rn + FROM candidate_rows LEFT JOIN {schema}.attempt_state AS attempt - ON attempt.job_id = claim_source.job_id - AND attempt.run_lease = claim_source.run_lease + ON attempt.job_id = candidate_rows.job_id + AND attempt.run_lease = candidate_rows.run_lease ), stale_candidates AS ( SELECT candidates.* @@ -13261,14 +13572,14 @@ impl QueueStorage { ))) .bind(slot) .bind(cutoff) - .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? .into_iter() .unzip(); - let rescued: Vec = sqlx::query_as(audited_sql(format!( + let (rescued, scan_full, cursor_at, cursor_job_id, cursor_run_lease): ReceiptRescueScanRow = sqlx::query_as(audited_sql(format!( r#" WITH {candidate_ctes}, stale_row_locked AS ( @@ -13464,11 +13775,9 @@ impl QueueStorage { rescue_cursor_run_lease = advance_target.run_lease FROM advance_target WHERE slots.slot = $1 - RETURNING slots.slot + RETURNING slots.rescue_cursor_claimed_at, slots.rescue_cursor_job_id, slots.rescue_cursor_run_lease ), - cursor_advance AS ( - SELECT count(*) FROM advance_cursor - ) + rescued_rows AS ( SELECT stale_locked.ready_slot, stale_locked.ready_generation, @@ -13487,19 +13796,33 @@ impl QueueStorage { ON closed_locked.claim_slot = stale_locked.claim_slot AND closed_locked.job_id = stale_locked.job_id AND closed_locked.run_lease = stale_locked.run_lease - CROSS JOIN cursor_advance + ) + SELECT + COALESCE((SELECT jsonb_agg(to_jsonb(rescued_rows)) FROM rescued_rows), '[]'::jsonb), + (SELECT count(*) >= $3 FROM candidates), + (SELECT rescue_cursor_claimed_at FROM advance_cursor), + (SELECT rescue_cursor_job_id FROM advance_cursor), + (SELECT rescue_cursor_run_lease FROM advance_cursor) "# ))) .bind(slot) .bind(cutoff) - .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .bind(&locked_job_ids) .bind(&locked_run_leases) - .fetch_all(tx.as_mut()) + .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; - Ok(rescued) + let scan_full = scan_full || rescued.0.len() as i64 >= rescue_limit; + Ok(( + rescued.0, + scan_full, + cursor_at + .zip(cursor_job_id) + .zip(cursor_run_lease) + .map(|((at, id), lease)| (at, id, lease)), + )) } /// Receipt-side counterpart to `rescue_expired_deadlines`: scans both @@ -13531,51 +13854,61 @@ impl QueueStorage { async fn rescue_expired_receipt_deadlines_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, - ) -> Result, AwaError> { - let mut rescued = Vec::new(); + options: &RescueBatchOptions, + batch: &mut ReceiptRescueBatch, + ) -> Result<(), AwaError> { let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; - let preferred_slot = self.oldest_initialized_claim_slot_tx(tx).await?; - - let mut slots: Vec = Vec::with_capacity(self.claim_slot_count() + 1); - if let Some(slot) = preferred_slot { - slots.push(slot); - } - for slot in 0..self.claim_slot_count() { - let slot = slot as i32; - if Some(slot) != preferred_slot { - slots.push(slot); - } - } - - for slot in slots { - let mut slot_rescued = self - .rescue_expired_batch_deadlines_for_slot_tx(tx, slot, remaining) + for slot in self.rescue_slot_order_tx(tx, options).await? { + let scan_limit = options + .cursor_scan_limit + .map_or(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT, |limit| { + limit.max(1) as i64 + }); + let (compact, compact_full, compact_cursor) = self + .rescue_expired_batch_deadlines_for_slot_tx(tx, slot, scan_limit, remaining) .await?; - remaining = remaining.saturating_sub(slot_rescued.len() as i64); - rescued.append(&mut slot_rescued); - if remaining == 0 { - break; - } - - let mut slot_rescued = self - .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining) + remaining = remaining.saturating_sub(compact.len() as i64); + let compact_found = !compact.is_empty(); + batch.rescued.extend(compact); + batch.saturated |= compact_full; + let (row_local, row_full, row_cursor) = self + .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, scan_limit, remaining) .await?; - remaining = remaining.saturating_sub(slot_rescued.len() as i64); - rescued.append(&mut slot_rescued); + batch.saturated |= row_full; + batch.scan_progress.push(RescueScanProgress { + slot, + cursors: vec![ + RescueScanCursor { + position: compact_cursor, + saturated: compact_full, + }, + RescueScanCursor { + position: row_cursor, + saturated: row_full, + }, + ], + }); + if !compact_found && row_local.is_empty() && !compact_full && !row_full { + batch.idle_slots.push(slot); + } + remaining = remaining.saturating_sub(row_local.len() as i64); + batch.rescued.extend(row_local); if remaining == 0 { + batch.saturated = true; break; } } - - Ok(rescued) + Ok(()) } - async fn oldest_initialized_claim_slot_tx<'a>( + /// The oldest initialized claim slot, and every slot ordered from it + /// around the ring to the current one, then any not yet initialized. + async fn claim_slots_by_age_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, - ) -> Result, AwaError> { + ) -> Result<(Option, Vec), AwaError> { let schema = self.schema(); - let preferred_slot = sqlx::query_as::<_, (i32, i64, i32)>(audited_sql(format!( + let ring_cursor = sqlx::query_as::<_, (i32, i64, i32)>(audited_sql(format!( r#" SELECT ledger.slot, ledger.generation, state.slot_count FROM {schema}.claim_ring_state AS state @@ -13590,24 +13923,40 @@ impl QueueStorage { ))) .fetch_optional(tx.as_mut()) .await - .map_err(map_sqlx_error)? - .and_then(|(current_slot, generation, slot_count)| { - oldest_initialized_ring_slot(current_slot, generation, slot_count) - .map(|(slot, _generation)| slot) - .filter(|slot| *slot >= 0 && (*slot as usize) < self.claim_slot_count()) - }); - - Ok(preferred_slot) + .map_err(map_sqlx_error)?; + let slot_count = self.claim_slot_count() as i32; + let Some((current_slot, generation, ring_slot_count)) = ring_cursor else { + return Ok((None, (0..slot_count).collect())); + }; + let oldest = oldest_initialized_ring_slot(current_slot, generation, ring_slot_count) + .map(|(slot, _generation)| slot) + .filter(|slot| (0..slot_count).contains(slot)); + let start = oldest.unwrap_or(current_slot).rem_euclid(slot_count); + let mut by_age: Vec = Vec::with_capacity(slot_count as usize); + let mut slot = start; + loop { + by_age.push(slot); + if slot == current_slot.rem_euclid(slot_count) || by_age.len() == slot_count as usize { + break; + } + slot = (slot + 1).rem_euclid(slot_count); + } + let uninitialized: Vec = (0..slot_count) + .filter(|slot| !by_age.contains(slot)) + .collect(); + by_age.extend(uninitialized); + Ok((oldest, by_age)) } async fn rescue_expired_receipt_deadlines_for_slot_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, + scan_limit: i64, rescue_limit: i64, - ) -> Result, AwaError> { + ) -> Result { if rescue_limit <= 0 { - return Ok(Vec::new()); + return Ok((Vec::new(), true, None)); } let schema = self.schema(); @@ -13737,14 +14086,14 @@ impl QueueStorage { "# ))) .bind(slot) - .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? .into_iter() .unzip(); - let rescued: Vec = sqlx::query_as(audited_sql(format!( + let (rescued, scan_full, cursor_at, cursor_job_id, cursor_run_lease): ReceiptRescueScanRow = sqlx::query_as(audited_sql(format!( r#" WITH {candidate_ctes}, expired_locked AS ( @@ -13822,11 +14171,9 @@ impl QueueStorage { deadline_cursor_run_lease = advance_target.run_lease FROM advance_target WHERE slots.slot = $1 - RETURNING slots.slot + RETURNING slots.deadline_cursor_deadline_at, slots.deadline_cursor_job_id, slots.deadline_cursor_run_lease ), - cursor_advance AS ( - SELECT count(*) FROM advance_cursor - ) + rescued_rows AS ( SELECT expired_locked.ready_slot, expired_locked.ready_generation, @@ -13845,18 +14192,32 @@ impl QueueStorage { ON inserted.claim_slot = expired_locked.claim_slot AND inserted.job_id = expired_locked.job_id AND inserted.run_lease = expired_locked.run_lease - CROSS JOIN cursor_advance + ) + SELECT + COALESCE((SELECT jsonb_agg(to_jsonb(rescued_rows)) FROM rescued_rows), '[]'::jsonb), + (SELECT count(*) >= $2 FROM candidates), + (SELECT deadline_cursor_deadline_at FROM advance_cursor), + (SELECT deadline_cursor_job_id FROM advance_cursor), + (SELECT deadline_cursor_run_lease FROM advance_cursor) "# ))) .bind(slot) - .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .bind(&locked_job_ids) .bind(&locked_run_leases) - .fetch_all(tx.as_mut()) + .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; - Ok(rescued) + let scan_full = scan_full || rescued.0.len() as i64 >= rescue_limit; + Ok(( + rescued.0, + scan_full, + cursor_at + .zip(cursor_job_id) + .zip(cursor_run_lease) + .map(|((at, id), lease)| (at, id, lease)), + )) } /// Compact-batch counterpart of @@ -13883,10 +14244,11 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, + scan_limit: i64, rescue_limit: i64, - ) -> Result, AwaError> { + ) -> Result { if rescue_limit <= 0 { - return Ok(Vec::new()); + return Ok((Vec::new(), true, None)); } let schema = self.schema(); @@ -14048,14 +14410,14 @@ impl QueueStorage { "# ))) .bind(slot) - .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? .into_iter() .unzip(); - let rescued: Vec = sqlx::query_as(audited_sql(format!( + let (rescued, scan_full, cursor_at, cursor_job_id, cursor_run_lease): ReceiptRescueScanRow = sqlx::query_as(audited_sql(format!( r#" WITH {candidate_ctes}, expired_locked AS ( @@ -14175,11 +14537,9 @@ impl QueueStorage { batch_deadline_cursor_batch_id = advance_target.batch_id FROM advance_target WHERE slots.slot = $1 - RETURNING slots.slot + RETURNING slots.batch_deadline_cursor_deadline_at, slots.batch_deadline_cursor_batch_id, 0::bigint AS batch_deadline_cursor_run_lease ), - cursor_advance AS ( - SELECT count(*) FROM advance_cursor - ) + rescued_rows AS ( SELECT expired_locked.ready_slot, expired_locked.ready_generation, @@ -14198,18 +14558,32 @@ impl QueueStorage { ON closed_locked.claim_slot = expired_locked.claim_slot AND closed_locked.job_id = expired_locked.job_id AND closed_locked.run_lease = expired_locked.run_lease - CROSS JOIN cursor_advance + ) + SELECT + COALESCE((SELECT jsonb_agg(to_jsonb(rescued_rows)) FROM rescued_rows), '[]'::jsonb), + (SELECT count(*) >= $2 FROM candidate_keys), + (SELECT batch_deadline_cursor_deadline_at FROM advance_cursor), + (SELECT batch_deadline_cursor_batch_id FROM advance_cursor), + (SELECT batch_deadline_cursor_run_lease FROM advance_cursor) "# ))) .bind(slot) - .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .bind(&locked_job_ids) .bind(&locked_run_leases) - .fetch_all(tx.as_mut()) + .fetch_one(tx.as_mut()) .await .map_err(map_sqlx_error)?; - Ok(rescued) + let scan_full = scan_full || rescued.0.len() as i64 >= rescue_limit; + Ok(( + rescued.0, + scan_full, + cursor_at + .zip(cursor_job_id) + .zip(cursor_run_lease) + .map(|((at, id), lease)| (at, id, lease)), + )) } pub async fn load_job(&self, pool: &PgPool, job_id: i64) -> Result, AwaError> { @@ -15628,7 +16002,8 @@ impl QueueStorage { /// Heartbeat compact receipt claims by their batch identity. Each claim /// batch is addressed by primary key and the attempt by item index, with /// index-only closure probes and no batch row lock, so the cost is - /// O(in-flight attempts) rather than O(open claim batches). + /// O(in-flight attempts) rather than O(open claim batches). Attempts that + /// a closer holds right now are skipped rather than waited for. pub async fn heartbeat_claims_batch( &self, pool: &PgPool, @@ -15643,7 +16018,19 @@ impl QueueStorage { .iter() .map(|claim| (claim.job_id, claim.run_lease)) .collect(); - self.lock_receipt_attempt_keys_tx(&mut tx, &pairs).await?; + let locked: HashSet<(i64, i64)> = self + .try_lock_receipt_attempt_keys_tx(&mut tx, &pairs) + .await? + .into_iter() + .collect(); + let claims: Vec<&ClaimHeartbeatRef> = claims + .iter() + .filter(|claim| locked.contains(&(claim.job_id, claim.run_lease))) + .collect(); + if claims.is_empty() { + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(0); + } let job_ids: Vec = claims.iter().map(|claim| claim.job_id).collect(); let run_leases: Vec = claims.iter().map(|claim| claim.run_lease).collect(); let claim_slots: Vec = claims.iter().map(|claim| claim.claim_slot).collect(); @@ -15715,6 +16102,38 @@ impl QueueStorage { Ok(updated as usize) } + /// Lock legacy lease rows without waiting behind a finalizer. Keep this + /// separate from the updating statement: the latter must use a fresh + /// snapshot and must not mix partition pruning with row locking (PG18). + async fn lock_live_leases_for_heartbeat_tx( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + jobs: &[(i64, i64)], + ) -> Result, AwaError> { + if jobs.is_empty() { + return Ok(Vec::new()); + } + let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); + let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); + sqlx::query_as::<_, (i64, i64)>(audited_sql(format!( + r#" + SELECT lease.job_id, lease.run_lease + FROM {table} AS lease + JOIN unnest($1::bigint[], $2::bigint[]) AS input(job_id, run_lease) + ON lease.job_id = input.job_id AND lease.run_lease = input.run_lease + WHERE lease.state = 'running' + ORDER BY lease.job_id, lease.run_lease + FOR UPDATE OF lease SKIP LOCKED + "#, + table = self.leases_table(), + ))) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error) + } + pub async fn heartbeat_batch( &self, pool: &PgPool, @@ -15735,14 +16154,18 @@ impl QueueStorage { // dead entries per write under sustained churn). Compat // reads in `LeaseJobRow` SELECTs and the `awa.jobs` view // COALESCE attempt_state.heartbeat_at first. + let locked = self.try_lock_receipt_attempt_keys_tx(&mut tx, jobs).await?; updated += self - .upsert_attempt_state_from_receipts_tx(&mut tx, jobs) + .upsert_attempt_state_from_receipts_tx(&mut tx, &locked) .await?; } else { // Legacy non-receipts mode (custom schemas with // `lease_claim_receipts=FALSE`): the `leases.heartbeat_at` // write is still the heartbeat home, since there is no // upsert_attempt_state path firing. + let jobs = self + .lock_live_leases_for_heartbeat_tx(&mut tx, jobs) + .await?; let job_ids: Vec = jobs.iter().map(|(job_id, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease)| *run_lease).collect(); let result = sqlx::query(audited_sql(format!( @@ -15793,9 +16216,19 @@ impl QueueStorage { self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs) .await? } else { - // Legacy non-receipts mode keeps the old CTE shape that - // updates leases.heartbeat_at + leases.progress - // (via attempt_state upsert) in a single round-trip. + // Prelock independently so a finalizer cannot stall the entire + // heartbeat pass. Skipped progress snapshots remain unacknowledged + // in the worker and are retried on the next pass. + let keys: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect(); + let locked: HashSet<_> = self + .lock_live_leases_for_heartbeat_tx(&mut tx, &keys) + .await? + .into_iter() + .collect(); + let jobs: Vec<_> = jobs + .iter() + .filter(|(id, lease, _)| locked.contains(&(*id, *lease))) + .collect(); let schema = self.schema(); let job_ids: Vec = jobs.iter().map(|(job_id, _, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect(); @@ -16174,59 +16607,198 @@ impl QueueStorage { .await } - /// Settle one attempt that maintenance rescued (stale heartbeat, expired + /// Settle attempts that maintenance rescued (stale heartbeat, expired /// deadline or callback timeout): back to `retryable` with backoff while /// attempts remain, `failed` once the rescued attempt was the last. A job /// that kills or hangs its worker every time therefore stops after /// `max_attempts`, like one whose handler returns an error. /// - /// An exhausted attempt on a queue whose DLQ is enabled (`dlq_reason` - /// is `Some`) goes straight to `dlq_entries` in this transaction, as - /// `fail_to_dlq` does for a handler error, so there is no window in - /// which it sits `failed` outside the DLQ. - async fn settle_rescued_attempt_tx( + /// An exhausted attempt on a queue whose DLQ is enabled goes straight to + /// `dlq_entries` in this transaction, as `fail_to_dlq` does for a handler + /// error, so there is no window in which it sits `failed` outside the DLQ. + /// + /// Each outcome is written in one statement per previous state. If that + /// statement hits a unique-claim conflict, the group is retried row by + /// row so only the conflicting attempts become cancellations. + async fn settle_rescued_attempts_tx( &self, tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, - row: LeaseTransitionRow, + rows: Vec, error: &str, - dlq_reason: Option<&str>, - ) -> Result { - let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state( - row.payload.clone(), - row.progress.clone(), - )?)?; - let exhausted = row.attempt >= row.max_attempts; - payload.push_error(lifecycle_error(error, row.attempt, exhausted)); - let duplicate_error = "rescued as duplicate: unique claim held by a newer job"; - if exhausted { - let finalized_at = Utc::now(); - if let Some(dlq_reason) = dlq_reason { - let dlq = row.clone().into_dlq_row( + dlq: RescueDlq<'_>, + retry_not_before: Option>, + ) -> Result, AwaError> { + if rows.is_empty() { + return Ok(Vec::new()); + } + + let (retry_attempts, retry_max_attempts): (Vec, Vec) = rows + .iter() + .filter(|row| row.attempt < row.max_attempts) + .map(|row| (row.attempt, row.max_attempts)) + .unzip(); + let mut retry_run_ats = self + .backoffs_at_tx(tx, &retry_attempts, &retry_max_attempts) + .await? + .into_iter(); + + let finalized_at = Utc::now(); + let mut retries: Vec<(LeaseTransitionRow, DeferredJobRow)> = Vec::new(); + let mut failures: Vec<(LeaseTransitionRow, DoneJobRow)> = Vec::new(); + let mut dead_letters: Vec<(LeaseTransitionRow, DlqJobRow)> = Vec::new(); + for row in rows { + let mut payload = RuntimePayload::from_json(Self::payload_with_attempt_state( + row.payload.clone(), + row.progress.clone(), + )?)?; + let exhausted = row.attempt >= row.max_attempts; + payload.push_error(lifecycle_error(error, row.attempt, exhausted)); + if !exhausted { + let backoff_at = retry_run_ats.next().ok_or_else(|| { + AwaError::Validation("rescue backoff missing for a retryable attempt".into()) + })?; + let run_at = retry_not_before.map_or(backoff_at, |floor| floor.max(backoff_at)); + let deferred = row.clone().into_deferred_row( + JobState::Retryable, + run_at, + Some(finalized_at), + payload.into_json(), + ); + retries.push((row, deferred)); + } else if (dlq.enabled)(&row.queue) { + let dead_letter = row.clone().into_dlq_row( finalized_at, payload.into_json(), - dlq_reason.to_string(), + dlq.reason.to_string(), finalized_at, ); - return self - .insert_rescued_dlq_or_cancel_tx(tx, row, dlq, duplicate_error) - .await; + dead_letters.push((row, dead_letter)); + } else { + let done = + row.clone() + .into_done_row(JobState::Failed, finalized_at, payload.into_json()); + failures.push((row, done)); } - let done = - row.clone() - .into_done_row(JobState::Failed, finalized_at, payload.into_json()); - self.insert_rescued_done_or_cancel_tx(tx, row, done, duplicate_error) + } + + let mut settled = Vec::with_capacity(retries.len() + failures.len() + dead_letters.len()); + for (old_state, group) in group_by_previous_state(retries) { + let deferred: Vec = group.iter().map(|(_, row)| row.clone()).collect(); + let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; + match self + .insert_deferred_rows_tx(&mut savepoint, deferred, Some(old_state)) .await - } else { - let deferred = row.clone().into_deferred_row( - JobState::Retryable, - self.backoff_at_tx(tx, row.attempt, row.max_attempts) - .await?, - Some(Utc::now()), - payload.into_json(), - ); - self.insert_rescued_deferred_or_cancel_tx(tx, row, deferred, duplicate_error) + { + Ok(_) => { + savepoint.commit().await.map_err(map_sqlx_error)?; + for (_, deferred) in group { + settled.push(deferred.into_job_row()?); + } + } + Err(AwaError::UniqueConflict { .. }) => { + savepoint.rollback().await.map_err(map_sqlx_error)?; + for (row, deferred) in group { + settled.push( + self.insert_rescued_deferred_or_cancel_tx( + tx, + row, + deferred, + RESCUED_DUPLICATE_ERROR, + ) + .await?, + ); + } + } + Err(err) => return Err(err), + } + } + for (old_state, group) in group_by_previous_state(failures) { + let done: Vec = group.iter().map(|(_, row)| row.clone()).collect(); + let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; + match self + .insert_done_rows_tx(&mut savepoint, &done, Some(old_state)) .await + { + Ok(_) => { + savepoint.commit().await.map_err(map_sqlx_error)?; + for done in done { + settled.push(done.into_job_row()?); + } + } + Err(AwaError::UniqueConflict { .. }) => { + savepoint.rollback().await.map_err(map_sqlx_error)?; + for (row, done) in group { + settled.push( + self.insert_rescued_done_or_cancel_tx( + tx, + row, + done, + RESCUED_DUPLICATE_ERROR, + ) + .await?, + ); + } + } + Err(err) => return Err(err), + } } + for (old_state, group) in group_by_previous_state(dead_letters) { + let dead_letters: Vec = group.iter().map(|(_, row)| row.clone()).collect(); + let mut savepoint = tx.begin().await.map_err(map_sqlx_error)?; + match self + .insert_dlq_rows_tx(&mut savepoint, &dead_letters, Some(old_state)) + .await + { + Ok(_) => { + savepoint.commit().await.map_err(map_sqlx_error)?; + for dead_letter in dead_letters { + settled.push(dead_letter.into_job_row()?); + } + } + Err(AwaError::UniqueConflict { .. }) => { + savepoint.rollback().await.map_err(map_sqlx_error)?; + for (row, dead_letter) in group { + settled.push( + self.insert_rescued_dlq_or_cancel_tx( + tx, + row, + dead_letter, + RESCUED_DUPLICATE_ERROR, + ) + .await?, + ); + } + } + Err(err) => return Err(err), + } + } + Ok(settled) + } + + /// `clock_timestamp() + awa.backoff_duration(..)` for each attempt, in + /// one round trip, in input order. + async fn backoffs_at_tx( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + attempts: &[i16], + max_attempts: &[i16], + ) -> Result>, AwaError> { + if attempts.is_empty() { + return Ok(Vec::new()); + } + sqlx::query_scalar( + r#" + SELECT clock_timestamp() + awa.backoff_duration(attempts.attempt, attempts.max_attempts) + FROM unnest($1::smallint[], $2::smallint[]) + WITH ORDINALITY AS attempts(attempt, max_attempts, ord) + ORDER BY attempts.ord + "#, + ) + .bind(attempts) + .bind(max_attempts) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error) } /// Same fallback for rescue arms that write a terminal row directly @@ -17106,16 +17678,36 @@ impl QueueStorage { /// [`Self::rescue_stale_heartbeats`], moving an exhausted attempt into /// `dlq_entries` (reason `heartbeat_stale`) in the rescue transaction /// when `dlq_enabled` returns true for its queue. + pub async fn rescue_stale_heartbeats_with_dlq( + &self, + pool: &PgPool, + staleness: Duration, + dlq_enabled: &(dyn Fn(&str) -> bool + Sync), + ) -> Result, AwaError> { + Ok(self + .rescue_stale_heartbeats_batch( + pool, + staleness, + dlq_enabled, + &RescueBatchOptions::default(), + ) + .await? + .rescued) + } + + /// One stale-heartbeat rescue batch in its own transaction: up to + /// 500 lease rows and 500 receipt claims. #[tracing::instrument( - skip(self, pool, dlq_enabled), + skip(self, pool, dlq_enabled, options), name = "queue_storage.rescue_stale_heartbeats" )] - pub async fn rescue_stale_heartbeats_with_dlq( + pub async fn rescue_stale_heartbeats_batch( &self, pool: &PgPool, staleness: Duration, dlq_enabled: &(dyn Fn(&str) -> bool + Sync), - ) -> Result, AwaError> { + options: &RescueBatchOptions, + ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let cutoff = Utc::now() @@ -17158,7 +17750,7 @@ impl QueueStorage { WHERE lease.state = 'running' AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1 ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC - LIMIT 500 + LIMIT $2 FOR UPDATE OF lease SKIP LOCKED ) RETURNING @@ -17181,19 +17773,30 @@ impl QueueStorage { "# ))) .bind(cutoff) + .bind(RECEIPT_RESCUE_BATCH_LIMIT) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; + let mut saturated = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; - let rescued_receipts = if self.lease_claim_receipts() { - self.rescue_stale_receipt_claims_tx(&mut tx, cutoff).await? - } else { - Vec::new() - }; + let mut receipts = ReceiptRescueBatch::default(); + if self.lease_claim_receipts() { + self.rescue_stale_receipt_claims_tx(&mut tx, cutoff, options, &mut receipts) + .await?; + } + saturated |= receipts.saturated; + let idle_slots = receipts.idle_slots; + let scan_progress = receipts.scan_progress; + let rescued_receipts = receipts.rescued; if deleted.is_empty() && rescued_receipts.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); + return Ok(RescueBatch { + idle_slots, + scan_progress, + saturated, + ..RescueBatch::default() + }); } let moved_leases = self @@ -17203,33 +17806,38 @@ impl QueueStorage { .hydrate_deleted_leases_tx(&mut tx, rescued_receipts, true) .await?; - 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 mut rescued = self + .settle_rescued_attempts_tx( + &mut tx, + moved_leases, + "heartbeat stale: worker presumed dead", + RescueDlq { + reason: "heartbeat_stale", + enabled: dlq_enabled, + }, + options.retry_not_before, + ) + .await?; + rescued.extend( + self.settle_rescued_attempts_tx( + &mut tx, + moved_receipts, + "receipt claim stale: worker presumed dead", + RescueDlq { + reason: "heartbeat_stale", + enabled: dlq_enabled, + }, + options.retry_not_before, + ) + .await?, + ); tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) + Ok(RescueBatch { + rescued, + saturated, + idle_slots, + scan_progress, + }) } /// Rescue attempts past their hard deadline. Exhausted attempts fail @@ -17243,15 +17851,29 @@ impl QueueStorage { /// [`Self::rescue_expired_deadlines`], moving an exhausted attempt into /// `dlq_entries` (reason `deadline_exceeded`) in the rescue transaction /// when `dlq_enabled` returns true for its queue. + pub async fn rescue_expired_deadlines_with_dlq( + &self, + pool: &PgPool, + dlq_enabled: &(dyn Fn(&str) -> bool + Sync), + ) -> Result, AwaError> { + Ok(self + .rescue_expired_deadlines_batch(pool, dlq_enabled, &RescueBatchOptions::default()) + .await? + .rescued) + } + + /// One deadline rescue batch in its own transaction: up to + /// 500 lease rows and 500 receipt claims. #[tracing::instrument( - skip(self, pool, dlq_enabled), + skip(self, pool, dlq_enabled, options), name = "queue_storage.rescue_expired_deadlines" )] - pub async fn rescue_expired_deadlines_with_dlq( + pub async fn rescue_expired_deadlines_batch( &self, pool: &PgPool, dlq_enabled: &(dyn Fn(&str) -> bool + Sync), - ) -> Result, AwaError> { + options: &RescueBatchOptions, + ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(audited_sql(format!( @@ -17264,7 +17886,7 @@ impl QueueStorage { AND deadline_at IS NOT NULL AND deadline_at < clock_timestamp() ORDER BY deadline_at ASC - LIMIT 500 + LIMIT $1 FOR UPDATE SKIP LOCKED ) RETURNING @@ -17286,9 +17908,11 @@ impl QueueStorage { callback_timeout_at "# ))) + .bind(RECEIPT_RESCUE_BATCH_LIMIT) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; + let mut saturated = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; // Receipts-mode short-path claims hold their deadline on // `lease_claim_batches.deadline_at` (compact batches, #246) or on @@ -17296,15 +17920,24 @@ impl QueueStorage { // `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() - }; + let mut receipts = ReceiptRescueBatch::default(); + if self.lease_claim_receipts() { + self.rescue_expired_receipt_deadlines_tx(&mut tx, options, &mut receipts) + .await?; + } + saturated |= receipts.saturated; + let idle_slots = receipts.idle_slots; + let scan_progress = receipts.scan_progress; + let receipt_deleted = receipts.rescued; if deleted.is_empty() && receipt_deleted.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); + return Ok(RescueBatch { + idle_slots, + scan_progress, + saturated, + ..RescueBatch::default() + }); } let mut moved = self @@ -17315,16 +17948,25 @@ impl QueueStorage { .await?, ); - 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 rescued = self + .settle_rescued_attempts_tx( + &mut tx, + moved, + "hard deadline exceeded", + RescueDlq { + reason: "deadline_exceeded", + enabled: dlq_enabled, + }, + options.retry_not_before, + ) + .await?; tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) + Ok(RescueBatch { + rescued, + saturated, + idle_slots, + scan_progress, + }) } /// Rescue attempts whose callback timed out. Exhausted attempts fail @@ -17338,15 +17980,30 @@ impl QueueStorage { /// [`Self::rescue_expired_callbacks`], moving an exhausted attempt into /// `dlq_entries` (reason `callback_timeout`) in the rescue transaction /// when `dlq_enabled` returns true for its queue. + pub async fn rescue_expired_callbacks_with_dlq( + &self, + pool: &PgPool, + dlq_enabled: &(dyn Fn(&str) -> bool + Sync), + ) -> Result, AwaError> { + Ok(self + .rescue_expired_callbacks_batch(pool, dlq_enabled, &RescueBatchOptions::default()) + .await? + .rescued) + } + + /// One callback-timeout rescue batch in its own transaction: up to + /// 500 attempts. `slot_order` does not apply; + /// callback waits live in `leases`, not in the claim ring. #[tracing::instrument( - skip(self, pool, dlq_enabled), + skip(self, pool, dlq_enabled, options), name = "queue_storage.rescue_expired_callbacks" )] - pub async fn rescue_expired_callbacks_with_dlq( + pub async fn rescue_expired_callbacks_batch( &self, pool: &PgPool, dlq_enabled: &(dyn Fn(&str) -> bool + Sync), - ) -> Result, AwaError> { + options: &RescueBatchOptions, + ) -> Result { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(audited_sql(format!( @@ -17359,7 +18016,7 @@ impl QueueStorage { AND callback_timeout_at IS NOT NULL AND callback_timeout_at < clock_timestamp() ORDER BY callback_timeout_at ASC - LIMIT 500 + LIMIT $1 FOR UPDATE SKIP LOCKED ) RETURNING @@ -17381,29 +18038,40 @@ impl QueueStorage { callback_timeout_at "# ))) + .bind(RECEIPT_RESCUE_BATCH_LIMIT) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; + let saturated = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; if deleted.is_empty() { tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); + return Ok(RescueBatch::default()); } let moved = self .hydrate_deleted_leases_tx(&mut tx, deleted, true) .await?; - 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) - .await?, - ); - } + let rescued = self + .settle_rescued_attempts_tx( + &mut tx, + moved, + "callback timed out", + RescueDlq { + reason: "callback_timeout", + enabled: dlq_enabled, + }, + options.retry_not_before, + ) + .await?; tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) + Ok(RescueBatch { + rescued, + saturated, + idle_slots: Vec::new(), + scan_progress: Vec::new(), + }) } pub async fn promote_due( diff --git a/awa-python/python/awa/_awa.pyi b/awa-python/python/awa/_awa.pyi index 2219bfc3b..34fcc4b78 100644 --- a/awa-python/python/awa/_awa.pyi +++ b/awa-python/python/awa/_awa.pyi @@ -604,6 +604,8 @@ class Client: heartbeat_staleness_ms: int | None = None, deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, + rescue_pass_budget_ms: int | None = None, + rescue_pass_max_rows: int | None = None, terminal_count_rollup_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = 16, diff --git a/awa-python/python/awa/client.py b/awa-python/python/awa/client.py index f734ca679..bd2f11828 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -921,6 +921,8 @@ async def start( heartbeat_staleness_ms: int | None = None, deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, + rescue_pass_budget_ms: int | None = None, + rescue_pass_max_rows: int | None = None, terminal_count_rollup_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = DEFAULT_QUEUE_STORAGE_QUEUE_SLOT_COUNT, @@ -967,6 +969,11 @@ async def start( counters. Exact reads include pending deltas, so the default favors fewer hot counter updates over immediate folding. + ``rescue_pass_budget_ms`` (default 1000) and ``rescue_pass_max_rows`` + (default 5000) bound how long one queue-storage rescue pass keeps + taking 500-attempt batches and how many attempts it rescues; see the + configuration guide's "Draining a backlog of abandoned claims". + Queue dicts also accept ``claimers`` and ``claim_batch_size`` for hot queue-storage queues. ``claimers`` adds dispatcher/claimer loops that share the queue's worker permits; benchmark ``2`` or ``4`` before @@ -987,6 +994,8 @@ async def start( heartbeat_staleness_ms=heartbeat_staleness_ms, deadline_rescue_interval_ms=deadline_rescue_interval_ms, callback_rescue_interval_ms=callback_rescue_interval_ms, + rescue_pass_budget_ms=rescue_pass_budget_ms, + rescue_pass_max_rows=rescue_pass_max_rows, terminal_count_rollup_interval_ms=terminal_count_rollup_interval_ms, queue_storage_schema=queue_storage_schema, queue_storage_queue_slot_count=queue_storage_queue_slot_count, diff --git a/awa-python/src/client.rs b/awa-python/src/client.rs index f141508a7..473453d4d 100644 --- a/awa-python/src/client.rs +++ b/awa-python/src/client.rs @@ -1796,7 +1796,7 @@ impl PyClient { }) } - #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, terminal_count_rollup_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] + #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, rescue_pass_budget_ms=None, rescue_pass_max_rows=None, terminal_count_rollup_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] #[allow(clippy::too_many_arguments)] fn start<'py>( &self, @@ -1815,6 +1815,8 @@ impl PyClient { heartbeat_staleness_ms: Option, deadline_rescue_interval_ms: Option, callback_rescue_interval_ms: Option, + rescue_pass_budget_ms: Option, + rescue_pass_max_rows: Option, terminal_count_rollup_interval_ms: Option, queue_storage_schema: Option, queue_storage_queue_slot_count: u32, @@ -1942,6 +1944,12 @@ impl PyClient { if let Some(ms) = callback_rescue_interval_ms { builder = builder.callback_rescue_interval(Duration::from_millis(ms)); } + if let Some(ms) = rescue_pass_budget_ms { + builder = builder.rescue_pass_budget(Duration::from_millis(ms)); + } + if let Some(max_rows) = rescue_pass_max_rows { + builder = builder.rescue_pass_max_rows(max_rows); + } if let Some(ms) = terminal_count_rollup_interval_ms { if ms == 0 { return Err(pyo3::exceptions::PyValueError::new_err( diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 253bbde7a..07556ed9b 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -156,6 +156,8 @@ pub struct ClientBuilder { heartbeat_staleness: Option, deadline_rescue_interval: Option, callback_rescue_interval: Option, + rescue_pass_budget: Option, + rescue_pass_max_rows: Option, periodic_jobs: Vec, periodic_reconciliation: Option, global_max_workers: Option, @@ -219,6 +221,8 @@ impl ClientBuilder { heartbeat_staleness: None, deadline_rescue_interval: None, callback_rescue_interval: None, + rescue_pass_budget: None, + rescue_pass_max_rows: None, periodic_jobs: Vec::new(), periodic_reconciliation: None, global_max_workers: None, @@ -815,6 +819,26 @@ impl ClientBuilder { self } + /// Set how long one queue-storage rescue pass may keep starting batches of + /// 500 attempts, each in its own transaction (default: 1s, capped at half + /// the rescue interval). While a backlog remains the next pass runs + /// a second later. `Duration::ZERO` keeps every pass to one batch. + pub fn rescue_pass_budget(mut self, budget: Duration) -> Self { + self.rescue_pass_budget = Some(budget); + self + } + + /// Set how many attempts one queue-storage rescue pass rescues before it + /// stops (default: 5000). With a pass a second while a backlog remains, + /// this bounds the number rescued per second per kind, allowing one + /// final batch beyond the cap. Zero disables follow-up batches and early + /// re-arms. Retried attempts retain their normal backoff; their due times + /// can cluster, so this does not cap the promotion rate. + pub fn rescue_pass_max_rows(mut self, max_rows: usize) -> Self { + self.rescue_pass_max_rows = Some(max_rows); + self + } + /// Set the leader election retry interval (default: 10s). /// /// Controls how often a non-leader instance retries acquiring the maintenance @@ -1235,6 +1259,8 @@ impl ClientBuilder { heartbeat_staleness: self.heartbeat_staleness, deadline_rescue_interval: self.deadline_rescue_interval, callback_rescue_interval: self.callback_rescue_interval, + rescue_pass_budget: self.rescue_pass_budget, + rescue_pass_max_rows: self.rescue_pass_max_rows, periodic_jobs: Arc::new(self.periodic_jobs), periodic_reconciliation: self.periodic_reconciliation, periodic_manifest, @@ -1335,6 +1361,8 @@ pub struct Client { heartbeat_staleness: Option, deadline_rescue_interval: Option, callback_rescue_interval: Option, + rescue_pass_budget: Option, + rescue_pass_max_rows: Option, periodic_jobs: Arc>, periodic_reconciliation: Option, periodic_manifest: Option>, @@ -1920,6 +1948,12 @@ impl Client { if let Some(interval) = self.callback_rescue_interval { maintenance = maintenance.callback_rescue_interval(interval); } + if let Some(budget) = self.rescue_pass_budget { + maintenance = maintenance.rescue_pass_budget(budget); + } + if let Some(max_rows) = self.rescue_pass_max_rows { + maintenance = maintenance.rescue_pass_max_rows(max_rows); + } if let Some(interval) = self.leader_election_interval { maintenance = maintenance.leader_election_interval(interval); } diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 894881c2e..6c3472936 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -409,6 +409,7 @@ impl JobExecutor { claim, ) }), + finalizing: false, }; in_flight.insert((job_id, job_run_lease), in_flight_state); if let Some(counter) = queue_in_flight.get(&job_queue) { @@ -514,7 +515,12 @@ impl JobExecutor { } } - // Keep worker/heartbeat ownership while the bounded finalizer queue is full. + // Keep heartbeating through the finalization queue and its durable write. + in_flight.mark_finalizing((job_id, job_run_lease)); + let heartbeat = HeartbeatUntilFinalized { + in_flight: in_flight.clone(), + key: (job_id, job_run_lease), + }; let finalize_permit = match result { Some(_) => tokio::select! { permit = finalize_capacity.acquire_owned() => Some(permit.expect("finalizer capacity is never closed")), @@ -523,7 +529,6 @@ impl JobExecutor { None => None, }; - in_flight.remove((job_id, job_run_lease)); if let Some(counter) = queue_in_flight.get(&job_queue) { counter.fetch_sub(1, Ordering::SeqCst); } @@ -543,6 +548,7 @@ impl JobExecutor { job.output = sealed_output; finalizers.spawn( async move { + let _heartbeat = heartbeat; let event_kind = job_kind.clone(); let abandon_output = had_output.then(|| (metrics.clone(), job_queue.clone())); let finalize = async move { @@ -3369,6 +3375,18 @@ async fn retry_backoff_canonical_with_followups( Ok(Some(updated_job)) } +/// Keep finalizing attempts live until their durable write or abandonment. +struct HeartbeatUntilFinalized { + in_flight: InFlightMap, + key: (i64, i64), +} + +impl Drop for HeartbeatUntilFinalized { + fn drop(&mut self) { + self.in_flight.remove(self.key); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/awa-worker/src/heartbeat.rs b/awa-worker/src/heartbeat.rs index 91b74f81b..bbe027d22 100644 --- a/awa-worker/src/heartbeat.rs +++ b/awa-worker/src/heartbeat.rs @@ -1,5 +1,6 @@ use crate::runtime::{InFlightMap, RunLease}; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; +use awa_model::queue_storage::ClaimHeartbeatRef; use sqlx::PgPool; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -140,7 +141,7 @@ impl HeartbeatService { .iter() .filter(|(key, _)| !progress_keys.contains(key)) .partition(|(_, claim)| claim.is_some()); - let claimed: Vec = + let claimed: Vec = claimed.into_iter().filter_map(|(_, claim)| claim).collect(); let plain: Vec<(i64, RunLease)> = plain.into_iter().map(|(key, _)| key).collect(); self.heartbeat_once_queue_storage(runtime, &claimed, &plain, &pending_progress) @@ -170,7 +171,7 @@ impl HeartbeatService { AND jobs.run_lease = inflight.run_lease WHERE jobs.state = 'running' ORDER BY jobs.id - FOR UPDATE OF jobs + FOR UPDATE OF jobs SKIP LOCKED ) UPDATE awa.jobs_hot AS jobs SET heartbeat_at = now() @@ -217,7 +218,7 @@ impl HeartbeatService { AND jobs.run_lease = inflight.run_lease WHERE jobs.state = 'running' ORDER BY jobs.id - FOR UPDATE OF jobs + FOR UPDATE OF jobs SKIP LOCKED ) UPDATE awa.jobs_hot AS jobs SET heartbeat_at = now(), @@ -243,7 +244,13 @@ impl HeartbeatService { .iter() .map(|(id, lease, gen, _)| (*id, *lease, *gen)) .collect(); - self.in_flight.ack_progress(&acked); + if result.rows_affected() as usize == chunk.len() { + self.in_flight.ack_progress(&acked); + } else { + // Re-send successful snapshots too: replay is safe, + // whereas acknowledging a skipped write loses progress. + self.in_flight.clear_in_flight_progress(&acked); + } } Err(err) => { warn!(error = %err, "Failed to send heartbeat+progress batch"); @@ -261,7 +268,7 @@ impl HeartbeatService { async fn heartbeat_once_queue_storage( &self, runtime: &QueueStorageRuntime, - claimed: &[awa_model::queue_storage::ClaimHeartbeatRef], + claimed: &[ClaimHeartbeatRef], heartbeat_only: &[(i64, RunLease)], pending_progress: &[(i64, i64, u64, serde_json::Value)], ) { @@ -314,7 +321,11 @@ impl HeartbeatService { .iter() .map(|(id, lease, gen, _)| (*id, *lease, *gen)) .collect(); - self.in_flight.ack_progress(&acked); + if updated == chunk.len() { + self.in_flight.ack_progress(&acked); + } else { + self.in_flight.clear_in_flight_progress(&acked); + } } Err(err) => { warn!(error = %err, "Failed to send heartbeat+progress batch"); diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index d84960d37..febcd7038 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -8,14 +8,15 @@ use awa_model::cron::{ #[cfg(test)] use awa_model::SkipReason; use awa_model::{ - JobRow, JobState, PeriodicJob, PruneOutcome, RotateOutcome, TerminalDeltaRollupOutcome, + JobRow, JobState, PeriodicJob, PruneOutcome, RescueBatch, RescueBatchOptions, + RescueScanProgress, RescueSlotOrder, RotateOutcome, TerminalDeltaRollupOutcome, }; -use chrono::Utc; +use chrono::{DateTime, Utc}; use croner::Cron; use sqlx::pool::PoolConnection; use sqlx::{PgPool, Postgres}; use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::task::JoinHandle; @@ -225,6 +226,268 @@ const CALLBACK_RESCUE_PER_ROW_SQL: &str = r#" /// How often the leader scans for attempts past their deadline by default. pub(crate) const DEFAULT_DEADLINE_RESCUE_INTERVAL: Duration = Duration::from_secs(30); +/// Default wall-clock budget for one queue-storage rescue pass. +pub(crate) const DEFAULT_RESCUE_PASS_BUDGET: Duration = Duration::from_secs(1); + +/// Default cap on attempts one queue-storage rescue pass rescues. +pub(crate) const DEFAULT_RESCUE_PASS_MAX_ROWS: usize = 5_000; + +/// While a rescue kind still has a backlog, its next pass runs after this +/// rather than after its full interval. With the per-pass row cap this bounds +/// a drain to about `rescue_pass_max_rows` attempts per second per kind. +const RESCUE_BACKLOG_INTERVAL: Duration = Duration::from_secs(1); + +/// Claims each slot's cursor walk examines during bulk heartbeat rescue and +/// follow-up batches. Other first batches keep the store's wider default to +/// walk past long closed prefixes when acceleration is unavailable or disabled. +const RESCUE_FOLLOW_UP_SCAN_LIMIT: usize = 2_000; + +/// One batch in this many visits the other claim slots before the oldest, so +/// a long drain of the oldest slot leaves the rest at least that share. +const RESCUE_OTHERS_FIRST_EVERY: usize = 4; + +/// Runtime snapshots the leader reads before accelerating stale-heartbeat rescue. A runtime is `fresh` while its snapshot is +/// within three effective snapshot intervals. An unconfirmed phase exists only +/// when 3 * max(snapshot_interval, 1s) < heartbeat staleness ($1, milliseconds). +/// With a larger interval a stalled live runtime can pass directly from fresh +/// to gone; this evidence is a bounded heuristic, not proof of process death. +/// `last_gone_heartbeat` estimates the last heartbeat of a gone runtime from +/// its final snapshot and configured interval. +const RESCUE_RUNTIME_EVIDENCE_SQL: &str = r#" + SELECT + count(*) FILTER (WHERE fresh AND NOT heartbeat_alive)::bigint AS fresh_without_heartbeat, + count(*) FILTER (WHERE NOT fresh AND NOT gone)::bigint AS unconfirmed, + max(last_seen_at + snapshot_interval_ms * 2 * interval '1 millisecond') + FILTER (WHERE gone) AS last_gone_heartbeat + FROM ( + SELECT + heartbeat_alive, + last_seen_at, + snapshot_interval_ms, + last_seen_at >= now() - GREATEST(snapshot_interval_ms, 1000) * 3 * interval '1 millisecond' + AS fresh, + last_seen_at < now() - GREATEST(snapshot_interval_ms, 1000) * 3 * interval '1 millisecond' + AND last_seen_at < now() - $1 * interval '1 millisecond' + AS gone + FROM awa.runtime_instances + ) AS runtimes +"#; + +/// The queue-storage rescue kinds a leader drains in passes of batches. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum RescueKind { + Heartbeat, + Deadline, + Callback, +} + +impl RescueKind { + /// The `awa.rescue.kind` attribute value, as on `awa.maintenance.rescues`. + fn as_str(self) -> &'static str { + match self { + Self::Heartbeat => "heartbeat", + Self::Deadline => "deadline", + Self::Callback => "callback_timeout", + } + } +} + +/// Why a queue-storage rescue pass stopped taking batches. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RescuePassEnd { + /// The last batch found fewer eligible attempts than its limit. + Drained, + /// The pass budget is spent. + TimeBudget, + /// The pass reached `rescue_pass_max_rows`. + RowLimit, + /// Stale-heartbeat only: the runtime snapshots do not yet show that the + /// owners of stale claims are gone, so the pass keeps to one batch. + OwnersUnconfirmed, + Error, +} + +impl RescuePassEnd { + fn as_str(self) -> &'static str { + match self { + Self::Drained => "drained", + Self::TimeBudget => "time_budget", + Self::RowLimit => "row_limit", + Self::OwnersUnconfirmed => "owners_unconfirmed", + Self::Error => "error", + } + } + + /// The pass stopped with eligible attempts left that it was allowed to + /// take, so the next pass should come sooner than the interval. + fn backlog_remains(self) -> bool { + matches!(self, Self::TimeBudget | Self::RowLimit) + } +} + +/// The leader's record of whether every runtime that could own a stale +/// claim is gone. A claim does not record its runtime, so the evidence is +/// fleet-wide: every runtime snapshot is either fresh with its heartbeat +/// loop alive, or silent for at least the heartbeat staleness. That must +/// hold at every sample for a full staleness window before a stale-heartbeat +/// rescue may run extra batches or early passes. This settled window may have +/// started before a runtime stalled. If freshness lasts at least the staleness, +/// that runtime can pass directly from fresh to gone without resetting it. +/// A gap between samples longer than half the staleness restarts the window, +/// because a stall the leader did not sample could have hidden behind it. +#[derive(Debug, Default)] +struct RescueOwnerEvidence { + steady_since: Option, + last_sample: Option, +} + +impl RescueOwnerEvidence { + /// Record one sample and return whether the evidence now allows a + /// stale-heartbeat pass to take more than one batch. + fn observe(&mut self, now: Instant, settled: bool, staleness: Duration) -> bool { + let sampled_recently = self + .last_sample + .is_some_and(|last| now.saturating_duration_since(last) <= staleness / 2); + self.last_sample = Some(now); + if !settled { + self.steady_since = None; + return false; + } + if !sampled_recently { + self.steady_since = Some(now); + } + let steady_since = *self.steady_since.get_or_insert(now); + now.saturating_duration_since(steady_since) >= staleness + } +} + +/// A normal interval tick is a baseline visit even if scheduler jitter makes +/// its actual start slightly earlier than one interval after the last start. +/// A reset tick is accelerated work and must retain the owner/cadence gate. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RescueTickOrigin { + Ordinary, + Backlog, +} + +struct RescueTimer { + interval: tokio::time::Interval, + next_origin: RescueTickOrigin, +} + +impl RescueTimer { + fn new(period: Duration) -> Self { + Self { + interval: leader_loop_interval(period), + next_origin: RescueTickOrigin::Ordinary, + } + } + + async fn tick(&mut self) -> RescueTickOrigin { + self.interval.tick().await; + // Consume even when the caller's branch cooldown skips this tick. + std::mem::replace(&mut self.next_origin, RescueTickOrigin::Ordinary) + } + + fn reset_after(&mut self, duration: Duration) { + self.interval.reset_after(duration); + self.next_origin = RescueTickOrigin::Backlog; + } +} + +/// A sweep survives early re-arms: closed history must not keep the leader +/// scanning once a cursor has made a full revolution. The first window is +/// already visited, so crossing its end after wrapping completes the sweep. +#[derive(Debug, Default)] +struct RescueSweep { + cursors: HashMap>, + idle_slots: Vec, + last_revisit: Option, + batches_started: usize, +} + +#[derive(Debug, Default)] +struct RescueSweepCursor { + first: Option<(DateTime, i64, i64)>, + previous: Option<(DateTime, i64, i64)>, + wrapped: bool, + done: bool, +} + +impl RescueSweep { + /// A long drain must revisit slots whose claims have become eligible + /// since their last scan. Keep unfinished cursors so a large history + /// sweep still converges across these ordinary-cadence revisits. + fn revisit(&mut self, now: Instant, interval: Duration, origin: RescueTickOrigin) { + if origin == RescueTickOrigin::Ordinary + || self + .last_revisit + .is_none_or(|last| now.saturating_duration_since(last) >= interval) + { + self.idle_slots.clear(); + for cursor in self.cursors.values_mut().flatten() { + if cursor.done { + *cursor = RescueSweepCursor::default(); + } + } + self.last_revisit = Some(now); + } + } + + fn observe(&mut self, progress: RescueScanProgress) { + let cursors = self.cursors.entry(progress.slot).or_default(); + cursors.resize_with(progress.cursors.len(), RescueSweepCursor::default); + for (cursor, progress) in cursors.iter_mut().zip(progress.cursors) { + if !progress.saturated { + cursor.done = true; + } + if let Some(position) = progress.position { + if let Some(first) = cursor.first { + let crossed = cursor.previous.is_some_and(|previous| position <= previous); + cursor.done |= (cursor.wrapped && crossed) + || ((cursor.wrapped || crossed) && position >= first); + cursor.wrapped |= crossed; + } else { + cursor.first = Some(position); + } + cursor.previous = Some(position); + } + } + if !cursors.is_empty() && cursors.iter().all(|cursor| cursor.done) { + self.idle_slots.push(progress.slot); + } + } +} + +/// An early re-arm may run after owner evidence disappears. Keep ordinary +/// first batches on their original cadence; every accelerated batch needs +/// the same owner cutoff as follow-up batches. +#[derive(Debug, Default)] +struct RescueHeartbeatCadence { + last_baseline: Option, +} + +impl RescueHeartbeatCadence { + fn select( + &mut self, + now: Instant, + interval: Duration, + cutoff: Option>, + origin: RescueTickOrigin, + ) -> Option>> { + if origin == RescueTickOrigin::Ordinary + || self + .last_baseline + .is_none_or(|last| now.saturating_duration_since(last) >= interval) + { + self.last_baseline = Some(now); + Some(None) + } else { + cutoff.map(Some) + } + } +} + /// Per-queue or global retention policy for completed and failed/cancelled jobs. #[derive(Debug, Clone)] pub struct RetentionPolicy { @@ -716,6 +979,13 @@ pub struct MaintenanceService { heartbeat_rescue_interval: Duration, deadline_rescue_interval: Duration, callback_rescue_interval: Duration, + rescue_pass_budget: Duration, + rescue_pass_max_rows: usize, + rescue_owner_evidence: std::sync::Mutex, + rescue_heartbeat_cadence: std::sync::Mutex, + rescue_sweeps: std::sync::Mutex>, + + rescue_others_first_offset: AtomicUsize, promote_interval: Duration, cleanup_interval: Duration, cron_sync_interval: Duration, @@ -797,6 +1067,12 @@ impl MaintenanceService { heartbeat_rescue_interval: Duration::from_secs(30), deadline_rescue_interval: DEFAULT_DEADLINE_RESCUE_INTERVAL, callback_rescue_interval: Duration::from_secs(30), + rescue_pass_budget: DEFAULT_RESCUE_PASS_BUDGET, + rescue_pass_max_rows: DEFAULT_RESCUE_PASS_MAX_ROWS, + rescue_owner_evidence: std::sync::Mutex::new(RescueOwnerEvidence::default()), + rescue_heartbeat_cadence: std::sync::Mutex::new(RescueHeartbeatCadence::default()), + rescue_sweeps: std::sync::Mutex::new(HashMap::new()), + rescue_others_first_offset: AtomicUsize::new(0), promote_interval: Duration::from_millis(250), cleanup_interval: Duration::from_secs(60), cron_sync_interval: Duration::from_secs(60), @@ -925,6 +1201,22 @@ impl MaintenanceService { self } + /// Set how long one queue-storage rescue pass may keep starting batches + /// (default: 1s, and never more than half the rescue interval). Zero + /// keeps every pass to one batch. + pub fn rescue_pass_budget(mut self, budget: Duration) -> Self { + self.rescue_pass_budget = budget; + self + } + + /// Set how many attempts one queue-storage rescue pass rescues before it + /// stops (default: 5000). A pass always finishes the batch it started; + /// zero keeps one batch per ordinary interval, as a zero pass budget does. + pub fn rescue_pass_max_rows(mut self, max_rows: usize) -> Self { + self.rescue_pass_max_rows = max_rows; + self + } + /// Set how long a heartbeat must be stale before the job is rescued (default: 90s). /// /// Should be at least 3× the heartbeat interval to avoid false rescues @@ -1035,9 +1327,9 @@ impl MaintenanceService { } // Run maintenance tasks as leader - let mut heartbeat_rescue_timer = leader_loop_interval(self.heartbeat_rescue_interval); - let mut deadline_rescue_timer = leader_loop_interval(self.deadline_rescue_interval); - let mut callback_rescue_timer = leader_loop_interval(self.callback_rescue_interval); + let mut heartbeat_rescue_timer = RescueTimer::new(self.heartbeat_rescue_interval); + let mut deadline_rescue_timer = RescueTimer::new(self.deadline_rescue_interval); + let mut callback_rescue_timer = RescueTimer::new(self.callback_rescue_interval); let mut promote_timer = leader_loop_interval(self.promote_interval); let mut cleanup_timer = leader_loop_interval(self.cleanup_interval); let mut cron_sync_timer = leader_loop_interval(self.cron_sync_interval); @@ -1081,6 +1373,20 @@ impl MaintenanceService { // Per-branch overrun tracker. Issue #242 — observability only; // the architectural split of this `tokio::select!` is deferred // to v0.7 conditional on the overrun counter showing fleet hits. + *self + .rescue_owner_evidence + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + RescueOwnerEvidence::default(); + *self + .rescue_heartbeat_cadence + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + RescueHeartbeatCadence::default(); + self.rescue_sweeps + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clear(); let branch_tracker = MaintenanceBranchTracker::new(); // Per-segment prune backoff (#169). Gates the prune step of // each rotate branch so a pinned MVCC snapshot doesn't make @@ -1133,22 +1439,31 @@ impl MaintenanceService { let _ = Self::release_leader(&mut leader_conn).await; return; } - _ = heartbeat_rescue_timer.tick() => { + origin = heartbeat_rescue_timer.tick() => { if let Some(timer) = branch_tracker.try_begin("rescue_stale_heartbeats", self.heartbeat_rescue_interval, &self.metrics) { - self.rescue_stale_heartbeats().await; + let backlog = self.rescue_stale_heartbeats(origin).await; timer.finish(); + if backlog { + heartbeat_rescue_timer.reset_after(RESCUE_BACKLOG_INTERVAL.min(self.heartbeat_rescue_interval)); + } } } - _ = deadline_rescue_timer.tick() => { + origin = deadline_rescue_timer.tick() => { if let Some(timer) = branch_tracker.try_begin("rescue_expired_deadlines", self.deadline_rescue_interval, &self.metrics) { - self.rescue_expired_deadlines().await; + let backlog = self.rescue_expired_deadlines(origin).await; timer.finish(); + if backlog { + deadline_rescue_timer.reset_after(RESCUE_BACKLOG_INTERVAL.min(self.deadline_rescue_interval)); + } } } - _ = callback_rescue_timer.tick() => { + origin = callback_rescue_timer.tick() => { if let Some(timer) = branch_tracker.try_begin("rescue_expired_callbacks", self.callback_rescue_interval, &self.metrics) { - self.rescue_expired_callbacks().await; + let backlog = self.rescue_expired_callbacks(origin).await; timer.finish(); + if backlog { + callback_rescue_timer.reset_after(RESCUE_BACKLOG_INTERVAL.min(self.callback_rescue_interval)); + } } } _ = promote_timer.tick() => { @@ -1626,42 +1941,301 @@ impl MaintenanceService { /// `canonical_live_backlog` above zero and wedge `finalize --wait`, /// and stuck queue-storage rows would wait out a drain-only leader. /// Both cross-plane sweeps are cheap no-ops once their tables are empty. + /// + /// Returns whether the queue-storage pass left a backlog it may keep + /// draining, so the next pass should run before the full interval. #[tracing::instrument(skip(self), name = "maintenance.rescue_stale")] - async fn rescue_stale_heartbeats(&self) { + async fn rescue_stale_heartbeats(&self, origin: RescueTickOrigin) -> bool { match &self.storage { RuntimeStorage::Canonical => { let outcome = self.rescue_canonical_stale_heartbeats().await; self.process_heartbeat_rescues(outcome, true, None).await; if let Some(runtime) = self.active_standby_queue_storage().await { - let outcome = runtime - .store - .rescue_stale_heartbeats_with_dlq( - &self.pool, - self.heartbeat_staleness, - &|queue| self.dlq_policy.enabled_for(queue), - ) - .await; - self.process_heartbeat_rescues(outcome, false, Some(runtime)) + return self + .rescue_queue_storage_pass(RescueKind::Heartbeat, runtime, false, origin) .await; } + false } RuntimeStorage::QueueStorage(runtime) => { - let outcome = runtime - .store - .rescue_stale_heartbeats_with_dlq( - &self.pool, - self.heartbeat_staleness, - &|queue| self.dlq_policy.enabled_for(queue), - ) - .await; - self.process_heartbeat_rescues(outcome, true, Some(runtime)) + let backlog = self + .rescue_queue_storage_pass(RescueKind::Heartbeat, runtime, true, origin) .await; let outcome = self.rescue_canonical_stale_heartbeats().await; self.process_heartbeat_rescues(outcome, false, None).await; + backlog } } } + /// One queue-storage rescue pass: batches of up to 500 attempts per + /// source, each in its own transaction, until a batch comes back short, + /// the pass budget is spent, or the pass reaches `rescue_pass_max_rows`. Most batches visit the oldest claim slot + /// first, since the ring cannot rotate back into it while it holds open + /// claims; every [`RESCUE_OTHERS_FIRST_EVERY`]th visits the others first. + /// A slot a batch found nothing in is left out for the rest of the pass. + /// + /// A stale heartbeat cannot tell a dead runtime from a stalled one, so a + /// stale-heartbeat pass takes more than one batch only when + /// [`RescueOwnerEvidence`] holds, and its later batches take only claims + /// last heartbeated before the most recently gone runtime stopped + /// reporting. This reduces the chance of draining a lagging live runtime, + /// but snapshot freshness and sampling cannot prove an owner is dead. + /// Deadline and callback rescues act on a deadline that has passed whoever + /// owns the attempt. + /// + /// Returns whether the pass stopped with a backlog it may keep draining, + /// so the next pass should run before the full interval. + async fn rescue_queue_storage_pass( + &self, + kind: RescueKind, + runtime: &QueueStorageRuntime, + signal_local: bool, + origin: RescueTickOrigin, + ) -> bool { + let started = Instant::now(); + let interval = match kind { + RescueKind::Heartbeat => self.heartbeat_rescue_interval, + RescueKind::Deadline => self.deadline_rescue_interval, + RescueKind::Callback => self.callback_rescue_interval, + }; + let budget = self.rescue_pass_budget.min(interval / 2); + let dlq_enabled = |queue: &str| self.dlq_policy.enabled_for(queue); + let bulk_heartbeat_cutoff = if kind == RescueKind::Heartbeat { + self.stale_claim_owners_gone().await + } else { + None + }; + let first_heartbeat_cutoff = if kind == RescueKind::Heartbeat { + let selected = self + .rescue_heartbeat_cadence + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .select(started, interval, bulk_heartbeat_cutoff, origin); + let Some(cutoff) = selected else { + self.rescue_sweeps + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&kind); + return false; + }; + cutoff + } else { + None + }; + let mut rescued_total = 0usize; + let mut oldest_attempt: Option> = None; + let mut batches = 0usize; + let mut sweep = self + .rescue_sweeps + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&kind) + .unwrap_or_default(); + sweep.revisit(started, interval, origin); + let pass_end = loop { + let slot_order = if sweep.batches_started % RESCUE_OTHERS_FIRST_EVERY + == RESCUE_OTHERS_FIRST_EVERY - 1 + { + RescueSlotOrder::OthersFirst( + self.rescue_others_first_offset + .fetch_add(1, Ordering::Relaxed), + ) + } else { + RescueSlotOrder::OldestFirst + }; + let options = RescueBatchOptions { + slot_order, + skip_slots: sweep.idle_slots.clone(), + retry_not_before: None, + cursor_scan_limit: (batches > 0 + || (kind == RescueKind::Heartbeat + && bulk_heartbeat_cutoff.is_some() + && !budget.is_zero() + && self.rescue_pass_max_rows > 0)) + .then_some(RESCUE_FOLLOW_UP_SCAN_LIMIT), + }; + let outcome = match kind { + RescueKind::Heartbeat => { + // Follow-up batches take only claims last heartbeated + // before a runtime known to be gone stopped reporting. + let cutoff = if batches == 0 { + first_heartbeat_cutoff + } else { + bulk_heartbeat_cutoff + }; + let staleness = cutoff + .and_then(|cutoff| (Utc::now() - cutoff).to_std().ok()) + .map_or(self.heartbeat_staleness, |since_cutoff| { + since_cutoff.max(self.heartbeat_staleness) + }); + runtime + .store + .rescue_stale_heartbeats_batch( + &self.pool, + staleness, + &dlq_enabled, + &options, + ) + .await + } + RescueKind::Deadline => { + runtime + .store + .rescue_expired_deadlines_batch(&self.pool, &dlq_enabled, &options) + .await + } + RescueKind::Callback => { + runtime + .store + .rescue_expired_callbacks_batch(&self.pool, &dlq_enabled, &options) + .await + } + }; + batches += 1; + sweep.batches_started = sweep.batches_started.wrapping_add(1); + let RescueBatch { + rescued, + saturated, + idle_slots: batch_idle_slots, + scan_progress, + } = match outcome { + Ok(batch) => batch, + Err(err) => { + self.process_queue_storage_rescues(kind, Err(err), signal_local, runtime) + .await; + break RescuePassEnd::Error; + } + }; + sweep.idle_slots.extend(batch_idle_slots); + for progress in scan_progress { + sweep.observe(progress); + } + rescued_total += rescued.len(); + if let Some(attempted_at) = rescued.iter().filter_map(|job| job.attempted_at).min() { + oldest_attempt = + Some(oldest_attempt.map_or(attempted_at, |oldest| oldest.min(attempted_at))); + } + self.process_queue_storage_rescues(kind, Ok(rescued), signal_local, runtime) + .await; + + if !saturated { + break RescuePassEnd::Drained; + } + if budget.is_zero() || self.rescue_pass_max_rows == 0 { + break RescuePassEnd::TimeBudget; + } + if kind == RescueKind::Heartbeat && bulk_heartbeat_cutoff.is_none() { + break RescuePassEnd::OwnersUnconfirmed; + } + if rescued_total >= self.rescue_pass_max_rows { + break RescuePassEnd::RowLimit; + } + if started.elapsed() >= budget { + break RescuePassEnd::TimeBudget; + } + }; + + let oldest_attempt_age = oldest_attempt + .and_then(|attempted_at| (Utc::now() - attempted_at).to_std().ok()) + .unwrap_or_default(); + self.metrics.record_rescue_pass( + kind.as_str(), + pass_end.as_str(), + rescued_total as u64, + oldest_attempt_age, + ); + if batches > 1 { + info!( + rescue_kind = kind.as_str(), + rescued = rescued_total, + batches, + pass_end = pass_end.as_str(), + elapsed_ms = started.elapsed().as_millis() as u64, + oldest_attempt_age_s = oldest_attempt_age.as_secs(), + "Queue-storage rescue pass drained a backlog" + ); + } + let rearm = pass_end.backlog_remains() + && !budget.is_zero() + && self.rescue_pass_max_rows > 0 + && (kind != RescueKind::Heartbeat || bulk_heartbeat_cutoff.is_some()); + if rearm { + self.rescue_sweeps + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(kind, sweep); + } + rearm + } + + async fn process_queue_storage_rescues( + &self, + kind: RescueKind, + outcome: Result, awa_model::AwaError>, + signal_local: bool, + runtime: &QueueStorageRuntime, + ) { + match kind { + RescueKind::Heartbeat => { + self.process_heartbeat_rescues(outcome, signal_local, Some(runtime)) + .await + } + RescueKind::Deadline => { + self.process_deadline_rescues(outcome, signal_local, Some(runtime)) + .await + } + RescueKind::Callback => { + self.process_callback_rescues(outcome, signal_local, Some(runtime)) + .await + } + } + } + + /// Sample the runtime snapshots and update [`RescueOwnerEvidence`]. + /// Returns the heartbeat cutoff for accelerated stale-heartbeat batches + /// when the evidence holds and some runtime is known to be gone; `None` + /// keeps the pass to one batch. A failed read counts as no evidence. + async fn stale_claim_owners_gone(&self) -> Option> { + let staleness_ms = i64::try_from(self.heartbeat_staleness.as_millis()).unwrap_or(i64::MAX); + let (settled, last_gone_heartbeat) = match sqlx::query_as::< + _, + (i64, i64, Option>), + >(RESCUE_RUNTIME_EVIDENCE_SQL) + .bind(staleness_ms) + .fetch_one(&self.pool) + .await + { + Ok((fresh_without_heartbeat, unconfirmed, last_gone_heartbeat)) => { + if fresh_without_heartbeat > 0 || unconfirmed > 0 { + debug!( + fresh_without_heartbeat, + unconfirmed, + "Stale-heartbeat rescue kept to one batch: runtime snapshots do not show the claim owners gone" + ); + } + ( + fresh_without_heartbeat == 0 && unconfirmed == 0, + last_gone_heartbeat, + ) + } + Err(err) => { + warn!(error = %err, "Failed to read runtime snapshots for stale-heartbeat rescue"); + (false, None) + } + }; + let steady = self + .rescue_owner_evidence + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .observe(Instant::now(), settled, self.heartbeat_staleness); + if steady { + last_gone_heartbeat + } else { + None + } + } + /// Batched canonical stale-heartbeat rescue with the #388 row-at-a-time /// fallback on unique-claim conflict. Exhausted attempts fail in place, /// as callback timeouts do, so a job that keeps killing or hanging its @@ -1825,33 +2399,25 @@ impl MaintenanceService { /// #456 mirror: sweeps both storage planes; see /// [`Self::rescue_stale_heartbeats`]. #[tracing::instrument(skip(self), name = "maintenance.rescue_deadline")] - async fn rescue_expired_deadlines(&self) { + async fn rescue_expired_deadlines(&self, origin: RescueTickOrigin) -> bool { match &self.storage { RuntimeStorage::Canonical => { let outcome = self.rescue_canonical_expired_deadlines().await; self.process_deadline_rescues(outcome, true, None).await; if let Some(runtime) = self.active_standby_queue_storage().await { - let outcome = runtime - .store - .rescue_expired_deadlines_with_dlq(&self.pool, &|queue| { - self.dlq_policy.enabled_for(queue) - }) - .await; - self.process_deadline_rescues(outcome, false, Some(runtime)) + return self + .rescue_queue_storage_pass(RescueKind::Deadline, runtime, false, origin) .await; } + false } RuntimeStorage::QueueStorage(runtime) => { - let outcome = runtime - .store - .rescue_expired_deadlines_with_dlq(&self.pool, &|queue| { - self.dlq_policy.enabled_for(queue) - }) - .await; - self.process_deadline_rescues(outcome, true, Some(runtime)) + let backlog = self + .rescue_queue_storage_pass(RescueKind::Deadline, runtime, true, origin) .await; let outcome = self.rescue_canonical_expired_deadlines().await; self.process_deadline_rescues(outcome, false, None).await; + backlog } } } @@ -1998,33 +2564,25 @@ impl MaintenanceService { /// doubly here because callback-carrying jobs never migrate off the /// canonical plane during a transition. #[tracing::instrument(skip(self), name = "maintenance.rescue_callback_timeout")] - async fn rescue_expired_callbacks(&self) { + async fn rescue_expired_callbacks(&self, origin: RescueTickOrigin) -> bool { match &self.storage { RuntimeStorage::Canonical => { let outcome = self.rescue_canonical_expired_callbacks().await; self.process_callback_rescues(outcome, true, None).await; if let Some(runtime) = self.active_standby_queue_storage().await { - let outcome = runtime - .store - .rescue_expired_callbacks_with_dlq(&self.pool, &|queue| { - self.dlq_policy.enabled_for(queue) - }) - .await; - self.process_callback_rescues(outcome, false, Some(runtime)) + return self + .rescue_queue_storage_pass(RescueKind::Callback, runtime, false, origin) .await; } + false } RuntimeStorage::QueueStorage(runtime) => { - let outcome = runtime - .store - .rescue_expired_callbacks_with_dlq(&self.pool, &|queue| { - self.dlq_policy.enabled_for(queue) - }) - .await; - self.process_callback_rescues(outcome, true, Some(runtime)) + let backlog = self + .rescue_queue_storage_pass(RescueKind::Callback, runtime, true, origin) .await; let outcome = self.rescue_canonical_expired_callbacks().await; self.process_callback_rescues(outcome, false, None).await; + backlog } } } @@ -4036,8 +4594,12 @@ mod tests { Arc::new(HashMap::new()), ); - service.rescue_stale_heartbeats().await; - service.rescue_expired_callbacks().await; + service + .rescue_stale_heartbeats(RescueTickOrigin::Backlog) + .await; + service + .rescue_expired_callbacks(RescueTickOrigin::Backlog) + .await; // Post-flip the `awa.jobs` view exposes only the queue-storage // plane, so assert against the physical canonical tables: both @@ -4290,6 +4852,645 @@ mod tests { ); } + #[tokio::test(start_paused = true)] + async fn ordinary_rescue_ticks_survive_jitter_and_revisit_finished_slots() { + let interval = Duration::from_secs(30); + let mut timer = RescueTimer::new(interval); + timer.tick().await; + let mut cadence = RescueHeartbeatCadence::default(); + let mut sweep = RescueSweep::default(); + tokio::time::advance(interval + Duration::from_millis(2)).await; + let first_origin = timer.tick().await; + assert_eq!(first_origin, RescueTickOrigin::Ordinary); + let first_start = tokio::time::Instant::now().into_std(); + assert_eq!( + cadence.select(first_start, interval, None, first_origin), + Some(None) + ); + sweep.revisit(first_start, interval, first_origin); + let stamp = Utc::now(); + for (slot, saturated) in [(0, false), (1, true)] { + sweep.observe(RescueScanProgress { + slot, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 100, 1)), + saturated, + }], + }); + } + assert_eq!(sweep.idle_slots, vec![0]); + tokio::time::advance(interval - Duration::from_millis(2)).await; + let next_origin = timer.tick().await; + let next_start = tokio::time::Instant::now().into_std(); + assert_eq!( + next_start.duration_since(first_start), + interval - Duration::from_millis(2) + ); + assert_eq!(next_origin, RescueTickOrigin::Ordinary); + assert_eq!( + cadence.select(next_start, interval, None, next_origin), + Some(None), + "an ordinary tick must not disappear because the prior tick started 2ms late" + ); + sweep.revisit(next_start, interval, next_origin); + assert!( + sweep.idle_slots.is_empty(), + "newly eligible finished slots must be revisited" + ); + assert!(!sweep.cursors[&0][0].done); + assert_eq!( + sweep.cursors[&1][0].first, + Some((stamp, 100, 1)), + "ordinary revisits must retain unfinished sweep progress" + ); + } + + #[test] + fn ordinary_rescue_revisit_preserves_unfinished_cursor_under_jitter() { + let interval = Duration::from_secs(30); + let now = Instant::now(); + let first_start = now + interval + Duration::from_millis(2); + let second_start = now + interval * 2; + let stamp = Utc::now(); + let mut sweep = RescueSweep::default(); + sweep.revisit(first_start, interval, RescueTickOrigin::Ordinary); + for (slot, saturated) in [(0, false), (1, true)] { + sweep.observe(RescueScanProgress { + slot, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 100, 1)), + saturated, + }], + }); + } + assert_eq!(sweep.idle_slots, vec![0]); + sweep.revisit(second_start, interval, RescueTickOrigin::Ordinary); + assert!( + sweep.idle_slots.is_empty(), + "finished slot must be revisited on the ordinary tick" + ); + assert!(!sweep.cursors[&0][0].done); + assert_eq!( + sweep.cursors[&1][0].first, + Some((stamp, 100, 1)), + "unfinished sweep progress must survive the same ordinary revisit" + ); + } + + #[tokio::test(start_paused = true)] + async fn delayed_rescue_baseline_does_not_turn_early_rearm_into_baseline() { + let interval = Duration::from_secs(30); + let mut timer = RescueTimer::new(interval); + timer.tick().await; + tokio::time::advance(Duration::from_secs(52)).await; + let first_origin = timer.tick().await; + let first_start = tokio::time::Instant::now().into_std(); + let mut cadence = RescueHeartbeatCadence::default(); + assert_eq!( + cadence.select(first_start, interval, None, first_origin), + Some(None) + ); + timer.reset_after(Duration::from_secs(8)); + let origin = timer.tick().await; + let now = tokio::time::Instant::now().into_std(); + assert_eq!(now.duration_since(first_start), Duration::from_secs(8)); + assert_eq!(origin, RescueTickOrigin::Backlog); + assert_eq!(cadence.select(now, interval, None, origin), None, + "a tick scheduled 30s but serviced at52s must not make backlog60s an unguarded baseline"); + let cutoff = Utc::now(); + assert_eq!( + cadence.select(now, interval, Some(cutoff), origin), + Some(Some(cutoff)), + "an early re-arm must retain the owner cutoff" + ); + let origin = timer.tick().await; + let now = tokio::time::Instant::now().into_std(); + assert_eq!(origin, RescueTickOrigin::Ordinary); + assert_eq!(now.duration_since(first_start), Duration::from_secs(38)); + assert_eq!(cadence.select(now, interval, None, origin), Some(None)); + } + + #[tokio::test(start_paused = true)] + async fn rescue_backlog_tick_is_consumed_even_when_branch_cooldown_skips_it() { + let (tracker, clock, cooldown_until) = tracker_flipped_by_250ms_bodies(); + let metrics = metrics_for_test(); + let mut timer = RescueTimer::new(Duration::from_secs(30)); + timer.tick().await; + timer.reset_after(TICK_100MS); + let origin = timer.tick().await; + assert_eq!(origin, RescueTickOrigin::Backlog); + let skipped_at = clock + Duration::from_nanos(1); + assert!(skipped_at < cooldown_until); + assert!(tracker + .try_begin_at("cleanup", TICK_100MS, &metrics, true, skipped_at) + .is_none()); + // tick() consumed the marker before try_begin could skip the work. + assert_eq!(timer.tick().await, RescueTickOrigin::Ordinary); + } + + #[test] + fn rescue_accelerated_first_batch_requires_owner_evidence() { + let mut cadence = RescueHeartbeatCadence::default(); + let now = Instant::now(); + let interval = Duration::from_secs(30); + let cutoff = Utc::now(); + assert_eq!( + cadence.select(now, interval, None, RescueTickOrigin::Backlog), + Some(None) + ); + assert_eq!( + cadence.select( + now + Duration::from_secs(1), + interval, + Some(cutoff), + RescueTickOrigin::Backlog + ), + Some(Some(cutoff)) + ); + // The preceding accelerated pass re-armed the timer, but owners + // became unconfirmed before it fired: no ordinary batch may slip in. + assert_eq!( + cadence.select( + now + Duration::from_secs(2), + interval, + None, + RescueTickOrigin::Backlog + ), + None + ); + assert_eq!( + cadence.select(now + interval, interval, None, RescueTickOrigin::Backlog), + Some(None) + ); + } + + #[test] + fn rescue_closed_history_sweep_stops_after_wrap() { + let mut sweep = RescueSweep::default(); + let stamp = Utc::now(); + // First window ends midway around a slot. Sweep state survives + // the pass's time budget and the next early timer fire. + for id in [20, 40, 60, 10] { + sweep.observe(RescueScanProgress { + slot: 1, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, id, 1)), + saturated: true, + }], + }); + assert!(sweep.idle_slots.is_empty()); + } + sweep.observe(RescueScanProgress { + slot: 1, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 20, 1)), + saturated: true, + }], + }); + assert_eq!(sweep.idle_slots, vec![1]); + + let mut one = RescueSweep::default(); + for _ in 0..2 { + one.observe(RescueScanProgress { + slot: 0, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 1, 1)), + saturated: true, + }], + }); + } + assert_eq!(one.idle_slots, vec![0]); + } + + #[test] + fn rescue_sweep_revisits_newly_eligible_slots_at_the_ordinary_interval() { + let mut sweep = RescueSweep::default(); + let now = Instant::now(); + let interval = Duration::from_secs(30); + let stamp = Utc::now(); + sweep.revisit(now, interval, RescueTickOrigin::Backlog); + // Slot A's deadlines are in the future; B keeps a large pass open. + sweep.observe(RescueScanProgress { + slot: 0, + cursors: vec![awa_model::RescueScanCursor { + position: None, + saturated: false, + }], + }); + sweep.observe(RescueScanProgress { + slot: 1, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 100, 1)), + saturated: true, + }], + }); + sweep.revisit(now + interval / 2, interval, RescueTickOrigin::Backlog); + assert_eq!(sweep.idle_slots, vec![0]); + // A can now be expired (or contain claims from a reused generation). + sweep.revisit(now + interval, interval, RescueTickOrigin::Backlog); + assert!(sweep.idle_slots.is_empty()); + assert!(!sweep.cursors[&0][0].done); + assert_eq!( + sweep.cursors[&1][0].first, + Some((stamp, 100, 1)), + "revisiting an idle slot must not restart an unfinished sweep" + ); + } + + #[test] + fn rescue_sweep_detects_a_window_crossing_its_origin() { + let mut sweep = RescueSweep::default(); + let stamp = Utc::now(); + // A window can cross the origin and wrap again before its last key. + // Its end alone need never equal or exceed the original last key. + for id in [10_000, 1_999, 3_999, 5_999, 7_999, 9_999, 1_998] { + sweep.observe(RescueScanProgress { + slot: 0, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, id, 1)), + saturated: true, + }], + }); + } + assert_eq!(sweep.idle_slots, vec![0]); + } + + struct SparseHeartbeatFixture { + pool: PgPool, + runtime: QueueStorageRuntime, + service: MaintenanceService, + stale_job: i64, + fresh_job: i64, + fresh_runtime: Uuid, + } + + impl SparseHeartbeatFixture { + async fn new(budget: Duration, max_rows: usize, confirmed: bool) -> Self { + let pool = setup_pool(4).await; + reset_schema(&pool).await; + migrations::run(&pool).await.expect("migrations"); + let config = QueueStorageConfig::default(); + let store = QueueStorage::new(config.clone()).expect("queue storage"); + store.install(&pool).await.expect("install"); + let jobs: Vec<_> = (0..2) + .map(|id| awa_model::InsertParams { + kind: "sparse_heartbeat".into(), + args: serde_json::json!({ "id": id }), + opts: awa_model::InsertOpts { + queue: "sparse_heartbeat".into(), + max_attempts: 3, + ..Default::default() + }, + }) + .collect(); + store + .enqueue_params_batch(&pool, &jobs) + .await + .expect("enqueue"); + let claims = store + .claim_runtime_batch(&pool, "sparse_heartbeat", 2, Duration::from_secs(3600)) + .await + .expect("claim real jobs"); + assert_eq!(claims.len(), 2); + let slot = claims[0].claim.claim_slot; + assert_eq!(claims[1].claim.claim_slot, slot); + let stamp = Utc::now() - chrono::Duration::hours(2); + + // Keep the old row-local representation valid on both versions. + // Closed history needs no payload; the two open claims retain their + // real ready rows so rescue exercises the actual disposition path. + sqlx::query("DELETE FROM awa.lease_claim_batches WHERE queue = 'sparse_heartbeat'") + .execute(&pool) + .await + .expect("replace compact claims with legacy rows"); + for (index, claim) in claims.iter().enumerate() { + sqlx::query( + "INSERT INTO awa.lease_claims + (claim_slot, receipt_id, job_id, run_lease, ready_slot, ready_generation, + queue, priority, attempt, max_attempts, lane_seq, enqueue_shard, claimed_at, deadline_at) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,clock_timestamp()+interval '1 hour') + ON CONFLICT (claim_slot,job_id,run_lease) DO UPDATE + SET claimed_at=EXCLUDED.claimed_at, deadline_at=EXCLUDED.deadline_at", + ) + .bind(slot).bind(claim.claim.receipt_id).bind(claim.job.id).bind(claim.job.run_lease) + .bind(claim.claim.ready_slot).bind(claim.claim.ready_generation).bind(&claim.claim.queue) + .bind(claim.claim.priority).bind(claim.job.attempt).bind(claim.job.max_attempts) + .bind(claim.claim.lane_seq).bind(claim.claim.enqueue_shard) + .bind(stamp + chrono::Duration::microseconds(if index == 0 { 3500 } else { 4500 })) + .execute(&pool).await.expect("legacy open receipt"); + } + sqlx::query( + "INSERT INTO awa.lease_claims + (claim_slot,job_id,run_lease,ready_slot,ready_generation,queue,priority, + attempt,max_attempts,lane_seq,claimed_at,closed_at) + SELECT $1,1000000+id,1,0,0,'sparse_history',0,1,3,id, + $2::timestamptz+id*interval '1 microsecond',clock_timestamp() + FROM generate_series(1,6000) AS id", + ) + .bind(slot) + .bind(stamp) + .execute(&pool) + .await + .expect("sparse closed prefix"); + let stale_job = claims[0].job.id; + let fresh_job = claims[1].job.id; + sqlx::query( + "INSERT INTO awa.attempt_state(job_id,run_lease,heartbeat_at,updated_at) + VALUES($1,$2,clock_timestamp(),clock_timestamp())", + ) + .bind(fresh_job) + .bind(claims[1].job.run_lease) + .execute(&pool) + .await + .expect("protect fresh claim"); + let fresh_runtime = Uuid::new_v4(); + for (id, seen, heartbeat_alive) in [ + (Uuid::new_v4(), stamp + chrono::Duration::hours(1), false), + (fresh_runtime, Utc::now(), true), + ] { + sqlx::query( + "INSERT INTO awa.runtime_instances + (instance_id,hostname,pid,version,started_at,last_seen_at,snapshot_interval_ms, + healthy,postgres_connected,poll_loop_alive,heartbeat_alive,shutting_down,leader) + VALUES($1,'sparse_test',1,'0.6.13',$2,$2,1000,true,true,true,$3,false,false)", + ) + .bind(id).bind(seen).bind(heartbeat_alive) + .execute(&pool).await.expect("runtime evidence"); + } + let runtime = + QueueStorageRuntime::new(config, Duration::from_secs(1), Duration::from_secs(1)) + .expect("runtime"); + let service = MaintenanceService::new( + pool.clone(), + metrics_for_test(), + Arc::new(AtomicBool::new(true)), + Arc::new(AtomicBool::new(true)), + CancellationToken::new(), + Arc::new(Vec::new()), + InFlightMap::default(), + RuntimeStorage::QueueStorage(runtime.clone()), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + ) + .rescue_pass_budget(budget) + .rescue_pass_max_rows(max_rows); + if confirmed { + let now = Instant::now(); + *service + .rescue_owner_evidence + .lock() + .expect("owner evidence") = RescueOwnerEvidence { + steady_since: Some(now - service.heartbeat_staleness * 2), + last_sample: Some(now), + }; + } + Self { + pool, + runtime, + service, + stale_job, + fresh_job, + fresh_runtime, + } + } + + async fn pass(&self) -> bool { + sqlx::query("UPDATE awa.runtime_instances SET last_seen_at=clock_timestamp() WHERE instance_id=$1") + .bind(self.fresh_runtime).execute(&self.pool).await.expect("fresh owner snapshot"); + self.service + .rescue_queue_storage_pass( + RescueKind::Heartbeat, + &self.runtime, + false, + RescueTickOrigin::Backlog, + ) + .await + } + + async fn assert_dispositions(&self) { + let stale = self + .runtime + .store + .load_job(&self.pool, self.stale_job) + .await + .expect("stale job lookup") + .expect("stale job"); + assert_eq!(stale.state, JobState::Retryable); + let fresh = self + .runtime + .store + .load_job(&self.pool, self.fresh_job) + .await + .expect("fresh job lookup") + .expect("fresh job"); + assert_eq!( + fresh.state, + JobState::Running, + "fresh heartbeat stays protected" + ); + let closures: i64 = + sqlx::query_scalar("SELECT count(*) FROM awa.lease_claim_closures WHERE job_id=$1") + .bind(self.stale_job) + .fetch_one(&self.pool) + .await + .expect("closure count"); + assert_eq!(closures, 1, "sparse stale attempt is rescued exactly once"); + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_confirmed_heartbeat_sparse_sweep_progresses_and_stops() { + let _guard = db_test_mutex().lock().await; + // A positive budget below any database round trip makes each pass + // finish its first batch and rearm; no elapsed-time assertion is used. + let run = SparseHeartbeatFixture::new(Duration::from_nanos(1), 5000, true).await; + assert!(run.pass().await, "closed prefix leaves a resumable sweep"); + let cursor: i64 = sqlx::query_scalar( + "SELECT rescue_cursor_job_id FROM awa.claim_ring_slots WHERE slot=0", + ) + .fetch_one(&run.pool) + .await + .expect("first persisted cursor"); + assert_eq!( + cursor, 1002000, + "confirmed bulk starts with one small window" + ); + let stale = run + .runtime + .store + .load_job(&run.pool, run.stale_job) + .await + .expect("unvisited job lookup") + .expect("unvisited job"); + assert_eq!( + stale.state, + JobState::Running, + "sparse job lies beyond the first window" + ); + let mut drained = false; + for _ in 0..12 { + if !run.pass().await { + drained = true; + break; + } + } + assert!( + drained, + "persisted sparse sweep converges instead of rearming forever" + ); + assert!(run.service.rescue_sweeps.lock().expect("sweeps").is_empty()); + run.assert_dispositions().await; + run.pool.close().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_disabled_acceleration_keeps_wide_first_heartbeat_window() { + let _guard = db_test_mutex().lock().await; + for (budget, max_rows) in [(Duration::ZERO, 5000), (Duration::from_secs(1), 0)] { + let run = SparseHeartbeatFixture::new(budget, max_rows, true).await; + assert!(!run.pass().await, "disabled acceleration never rearms"); + run.assert_dispositions().await; + run.pool.close().await; + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_unconfirmed_heartbeat_keeps_wide_first_window() { + let _guard = db_test_mutex().lock().await; + let run = SparseHeartbeatFixture::new(Duration::from_nanos(1), 5000, false).await; + assert!( + !run.pass().await, + "unconfirmed owners keep the ordinary one-batch cadence" + ); + run.assert_dispositions().await; + run.pool.close().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_closed_history_stops_rearming() { + let _guard = db_test_mutex().lock().await; + let pool = setup_pool(4).await; + reset_schema(&pool).await; + migrations::run(&pool).await.expect("migrations"); + let config = QueueStorageConfig::default(); + let store = QueueStorage::new(config.clone()).expect("queue storage"); + store.install(&pool).await.expect("install"); + sqlx::query( + "INSERT INTO awa.lease_claims + (claim_slot, job_id, run_lease, ready_slot, ready_generation, + queue, priority, attempt, max_attempts, lane_seq, claimed_at, closed_at, deadline_at) + SELECT 0, id, 1, 0, 0, 'closed_history', 0, 1, 3, id, + now() - interval '1 hour', now(), now() - interval '1 minute' + FROM generate_series(1, 10001) AS id", + ) + .execute(&pool) + .await + .expect("seed more closed history than a scan window"); + let runtime = + QueueStorageRuntime::new(config, Duration::from_secs(1), Duration::from_secs(1)) + .expect("runtime"); + let service = MaintenanceService::new( + pool.clone(), + metrics_for_test(), + Arc::new(AtomicBool::new(true)), + Arc::new(AtomicBool::new(true)), + CancellationToken::new(), + Arc::new(Vec::new()), + InFlightMap::default(), + RuntimeStorage::QueueStorage(runtime.clone()), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + ) + .rescue_pass_budget(Duration::from_millis(200)); + let mut drained = false; + for _ in 0..20 { + if !service + .rescue_queue_storage_pass( + RescueKind::Deadline, + &runtime, + false, + RescueTickOrigin::Backlog, + ) + .await + { + drained = true; + break; + } + } + assert!( + drained, + "closed history must finish a bounded sweep instead of rearming forever" + ); + assert!(service.rescue_sweeps.lock().expect("sweeps").is_empty()); + // Zero is a documented way to retain the ordinary one-batch cadence. + let service = service.rescue_pass_max_rows(0); + assert!( + !service + .rescue_queue_storage_pass( + RescueKind::Deadline, + &runtime, + false, + RescueTickOrigin::Backlog + ) + .await + ); + assert!(service.rescue_sweeps.lock().expect("sweeps").is_empty()); + pool.close().await; + } + + #[test] + fn rescue_owner_evidence_needs_a_full_settled_staleness_window() { + let staleness = Duration::from_secs(90); + let sample_every = Duration::from_secs(30); + let mut evidence = RescueOwnerEvidence::default(); + let mut clock = Instant::now(); + + for _ in 0..3 { + assert!(!evidence.observe(clock, true, staleness)); + clock += sample_every; + } + assert!( + evidence.observe(clock, true, staleness), + "settled at every sample across the staleness window" + ); + + clock += sample_every; + assert!( + !evidence.observe(clock, false, staleness), + "an unsettled sample closes the gate at once" + ); + for _ in 0..3 { + clock += sample_every; + assert!( + !evidence.observe(clock, true, staleness), + "the window restarts after an unsettled sample" + ); + } + clock += sample_every; + assert!(evidence.observe(clock, true, staleness)); + } + + #[test] + fn rescue_owner_evidence_restarts_after_an_unsampled_gap() { + let staleness = Duration::from_secs(90); + let mut evidence = RescueOwnerEvidence::default(); + let mut clock = Instant::now(); + for _ in 0..4 { + evidence.observe(clock, true, staleness); + clock += Duration::from_secs(30); + } + assert!(evidence.observe(clock, true, staleness)); + + // The leader did not sample for longer than half the staleness: a + // stall in that gap would have been invisible, so start over. + clock += Duration::from_secs(46); + assert!(!evidence.observe(clock, true, staleness)); + clock += Duration::from_secs(45); + assert!(!evidence.observe(clock, true, staleness)); + clock += Duration::from_secs(45); + assert!(evidence.observe(clock, true, staleness)); + } + #[test] fn branch_tracker_without_cooldown_tracks_overrun_but_keeps_running() { let tracker = MaintenanceBranchTracker::new(); diff --git a/awa-worker/src/runtime.rs b/awa-worker/src/runtime.rs index 2332a9ca8..c6e01631a 100644 --- a/awa-worker/src/runtime.rs +++ b/awa-worker/src/runtime.rs @@ -23,6 +23,9 @@ pub(crate) struct InFlightState { /// queue storage with one; lets the heartbeat address the claim batch /// directly. pub claim: Option, + /// The handler has returned and the attempt waits for its finalize. It is + /// still heartbeated, but its progress is settled by the finalize. + pub finalizing: bool, } /// Mutable progress buffer shared between handler and heartbeat service. @@ -163,7 +166,7 @@ impl InFlightRegistry { guard.get(&key).map(|s| s.cancel.clone()) } - /// Every in-flight attempt with its compact claim identity, if any. + /// Every attempt to heartbeat, with its compact claim identity, if any. pub fn claim_refs(&self) -> Vec<(InFlightKey, Option)> { let mut refs = Vec::new(); for shard in &*self.shards { @@ -173,6 +176,18 @@ impl InFlightRegistry { refs } + /// The handler has returned: keep heartbeating the attempt until its + /// finalize is done, and stop flushing its progress. + pub fn mark_finalizing(&self, key: InFlightKey) { + let mut guard = self + .shard_for(&key) + .write() + .expect("in_flight shard poisoned"); + if let Some(state) = guard.get_mut(&key) { + state.finalizing = true; + } + } + pub fn flags(&self) -> Vec { let mut flags = Vec::new(); for shard in &*self.shards { @@ -192,6 +207,9 @@ impl InFlightRegistry { for shard in &*self.shards { let guard = shard.read().expect("in_flight shard poisoned"); for (&(job_id, run_lease), state) in guard.iter() { + if state.finalizing { + continue; + } let mut progress = state.progress.lock().expect("progress lock poisoned"); if progress.has_pending() && progress.in_flight.is_none() { if let Some(ref value) = progress.latest { @@ -278,6 +296,7 @@ mod tests { cancel: Arc::new(AtomicBool::new(false)), progress: progress.clone(), claim: None, + finalizing: false, }, ); diff --git a/awa/tests/bulk_rescue_test.rs b/awa/tests/bulk_rescue_test.rs new file mode 100644 index 000000000..69803186d --- /dev/null +++ b/awa/tests/bulk_rescue_test.rs @@ -0,0 +1,908 @@ +//! A runtime holding tens of thousands of receipt claims dies without a +//! graceful shutdown. Maintenance must rescue the whole backlog in bounded +//! time, the claim ring must resume rotating, a runtime that stays alive +//! must never have a claim rescued, and every rescued attempt must count +//! (#627). + +#[path = "ci_timing.rs"] +mod ci_timing; + +use awa::audited_sql; +use awa::model::{QueueStorage, QueueStorageConfig, RingRotationCadence}; +use awa::{Client, InsertOpts, JobArgs, JobEvent, JobResult, QueueConfig}; +use awa_testing::setup::TestDatabase; +use ci_timing::{scaled_staleness, scaled_timeout}; +use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData}; +use opentelemetry_sdk::metrics::{InMemoryMetricExporter, SdkMeterProvider}; +use serde::{Deserialize, Serialize}; +use sqlx::postgres::PgPoolOptions; +use sqlx::PgPool; +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +const ABANDONED: i64 = 50_000; +const ABANDONED_CHUNK: i64 = 10_000; +/// Every tenth abandoned job has one attempt, so its rescue exhausts it. +const EXHAUSTED_EVERY: i64 = 10; +const ABANDONED_QUEUE: &str = "bulk_rescue_abandoned"; +const LIVE_QUEUE: &str = "bulk_rescue_live"; + +/// What each run of a job saw, by job id. +type RunsByJob = HashMap>; + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct AbandonedJob { + n: i64, +} + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct LiveJob { + n: i64, + hold_ms: u64, +} + +fn queue_storage_builder( + pool: &PgPool, + staleness: Duration, + heartbeat_interval: Duration, +) -> awa::ClientBuilder { + Client::builder(pool.clone()) + .queue_storage( + QueueStorageConfig::default(), + Duration::from_secs(1), + Duration::from_secs(1), + ) + // Rotate the claim ring every tick so the abandoned claims spread + // across slots and the ring wraps onto the oldest of them. + .claim_rotate_interval(Duration::from_millis(200)) + .ring_rotation_cadence(RingRotationCadence::EAGER) + .heartbeat_interval(heartbeat_interval) + .heartbeat_staleness(staleness) + .runtime_snapshot_interval(Duration::from_millis(250)) + .leader_check_interval(Duration::from_millis(200)) + .promote_interval(Duration::from_millis(100)) + .dlq_enabled_by_default(true) +} + +async fn count(pool: &PgPool, sql: String) -> i64 { + sqlx::query_scalar(audited_sql(sql.clone())) + .fetch_one(pool) + .await + .unwrap_or_else(|err| panic!("{sql}: {err}")) +} + +/// Retry promotion transfers a row between these states in one transaction. +/// Read both counts from one statement so a commit between reads cannot count +/// the same job once in deferred storage and again in ready storage. +async fn waiting_retry_counts(pool: &PgPool, queue: &str) -> (i64, i64) { + let overview = awa::model::admin::queue_overview(pool, queue) + .await + .expect("retry accounting overview") + .expect("retry accounting queue"); + (overview.retryable, overview.available) +} + +#[tokio::test] +async fn retry_accounting_uses_one_snapshot_during_promotion() { + let db = TestDatabase::queue_storage().await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let queue = "bulk_rescue_retry_snapshot"; + let params = awa::model::insert::params_with( + &AbandonedJob { n: 0 }, + InsertOpts { + queue: queue.into(), + max_attempts: 3, + ..Default::default() + }, + ) + .expect("retry accounting job"); + store + .enqueue_params_batch(db.pool(), &[params]) + .await + .expect("enqueue retry accounting job"); + let claims = store + .claim_runtime_batch(db.pool(), queue, 1, Duration::ZERO) + .await + .expect("claim retry accounting job"); + assert_eq!(claims.len(), 1); + store + .retry_after( + db.pool(), + claims[0].job.id, + claims[0].job.run_lease, + Duration::ZERO, + None, + ) + .await + .expect("retry accounting job") + .expect("claimed job becomes retryable"); + assert_eq!(waiting_retry_counts(db.pool(), queue).await, (1, 0)); + + // Hold the observer's only connection after its first statement. A real + // promotion commits before any second observer statement can start. The + // old split read therefore counts this one job twice; the overview's two + // counts always describe the same statement snapshot. + let armed = Arc::new(AtomicBool::new(false)); + let read_finished = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let observer = PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(scaled_timeout(Duration::from_secs(10))) + .after_release({ + let armed = armed.clone(); + let read_finished = read_finished.clone(); + let release = release.clone(); + move |_, _| { + let gate = armed.swap(false, Ordering::SeqCst); + let read_finished = read_finished.clone(); + let release = release.clone(); + Box::pin(async move { + if gate { + read_finished.notify_one(); + release.notified().await; + } + Ok(true) + }) + } + }) + .connect(&db.url()) + .await + .expect("retry accounting observer pool"); + let primed = observer.acquire().await.expect("prime observer connection"); + drop(primed); + tokio::time::timeout(scaled_timeout(Duration::from_secs(10)), async { + while observer.num_idle() != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("observer's initial connection is idle before arming"); + armed.store(true, Ordering::SeqCst); + let promote = async { + let observed = tokio::time::timeout( + scaled_timeout(Duration::from_secs(10)), + read_finished.notified(), + ) + .await; + if observed.is_err() { + release.notify_one(); + panic!("retry accounting observer did not finish its first statement"); + } + let promoted = store + .promote_due(db.pool(), awa::JobState::Retryable, 1) + .await; + release.notify_one(); + assert_eq!(promoted.expect("promote between accounting reads"), 1); + }; + let (counts, ()) = tokio::time::timeout(scaled_timeout(Duration::from_secs(10)), async { + tokio::join!(waiting_retry_counts(&observer, queue), promote) + }) + .await + .expect("bounded retry accounting interleaving"); + observer.close().await; + assert_eq!(counts.0 + counts.1, 1, "one job is waiting to retry"); + assert_eq!(waiting_retry_counts(db.pool(), queue).await, (0, 1)); +} + +async fn running(pool: &PgPool, store: &QueueStorage, queue: &str) -> i64 { + store + .queue_counts(pool, queue) + .await + .expect("queue counts") + .running +} + +async fn claim_ring_generation(pool: &PgPool) -> i64 { + count( + pool, + "SELECT generation FROM awa.claim_ring_rotations ORDER BY generation DESC LIMIT 1".into(), + ) + .await +} + +async fn abandoned_claim_slots(pool: &PgPool) -> i64 { + count( + pool, + format!( + "SELECT count(DISTINCT claim_slot) FROM ( + SELECT claim_slot FROM awa.lease_claim_batches WHERE queue = '{ABANDONED_QUEUE}' + UNION ALL + SELECT claim_slot FROM awa.lease_claims WHERE queue = '{ABANDONED_QUEUE}' + ) AS claims" + ), + ) + .await +} + +/// Rescued and dead-lettered rows for a queue: what a false rescue would leave. +async fn rescued_rows(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT (SELECT count(*) FROM awa.deferred_jobs WHERE queue = $1 AND state = 'retryable') + + (SELECT count(*) FROM awa.dlq_entries WHERE queue = $1)", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("count rescued rows") +} + +/// Failure-only evidence: never delay a passing test or replace its assertion. +/// Row caps bound output; statement and outer deadlines also bound the scans. +async fn live_rescue_diagnostic( + pool: &PgPool, + staleness: Duration, + heartbeat_interval: Duration, +) -> serde_json::Value { + let capture = async { + let mut tx = pool.begin().await?; + sqlx::query("SET TRANSACTION READ ONLY") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL statement_timeout = '250ms'") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL lock_timeout = '100ms'") + .execute(tx.as_mut()) + .await?; + let result = sqlx::query_scalar::<_, serde_json::Value>( + r#" + WITH stamp AS MATERIALIZED (SELECT clock_timestamp() AS at), + rescued AS MATERIALIZED ( + SELECT * FROM ( + SELECT 'retryable' AS plane, job_id, run_lease, attempt, + attempted_at, finalized_at, NULL::text AS dlq_reason, + NULL::bigint AS original_run_lease, payload + FROM awa.deferred_jobs WHERE queue = $1 AND state = 'retryable' + UNION ALL + SELECT 'dlq', job_id, run_lease, attempt, + attempted_at, finalized_at, dlq_reason, original_run_lease, payload + FROM awa.dlq_entries WHERE queue = $1 + ) retained ORDER BY job_id LIMIT 16 + ), + claims AS ( + SELECT batch.claim_slot, batch.batch_id, batch.ready_slot, + batch.ready_generation, batch.claimed_at, item.job_id, + item.run_lease, item.receipt_id, attempt.heartbeat_at, + extract(epoch FROM stamp.at - COALESCE(attempt.heartbeat_at, batch.claimed_at)) AS age_s + FROM awa.lease_claim_batches batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases, batch.receipt_ids) + AS item(job_id, run_lease, receipt_id) + CROSS JOIN stamp + LEFT JOIN awa.attempt_state attempt + ON attempt.job_id = item.job_id AND attempt.run_lease = item.run_lease + WHERE batch.queue = $1 AND item.job_id IN (SELECT job_id FROM rescued) + LIMIT 16 + ), + runtimes AS ( + SELECT instance_id, last_seen_at, heartbeat_alive, snapshot_interval_ms, + extract(epoch FROM stamp.at - last_seen_at) AS age_s + FROM awa.runtime_instances CROSS JOIN stamp LIMIT 8 + ), + activity AS ( + SELECT pid, state, wait_event_type, wait_event, + extract(epoch FROM stamp.at - query_start) AS query_age_s, + pg_blocking_pids(pid) AS blocking_pids, left(query, 256) AS query + FROM pg_stat_activity CROSS JOIN stamp + WHERE datname = current_database() AND pid <> pg_backend_pid() AND state <> 'idle' + ORDER BY query_start LIMIT 24 + ), + waiting AS ( + SELECT locks.pid, locks.locktype, locks.mode, relation.relname + FROM pg_locks locks LEFT JOIN pg_class relation ON relation.oid = locks.relation + WHERE NOT locks.granted + AND (locks.database IS NULL OR locks.database = (SELECT oid FROM pg_database WHERE datname = current_database())) + LIMIT 32 + ) + SELECT jsonb_build_object( + 'at', (SELECT at FROM stamp), 'database', current_database(), + 'rescued', (SELECT COALESCE(jsonb_agg(rescued), '[]') FROM rescued), + 'claims', (SELECT COALESCE(jsonb_agg(claims), '[]') FROM claims), + 'runtimes', (SELECT COALESCE(jsonb_agg(runtimes), '[]') FROM runtimes), + 'claim_slots', (SELECT COALESCE(jsonb_agg(ring_slot), '[]') FROM (SELECT * FROM awa.claim_ring_slots ORDER BY slot LIMIT 8) ring_slot), + 'activity', (SELECT COALESCE(jsonb_agg(activity), '[]') FROM activity), + 'waiting_locks', (SELECT COALESCE(jsonb_agg(waiting), '[]') FROM waiting), + 'wal', (SELECT to_jsonb(wal) FROM pg_stat_wal wal)) + "#, + ).bind(ABANDONED_QUEUE).fetch_one(tx.as_mut()).await; + tx.rollback().await?; + result + }; + let evidence = match tokio::time::timeout(Duration::from_secs(3), capture).await { + Ok(Ok(data)) => serde_json::json!({"status": "captured", "data": data}), + Ok(Err(error)) => serde_json::json!({"status": "incomplete", "error": error.to_string()}), + Err(_) => serde_json::json!({"status": "incomplete", "error": "3s diagnostic deadline"}), + }; + serde_json::json!({ + "staleness_ms": staleness.as_millis(), + "heartbeat_interval_ms": heartbeat_interval.as_millis(), + "CI": std::env::var("CI").ok(), + "multiplier_override": std::env::var("AWA_CHAOS_TIMEOUT_MULTIPLIER").ok(), + "evidence": evidence, + "caveat": "current bounded snapshot, not historical heartbeat proof; missing attempt_state may follow rescue" + }) +} + +#[tokio::test] +async fn live_rescue_diagnostic_captures_bounded_compact_evidence() { + let db = TestDatabase::queue_storage().await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let params: Vec<_> = (0..20) + .map(|n| { + awa::model::insert::params_with( + &AbandonedJob { n }, + InsertOpts { + queue: ABANDONED_QUEUE.into(), + max_attempts: 1, + ..Default::default() + }, + ) + .expect("diagnostic job") + }) + .collect(); + store + .enqueue_params_batch(db.pool(), ¶ms) + .await + .expect("enqueue"); + let claims = store + .claim_runtime_batch(db.pool(), ABANDONED_QUEUE, 20, Duration::ZERO) + .await + .expect("claim"); + assert_eq!(claims.len(), 20); + assert!(claims + .iter() + .all(|claim| claim.claim.claim_batch_id.is_some())); + let rescued = store + .rescue_stale_heartbeats_with_dlq(db.pool(), Duration::ZERO, &|_| true) + .await + .expect("force diagnostic evidence"); + assert_eq!(rescued.len(), 20); + let report = live_rescue_diagnostic( + db.pool(), + Duration::from_secs(18), + Duration::from_millis(2250), + ) + .await; + assert_eq!(report["evidence"]["status"], "captured", "{report}"); + let data = &report["evidence"]["data"]; + assert_eq!(data["rescued"].as_array().unwrap().len(), 16); + assert_eq!(data["rescued"][0]["dlq_reason"], "heartbeat_stale"); + assert_eq!(data["claims"].as_array().unwrap().len(), 16); + assert!(data["claim_slots"][0]["generation"].is_number()); + for item in data["claims"].as_array().unwrap() { + assert!(claims.iter().any( + |claim| item["job_id"] == claim.job.id && item["run_lease"] == claim.job.run_lease + )); + } + assert_eq!( + rescued_rows(db.pool(), ABANDONED_QUEUE).await, + 20, + "diagnostic is read-only" + ); +} + +#[tokio::test] +async fn live_rescue_diagnostic_query_error_is_incomplete() { + let db = TestDatabase::queue_storage().await; + let mut gate = db.pool().begin().await.expect("begin diagnostic lock gate"); + sqlx::query("LOCK TABLE awa.lease_claim_batches IN ACCESS EXCLUSIVE MODE") + .execute(gate.as_mut()) + .await + .expect("hold the diagnostic claim relation"); + let report = live_rescue_diagnostic( + db.pool(), + Duration::from_secs(18), + Duration::from_millis(2250), + ) + .await; + gate.rollback().await.expect("release diagnostic lock gate"); + assert_eq!(report["evidence"]["status"], "incomplete"); + assert!(report["evidence"]["error"] + .as_str() + .unwrap() + .contains("lock timeout")); +} + +async fn wait_until(timeout: Duration, what: &str, mut done: F) -> Duration +where + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + let started = Instant::now(); + loop { + if done().await { + return started.elapsed(); + } + assert!( + started.elapsed() < timeout, + "timed out after {timeout:?} waiting for {what}" + ); + tokio::time::sleep(Duration::from_millis(250)).await; + } +} + +fn largest_heartbeat_rescue_pass(exporter: &InMemoryMetricExporter) -> u64 { + let mut largest = 0; + for resource in exporter.get_finished_metrics().expect("read metrics") { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + if metric.name() != "awa.maintenance.rescue.pass_size" { + continue; + } + if let AggregatedMetrics::U64(MetricData::Histogram(histogram)) = metric.data() { + for point in histogram.data_points() { + let heartbeat = point.attributes().any(|attribute| { + attribute.key.as_str() == "awa.rescue.kind" + && attribute.value.as_str() == "heartbeat" + }); + if heartbeat { + largest = largest.max(point.max().unwrap_or(0)); + } + } + } + } + } + } + largest +} + +/// A client on its own Tokio runtime, so it can be killed the way SIGKILL +/// would: every task is dropped mid-flight, no shutdown runs, and its +/// connections close with their transactions rolled back. +struct DoomedClient { + kill: Option>, + thread: Option>, +} + +impl DoomedClient { + fn start(url: String, staleness: Duration, heartbeat_interval: Duration) -> Self { + let (kill, killed) = tokio::sync::oneshot::channel::<()>(); + let (started_tx, started_rx) = std::sync::mpsc::channel::<()>(); + let thread = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .expect("doomed client runtime"); + let client = runtime.block_on(async { + let pool = PgPoolOptions::new() + .max_connections(20) + .connect(&url) + .await + .expect("doomed client pool"); + let client = queue_storage_builder(&pool, staleness, heartbeat_interval) + .queue( + ABANDONED_QUEUE, + QueueConfig { + max_workers: ABANDONED as u32 + 1_000, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claimers: 4, + claim_batch_size: 1_000, + ..QueueConfig::default() + }, + ) + .register::(|_, _| async { + std::future::pending::>().await + }) + .build() + .expect("doomed client"); + client.start().await.expect("start doomed client"); + client + }); + started_tx.send(()).expect("report doomed client start"); + runtime.block_on(async { + let _ = killed.await; + }); + runtime.shutdown_background(); + drop(client); + }); + started_rx.recv().expect("doomed client started"); + Self { + kill: Some(kill), + thread: Some(thread), + } + } + + fn kill(mut self) { + if let Some(kill) = self.kill.take() { + let _ = kill.send(()); + } + if let Some(thread) = self.thread.take() { + thread.join().expect("doomed client thread"); + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn rescue_drains_an_abandoned_backlog_without_rescuing_live_claims() { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(); + let exporter = InMemoryMetricExporter::default(); + let meter_provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + opentelemetry::global::set_meter_provider(meter_provider.clone()); + + let db = TestDatabase::queue_storage().await; + let pool = db.pool().clone(); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let claim_slots = store.claim_slot_count() as i64; + + // Once the doomed runtime has claimed its backlog, live jobs run for + // three staleness windows, so a rescue that ignored their heartbeats + // would take them. Until then they finish at once, so their claims do + // not hold up the claim ring while the doomed runtime fills it. + let staleness = scaled_staleness(Duration::from_secs(6)); + let heartbeat_interval = staleness / 8; + let live_job_duration = staleness * 3; + + let live_pool = PgPoolOptions::new() + .max_connections(30) + .connect(&db.url()) + .await + .expect("live pool"); + let live_runs: Arc>> = Arc::default(); + // The live runtime leads maintenance, so it sees every rescue. + let rescues: Arc>> = Arc::default(); + let live = { + let live_runs = live_runs.clone(); + let live_rescues = rescues.clone(); + let abandoned_rescues = rescues.clone(); + queue_storage_builder(&live_pool, staleness, heartbeat_interval) + // Under half the staleness, so the leader samples runtime + // snapshots often enough to confirm the dead owner. + .heartbeat_rescue_interval(Duration::from_secs(2)) + .deadline_rescue_interval(Duration::from_secs(2)) + .leader_election_interval(Duration::from_millis(100)) + .queue( + LIVE_QUEUE, + QueueConfig { + max_workers: 1_000, + poll_interval: Duration::from_millis(25), + deadline_duration: Duration::ZERO, + ..QueueConfig::default() + }, + ) + .on_event::(move |event| { + if let JobEvent::Rescued { job, .. } = event { + live_rescues + .lock() + .expect("rescues") + .entry(job.id) + .or_default() + .push((LIVE_QUEUE, job.attempt)); + } + async {} + }) + .on_event::(move |event| { + if let JobEvent::Rescued { job, .. } = event { + abandoned_rescues + .lock() + .expect("rescues") + .entry(job.id) + .or_default() + .push((ABANDONED_QUEUE, job.attempt)); + } + async {} + }) + .register::(move |job, ctx| { + live_runs + .lock() + .expect("live runs") + .entry(ctx.job.id) + .or_default() + .push(ctx.job.attempt); + async move { + tokio::time::sleep(Duration::from_millis(job.hold_ms)).await; + Ok(JobResult::Completed) + } + }) + .build() + .expect("live client") + }; + live.start().await.expect("start live client"); + wait_until( + scaled_timeout(Duration::from_secs(10)), + "live leader", + || async { live.health_check().await.leader }, + ) + .await; + + // The live runtime keeps claiming new work for the whole test. + let producing = Arc::new(AtomicBool::new(true)); + let live_hold_ms = Arc::new(AtomicU64::new(0)); + let producer = { + let pool = pool.clone(); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let producing = producing.clone(); + let live_hold_ms = live_hold_ms.clone(); + tokio::spawn(async move { + let mut n = 0; + let mut produced = 0usize; + while producing.load(Ordering::SeqCst) { + let params: Vec<_> = (0..5) + .map(|i| { + awa::model::insert::params_with( + &LiveJob { + n: n + i, + hold_ms: live_hold_ms.load(Ordering::SeqCst), + }, + InsertOpts { + queue: LIVE_QUEUE.into(), + ..Default::default() + }, + ) + .expect("live params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue live jobs"); + n += 5; + produced += params.len(); + tokio::time::sleep(Duration::from_millis(500)).await; + } + produced + }) + }; + + // The doomed runtime claims the backlog in chunks, each after the claim + // ring has rotated, so the abandoned claims span several slots. + let generation_before = claim_ring_generation(&pool).await; + let doomed = DoomedClient::start(db.url(), staleness, heartbeat_interval); + let claim_started = Instant::now(); + for chunk_start in (0..ABANDONED).step_by(ABANDONED_CHUNK as usize) { + let chunk_end = (chunk_start + ABANDONED_CHUNK).min(ABANDONED); + let params: Vec<_> = (chunk_start..chunk_end) + .map(|n| { + awa::model::insert::params_with( + &AbandonedJob { n }, + InsertOpts { + queue: ABANDONED_QUEUE.into(), + max_attempts: if n % EXHAUSTED_EVERY == 0 { 1 } else { 3 }, + ..Default::default() + }, + ) + .expect("abandoned params") + }) + .collect(); + let generation = claim_ring_generation(&pool).await; + for batch in params.chunks(5_000) { + store + .enqueue_params_batch(&pool, batch) + .await + .expect("enqueue abandoned jobs"); + } + wait_until( + scaled_timeout(Duration::from_secs(60)), + "the doomed runtime to claim the chunk", + || async { running(&pool, &store, ABANDONED_QUEUE).await >= chunk_end }, + ) + .await; + // Give the ring a chance to move on before the next chunk; once it + // wraps onto a slot of abandoned claims it stays there. + let rotate_deadline = Instant::now() + scaled_timeout(Duration::from_secs(3)); + while claim_ring_generation(&pool).await <= generation && Instant::now() < rotate_deadline { + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + let claimed_in = claim_started.elapsed(); + let rescued_while_alive = rescued_rows(&pool, ABANDONED_QUEUE).await; + if rescued_while_alive != 0 { + eprintln!( + "live-rescue diagnostic: {}", + live_rescue_diagnostic(&pool, staleness, heartbeat_interval).await + ); + } + assert_eq!( + rescued_while_alive, 0, + "rescue took a claim while its runtime was alive and heartbeating" + ); + let generation_at_kill = claim_ring_generation(&pool).await; + live_hold_ms.store(live_job_duration.as_millis() as u64, Ordering::SeqCst); + doomed.kill(); + let slots = abandoned_claim_slots(&pool).await; + assert!( + slots >= 2, + "the abandoned claims should span several claim slots: {slots} \ + (ring generation {generation_before} -> {generation_at_kill})" + ); + + // Before the drain: the claims become stale one staleness after the + // kill, and the owner evidence needs one more staleness window. After + // that the backlog must drain at the multi-batch rate, not 500 claims + // per heartbeat-rescue interval. + let drained_in = wait_until( + staleness * 2 + scaled_timeout(Duration::from_secs(90)), + "rescue to drain the abandoned claims", + || async { running(&pool, &store, ABANDONED_QUEUE).await == 0 }, + ) + .await; + let generation_drained = claim_ring_generation(&pool).await; + eprintln!( + "claimed {ABANDONED} in {claimed_in:?}; drained {drained_in:?} after the kill \ + (claims stale after {staleness:?})" + ); + + wait_until( + scaled_timeout(Duration::from_secs(10)), + "a rescue event for every abandoned job", + || async { rescues.lock().expect("rescues").len() as i64 >= ABANDONED }, + ) + .await; + type Rescued = (&'static str, i16); + let rescued_twice = |rescues: &RunsByJob| -> Vec<(i64, Vec)> { + rescues + .iter() + .filter(|(_, rescued)| rescued.as_slice() != [(ABANDONED_QUEUE, 1)]) + .take(5) + .map(|(id, rescued)| (*id, rescued.clone())) + .collect() + }; + { + let rescues = rescues.lock().expect("rescues"); + assert_eq!(rescues.len() as i64, ABANDONED); + let unexpected = rescued_twice(&rescues); + assert!( + unexpected.is_empty(), + "each abandoned job is rescued once and nothing else is: {unexpected:?}" + ); + } + + let exhausted = ABANDONED / EXHAUSTED_EVERY; + assert_eq!( + count( + &pool, + format!( + "SELECT count(*) FROM awa.dlq_entries + WHERE queue = '{ABANDONED_QUEUE}' AND dlq_reason = 'heartbeat_stale' AND attempt = 1" + ), + ) + .await, + exhausted, + "every one-attempt job is dead-lettered by its rescue" + ); + let (retryable, promoted) = waiting_retry_counts(&pool, ABANDONED_QUEUE).await; + assert_eq!( + retryable + promoted, + ABANDONED - exhausted, + "every other abandoned job is waiting to retry" + ); + meter_provider.force_flush().expect("flush metrics"); + let largest_pass = largest_heartbeat_rescue_pass(&exporter); + assert!( + largest_pass > 500, + "a stale-heartbeat pass should take more than one 500-claim batch once the \ + owner is confirmed gone; the largest took {largest_pass}" + ); + + // Live claims can hold a slot for a whole live job. + let rotated_in = wait_until( + live_job_duration * 2 + scaled_timeout(Duration::from_secs(30)), + "the claim ring to rotate through every slot after the drain", + || async { claim_ring_generation(&pool).await >= generation_drained + claim_slots }, + ) + .await; + eprintln!("claim ring rotated a full revolution {rotated_in:?} after the drain"); + + // A fresh runtime runs every retried job exactly once, on its second attempt. + let recovered: Arc>> = Arc::default(); + let recovery_pool = PgPoolOptions::new() + .max_connections(20) + .connect(&db.url()) + .await + .expect("recovery pool"); + let recovery = { + let recovered = recovered.clone(); + queue_storage_builder(&recovery_pool, staleness, heartbeat_interval) + .queue( + ABANDONED_QUEUE, + // Small claims, so finished attempts never wait on a long + // completion backlog while the recovery runtime is under test. + QueueConfig { + max_workers: 100, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: 100, + ..QueueConfig::default() + }, + ) + .register::(move |_, ctx| { + let rescue_recorded = ctx + .job + .errors + .as_ref() + .and_then(|errors| errors.last()) + .and_then(|error| error.get("error")) + .and_then(|error| error.as_str()) + .is_some_and(|error| error.contains("presumed dead")); + recovered + .lock() + .expect("recovered") + .entry(ctx.job.id) + .or_default() + .push((ctx.job.attempt, rescue_recorded)); + async { Ok(JobResult::Completed) } + }) + .build() + .expect("recovery client") + }; + recovery.start().await.expect("start recovery client"); + wait_until( + scaled_timeout(Duration::from_secs(120)), + "the retried jobs to complete", + || async { + count(&pool, format!( + "SELECT count(*) FROM awa.terminal_jobs WHERE queue = '{ABANDONED_QUEUE}' AND state = 'completed'" + )).await == ABANDONED - exhausted + }, + ) + .await; + recovery.shutdown(Duration::from_secs(10)).await; + { + let recovered = recovered.lock().expect("recovered"); + assert_eq!(recovered.len() as i64, ABANDONED - exhausted); + let unexpected: Vec<_> = recovered + .iter() + .filter(|(_, runs)| runs.as_slice() != [(2, true)]) + .take(5) + .collect(); + assert!( + unexpected.is_empty(), + "each retried job runs once, on attempt 2, with the rescue recorded: {unexpected:?}" + ); + } + + producing.store(false, Ordering::SeqCst); + let produced = producer.await.expect("producer"); + wait_until( + live_job_duration + scaled_timeout(Duration::from_secs(30)), + "the live jobs to finish", + || async { + running(&pool, &store, LIVE_QUEUE).await == 0 + && store + .queue_counts(&pool, LIVE_QUEUE) + .await + .expect("live counts") + .available + == 0 + }, + ) + .await; + live.shutdown(Duration::from_secs(10)).await; + + assert_eq!( + rescued_rows(&pool, LIVE_QUEUE).await, + 0, + "rescue took a claim from the live runtime" + ); + { + let rescues = rescues.lock().expect("rescues"); + let unexpected = rescued_twice(&rescues); + assert!( + rescues.len() as i64 == ABANDONED && unexpected.is_empty(), + "no claim of a live runtime was rescued: {} rescued jobs, e.g. {unexpected:?}", + rescues.len() + ); + } + let live_runs = live_runs.lock().expect("live runs"); + assert_eq!(live_runs.len(), produced, "every live job ran"); + let rerun: Vec<_> = live_runs + .iter() + .filter(|(_, attempts)| attempts.as_slice() != [1]) + .take(5) + .collect(); + assert!( + rerun.is_empty(), + "a live job ran more than once or on a later attempt: {rerun:?}" + ); + eprintln!( + "live runtime ran {} jobs of {live_job_duration:?} each with no rescue", + live_runs.len() + ); +} diff --git a/awa/tests/heartbeat_pass_benchmark_test.rs b/awa/tests/heartbeat_pass_benchmark_test.rs new file mode 100644 index 000000000..5a37eb193 --- /dev/null +++ b/awa/tests/heartbeat_pass_benchmark_test.rs @@ -0,0 +1,193 @@ +//! Heartbeat pass cost against in-flight claim count. Ignored by default: +//! +//! ```text +//! DATABASE_URL=postgres://postgres:test@localhost:15432/awa_test \ +//! cargo test --release -p awa --test heartbeat_pass_benchmark_test -- --ignored --nocapture +//! ``` +//! +//! Each size claims that many jobs through the claim path the dispatcher uses, +//! in claims of `AWA_HEARTBEAT_BENCH_CLAIM_BATCH` jobs (default 16, as a busy +//! runtime refills a few freed workers at a time), and times a heartbeat pass +//! over them in the heartbeat service's 500-attempt chunks. +//! `AWA_HEARTBEAT_BENCH_SIZES` overrides the sizes. + +use awa::model::queue_storage::ClaimHeartbeatRef; +use awa::model::{insert, migrations, QueueStorage, QueueStorageConfig}; +use awa::{audited_sql, InsertOpts, JobArgs}; +use serde::{Deserialize, Serialize}; +use sqlx::postgres::{PgConnection, PgPoolOptions}; +use sqlx::{Connection, PgPool}; +use std::time::{Duration, Instant}; + +const HEARTBEAT_CHUNK: usize = 500; +const PASSES: usize = 6; + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct HeartbeatBenchJob { + n: i64, +} + +fn database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) +} + +fn with_database(url: &str, database: &str) -> String { + let (base, _) = url + .rsplit_once('/') + .expect("database URL has a database name"); + format!("{base}/{database}") +} + +async fn fresh_database(name: &str) -> PgPool { + // Database names come only from the fixed prefix and this process ID. + let mut admin = PgConnection::connect(&with_database(&database_url(), "postgres")) + .await + .expect("admin connection"); + sqlx::raw_sql(audited_sql(format!( + "DROP DATABASE IF EXISTS {name} WITH (FORCE)" + ))) + .execute(&mut admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(audited_sql(format!("CREATE DATABASE {name}"))) + .execute(&mut admin) + .await + .expect("create database"); + let pool = PgPoolOptions::new() + .max_connections(4) + .connect(&with_database(&database_url(), name)) + .await + .expect("bench pool"); + migrations::run(&pool).await.expect("migrate"); + pool +} + +async fn drop_database(name: &str) { + // Database names come only from the fixed prefix and this process ID. + let mut admin = PgConnection::connect(&with_database(&database_url(), "postgres")) + .await + .expect("admin connection"); + let _ = sqlx::raw_sql(audited_sql(format!( + "DROP DATABASE IF EXISTS {name} WITH (FORCE)" + ))) + .execute(&mut admin) + .await; +} + +#[tokio::test] +#[ignore = "benchmark; creates and drops its own database"] +async fn heartbeat_pass_time_by_in_flight_claims() { + let sizes: Vec = std::env::var("AWA_HEARTBEAT_BENCH_SIZES") + .ok() + .map(|sizes| { + sizes + .split(',') + .filter_map(|size| size.trim().parse().ok()) + .collect() + }) + .unwrap_or_else(|| vec![1_000, 10_000, 25_000]); + let claim_batch: i64 = std::env::var("AWA_HEARTBEAT_BENCH_CLAIM_BATCH") + .ok() + .and_then(|size| size.parse().ok()) + .unwrap_or(16); + let server: String = { + let pool = fresh_database(&format!("awa_hb_bench_{}", std::process::id())).await; + let version = sqlx::query_scalar("SHOW server_version") + .fetch_one(&pool) + .await + .expect("server version"); + pool.close().await; + version + }; + for size in sizes { + let name = format!("awa_hb_bench_{}", std::process::id()); + let pool = fresh_database(&name).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + store.install(&pool).await.expect("install queue storage"); + awa_testing::setup::activate_queue_storage(&pool).await; + let queue = "heartbeat_bench"; + for first in (0..size as i64).step_by(1_000) { + let params: Vec<_> = (first..(first + 1_000).min(size as i64)) + .map(|n| { + insert::params_with( + &HeartbeatBenchJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue"); + } + let mut claimed = Vec::with_capacity(size); + while claimed.len() < size { + let batch = store + .claim_runtime_batch(&pool, queue, claim_batch, Duration::ZERO) + .await + .expect("claim"); + assert!(!batch.is_empty(), "claimed {} of {size}", claimed.len()); + claimed.extend(batch); + } + sqlx::query("ANALYZE") + .execute(&pool) + .await + .expect("analyze"); + + // The first pass creates every attempt's heartbeat row; later passes + // refresh them. + let mut passes = Vec::with_capacity(PASSES); + let mut refreshed = 0; + for _ in 0..PASSES { + let started = Instant::now(); + refreshed = heartbeat_pass(&store, &pool, &claimed).await; + passes.push(started.elapsed()); + } + let first = passes[0]; + let mut later = passes[1..].to_vec(); + later.sort_unstable(); + eprintln!( + "heartbeat pass: server={server} in_flight={size} claim_batch={claim_batch} \ + refreshed={refreshed} first_ms={} median_ms={} max_ms={}", + first.as_millis(), + later[later.len() / 2].as_millis(), + later[later.len() - 1].as_millis() + ); + assert_eq!(refreshed, size, "every in-flight attempt is refreshed"); + pool.close().await; + drop_database(&name).await; + } +} + +async fn heartbeat_pass( + store: &QueueStorage, + pool: &PgPool, + claimed: &[awa::model::ClaimedRuntimeJob], +) -> usize { + // As the heartbeat service splits them: compact receipt claims by their + // batch identity, anything else by attempt key. + let mut claims = Vec::new(); + let mut plain = Vec::new(); + for job in claimed { + match ClaimHeartbeatRef::from_claim(job.job.id, job.job.run_lease, &job.claim) { + Some(claim) => claims.push(claim), + None => plain.push((job.job.id, job.job.run_lease)), + } + } + let mut refreshed = 0; + for chunk in claims.chunks(HEARTBEAT_CHUNK) { + refreshed += store + .heartbeat_claims_batch(pool, chunk) + .await + .expect("heartbeat claims"); + } + for chunk in plain.chunks(HEARTBEAT_CHUNK) { + refreshed += store.heartbeat_batch(pool, chunk).await.expect("heartbeat"); + } + refreshed +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index bf51546fe..b77972d0f 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -6,6 +6,8 @@ #[path = "../../awa-model/src/storage_footprint_sql.rs"] mod storage_footprint_sql; +mod ci_timing; + use awa::audited_sql; use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, @@ -2043,6 +2045,316 @@ impl Worker for MultiClientTrackingWorker { } } +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct FinalizeBacklogJob { + n: i64, +} + +/// A runtime's finished attempts stop heartbeating once they are handed to +/// finalization, so it must stop claiming when finalization falls behind. +/// Otherwise the backlog of finished, unheartbeated attempts grows with the +/// claim rate, outlives `heartbeat_staleness`, and rescue runs those jobs +/// again. Holding the first jobs' receipt-completion locks stalls the +/// completion batcher, standing in for a batcher that cannot keep up. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_stops_claiming_while_finalization_is_stalled() { + // The default completion pipeline: one shard of 512-job batches, one + // batch in flight and one queued. + const FINALIZER_CAPACITY: usize = 1_024; + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_finalize_backlog"; + let store_config = QueueStorageConfig { + schema: "awa_qs_finalize_backlog".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + for first in (0..JOBS).step_by(500) { + let params: Vec<_> = (first..(first + 500).min(JOBS)) + .map(|n| { + insert::params_with( + &FinalizeBacklogJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue backlog"); + } + + let first_jobs: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT job_id FROM {}.ready_entries WHERE queue = $1 ORDER BY lane_seq LIMIT $2", + store.schema() + ))) + .bind(queue) + .bind(i64::from(WORKERS) * 2) + .fetch_all(&pool) + .await + .expect("first jobs"); + let mut lock_holder = pool.acquire().await.expect("lock connection"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_lock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', job_id, 1), 0)) + FROM unnest($1::bigint[]) AS job_id", + ) + .bind(&first_jobs) + .execute(&mut *lock_holder) + .await + .expect("stall the first jobs' finalization"); + + let runs: Arc>> = Arc::default(); + let recorded = runs.clone(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: WORKERS, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: WORKERS as usize, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register::(move |_, ctx| { + *recorded + .lock() + .expect("runs") + .entry(ctx.job.id) + .or_insert(0) += 1; + async { Ok(JobResult::Completed) } + }) + .heartbeat_staleness(Duration::from_secs(300)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build finalize backlog client"); + client.start().await.expect("start finalize backlog client"); + + // Wait for claiming to stop, or to run past the bound. A fixed runtime + // can never pass it, so a pause in claiming only weakens the check. + let limit = FINALIZER_CAPACITY + WORKERS as usize; + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + let mut last = (0, Instant::now()); + let started_while_stalled = loop { + let started = runs.lock().expect("runs").len(); + if started > limit { + break started; + } + if started != last.0 { + last = (started, Instant::now()); + } else if started > 0 && last.1.elapsed() >= Duration::from_secs(5) { + break started; + } + assert!( + Instant::now() < deadline, + "claiming neither stopped nor passed {limit}" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + }; + + sqlx::query("SELECT pg_catalog.pg_advisory_unlock_all()") + .execute(&mut *lock_holder) + .await + .expect("release the stalled finalization"); + drop(lock_holder); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(60)); + while completed_terminal_count(&pool, &store, queue).await < JOBS { + assert!(Instant::now() < deadline, "the backlog did not drain"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + client.shutdown(Duration::from_secs(10)).await; + + assert!( + started_while_stalled <= limit, + "{started_while_stalled} jobs started while finalization was stalled; \ + at most {limit} may wait without a heartbeat or hold a worker" + ); + let runs = runs.lock().expect("runs"); + assert_eq!(runs.len() as i64, JOBS); + assert!( + runs.values().all(|count| *count == 1), + "a job ran more than once" + ); +} + +/// A finished attempt keeps heartbeating until its finalize is over, so a +/// finalize that waits behind a stalled completion batcher for several +/// staleness windows is not mistaken for a dead worker and run again. +/// An independent trigger gate stalls a completion after it has acquired all +/// of its attempts' locks. That batch remains protected from rescue while +/// every other finished attempt must still be heartbeated. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_finalizing_attempts_keep_heartbeating_through_a_stalled_finalize() { + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_finalize_heartbeat"; + let store_config = QueueStorageConfig { + schema: "awa_qs_finalize_heartbeat".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + for first in (0..JOBS).step_by(500) { + let params: Vec<_> = (first..(first + 500).min(JOBS)) + .map(|n| { + insert::params_with( + &FinalizeBacklogJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue backlog"); + } + + // The trigger runs after the completer owns every attempt lock in its + // flush. Holding an attempt lock directly would suppress its heartbeat + // before completion owns it, racing rescue when the test unlocks it. + sqlx::raw_sql(audited_sql(format!( + r#" + CREATE FUNCTION {schema}.hold_completion_for_heartbeat_test() RETURNS trigger + LANGUAGE plpgsql AS $$ + BEGIN + IF NEW.outcome = 'completed' THEN + PERFORM pg_catalog.pg_advisory_xact_lock( + pg_catalog.hashtextextended('awa.test.finalize-heartbeat-gate', 0) + ); + END IF; + RETURN NEW; + END + $$; + CREATE TRIGGER hold_completion_for_heartbeat_test + BEFORE INSERT ON {schema}.lease_claim_closure_batches + FOR EACH ROW EXECUTE FUNCTION {schema}.hold_completion_for_heartbeat_test(); + "#, + schema = store.schema(), + ))) + .execute(&pool) + .await + .expect("install the completion gate"); + let mut lock_holder = pool.acquire().await.expect("lock connection"); + let lock_holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *lock_holder) + .await + .expect("gate holder PID"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_lock( + pg_catalog.hashtextextended('awa.test.finalize-heartbeat-gate', 0))", + ) + .execute(&mut *lock_holder) + .await + .expect("hold the independent completion gate"); + + // Sequential heartbeat chunks can need several WAL-sync round trips + // under contention. Allow margin for those writes while blocking the + // completion for four full staleness windows below. + let staleness = ci_timing::scaled_staleness(Duration::from_secs(5)); + let runs: Arc>> = Arc::default(); + let recorded = runs.clone(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: WORKERS, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: WORKERS as usize, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register::(move |_, ctx| { + *recorded + .lock() + .expect("runs") + .entry(ctx.job.id) + .or_insert(0) += 1; + async { Ok(JobResult::Completed) } + }) + .heartbeat_interval(staleness / 5) + .heartbeat_staleness(staleness) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build finalize heartbeat client"); + client + .start() + .await + .expect("start finalize heartbeat client"); + + // Hold the stall for several staleness windows once claiming has stopped. + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + loop { + let completion_at_gate: bool = sqlx::query_scalar( + "SELECT EXISTS ( + SELECT 1 FROM pg_stat_activity + WHERE datname = current_database() + AND wait_event = 'advisory' + AND $1 = ANY(pg_catalog.pg_blocking_pids(pid)) + )", + ) + .bind(lock_holder_pid) + .fetch_one(&pool) + .await + .expect("observe the completion gate"); + if completion_at_gate && runs.lock().expect("runs").len() >= WORKERS as usize * 2 { + break; + } + assert!( + Instant::now() < deadline, + "the runtime did not queue attempts behind the completion gate" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + } + tokio::time::sleep(staleness * 4).await; + + sqlx::query("SELECT pg_catalog.pg_advisory_unlock_all()") + .execute(&mut *lock_holder) + .await + .expect("release the stalled finalization"); + drop(lock_holder); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(60)); + while completed_terminal_count(&pool, &store, queue).await < JOBS { + assert!(Instant::now() < deadline, "the backlog did not drain"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + client.shutdown(Duration::from_secs(10)).await; + + let runs = runs.lock().expect("runs"); + assert_eq!(runs.len() as i64, JOBS); + let rerun: Vec<_> = runs + .iter() + .filter(|(_, count)| **count > 1) + .take(5) + .collect(); + assert!(rerun.is_empty(), "jobs ran more than once: {rerun:?}"); +} + /// ADR-023 claim-ring control-plane smoke test. Exercises /// `rotate_claims`, the #371 idle gate, and `prune_oldest_claims` on an /// empty schema: rotate is idle-skipped when nothing has been written @@ -8522,6 +8834,193 @@ async fn test_queue_storage_expired_deadline_batch_blocks_prune_until_rescued() ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_rescue_batch_keeps_a_partially_rescued_scan_active() { + let (_db_guard, pool) = setup_pool(10).await; + let schema = "awa_rescue_partial_window"; + let queue = "rescue_partial_window"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await; + for id in 0..501 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await; + } + let claims = store + .claim_runtime_batch(&pool, queue, 501, Duration::from_secs(30)) + .await + .expect("claim a backlog larger than one rescue batch"); + assert_eq!(claims.len(), 501); + for table in ["lease_claims", "lease_claim_batches"] { + sqlx::query(audited_sql(format!("UPDATE {schema}.{table} SET claimed_at = now() - interval '1 hour', deadline_at = now() - interval '1 minute'"))) + .execute(&pool).await.expect("expire attempts"); + } + let batch = store + .rescue_expired_deadlines_batch(&pool, &|_| false, &Default::default()) + .await + .expect("first rescue batch"); + assert_eq!(batch.rescued.len(), 500); + assert!(batch.saturated); + let slot = claims[0].claim.claim_slot; + assert!(!batch.idle_slots.contains(&slot)); + assert!( + batch + .scan_progress + .iter() + .filter(|scan| scan.slot == slot) + .flat_map(|scan| &scan.cursors) + .any(|cursor| cursor.saturated), + "a partially rescued candidate window must not complete that slot's sweep" + ); + let tail = store + .rescue_expired_deadlines_batch(&pool, &|_| false, &Default::default()) + .await + .expect("rescue the remaining attempt"); + assert_eq!(tail.rescued.len(), 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_rescue_batch_continues_past_closed_scan_windows() { + for (deadline, fresh_history) in [(false, false), (true, false), (false, true)] { + let (_db_guard, pool) = setup_pool(10).await; + let suffix = if deadline { + "deadline" + } else if fresh_history { + "fresh" + } else { + "heartbeat" + }; + let schema = format!("awa_rescue_sparse_{suffix}"); + let queue = format!("rescue_sparse_{suffix}"); + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.clone(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await; + for id in 0..3 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.clone(), + ..Default::default() + }, + ) + .await; + } + let mut claims = Vec::new(); + for _ in 0..3 { + claims.extend( + store + .claim_runtime_batch(&pool, &queue, 1, Duration::from_secs(30)) + .await + .expect("claim sparse backlog"), + ); + } + assert_eq!(claims.len(), 3); + for claim in &claims { + age_receipt_claim( + &pool, + &store, + claim.job.id, + claim.job.run_lease, + Duration::from_secs(300), + ) + .await; + } + if deadline { + sqlx::query(audited_sql(format!( + "UPDATE {schema}.lease_claim_batches SET deadline_at = now() - interval '1 second'" + ))) + .execute(&pool) + .await + .expect("expire deadlines"); + } + if fresh_history { + for claim in &claims[..2] { + sqlx::query(audited_sql(format!("INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at) VALUES ($1, $2, now(), now())"))) + .bind(claim.job.id).bind(claim.job.run_lease).execute(&pool).await.expect("fresh heartbeat history"); + } + } else { + store + .complete_runtime_batch(&pool, &claims[..2]) + .await + .expect("complete history"); + } + let options = awa::model::RescueBatchOptions { + cursor_scan_limit: Some(1), + ..Default::default() + }; + let batch = if deadline { + store + .rescue_expired_deadlines_batch(&pool, &|_| false, &options) + .await + } else { + store + .rescue_stale_heartbeats_batch(&pool, Duration::from_secs(90), &|_| false, &options) + .await + } + .expect("walk first closed window"); + assert!(batch.rescued.is_empty()); + assert!( + batch.saturated, + "a full closed scan window must continue toward unseen abandoned claims ({suffix})" + ); + assert!( + !batch.idle_slots.contains(&claims[0].claim.claim_slot), + "cursor progress is not an idle slot" + ); + let mut rescued = Vec::new(); + for _ in 0..3 { + let batch = if deadline { + store + .rescue_expired_deadlines_batch(&pool, &|_| false, &options) + .await + } else { + store + .rescue_stale_heartbeats_batch( + &pool, + Duration::from_secs(90), + &|_| false, + &options, + ) + .await + } + .expect("continue sparse scan"); + rescued.extend(batch.rescued); + } + assert_eq!( + rescued.iter().map(|job| job.id).collect::>(), + vec![claims[2].job.id] + ); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_receipt_rescue_cursor_sweeps_past_fresh_claims_and_wraps() { let (_db_guard, pool) = setup_pool(10).await; @@ -21093,3 +21592,206 @@ async fn test_queue_storage_legacy_row_claims_are_locked_before_the_heartbeat_st .expect("join heartbeat") .expect("heartbeat after the rewrite"); } + +/// A finalizing attempt must not hold up unrelated heartbeats, and an old +/// heartbeat snapshot must not recreate state after durable completion. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_heartbeat_skips_closers_and_never_revives_closed_claims() { + let (_db_guard, pool) = setup_pool(10).await; + let config = QueueStorageConfig { + schema: "awa_qs_heartbeat_closure".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, config).await; + let queue = "heartbeat_closure"; + for id in 0..2 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + } + let jobs = store + .claim_runtime_batch(&pool, queue, 2, Duration::ZERO) + .await + .expect("claim two compact receipts"); + assert_eq!(jobs.len(), 2); + let claims: Vec<_> = jobs + .iter() + .map(|job| { + awa::model::queue_storage::ClaimHeartbeatRef::from_claim( + job.job.id, + job.job.run_lease, + &job.claim, + ) + .expect("compact claim identity") + }) + .collect(); + let mut closer = pool.begin().await.expect("begin closer"); + sqlx::query( + "SELECT pg_advisory_xact_lock(hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(jobs[0].job.id) + .bind(jobs[0].job.run_lease) + .execute(closer.as_mut()) + .await + .expect("hold one attempt's closer lock"); + let refreshed = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.heartbeat_claims_batch(&pool, &claims), + ) + .await + .expect("a closer cannot stall the heartbeat pass") + .expect("heartbeat the unlocked claim"); + assert_eq!(refreshed, 1); + closer.rollback().await.expect("release closer"); + + let completed = store + .complete_runtime_batch(&pool, &jobs) + .await + .expect("durable completion"); + assert_eq!(completed.len(), 2); + assert_eq!( + store + .heartbeat_claims_batch(&pool, &claims) + .await + .expect("a cached heartbeat snapshot after completion"), + 0, + ); + let remaining: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {}.attempt_state WHERE job_id = ANY($1)", + store.schema() + ))) + .bind(jobs.iter().map(|job| job.job.id).collect::>()) + .fetch_one(&pool) + .await + .expect("read attempt state"); + assert_eq!( + remaining, 0, + "a late heartbeat must not revive a closed attempt" + ); +} + +/// A progress snapshot captured before a handler returns must not wait behind +/// its finalizer and prevent unrelated attempts from heartbeating. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_pending_progress_skips_a_closer_and_can_retry() { + progress_snapshot_skips_closer(true).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_legacy_pending_progress_skips_a_closer_and_can_retry() { + progress_snapshot_skips_closer(false).await; +} + +async fn progress_snapshot_skips_closer(receipts: bool) { + let (_db_guard, pool) = setup_pool(10).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: if receipts { + "awa_qs_pending_progress_closer".to_string() + } else { + "awa_qs_pending_progress_closer_legacy".to_string() + }, + lease_claim_receipts: receipts, + ..Default::default() + }, + ) + .await; + let queue = "pending_progress_closer"; + for id in 0..2 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + } + let jobs = store + .claim_runtime_batch(&pool, queue, 2, Duration::ZERO) + .await + .expect("claim compact receipts"); + assert_eq!(jobs.len(), 2); + let pending: Vec<_> = jobs + .iter() + .map(|job| { + ( + job.job.id, + job.job.run_lease, + serde_json::json!({"step": 2}), + ) + }) + .collect(); + let mut closer = pool.begin().await.expect("begin closer"); + if receipts { + sqlx::query( + "SELECT pg_advisory_xact_lock(hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(jobs[0].job.id) + .bind(jobs[0].job.run_lease) + .execute(closer.as_mut()) + .await + .expect("hold a closer lock after the progress snapshot"); + } else { + sqlx::query(audited_sql(format!( + "SELECT 1 FROM {}.leases WHERE job_id = $1 AND run_lease = $2 FOR UPDATE", + store.schema() + ))) + .bind(jobs[0].job.id) + .bind(jobs[0].job.run_lease) + .execute(closer.as_mut()) + .await + .expect("hold a legacy finalizer row lock"); + } + let updated = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.heartbeat_progress_batch(&pool, &pending), + ) + .await + .expect("a stale progress snapshot cannot block the heartbeat pass") + .expect("write the unlocked snapshot"); + assert_eq!(updated, 1); + let pairs: Vec<_> = jobs + .iter() + .map(|job| (job.job.id, job.job.run_lease)) + .collect(); + assert_eq!( + tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.heartbeat_batch(&pool, &pairs), + ) + .await + .expect("heartbeat-only passes must also skip a closer") + .expect("heartbeat unlocked claims"), + 1, + ); + closer.rollback().await.expect("release closer"); + assert_eq!( + store + .heartbeat_progress_batch(&pool, &pending) + .await + .expect("retry the unacknowledged snapshots"), + 2, + ); + let progress: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT progress FROM {}.attempt_state WHERE job_id = ANY($1) ORDER BY job_id", + store.schema() + ))) + .bind(jobs.iter().map(|job| job.job.id).collect::>()) + .fetch_all(&pool) + .await + .expect("read durable progress"); + assert_eq!(progress, vec![serde_json::json!({"step": 2}); 2]); +} diff --git a/docs/configuration.md b/docs/configuration.md index 930bc6bde..2a87840ce 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -484,6 +484,27 @@ If you suspend a handler with `wait_for_callback()` and the external system neve | --- | --- | --- | | `ClientBuilder::callback_rescue_interval` | `30s` | How often the maintenance leader scans for `callback_timeout_at < now()`. The per-callback timeout itself is set when registering the callback in the handler. | +### Draining a backlog of abandoned claims { #rescue-backlog } + +A runtime that is killed (OOM, `SIGKILL`, a lost node) leaves every attempt it had claimed open. On queue storage those open claims also stop the claim ring from rotating back into their slot, so a large backlog has to be rescued quickly, not 500 attempts per rescue interval. + +Each stale-heartbeat, deadline or callback-timeout rescue therefore runs as a pass of batches. A batch scans a bounded claim window and rescues up to 500 attempts per source (lease rows and receipt claims) in its own short transaction. The pass keeps starting batches while a scan window or rescue batch came back full, until its time budget is spent or it has rescued `rescue_pass_max_rows`. While a backlog remains, the next pass of that kind runs one second later instead of after its interval, so a drain proceeds at up to about `rescue_pass_max_rows` attempts per second per rescue kind. + +| Knob | Default | What it does | +| --- | --- | --- | +| `ClientBuilder::rescue_pass_budget` (Rust) / `rescue_pass_budget_ms` (Python kwarg) | `1s` | How long one pass keeps starting batches. Capped at half the rescue interval. `Duration::ZERO` / `0` keeps every pass to one batch. | +| `ClientBuilder::rescue_pass_max_rows` (Rust) / `rescue_pass_max_rows` (Python kwarg) | `5000` | How many attempts one pass rescues before it stops. Bounds the write burst (closures, retry rows, WAL, replica and CDC traffic), with up to one final batch beyond the cap. `0` keeps one batch per ordinary interval without early re-arms. | + +How a pass chooses and paces its work: + +- **Oldest claims first.** Batches visit claim slots from the oldest to the current one, because the ring cannot rotate back into the oldest slot until its claims are closed. Every fourth batch visits the other slots first, so a long drain of one slot does not starve the rest. A slot a batch found nothing in is skipped for the rest of the pass. +- **Stale-heartbeat acceleration uses runtime evidence.** A stale heartbeat cannot distinguish a dead runtime from a stalled one. Extra batches and early passes require every snapshot in `awa.runtime_instances` to be either fresh with its heartbeat loop alive or silent for at least `heartbeat_staleness`, with that condition observed for a full staleness window. Accelerated batches also require the claim's last heartbeat to precede the latest gone-runtime cutoff. The per-attempt staleness condition is unchanged. + This is a bounded heuristic, not proof that an owner has exited. To expose an unconfirmed phase between fresh and gone, each runtime needs `3 × max(runtime_snapshot_interval, 1s) < heartbeat_staleness`. With 5-second snapshots and 12-second staleness, a stalled live runtime goes directly from fresh to gone; steady evidence may predate the stall, so acceleration can rescue its stale claims faster than the ordinary one-batch cadence. The leader must also sample frequently enough: gaps longer than half the staleness reset the evidence window, and short stalls between samples can be missed. The defaults (10-second snapshots, 90-second staleness, 30-second rescue interval) provide an unconfirmed phase; after a runtime stops reporting, acceleration usually begins after about 3–3.5 minutes. A new maintenance leader resets the evidence window. Deadline and callback rescues use their expired deadline and do not need owner evidence. +- **Retries retain their backoff.** Rescue bounds closures and retry writes per pass. Backoffs and promotion delays can make retries become due together; `rescue_pass_max_rows` does not cap the rate at which promotion admits them to the ready lanes. +- **Exhausted attempts are dead-lettered in the same batches.** On a DLQ-enabled queue an attempt rescued on its last try goes to the DLQ in its batch's transaction (reason `heartbeat_stale`, `deadline_exceeded` or `callback_timeout`), so a drain can move many jobs to the DLQ at once. `awa.job.dlq_moved` counts them by reason. + +A pass holds the maintenance leader's loop for up to its budget plus one batch, so promotion and the other maintenance branches can wait about that long while a backlog drains. Watch a drain with `awa.maintenance.rescue.pass_size` (attempts per pass, by `awa.rescue.kind` and `awa.rescue.pass_end`: `drained`, `time_budget`, `row_limit`, `owners_unconfirmed` or `error`) and `awa.maintenance.rescue.oldest_attempt_age` (the age of the oldest attempt the latest pass rescued, `0` when it rescued nothing). An `oldest_attempt_age` far above `heartbeat_staleness` or the queue's deadline means rescue is behind. + ### Retention and cleanup Retention depends on the storage path. diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 68c0e0ae6..41b7f6d9d 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -255,6 +255,8 @@ In Prometheus they appear under the OpenTelemetry [OTLP-to-Prometheus translatio | `awa.maintenance.promote_batch_size` | Histogram | state | Jobs promoted | | `awa.maintenance.promote_duration` | Histogram (s) | state | Promotion time | | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | +| `awa.maintenance.rescue.pass_size` | Histogram | rescue_kind, rescue_pass_end | Attempts rescued per queue-storage rescue pass, and why the pass stopped (`drained`, `time_budget`, `row_limit`, `owners_unconfirmed`, `error`) | +| `awa.maintenance.rescue.oldest_attempt_age` | Gauge (s) | rescue_kind | Age of the oldest attempt the latest queue-storage rescue pass on this process rescued; `0` when it rescued none. A former leader retains its last value until it runs another pass; filter to the current leader | | `awa.maintenance.prune.attempts` | Counter | ring, outcome, reason | Ring prune outcomes, including destructive `pruned` and no-DDL `already_pruned` | | `awa.maintenance.prune.duration` | Histogram (s) | ring, phase | Successful destructive prune time split into `lock`, `truncate`, and `commit` | diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 1672c5865..4ac4f307c 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -285,6 +285,31 @@ If long-running jobs are expected, remember: If a job really needs `20m`, set a longer deadline for that queue. +## Claim Ring Stops Rotating After A Runtime Is Killed + +Symptoms on queue storage, usually after a high-throughput runtime was OOM-killed or lost without a graceful shutdown: + +- `awa.maintenance.rotate.attempts` for the claim ring reports `skipped_busy` with blocker `claim.rows`, and `awa.ring.generation` for the claim ring stops advancing +- the current claim slot keeps growing +- jobs the dead runtime held wait well past `heartbeat_staleness` before they run again + +The dead runtime's claims stay open until rescue closes them, and the ring cannot rotate back into a slot that still holds open claims. Rescue drains such a backlog in passes of 500-attempt batches ([Draining a backlog of abandoned claims](configuration.md#rescue-backlog)). To see where it is: + +- `awa.maintenance.rescue.oldest_attempt_age` by `awa.rescue.kind` shows how old the attempts rescue is closing are. It should fall towards `heartbeat_staleness` (or the queue's deadline) as the backlog drains. +- `awa.maintenance.rescue.pass_size` by `awa.rescue.pass_end` shows why passes stop. `time_budget` or `row_limit` means a drain is running at its configured rate. `owners_unconfirmed` means stale-heartbeat rescue is waiting for evidence that the claims' runtime is gone and is taking one batch per ordinary rescue interval. + +If stale-heartbeat passes stay at `owners_unconfirmed`, look at the runtime snapshots: + +```sql +SELECT instance_id, hostname, last_seen_at, snapshot_interval_ms, heartbeat_alive +FROM awa.runtime_instances +ORDER BY last_seen_at DESC; +``` + +A snapshot that is no longer fresh but is younger than `heartbeat_staleness` holds acceleration back, as does a fresh snapshot with `heartbeat_alive = false`. The unconfirmed phase requires `3 × max(runtime_snapshot_interval, 1s) < heartbeat_staleness` for every runtime; otherwise a stalled live runtime can pass directly from fresh to gone ([owner-evidence limits](configuration.md#rescue-backlog)). The drain also needs `heartbeat_rescue_interval` below half of `heartbeat_staleness`. With the defaults, acceleration usually starts about 3–3.5 minutes after the last runtime snapshot, and a leadership change restarts the evidence window. Deadline rescue is not held back this way, so queues with a `deadline_duration` drain through it regardless. + +To drain faster, raise `rescue_pass_max_rows` and `rescue_pass_budget`, keeping in mind that a pass holds the maintenance leader for its budget and can exceed `rescue_pass_max_rows` by one final batch of closures and retry rows. Retry due times can cluster; this cap does not bound promotion throughput. + ## Dead Tuples Growing In Queue Storage ### What It Usually Means diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index 4cfb039b0..a62e48761 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -146,6 +146,25 @@ Late completions from a rescued attempt lose against the newer attempt via the (default 90s) to at least 3× the heartbeat interval; detection latency is roughly `heartbeat_staleness + heartbeat_rescue_interval`. +On queue storage a killed runtime's open claims stop the claim ring rotating +until rescue closes them. Rescue drains such a backlog in passes of 500-attempt +batches, bounded by `rescue_pass_budget` (1s) and `rescue_pass_max_rows` +(5000), with the next pass a second later while a backlog remains. +Stale-heartbeat rescue runs extra batches and early passes only after runtime +snapshots support a full settled staleness window. Steady evidence may predate a +stall; it is a sampled heuristic, not proof that an owner has exited. An +unconfirmed phase requires `3 × max(runtime_snapshot_interval, 1s) < heartbeat_staleness`; +5-second snapshots with 12-second staleness skip that phase and can accelerate +rescue of stale live claims. See [rescue backlog](../../docs/configuration.md#rescue-backlog) +for cutoff and sampling limits. The ordinary first batch keeps its configured interval. This evidence needs `heartbeat_rescue_interval` +below half of `heartbeat_staleness`; at the defaults acceleration usually +starts 3–3.5 minutes after a runtime's last snapshot. Leadership changes reset +the evidence window. Zero for either rescue pass control keeps one batch per +ordinary interval. Retried jobs retain their backoff, and due times can cluster; +the pass row cap does not cap promotion throughput. Watch +`awa.maintenance.rescue.pass_size` (by `awa.rescue.pass_end`) and +`awa.maintenance.rescue.oldest_attempt_age`, filtering to the current leader. + Finished handlers retry durable finalization for up to 120s from the first finalize attempt when database writes are unavailable, including connection, pool-acquire and read-only-session errors. Recovery requires a writable endpoint