diff --git a/.config/nextest.toml b/.config/nextest.toml index 8b0ad8e3e..04b07483c 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -15,7 +15,15 @@ slow-timeout = { period = "60s" } [test-groups] migrations = { max-threads = 1 } +bulk-rescue-backlog = { max-threads = 1 } [[profile.default.overrides]] filter = 'binary(migration_test)' test-group = 'migrations' + +# The two bulk rescue backlog tests each load the same database heavily; run +# concurrently, the live runtime in one cannot heartbeat its claims inside the +# staleness window. +[[profile.default.overrides]] +filter = 'binary(bulk_rescue_test) and test(/^rescue_drains_/)' +test-group = 'bulk-rescue-backlog' diff --git a/CHANGELOG.md b/CHANGELOG.md index 878971946..0555517d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,10 @@ Notable changes between releases. Detailed migration notes for storage transitio - Everything [0.6.15](https://github.com/hardbyte/awa/releases/tag/v0.6.15) shipped that this branch did not already have; see its release notes for the details. Lifecycle hooks hold their attempt's finalizer slot until they return, so hooks slower than job throughput slow claiming instead of growing worker memory, and maintenance keeps at most 1,024 detached `Rescued` and duplicate `Cancelled` hook tasks per worker. By-job lookups of compact receipt claims use a GIN index on `lease_claim_batches.job_ids` that claim prune builds on each partition as it empties it (migration v054). Compatibility COPY inserts go through `awa.insert_jobs_compat_batch` (migration v055), which takes each queue lane's enqueue lock only when it writes that lane rather than at the batch's first ready row (the lock is still held until the transaction ends), and the new `insert_many_skip_unique_conflicts` inserts a batch through it while skipping claimed unique keys. Rescue reaches stale attempts behind a full batch of ambiguous ones, `Completed` events and follow-ups carry the completed job without the attempt's heartbeat and deadline, slot-scoped admin receipt reads stay off other slots' closure partitions, and unfiltered receipt job lists select candidate IDs before loading bodies. v054 and v055 add functions only and take no table locks; a database coming from 0.6.15 already has both functions from its schema patches. See [v054](docs/upgrade-0.6-to-0.7.md#v054-claim-batch-job-index) and [v055](docs/upgrade-0.6-to-0.7.md#v055-batch-compatibility-insert) in the upgrade guide. +### Forward-ported from 0.6.16 + +- Queue-storage rescue drains a large backlog of abandoned claims in bounded time. A killed runtime leaves every attempt it claimed open, and until they close, queue prune skips the queue slots they reference and claim prune skips their claim slots, so neither ring can wrap. A rescue pass takes 500-attempt batches, each in its own transaction, until a batch comes back short, `rescue_pass_budget` (default 1 s, at most half the rescue interval) is spent, or it has rescued `rescue_pass_max_rows` (default 5,000), and while a backlog remains the next pass runs a second later. `ClientBuilder::rescue_pass_budget` and `rescue_pass_max_rows` (Python `rescue_pass_budget_ms` and `rescue_pass_max_rows`) set both; zero for either keeps one batch per interval. Deadline and callback rescue act on their expired deadlines; stale-heartbeat rescue takes more than one batch per pass only once runtime snapshots show the claims' owners gone. Rescue's cursor scans seek to their cursor through the index, including the compact deadline-batch scan that 0.6 does not have, and rescue transactions run without JIT. `awa.maintenance.rescue.pass_size` (by `awa.rescue.kind` and `awa.rescue.pass_end`) and `awa.maintenance.rescue.oldest_attempt_age` show a drain in progress. Heartbeats no longer queue behind a completion: an attempt whose closer holds its lock is skipped, lease and progress writes skip rows a finalizer has locked and retry them on the next pass, and a finished attempt stops flushing progress once its handler returns. A point read of a failed job that starts while queue prune moves its row out of the slot being pruned is covered by a test; on this branch the read already locks the ring before it takes its snapshot. See [draining a backlog of abandoned claims](docs/configuration.md#rescue-backlog). + ### Added - **The queue-storage prune `lock_timeout` is configurable, and prune diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 9dc0cfd8e..2998719c0 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"; @@ -204,6 +206,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. @@ -508,6 +518,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") @@ -908,6 +931,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( @@ -1701,7 +1746,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, @@ -1716,6 +1761,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 fe279c4a9..15bb9b1e7 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -56,8 +56,10 @@ pub use partitioned_queue::{ }; pub use queue_storage::{ ClaimedEntry, ClaimedRuntimeJob, PromoteOutcome, PromotionLane, PruneDurations, PruneOutcome, - QueueCounts, QueueStorage, QueueStorageConfig, RingLedgerFoldOutcome, RingRotationCadence, - RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, TerminalRollupFoldOutcome, + QueueCounts, QueueStorage, 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 1c93ab888..4b8d285b3 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -39,11 +39,15 @@ 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. const DRAIN_TRANSACTION_BEGIN: &str = "BEGIN; SET LOCAL jit = off"; const DRAIN_CLAIM_CHECK_ROWS: i64 = 64; +// Rescue cursor scans are estimated far above their real cost, enough to +// JIT-compile them, which takes seconds against about 100 ms of execution. +const RESCUE_TRANSACTION_BEGIN: &str = "BEGIN; SET LOCAL jit = off"; const RECEIPT_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000; const RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT: i64 = 10_000; const DEFAULT_PRUNE_LOCK_TIMEOUT: Duration = Duration::from_millis(10); @@ -599,6 +603,116 @@ 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, + /// How many claims past its cursor each slot's stale-heartbeat or + /// row-local deadline walk examines (default 10,000). Every claim + /// examined costs its closure checks, so a batch that follows one that + /// filled up on a dense backlog can examine far fewer. + pub cursor_scan_limit: Option, + /// A stale-heartbeat batch takes only attempts last heartbeated before + /// this instant as well as before its staleness cutoff. The other + /// rescue kinds ignore it. + pub heartbeated_before: 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, @@ -3472,7 +3586,7 @@ enum DueLaneRows<'a> { Oldest(i64), } -#[derive(Debug, Clone, sqlx::FromRow)] +#[derive(Debug, Clone, serde::Deserialize, sqlx::FromRow)] struct DeletedLeaseRow { ready_slot: i32, ready_generation: i64, @@ -12261,6 +12375,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>, @@ -12566,7 +12721,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 @@ -12602,10 +12756,14 @@ impl QueueStorage { Ok(updated as usize) } + /// With `skip_locked`, attempts whose lock another transaction holds are + /// left out, as a heartbeat batch wants. Without it the call waits for + /// each attempt's lock, so an explicit flush is never dropped. async fn upsert_attempt_state_progress_from_receipts_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, jobs: &[(i64, i64, serde_json::Value)], + skip_locked: bool, ) -> Result { if jobs.is_empty() { return Ok(0); @@ -12615,6 +12773,28 @@ impl QueueStorage { .iter() .map(|(job_id, run_lease, _)| (*job_id, *run_lease)) .collect(); + let locked: Option> = if skip_locked { + Some( + self.try_lock_receipt_attempt_keys_tx(tx, &receipt_jobs) + .await? + .into_iter() + .collect(), + ) + } else { + None + }; + let jobs: Vec<_> = jobs + .iter() + .filter(|(job_id, run_lease, _)| { + locked + .as_ref() + .is_none_or(|locked| 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(); @@ -12684,7 +12864,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 @@ -13904,57 +14083,90 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, cutoff: DateTime, + options: &RescueBatchOptions, excluded: &BTreeSet<(i64, i64)>, selected: Option<&BTreeSet<(i64, i64)>>, - ) -> Result, AwaError> { - let mut rescued = Vec::new(); + 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 + 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, remaining, excluded, selected, + tx, + slot, + cutoff, + options + .cursor_scan_limit + .map_or(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT, |limit| { + limit.max(1) as i64 + }), + remaining, + excluded, + selected, ) .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, excluded, selected, - ) - .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) } + #[allow(clippy::too_many_arguments)] async fn rescue_stale_receipt_claims_for_slot_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, cutoff: DateTime, + scan_limit: i64, rescue_limit: i64, excluded: &BTreeSet<(i64, i64)>, selected: Option<&BTreeSet<(i64, i64)>>, - ) -> Result, AwaError> { + ) -> Result { let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = excluded.iter().copied().unzip(); let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected @@ -13976,8 +14188,32 @@ 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 > (SELECT rescue_cursor_claimed_at FROM cursor_row) + ) OVER ( + ORDER BY claim_batches.claimed_at, claim_batches.batch_id + ) AS covered + FROM {claim_batch_child} AS claim_batches + WHERE claim_batches.claim_slot = $1 + AND claim_batches.claimed_at < $2 + AND claim_batches.claimed_at >= (SELECT rescue_cursor_claimed_at FROM cursor_row) + ) AS running + WHERE running.covered >= $3 + ), + after_source AS ( + ( + SELECT claims.claim_slot, NULL::bigint AS batch_id, claims.ready_slot, @@ -13994,8 +14230,18 @@ 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 + WHERE claims.claim_slot = $1 + AND claims.claimed_at < $2 + AND (claims.claimed_at, claims.job_id, claims.run_lease) + > ( + (SELECT rescue_cursor_claimed_at FROM cursor_row), + (SELECT rescue_cursor_job_id FROM cursor_row), + (SELECT rescue_cursor_run_lease FROM cursor_row) + ) + ORDER BY claims.claimed_at, claims.job_id, claims.run_lease + LIMIT $3 + ) UNION ALL SELECT claim_batches.claim_slot, @@ -14015,6 +14261,7 @@ impl QueueStorage { NULL::timestamptz AS closed_at, true AS compact_batch FROM {claim_batch_child} AS claim_batches + CROSS JOIN after_batch_bound CROSS JOIN LATERAL unnest( claim_batches.job_ids, claim_batches.run_leases, @@ -14024,119 +14271,194 @@ 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 + AND claim_batches.claimed_at < $2 + AND claim_batches.claimed_at >= (SELECT rescue_cursor_claimed_at FROM cursor_row) + AND claim_batches.claimed_at + <= COALESCE(after_batch_bound.claimed_at, 'infinity'::timestamptz) + AND (claim_batches.claimed_at, items.job_id, items.run_lease) + > ( + (SELECT rescue_cursor_claimed_at FROM cursor_row), + (SELECT rescue_cursor_job_id FROM cursor_row), + (SELECT rescue_cursor_run_lease FROM cursor_row) + ) ), after_cursor AS MATERIALIZED ( SELECT - claim_source.claim_slot, - claim_source.job_id, - claim_source.run_lease, + after_source.*, row_number() OVER ( - ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease + ORDER BY after_source.claimed_at, after_source.job_id, after_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) - > ( - 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 + 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_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 < (SELECT rescue_cursor_claimed_at FROM cursor_row) + ) OVER ( + ORDER BY claim_batches.claimed_at, claim_batches.batch_id + ) AS covered + FROM {claim_batch_child} AS claim_batches + WHERE claim_batches.claim_slot = $1 + AND claim_batches.claimed_at < $2 + AND claim_batches.claimed_at <= (SELECT rescue_cursor_claimed_at FROM cursor_row) + ) 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 + 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) + <= ( + (SELECT rescue_cursor_claimed_at FROM cursor_row), + (SELECT rescue_cursor_job_id FROM cursor_row), + (SELECT rescue_cursor_run_lease FROM cursor_row) + ) + ORDER BY claims.claimed_at, claims.job_id, claims.run_lease + LIMIT $3 + ) + UNION ALL + SELECT + 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 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 <= (SELECT rescue_cursor_claimed_at FROM cursor_row) + AND claim_batches.claimed_at + <= COALESCE(before_batch_bound.claimed_at, 'infinity'::timestamptz) + AND (claim_batches.claimed_at, items.job_id, items.run_lease) + <= ( + (SELECT rescue_cursor_claimed_at FROM cursor_row), + (SELECT rescue_cursor_job_id FROM cursor_row), + (SELECT rescue_cursor_run_lease FROM cursor_row) + ) + ), before_cursor AS MATERIALIZED ( SELECT - claim_source.claim_slot, - claim_source.job_id, - claim_source.run_lease, + before_source.*, after_count.count + row_number() OVER ( - ORDER BY claim_source.claimed_at, claim_source.job_id, claim_source.run_lease + ORDER BY before_source.claimed_at, before_source.job_id, before_source.run_lease ) AS rn - FROM claim_source - CROSS JOIN cursor_row + FROM before_source 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) - <= ( - 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 + 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.* @@ -14177,7 +14499,7 @@ impl QueueStorage { ))) .bind(slot) .bind(cutoff) - .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .bind(&excluded_job_ids) .bind(&excluded_run_leases) @@ -14189,7 +14511,7 @@ impl QueueStorage { .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 ( @@ -14385,11 +14707,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, @@ -14408,12 +14728,18 @@ 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(&excluded_job_ids) .bind(&excluded_run_leases) @@ -14422,10 +14748,18 @@ impl QueueStorage { .bind(&selected_run_leases) .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 @@ -14457,55 +14791,67 @@ impl QueueStorage { async fn rescue_expired_receipt_deadlines_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + options: &RescueBatchOptions, excluded: &BTreeSet<(i64, i64)>, selected: Option<&BTreeSet<(i64, i64)>>, - ) -> Result, AwaError> { - let mut rescued = Vec::new(); + 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, excluded, selected) + 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, excluded, selected, + ) .await?; - remaining = remaining.saturating_sub(slot_rescued.len() as i64); - rescued.append(&mut slot_rescued); - if remaining == 0 { - break; - } - - let mut slot_rescued = self + 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, remaining, excluded, selected, + tx, slot, scan_limit, remaining, excluded, selected, ) .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 @@ -14520,26 +14866,42 @@ 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, excluded: &BTreeSet<(i64, i64)>, selected: Option<&BTreeSet<(i64, i64)>>, - ) -> Result, AwaError> { + ) -> Result { if rescue_limit <= 0 { - return Ok(Vec::new()); + return Ok((Vec::new(), true, None)); } let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = @@ -14573,15 +14935,14 @@ impl QueueStorage { ORDER BY claims.deadline_at, claims.job_id, claims.run_lease ) AS rn FROM {claim_child} AS claims - CROSS JOIN cursor_row WHERE claims.claim_slot = $1 AND claims.deadline_at IS NOT NULL AND claims.deadline_at < clock_timestamp() AND (claims.deadline_at, claims.job_id, claims.run_lease) > ( - cursor_row.deadline_cursor_deadline_at, - cursor_row.deadline_cursor_job_id, - cursor_row.deadline_cursor_run_lease + (SELECT deadline_cursor_deadline_at FROM cursor_row), + (SELECT deadline_cursor_job_id FROM cursor_row), + (SELECT deadline_cursor_run_lease FROM cursor_row) ) ORDER BY claims.deadline_at, claims.job_id, claims.run_lease LIMIT $2 @@ -14598,7 +14959,6 @@ impl QueueStorage { ORDER BY claims.deadline_at, claims.job_id, claims.run_lease ) AS rn FROM {claim_child} AS claims - CROSS JOIN cursor_row CROSS JOIN after_count WHERE after_count.count < $2 AND claims.claim_slot = $1 @@ -14606,9 +14966,9 @@ impl QueueStorage { AND claims.deadline_at < clock_timestamp() AND (claims.deadline_at, claims.job_id, claims.run_lease) <= ( - cursor_row.deadline_cursor_deadline_at, - cursor_row.deadline_cursor_job_id, - cursor_row.deadline_cursor_run_lease + (SELECT deadline_cursor_deadline_at FROM cursor_row), + (SELECT deadline_cursor_job_id FROM cursor_row), + (SELECT deadline_cursor_run_lease FROM cursor_row) ) ORDER BY claims.deadline_at, claims.job_id, claims.run_lease LIMIT (SELECT GREATEST($2 - count, 0) FROM after_count) @@ -14682,7 +15042,7 @@ impl QueueStorage { "# ))) .bind(slot) - .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .bind(&excluded_job_ids) .bind(&excluded_run_leases) @@ -14694,7 +15054,7 @@ impl QueueStorage { .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 ( @@ -14772,11 +15132,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, @@ -14795,11 +15153,17 @@ 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(&excluded_job_ids) .bind(&excluded_run_leases) @@ -14808,10 +15172,18 @@ impl QueueStorage { .bind(&selected_run_leases) .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 @@ -14838,12 +15210,13 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, + scan_limit: i64, rescue_limit: i64, excluded: &BTreeSet<(i64, i64)>, selected: Option<&BTreeSet<(i64, i64)>>, - ) -> Result, AwaError> { + ) -> Result { if rescue_limit <= 0 { - return Ok(Vec::new()); + return Ok((Vec::new(), true, None)); } let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = @@ -14865,55 +15238,95 @@ impl QueueStorage { WHERE slot = $1 FOR UPDATE ), + -- The window is counted in members, not batches: it takes whole + -- batches after the cursor in (deadline_at, batch_id) order until + -- they hold $2 members, then wraps to the first batches at or + -- before the cursor for the rest. Every batch has at least one + -- member, so $2 batches are always enough to read. after_cursor AS MATERIALIZED ( SELECT - claim_batches.claim_slot, - claim_batches.batch_id, + counted.claim_slot, + counted.batch_id, + counted.claimed_count, row_number() OVER ( - ORDER BY claim_batches.deadline_at, claim_batches.batch_id + ORDER BY counted.deadline_at, counted.batch_id ) AS rn - FROM {claim_batch_child} AS claim_batches - CROSS JOIN cursor_row - WHERE claim_batches.claim_slot = $1 - AND claim_batches.deadline_at IS NOT NULL - AND claim_batches.deadline_at < clock_timestamp() - AND (claim_batches.deadline_at, claim_batches.batch_id) - > ( - cursor_row.batch_deadline_cursor_deadline_at, - cursor_row.batch_deadline_cursor_batch_id - ) - ORDER BY claim_batches.deadline_at, claim_batches.batch_id - LIMIT $2 + FROM ( + SELECT + nearest.*, + COALESCE(sum(nearest.claimed_count) OVER ( + ORDER BY nearest.deadline_at, nearest.batch_id + ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING + ), 0) AS members_before + FROM ( + SELECT + claim_batches.claim_slot, + claim_batches.batch_id, + claim_batches.deadline_at, + claim_batches.claimed_count + FROM {claim_batch_child} AS claim_batches + WHERE claim_batches.claim_slot = $1 + AND claim_batches.deadline_at IS NOT NULL + AND claim_batches.deadline_at < clock_timestamp() + AND (claim_batches.deadline_at, claim_batches.batch_id) + > ( + (SELECT batch_deadline_cursor_deadline_at FROM cursor_row), + (SELECT batch_deadline_cursor_batch_id FROM cursor_row) + ) + ORDER BY claim_batches.deadline_at, claim_batches.batch_id + LIMIT $2 + ) AS nearest + ) AS counted + WHERE counted.members_before < $2 ), - after_count AS ( - SELECT count(*)::bigint AS count FROM after_cursor + after_count AS MATERIALIZED ( + SELECT + count(*)::bigint AS count, + COALESCE(sum(claimed_count), 0)::bigint AS members + FROM after_cursor ), before_cursor AS MATERIALIZED ( SELECT - claim_batches.claim_slot, - claim_batches.batch_id, + counted.claim_slot, + counted.batch_id, + counted.claimed_count, after_count.count + row_number() OVER ( - ORDER BY claim_batches.deadline_at, claim_batches.batch_id + ORDER BY counted.deadline_at, counted.batch_id ) AS rn - FROM {claim_batch_child} AS claim_batches - CROSS JOIN cursor_row + FROM ( + SELECT + nearest.*, + COALESCE(sum(nearest.claimed_count) OVER ( + ORDER BY nearest.deadline_at, nearest.batch_id + ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING + ), 0) AS members_before + FROM ( + SELECT + claim_batches.claim_slot, + claim_batches.batch_id, + claim_batches.deadline_at, + claim_batches.claimed_count + FROM {claim_batch_child} AS claim_batches + WHERE (SELECT members FROM after_count) < $2 + AND claim_batches.claim_slot = $1 + AND claim_batches.deadline_at IS NOT NULL + AND claim_batches.deadline_at < clock_timestamp() + AND (claim_batches.deadline_at, claim_batches.batch_id) + <= ( + (SELECT batch_deadline_cursor_deadline_at FROM cursor_row), + (SELECT batch_deadline_cursor_batch_id FROM cursor_row) + ) + ORDER BY claim_batches.deadline_at, claim_batches.batch_id + LIMIT (SELECT GREATEST($2 - members, 0) FROM after_count) + ) AS nearest + ) AS counted CROSS JOIN after_count - WHERE after_count.count < $2 - AND claim_batches.claim_slot = $1 - AND claim_batches.deadline_at IS NOT NULL - AND claim_batches.deadline_at < clock_timestamp() - AND (claim_batches.deadline_at, claim_batches.batch_id) - <= ( - cursor_row.batch_deadline_cursor_deadline_at, - cursor_row.batch_deadline_cursor_batch_id - ) - ORDER BY claim_batches.deadline_at, claim_batches.batch_id - LIMIT (SELECT GREATEST($2 - count, 0) FROM after_count) + WHERE counted.members_before < $2 - after_count.members ), candidate_keys AS MATERIALIZED ( - SELECT claim_slot, batch_id, rn FROM after_cursor + SELECT claim_slot, batch_id, claimed_count, rn FROM after_cursor UNION ALL - SELECT claim_slot, batch_id, rn FROM before_cursor + SELECT claim_slot, batch_id, claimed_count, rn FROM before_cursor ), -- Every member of a candidate batch, annotated with the same -- durable closure evidence and materialized-lease anti-joins the @@ -15018,7 +15431,7 @@ impl QueueStorage { "# ))) .bind(slot) - .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .bind(&excluded_job_ids) .bind(&excluded_run_leases) @@ -15030,7 +15443,7 @@ impl QueueStorage { .map_err(map_sqlx_error)? .into_iter() .unzip(); - let rescued: Vec = sqlx::query_as(audited_sql(format!( + let (rescued, scan_full, cursor_at, cursor_batch_id, cursor_run_lease): ReceiptRescueScanRow = sqlx::query_as(audited_sql(format!( r#" WITH {candidate_ctes}, expired_locked AS ( @@ -15150,11 +15563,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 ), - cursor_advance AS ( - SELECT count(*) FROM advance_cursor - ) + rescued_rows AS ( SELECT expired_locked.ready_slot, expired_locked.ready_generation, @@ -15173,11 +15584,17 @@ 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 COALESCE(sum(claimed_count), 0) >= $2 FROM candidate_keys), + (SELECT batch_deadline_cursor_deadline_at FROM advance_cursor), + (SELECT batch_deadline_cursor_batch_id FROM advance_cursor), + (SELECT 0::bigint FROM advance_cursor) "# ))) .bind(slot) - .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) + .bind(scan_limit) .bind(rescue_limit) .bind(&excluded_job_ids) .bind(&excluded_run_leases) @@ -15186,10 +15603,18 @@ impl QueueStorage { .bind(&selected_run_leases) .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_batch_id) + .zip(cursor_run_lease) + .map(|((at, batch_id), lease)| (at, batch_id, lease)), + )) } async fn active_job_read_row_tx( @@ -16129,14 +16554,87 @@ impl QueueStorage { run_lease: i64, callback_id: Uuid, ) -> Result { - let mut tx = pool.begin().await.map_err(map_sqlx_error)?; + self.enter_callback_wait_with_progress(pool, job_id, run_lease, callback_id, None) + .await + } + + /// [`Self::enter_callback_wait`], storing `progress` as the parked + /// attempt's progress in the same transaction. + pub async fn enter_callback_wait_with_progress( + &self, + pool: &PgPool, + job_id: i64, + run_lease: i64, + callback_id: Uuid, + progress: Option<&serde_json::Value>, + ) -> Result { + let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let entered = self - .enter_callback_wait_in_tx(&mut tx, job_id, run_lease, callback_id) + .enter_callback_wait_with_progress_in_tx( + &mut tx, + job_id, + run_lease, + callback_id, + progress, + ) .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(entered) } + /// Transaction-aware variant of + /// [`Self::enter_callback_wait_with_progress`]. + pub async fn enter_callback_wait_with_progress_in_tx( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + job_id: i64, + run_lease: i64, + callback_id: Uuid, + progress: Option<&serde_json::Value>, + ) -> Result { + let entered = self + .enter_callback_wait_in_tx(tx, job_id, run_lease, callback_id) + .await?; + if let (true, Some(progress)) = (entered, progress) { + self.write_lease_progress_tx(tx, job_id, run_lease, progress) + .await?; + } + Ok(entered) + } + + /// Store `progress` for an attempt that has a lease row. + async fn write_lease_progress_tx( + &self, + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + job_id: i64, + run_lease: i64, + progress: &serde_json::Value, + ) -> Result<(), AwaError> { + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {} (job_id, run_lease, progress, updated_at) + SELECT lease.job_id, lease.run_lease, $3, clock_timestamp() + FROM {} AS lease + WHERE lease.job_id = $1 + AND lease.run_lease = $2 + AND lease.state IN ('running', 'waiting_external') + ON CONFLICT (job_id, run_lease) + DO UPDATE SET + progress = EXCLUDED.progress, + updated_at = clock_timestamp() + "#, + self.attempt_state_table(), + self.leases_table() + ))) + .bind(job_id) + .bind(run_lease) + .bind(progress) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(()) + } + /// Transaction-aware variant of [`Self::enter_callback_wait`] (ADR-029). /// Returns whether the row transitioned to `waiting_external`. pub async fn enter_callback_wait_in_tx( @@ -16788,31 +17286,12 @@ impl QueueStorage { self.upsert_attempt_state_progress_from_receipts_tx( &mut tx, &[(job_id, run_lease, progress.clone())], + false, ) .await?; } - sqlx::query(audited_sql(format!( - r#" - INSERT INTO {} (job_id, run_lease, progress, updated_at) - SELECT lease.job_id, lease.run_lease, $3, clock_timestamp() - FROM {} AS lease - WHERE lease.job_id = $1 - AND lease.run_lease = $2 - AND lease.state IN ('running', 'waiting_external') - ON CONFLICT (job_id, run_lease) - DO UPDATE SET - progress = EXCLUDED.progress, - updated_at = clock_timestamp() - "#, - self.attempt_state_table(), - self.leases_table() - ))) - .bind(job_id) - .bind(run_lease) - .bind(progress) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + self.write_lease_progress_tx(&mut tx, job_id, run_lease, &progress) + .await?; tx.commit().await.map_err(map_sqlx_error)?; Ok(()) } @@ -16820,7 +17299,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, @@ -16835,7 +17315,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(); @@ -16907,6 +17399,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, @@ -16927,14 +17451,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!( @@ -16982,12 +17510,22 @@ impl QueueStorage { // anti-joined against durable closure evidence so it picks // up every open claim regardless of whether // `materialize_claims` has fanned it out to `leases` yet. - self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs) + self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs, true) .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(); @@ -17366,59 +17904,196 @@ 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<'_>, + ) -> 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 run_at = retry_run_ats.next().ok_or_else(|| { + AwaError::Validation("rescue backoff missing for a retryable attempt".into()) + })?; + 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 @@ -18308,21 +18983,47 @@ 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 mut tx = pool + .begin_with(RESCUE_TRANSACTION_BEGIN) + .await + .map_err(map_sqlx_error)?; let cutoff = Utc::now() - TimeDelta::from_std(staleness) .map_err(|err| AwaError::Validation(format!("invalid staleness: {err}")))?; + let cutoff = options + .heartbeated_before + .map_or(cutoff, |before| cutoff.min(before)); // #169 B1: the staleness predicate prefers // `attempt_state.heartbeat_at` (receipts-mode source of truth, // where heartbeat_batch writes go) and falls back to @@ -18350,6 +19051,8 @@ impl QueueStorage { // called at 30s cadence — cheap. let mut excluded = BTreeSet::new(); let mut selected: Option> = None; + let mut first_selection_full = false; + let mut first_receipt_scan: Option<(bool, Vec)> = None; loop { let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = excluded.iter().copied().unzip(); @@ -18359,73 +19062,97 @@ impl QueueStorage { .unwrap_or_default(); let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(audited_sql(format!( - r#" - DELETE FROM {schema}.leases AS target - WHERE (target.job_id, target.run_lease) IN ( - SELECT lease.job_id, lease.run_lease - FROM {schema}.leases AS lease - LEFT JOIN {schema}.attempt_state AS attempt - ON attempt.job_id = lease.job_id - AND attempt.run_lease = lease.run_lease - WHERE lease.state = 'running' - AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1 - AND NOT EXISTS ( - SELECT 1 FROM unnest($2::bigint[], $3::bigint[]) AS excluded(job_id, run_lease) - WHERE excluded.job_id = lease.job_id AND excluded.run_lease = lease.run_lease - ) - AND ($4::boolean OR EXISTS ( - SELECT 1 FROM unnest($5::bigint[], $6::bigint[]) AS selected(job_id, run_lease) - WHERE selected.job_id = lease.job_id AND selected.run_lease = lease.run_lease - )) - ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC - LIMIT 500 - FOR UPDATE OF lease SKIP LOCKED - ) - RETURNING - ready_slot, - ready_generation, - job_id, - queue, - state, - priority, - attempt, - run_lease, - max_attempts, - lane_seq, - enqueue_shard, - heartbeat_at, - deadline_at, - attempted_at, - callback_id, - callback_timeout_at - "# - ))) - .bind(cutoff) - .bind(&excluded_job_ids) - .bind(&excluded_run_leases) - .bind(selected.is_none()) - .bind(&selected_job_ids) - .bind(&selected_run_leases) - .fetch_all(attempt.as_mut()) - .await - .map_err(map_sqlx_error)?; + r#" + DELETE FROM {schema}.leases AS target + WHERE (target.job_id, target.run_lease) IN ( + SELECT lease.job_id, lease.run_lease + FROM {schema}.leases AS lease + LEFT JOIN {schema}.attempt_state AS attempt + ON attempt.job_id = lease.job_id + AND attempt.run_lease = lease.run_lease + WHERE lease.state = 'running' + AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1 + AND NOT EXISTS ( + SELECT 1 FROM unnest($2::bigint[], $3::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = lease.job_id AND excluded.run_lease = lease.run_lease + ) + AND ($4::boolean OR EXISTS ( + SELECT 1 FROM unnest($5::bigint[], $6::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = lease.job_id AND selected.run_lease = lease.run_lease + )) + ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC + LIMIT $7 + FOR UPDATE OF lease SKIP LOCKED + ) + RETURNING + ready_slot, + ready_generation, + job_id, + queue, + state, + priority, + attempt, + run_lease, + max_attempts, + lane_seq, + enqueue_shard, + heartbeat_at, + deadline_at, + attempted_at, + callback_id, + callback_timeout_at + "# + ))) + .bind(cutoff) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .bind(RECEIPT_RESCUE_BATCH_LIMIT) + .fetch_all(attempt.as_mut()) + .await + .map_err(map_sqlx_error)?; + if selected.is_none() { + first_selection_full = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; + } + let mut saturated = + first_selection_full || deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; let receipt_excluded: BTreeSet<_> = excluded .iter() .copied() .chain(deleted.iter().map(|row| (row.job_id, row.run_lease))) .collect(); - let rescued_receipts = if self.lease_claim_receipts() { + let mut receipts = ReceiptRescueBatch::default(); + if self.lease_claim_receipts() { self.rescue_stale_receipt_claims_tx( &mut attempt, cutoff, + options, &receipt_excluded, selected.as_ref(), + &mut receipts, ) - .await? - } else { - Vec::new() - }; + .await?; + } + saturated |= receipts.saturated; + let ReceiptRescueBatch { + rescued: rescued_receipts, + mut idle_slots, + scan_progress, + .. + } = receipts; + // A retry scans only the attempts the first selection found, so + // the window's fullness and the slots it found empty come from + // the first selection. + match &first_receipt_scan { + None => first_receipt_scan = Some((saturated, idle_slots.clone())), + Some((first_saturated, first_idle_slots)) => { + saturated |= *first_saturated; + idle_slots.clone_from(first_idle_slots); + } + } selected.get_or_insert_with(|| { deleted @@ -18438,7 +19165,12 @@ impl QueueStorage { if deleted.is_empty() && rescued_receipts.is_empty() { attempt.commit().await.map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); + return Ok(RescueBatch { + idle_slots, + scan_progress, + saturated, + ..RescueBatch::default() + }); } let moved_leases = match self @@ -18464,34 +19196,37 @@ impl QueueStorage { } }; - let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len()); - for row in moved_leases { - let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut attempt, - row, - "heartbeat stale: worker presumed dead", - dlq_reason, - ) - .await?, - ); - } - for row in moved_receipts { - let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut attempt, - row, - "receipt claim stale: worker presumed dead", - dlq_reason, - ) - .await?, - ); - } + let mut rescued = self + .settle_rescued_attempts_tx( + &mut attempt, + moved_leases, + "heartbeat stale: worker presumed dead", + RescueDlq { + reason: "heartbeat_stale", + enabled: dlq_enabled, + }, + ) + .await?; + rescued.extend( + self.settle_rescued_attempts_tx( + &mut attempt, + moved_receipts, + "receipt claim stale: worker presumed dead", + RescueDlq { + reason: "heartbeat_stale", + enabled: dlq_enabled, + }, + ) + .await?, + ); attempt.commit().await.map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; - return Ok(rescued); + return Ok(RescueBatch { + rescued, + saturated, + idle_slots, + scan_progress, + }); } } @@ -18506,19 +19241,38 @@ 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 mut tx = pool + .begin_with(RESCUE_TRANSACTION_BEGIN) + .await + .map_err(map_sqlx_error)?; let mut excluded = BTreeSet::new(); let mut selected: Option> = None; + let mut first_selection_full = false; + let mut first_receipt_scan: Option<(bool, Vec)> = None; loop { let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = excluded.iter().copied().unzip(); @@ -18528,53 +19282,59 @@ impl QueueStorage { .unwrap_or_default(); let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(audited_sql(format!( - r#" - DELETE FROM {schema}.leases - WHERE (job_id, run_lease) IN ( - SELECT job_id, run_lease - FROM {schema}.leases - WHERE state = 'running' - AND deadline_at IS NOT NULL - AND deadline_at < clock_timestamp() - AND NOT EXISTS ( - SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) - WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease - ) - AND ($3::boolean OR EXISTS ( - SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) - WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease - )) - ORDER BY deadline_at ASC - LIMIT 500 - FOR UPDATE SKIP LOCKED - ) - RETURNING - ready_slot, - ready_generation, - job_id, - queue, - state, - priority, - attempt, - run_lease, - max_attempts, - lane_seq, - enqueue_shard, - heartbeat_at, - deadline_at, - attempted_at, - callback_id, - callback_timeout_at - "# - ))) - .bind(&excluded_job_ids) - .bind(&excluded_run_leases) - .bind(selected.is_none()) - .bind(&selected_job_ids) - .bind(&selected_run_leases) - .fetch_all(attempt.as_mut()) - .await - .map_err(map_sqlx_error)?; + r#" + DELETE FROM {schema}.leases + WHERE (job_id, run_lease) IN ( + SELECT job_id, run_lease + FROM {schema}.leases + WHERE state = 'running' + AND deadline_at IS NOT NULL + AND deadline_at < clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease + ) + AND ($3::boolean OR EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease + )) + ORDER BY deadline_at ASC + LIMIT $6 + FOR UPDATE SKIP LOCKED + ) + RETURNING + ready_slot, + ready_generation, + job_id, + queue, + state, + priority, + attempt, + run_lease, + max_attempts, + lane_seq, + enqueue_shard, + heartbeat_at, + deadline_at, + attempted_at, + callback_id, + callback_timeout_at + "# + ))) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .bind(RECEIPT_RESCUE_BATCH_LIMIT) + .fetch_all(attempt.as_mut()) + .await + .map_err(map_sqlx_error)?; + if selected.is_none() { + first_selection_full = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; + } + let mut saturated = + first_selection_full || 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 @@ -18587,16 +19347,34 @@ impl QueueStorage { .copied() .chain(deleted.iter().map(|row| (row.job_id, row.run_lease))) .collect(); - let receipt_deleted = if self.lease_claim_receipts() { + let mut receipts = ReceiptRescueBatch::default(); + if self.lease_claim_receipts() { self.rescue_expired_receipt_deadlines_tx( &mut attempt, + options, &receipt_excluded, selected.as_ref(), + &mut receipts, ) - .await? - } else { - Vec::new() - }; + .await?; + } + saturated |= receipts.saturated; + let ReceiptRescueBatch { + rescued: receipt_deleted, + mut idle_slots, + scan_progress, + .. + } = receipts; + // A retry scans only the attempts the first selection found, so + // the window's fullness and the slots it found empty come from + // the first selection. + match &first_receipt_scan { + None => first_receipt_scan = Some((saturated, idle_slots.clone())), + Some((first_saturated, first_idle_slots)) => { + saturated |= *first_saturated; + idle_slots.clone_from(first_idle_slots); + } + } selected.get_or_insert_with(|| { deleted @@ -18609,7 +19387,12 @@ impl QueueStorage { if deleted.is_empty() && receipt_deleted.is_empty() { attempt.commit().await.map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); + return Ok(RescueBatch { + idle_slots, + scan_progress, + saturated, + ..RescueBatch::default() + }); } let mut moved = match self @@ -18623,36 +19406,37 @@ impl QueueStorage { continue; } }; - moved.extend( - match self - .try_hydrate_deleted_leases_tx(&mut attempt, receipt_deleted, true) - .await? - { - Ok(moved) => moved, - Err(busy) => { - attempt.rollback().await.map_err(map_sqlx_error)?; - busy.exclude(&mut excluded, &mut selected)?; - continue; - } - }, - ); - - let mut rescued = Vec::with_capacity(moved.len()); - for row in moved { - let dlq_reason = dlq_enabled(&row.queue).then_some("deadline_exceeded"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut attempt, - row, - "hard deadline exceeded", - dlq_reason, - ) - .await?, - ); + match self + .try_hydrate_deleted_leases_tx(&mut attempt, receipt_deleted, true) + .await? + { + Ok(receipts_moved) => moved.extend(receipts_moved), + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded, &mut selected)?; + continue; + } } + + let rescued = self + .settle_rescued_attempts_tx( + &mut attempt, + moved, + "hard deadline exceeded", + RescueDlq { + reason: "deadline_exceeded", + enabled: dlq_enabled, + }, + ) + .await?; attempt.commit().await.map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; - return Ok(rescued); + return Ok(RescueBatch { + rescued, + saturated, + idle_slots, + scan_progress, + }); } } @@ -18667,19 +19451,38 @@ 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. The options do 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 mut tx = pool + .begin_with(RESCUE_TRANSACTION_BEGIN) + .await + .map_err(map_sqlx_error)?; let mut excluded = BTreeSet::new(); let mut selected: Option> = None; + let mut first_selection_full = false; loop { let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = excluded.iter().copied().unzip(); @@ -18689,53 +19492,59 @@ impl QueueStorage { .unwrap_or_default(); let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(audited_sql(format!( - r#" - DELETE FROM {schema}.leases - WHERE (job_id, run_lease) IN ( - SELECT job_id, run_lease - FROM {schema}.leases - WHERE state = 'waiting_external' - AND callback_timeout_at IS NOT NULL - AND callback_timeout_at < clock_timestamp() - AND NOT EXISTS ( - SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) - WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease - ) - AND ($3::boolean OR EXISTS ( - SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) - WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease - )) - ORDER BY callback_timeout_at ASC - LIMIT 500 - FOR UPDATE SKIP LOCKED - ) - RETURNING - ready_slot, - ready_generation, - job_id, - queue, - state, - priority, - attempt, - run_lease, - max_attempts, - lane_seq, - enqueue_shard, - heartbeat_at, - deadline_at, - attempted_at, - callback_id, - callback_timeout_at - "# - ))) - .bind(&excluded_job_ids) - .bind(&excluded_run_leases) - .bind(selected.is_none()) - .bind(&selected_job_ids) - .bind(&selected_run_leases) - .fetch_all(attempt.as_mut()) - .await - .map_err(map_sqlx_error)?; + r#" + DELETE FROM {schema}.leases + WHERE (job_id, run_lease) IN ( + SELECT job_id, run_lease + FROM {schema}.leases + WHERE state = 'waiting_external' + AND callback_timeout_at IS NOT NULL + AND callback_timeout_at < clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease + ) + AND ($3::boolean OR EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease + )) + ORDER BY callback_timeout_at ASC + LIMIT $6 + FOR UPDATE SKIP LOCKED + ) + RETURNING + ready_slot, + ready_generation, + job_id, + queue, + state, + priority, + attempt, + run_lease, + max_attempts, + lane_seq, + enqueue_shard, + heartbeat_at, + deadline_at, + attempted_at, + callback_id, + callback_timeout_at + "# + ))) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .bind(RECEIPT_RESCUE_BATCH_LIMIT) + .fetch_all(attempt.as_mut()) + .await + .map_err(map_sqlx_error)?; + if selected.is_none() { + first_selection_full = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; + } + let saturated = + first_selection_full || deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; selected.get_or_insert_with(|| { deleted @@ -18747,7 +19556,7 @@ impl QueueStorage { if deleted.is_empty() { attempt.commit().await.map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); + return Ok(RescueBatch::default()); } let moved = match self @@ -18762,22 +19571,25 @@ impl QueueStorage { } }; - let mut rescued = Vec::with_capacity(moved.len()); - for row in moved { - let dlq_reason = dlq_enabled(&row.queue).then_some("callback_timeout"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut attempt, - row, - "callback timed out", - dlq_reason, - ) - .await?, - ); - } + let rescued = self + .settle_rescued_attempts_tx( + &mut attempt, + moved, + "callback timed out", + RescueDlq { + reason: "callback_timeout", + enabled: dlq_enabled, + }, + ) + .await?; attempt.commit().await.map_err(map_sqlx_error)?; tx.commit().await.map_err(map_sqlx_error)?; - return Ok(rescued); + return Ok(RescueBatch { + rescued, + saturated, + idle_slots: Vec::new(), + scan_progress: Vec::new(), + }); } } diff --git a/awa-python/python/awa/_awa.pyi b/awa-python/python/awa/_awa.pyi index aa1cd7872..ca7d5a50e 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 733114388..b2f4a3e5c 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, @@ -968,6 +970,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_storage_prune_lock_timeout_ms`` sets the PostgreSQL ``lock_timeout`` the maintenance leader uses when pruning the queue, lease and claim rings (default 10 ms). Each lock acquisition gets its @@ -1000,6 +1007,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 cf3bc70fd..e476a06ed 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_prune_lock_timeout_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_prune_lock_timeout_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, @@ -1943,6 +1945,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 b888310e6..3e76e680c 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, @@ -220,6 +222,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, @@ -818,6 +822,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 @@ -1272,6 +1296,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, @@ -1372,6 +1398,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>, @@ -1990,6 +2018,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 b37ab3122..8ecb5d2b7 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); let heartbeat_guard = storage.queue_storage().map(|_| FinalizationHeartbeatGuard { @@ -519,6 +520,9 @@ impl JobExecutor { } } + // The handler's progress snapshot is final: the finalize settles it, + // so heartbeats stop flushing progress for this attempt. + in_flight.mark_finalizing((job_id, job_run_lease)); // Keep worker/heartbeat ownership while the bounded finalizer queue is full. let finalize_permit = match result { Some(_) => tokio::select! { @@ -2384,9 +2388,15 @@ async fn complete_job_queue_storage( .and_then(|by_kind| by_kind.get(&job.kind)) .cloned(); if let Some(specs) = kind_specs.filter(|s| !s.is_empty()) { - let parked = - park_queue_storage_with_followups(runtime, pool, job, guard.id(), &specs) - .await?; + let parked = park_queue_storage_with_followups( + runtime, + pool, + job, + guard.id(), + progress_snapshot.as_ref(), + &specs, + ) + .await?; if let Some(parked_job) = parked { let event = if needs_event { Some(UntypedJobEvent::WaitingForCallback { job: parked_job }) @@ -2403,7 +2413,13 @@ async fn complete_job_queue_storage( let entered = runtime .store - .enter_callback_wait(pool, job.id, job.run_lease, guard.id()) + .enter_callback_wait_with_progress( + pool, + job.id, + job.run_lease, + guard.id(), + progress_snapshot.as_ref(), + ) .await?; if !entered { let current = runtime.store.load_job(pool, job.id).await?; @@ -2925,12 +2941,19 @@ async fn park_queue_storage_with_followups( pool: &PgPool, job: &JobRow, callback_id: uuid::Uuid, + progress_snapshot: Option<&serde_json::Value>, specs: &[crate::enqueue_specs::BoxedEnqueueSpec], ) -> Result, AwaError> { let mut tx = pool.begin().await?; let entered = runtime .store - .enter_callback_wait_in_tx(&mut tx, job.id, job.run_lease, callback_id) + .enter_callback_wait_with_progress_in_tx( + &mut tx, + job.id, + job.run_lease, + callback_id, + progress_snapshot, + ) .await?; if !entered { tx.rollback().await?; diff --git a/awa-worker/src/heartbeat.rs b/awa-worker/src/heartbeat.rs index 91b74f81b..49dd37706 100644 --- a/awa-worker/src/heartbeat.rs +++ b/awa-worker/src/heartbeat.rs @@ -170,7 +170,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 +217,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 +243,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"); @@ -314,7 +320,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 0a7136b7a..1993dd191 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -10,14 +10,15 @@ use awa_model::cron::{ use awa_model::SkipReason; use awa_model::{ AwaError, JobRow, JobState, PeriodicJob, PromotionLane, PruneOutcome, QueueStorage, - RotateOutcome, TerminalDeltaRollupOutcome, + 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, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::task::{JoinHandle, JoinSet}; @@ -230,6 +231,276 @@ 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 in every batch of a pass that may +/// take more than one: deadline rescue, and stale-heartbeat rescue once its +/// owners are confirmed gone. Every claim examined costs its closure checks, +/// so a narrow window keeps each batch's transaction short; the pass continues +/// while windows come back full. A pass limited to one batch keeps the store's +/// wider default so it can walk past a long closed prefix. +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, +} + +#[derive(Debug, Default)] +struct RescueSweepCursor { + first: Option<(DateTime, i64, i64)>, + previous: Option<(DateTime, i64, i64)>, + wrapped: bool, + /// The latest scan left its window unfilled. + drained: 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.drained = true; + cursor.done = true; + } else if cursor.drained { + // A drained cursor whose window fills again starts a new sweep. + *cursor = RescueSweepCursor::default(); + } + 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 { @@ -725,6 +996,15 @@ 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>, + /// Batches each rescue kind has started under this service, across + /// passes, for the slot-order rotation. + rescue_batches_started: std::sync::Mutex>, + rescue_others_first_offset: AtomicUsize, promote_interval: Duration, cleanup_interval: Duration, cron_sync_interval: Duration, @@ -1036,6 +1316,13 @@ 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_batches_started: 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), @@ -1164,6 +1451,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 @@ -1274,9 +1577,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); @@ -1320,6 +1623,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 @@ -1388,22 +1705,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() => { @@ -1881,42 +2207,306 @@ 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 batch_number = { + let mut started = self + .rescue_batches_started + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let started = started.entry(kind).or_default(); + let batch_number = *started; + *started = started.wrapping_add(1); + batch_number + }; + let slot_order = + if batch_number % 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(), + cursor_scan_limit: (batches > 0 + || (!budget.is_zero() + && self.rescue_pass_max_rows > 0 + && match kind { + RescueKind::Heartbeat => bulk_heartbeat_cutoff.is_some(), + RescueKind::Deadline => true, + RescueKind::Callback => false, + })) + .then_some(RESCUE_FOLLOW_UP_SCAN_LIMIT), + // Follow-up batches take only claims last heartbeated + // before a runtime known to be gone stopped reporting. + heartbeated_before: if batches == 0 { + first_heartbeat_cutoff + } else { + bulk_heartbeat_cutoff + }, + }; + let outcome = match kind { + RescueKind::Heartbeat => { + runtime + .store + .rescue_stale_heartbeats_batch( + &self.pool, + self.heartbeat_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; + 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 @@ -2080,33 +2670,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 } } } @@ -2253,33 +2835,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 } } } @@ -4296,8 +4870,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 @@ -4550,6 +5128,729 @@ 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]); + } + + /// A slot with two cursors is idle only while both are finished. One + /// that drained and then filled its window again is not finished when + /// the other drains. + #[test] + fn rescue_sweep_keeps_a_slot_whose_drained_cursor_fills_again() { + let stamp = Utc::now(); + let observe = |sweep: &mut RescueSweep, first: (i64, bool), second: (i64, bool)| { + sweep.observe(RescueScanProgress { + slot: 0, + cursors: [first, second] + .into_iter() + .map(|(id, saturated)| awa_model::RescueScanCursor { + position: Some((stamp, id, 0)), + saturated, + }) + .collect(), + }); + }; + let mut sweep = RescueSweep::default(); + observe(&mut sweep, (10, false), (100, true)); + assert!(sweep.cursors[&0][0].done); + assert!(sweep.idle_slots.is_empty()); + observe(&mut sweep, (20, true), (200, true)); + assert!( + !sweep.cursors[&0][0].done, + "a cursor whose window fills again has work left" + ); + observe(&mut sweep, (30, true), (300, false)); + assert!( + sweep.idle_slots.is_empty(), + "the slot stays active while its first cursor is saturated" + ); + observe(&mut sweep, (40, false), (300, false)); + 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; + } + + /// Passes that do not re-arm still share the slot-order rotation: every + /// fourth batch of a kind visits the other claim slots first. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_slot_order_rotates_across_passes_that_do_not_rearm() { + 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"); + 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_max_rows(0); + for pass in 1..=2 * RESCUE_OTHERS_FIRST_EVERY { + assert!( + !service + .rescue_queue_storage_pass( + RescueKind::Deadline, + &runtime, + false, + RescueTickOrigin::Ordinary, + ) + .await + ); + assert_eq!( + service.rescue_others_first_offset.load(Ordering::Relaxed), + pass / RESCUE_OTHERS_FIRST_EVERY, + "others-first batches after {pass} one-batch passes" + ); + } + 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..17fd2bd33 --- /dev/null +++ b/awa/tests/bulk_rescue_test.rs @@ -0,0 +1,1075 @@ +//! 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"; + +/// The backlog tests each load the database heavily enough to slow the other's +/// heartbeat passes past their staleness window, so they run one at a time: +/// this lock serialises them under `cargo test`, and a nextest test group in +/// `.config/nextest.toml` does so across processes. +static BACKLOG_TEST: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(()); + +/// 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") +} + +/// Record, as each compact claim is closed as rescued, how long it had gone +/// without a heartbeat. The trigger runs inside the closing statement, after +/// rescue holds the attempt's lock, so no later heartbeat can land. +async fn record_compact_rescue_heartbeats(pool: &PgPool) { + sqlx::raw_sql( + r#" + CREATE TABLE public.rescued_claim_heartbeats ( + job_id bigint NOT NULL, + queue text NOT NULL, + unheartbeated_for interval NOT NULL + ); + CREATE FUNCTION public.record_rescued_claim_heartbeat() RETURNS trigger + LANGUAGE plpgsql AS $$ + BEGIN + INSERT INTO public.rescued_claim_heartbeats + SELECT items.job_id, batches.queue, + NEW.closed_at - COALESCE(attempt.heartbeat_at, batches.claimed_at) + FROM awa.lease_claim_batches AS batches + CROSS JOIN LATERAL unnest(batches.receipt_ids, batches.job_ids, batches.run_leases) + AS items(receipt_id, job_id, run_lease) + LEFT JOIN awa.attempt_state AS attempt + ON attempt.job_id = items.job_id AND attempt.run_lease = items.run_lease + WHERE batches.claim_slot = NEW.claim_slot + AND items.receipt_id = ANY(NEW.receipt_ids); + RETURN NEW; + END + $$; + CREATE TRIGGER record_rescued_claim_heartbeat + AFTER INSERT ON awa.lease_claim_closure_batches + FOR EACH ROW WHEN (NEW.outcome = 'rescued') + EXECUTE FUNCTION public.record_rescued_claim_heartbeat(); + "#, + ) + .execute(pool) + .await + .expect("record rescued claim heartbeats"); +} + +/// Abandoned-queue claims rescued so far: how many had a heartbeat (or claim) +/// younger than `staleness` when closed, how many there were, and the shortest +/// time any of them had gone without a heartbeat. +async fn rescued_heartbeat_ages(pool: &PgPool, staleness: Duration) -> (i64, i64, String) { + sqlx::query_as( + "SELECT count(*) FILTER (WHERE unheartbeated_for < $2 * interval '1 millisecond'), + count(*), + COALESCE(min(unheartbeated_for)::text, 'n/a') + FROM public.rescued_claim_heartbeats + WHERE queue = $1", + ) + .bind(ABANDONED_QUEUE) + .bind(staleness.as_millis() as f64) + .fetch_one(pool) + .await + .expect("rescued claim heartbeat ages") +} + +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 _backlog_test = BACKLOG_TEST.lock().await; + 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; + record_compact_rescue_heartbeats(&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 (fresh, rescued, freshest) = rescued_heartbeat_ages(&pool, staleness).await; + assert_eq!( + fresh, 0, + "rescue closed {fresh} claims last heartbeated within the {staleness:?} staleness window" + ); + assert_eq!( + rescued_rows(&pool, ABANDONED_QUEUE).await, + 0, + "rescue took a claim while its runtime was alive: {rescued} claims had gone {freshest} \ + or longer without a heartbeat, beyond the {staleness:?} staleness, so the live \ + runtime's heartbeat passes fell behind" + ); + 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:?})" + ); + let (fresh, rescued, freshest) = rescued_heartbeat_ages(&pool, staleness).await; + assert_eq!( + (fresh, rescued), + (0, ABANDONED), + "every abandoned claim is rescued once it has gone {staleness:?} without a heartbeat \ + (the freshest had gone {freshest})" + ); + + 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"); + // Completed rows leave the ring once queue prune wraps past them, so the + // handler counts the recovery runs. + wait_until( + scaled_timeout(Duration::from_secs(120)), + "every retried job to run", + || async { recovered.lock().expect("recovered").len() as i64 >= ABANDONED - exhausted }, + ) + .await; + wait_until( + scaled_timeout(Duration::from_secs(30)), + "the retried jobs to finish", + || async { + running(&pool, &store, ABANDONED_QUEUE).await == 0 + && waiting_retry_counts(&pool, ABANDONED_QUEUE).await == (0, 0) + }, + ) + .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 exactly 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"); + assert_eq!(rescues.len() as i64, ABANDONED); + let unexpected = rescued_twice(&rescues); + assert!( + unexpected.is_empty(), + "neither live nor recovered attempts may be rescued: {unexpected:?}" + ); + } + 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() + ); +} + +const DEADLINE_QUEUE: &str = "bulk_rescue_deadline"; + +fn env_u64(name: &str) -> Option { + std::env::var(name).ok().map(|value| { + value + .parse() + .unwrap_or_else(|err| panic!("{name}={value}: {err}")) + }) +} + +async fn ring_generation(pool: &PgPool, ring: &str) -> i64 { + count( + pool, + format!( + "SELECT generation FROM awa.{ring}_ring_rotations ORDER BY generation DESC LIMIT 1" + ), + ) + .await +} + +/// First-attempt claims still open: a dead runtime's attempts that no rescue +/// has closed. A pruned claim slot takes its closed claims with it. +async fn open_first_attempt_claims(pool: &PgPool) -> i64 { + count( + pool, + format!( + "SELECT count(*) + FROM awa.lease_claim_batches AS batch + CROSS JOIN LATERAL unnest(batch.run_leases, batch.receipt_ids) + AS item(run_lease, receipt_id) + WHERE batch.queue = '{DEADLINE_QUEUE}' + AND item.run_lease = 1 + AND NOT EXISTS ( + SELECT 1 FROM awa.lease_claim_closure_batches AS closed + WHERE closed.claim_slot = batch.claim_slot + AND closed.receipt_ranges @> item.receipt_id + )" + ), + ) + .await +} + +/// A killed runtime left hundreds of thousands of claims with a deadline, in +/// compact claim batches that each share one deadline. They span two claim +/// slots, all reference one queue slot, were last heartbeated hours ago and +/// are past their deadlines. Until they close, queue prune skips that queue +/// slot and claim prune skips both claim slots, so neither ring can wrap. +/// With default rescue intervals, rescue must close every claim and both +/// rings must move past the pinned slots. +/// +/// `AWA_DEADLINE_RESCUE_CLAIMS` sets the backlog (default 20,000); +/// `AWA_DEADLINE_RESCUE_PASS_BUDGET_MS` and `AWA_DEADLINE_RESCUE_PASS_MAX_ROWS` +/// override the pass controls, and `AWA_DEADLINE_RESCUE_OBSERVE_SECS` reports +/// the drain rate after that long instead of requiring a full drain. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn rescue_drains_abandoned_deadline_claims_and_unpins_both_rings() { + let _backlog_test = BACKLOG_TEST.lock().await; + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(); + let abandoned = env_u64("AWA_DEADLINE_RESCUE_CLAIMS").unwrap_or(20_000) as i64; + let observe_for = env_u64("AWA_DEADLINE_RESCUE_OBSERVE_SECS").map(Duration::from_secs); + + let db = TestDatabase::queue_storage().await; + let pool = db.pool().clone(); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let queue_slots = store.queue_slot_count() as i64; + let claim_slots = store.claim_slot_count() as i64; + + let setup_started = Instant::now(); + for first in (0..abandoned).step_by(5_000) { + let params: Vec<_> = (first..(first + 5_000).min(abandoned)) + .map(|n| { + awa::model::insert::params_with( + &AbandonedJob { n }, + InsertOpts { + queue: DEADLINE_QUEUE.into(), + ..Default::default() + }, + ) + .expect("abandoned params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue abandoned jobs"); + } + // Half the backlog in each of two claim slots, as a runtime that kept + // claiming across a claim rotation leaves it. + let mut claimed = 0; + let mut rotated = false; + while claimed < abandoned { + if !rotated && claimed >= abandoned / 2 { + store.rotate_claims(&pool).await.expect("rotate claims"); + rotated = true; + } + let batch = store + .claim_runtime_batch(&pool, DEADLINE_QUEUE, 5_000, Duration::from_secs(300)) + .await + .expect("claim abandoned jobs"); + assert!(!batch.is_empty(), "the backlog should still be claimable"); + claimed += batch.len() as i64; + } + // The runtime died hours ago: its claims, deadlines and last heartbeats + // all predate the staleness window, and its snapshot stopped then too. + sqlx::raw_sql(audited_sql(format!( + "UPDATE awa.lease_claim_batches + SET claimed_at = claimed_at - interval '20 hours', + deadline_at = deadline_at - interval '20 hours' + WHERE queue = '{DEADLINE_QUEUE}'; + INSERT INTO awa.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + SELECT item.job_id, item.run_lease, + batch.claimed_at + interval '1 minute', batch.claimed_at + interval '1 minute' + FROM awa.lease_claim_batches AS batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + WHERE batch.queue = '{DEADLINE_QUEUE}' + ON CONFLICT (job_id, run_lease) DO UPDATE SET heartbeat_at = EXCLUDED.heartbeat_at; + INSERT INTO awa.runtime_instances ( + instance_id, pid, version, started_at, last_seen_at, snapshot_interval_ms, + healthy, postgres_connected, poll_loop_alive, heartbeat_alive, + maintenance_alive, shutting_down, leader, storage_capability, transition_role + ) VALUES ( + gen_random_uuid(), 1, 'killed', now() - interval '21 hours', + now() - interval '19 hours', 10000, true, true, true, true, true, false, false, + 'queue_storage', 'queue_storage_target' + ); + ANALYZE awa.ready_entries; + ANALYZE awa.lease_claim_batches; + ANALYZE awa.attempt_state;" + ))) + .execute(&pool) + .await + .expect("age the abandoned claims"); + let shape: (i64, i64, i64, i64, i64) = sqlx::query_as(audited_sql(format!( + "SELECT COALESCE(sum(claimed_count), 0)::bigint, count(DISTINCT claim_slot), + count(DISTINCT ready_slot), count(*) FILTER (WHERE deadline_at IS NULL), + (SELECT count(*) FROM awa.lease_claims) + FROM awa.lease_claim_batches WHERE queue = '{DEADLINE_QUEUE}'" + ))) + .fetch_one(&pool) + .await + .expect("backlog shape"); + assert_eq!( + shape, + (abandoned, 2, 1, 0, 0), + "deadline claim batches over two claim slots and one queue slot" + ); + let queue_generation = ring_generation(&pool, "queue").await; + let claim_generation = ring_generation(&pool, "claim").await; + eprintln!( + "{abandoned} abandoned deadline claims built in {:?}", + setup_started.elapsed() + ); + + // Completed rows leave the ring once queue prune wraps past them, so the + // handler counts the recovery runs. + let second_attempts = Arc::new(AtomicU64::new(0)); + let other_attempts = Arc::new(AtomicU64::new(0)); + let mut builder = Client::builder(pool.clone()) + .queue_storage( + QueueStorageConfig::default(), + Duration::from_secs(1), + Duration::from_secs(1), + ) + .claim_rotate_interval(Duration::from_secs(1)) + .queue( + DEADLINE_QUEUE, + QueueConfig { + max_workers: 200, + ..QueueConfig::default() + }, + ) + .register::({ + let second_attempts = second_attempts.clone(); + let other_attempts = other_attempts.clone(); + move |_, ctx| { + let runs = if ctx.job.attempt == 2 { + &second_attempts + } else { + &other_attempts + }; + runs.fetch_add(1, Ordering::SeqCst); + async { Ok(JobResult::Completed) } + } + }); + if let Some(budget) = env_u64("AWA_DEADLINE_RESCUE_PASS_BUDGET_MS") { + builder = builder.rescue_pass_budget(Duration::from_millis(budget)); + } + if let Some(max_rows) = env_u64("AWA_DEADLINE_RESCUE_PASS_MAX_ROWS") { + builder = builder.rescue_pass_max_rows(max_rows as usize); + } + let client = builder.build().expect("client"); + let started = Instant::now(); + client.start().await.expect("start client"); + + let mut first_rescue = None; + let mut drained = None; + loop { + let open = open_first_attempt_claims(&pool).await; + if open < abandoned && first_rescue.is_none() { + first_rescue = Some(started.elapsed()); + } + if open == 0 { + drained = Some(started.elapsed()); + break; + } + if let Some(observe_for) = observe_for { + if started.elapsed() >= observe_for { + let rescued = abandoned - open; + let since_first = started.elapsed() - first_rescue.unwrap_or_default(); + eprintln!( + "observed {rescued} of {abandoned} rescued in {:?} ({:?} after the first \ + rescue): {:.1} claims/s", + started.elapsed(), + since_first, + rescued as f64 / since_first.as_secs_f64().max(f64::EPSILON) + ); + break; + } + } + assert!( + started.elapsed() < scaled_timeout(Duration::from_secs(600 + abandoned as u64 / 200)), + "{open} of {abandoned} abandoned claims are still open" + ); + tokio::time::sleep(Duration::from_secs(1)).await; + } + + if let Some(drained) = drained { + let first_rescue = first_rescue.unwrap_or_default(); + eprintln!( + "first rescue {first_rescue:?} after start; all {abandoned} closed {drained:?} after \ + start: {:.1} claims/s from the first rescue", + abandoned as f64 / (drained - first_rescue).as_secs_f64().max(f64::EPSILON) + ); + let unpinned = wait_until( + scaled_timeout(Duration::from_secs(120 + abandoned as u64 / 10_000)), + "both rings to rotate past the slots the backlog pinned", + || async { + ring_generation(&pool, "queue").await >= queue_generation + queue_slots + && ring_generation(&pool, "claim").await >= claim_generation + claim_slots + }, + ) + .await; + eprintln!( + "queue and claim rings rotated past the pinned slots {unpinned:?} after the drain" + ); + wait_until( + scaled_timeout(Duration::from_secs(300 + abandoned as u64 / 500)), + "every rescued job to run on its second attempt", + || async { second_attempts.load(Ordering::SeqCst) >= abandoned as u64 }, + ) + .await; + assert_eq!(second_attempts.load(Ordering::SeqCst), abandoned as u64); + assert_eq!( + other_attempts.load(Ordering::SeqCst), + 0, + "only the abandoned first attempts are rescued" + ); + } + client.shutdown(Duration::from_secs(10)).await; +} diff --git a/awa/tests/heartbeat_pass_benchmark_test.rs b/awa/tests/heartbeat_pass_benchmark_test.rs new file mode 100644 index 000000000..8e2b123b6 --- /dev/null +++ b/awa/tests/heartbeat_pass_benchmark_test.rs @@ -0,0 +1,200 @@ +//! 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 (without_query, query) = match url.split_once('?') { + Some((prefix, query)) => (prefix, Some(query)), + None => (url, None), + }; + let (base, _) = without_query + .rsplit_once('/') + .expect("database URL has a database name"); + match query { + Some(query) => format!("{base}/{database}?{query}"), + None => 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/integration_test.rs b/awa/tests/integration_test.rs index a4b9c6567..33326c38e 100644 --- a/awa/tests/integration_test.rs +++ b/awa/tests/integration_test.rs @@ -2169,6 +2169,7 @@ async fn test_cleanup_runtime_snapshots_preserves_fresh() { use chrono::{TimeDelta, Utc}; use uuid::Uuid; + let _guard = test_lock().lock().await; let client = setup().await; let fresh_id = Uuid::new_v4(); diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index b89177f18..656c90753 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -2557,6 +2557,171 @@ async fn assert_claiming_stops_while_finalization_is_stalled( ); } +/// 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 @@ -9037,6 +9202,464 @@ 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); +} + +/// Out-of-line array reads from one claim slot's batch partition so far, as +/// seen by the only connection of `pool`. A large claim batch stores its +/// member arrays out of line, and a statement that expands the batch's members +/// reads each array it uses once. +async fn claim_batch_array_reads(pool: &sqlx::PgPool, schema: &str, slot: i32) -> i64 { + sqlx::query("SELECT pg_stat_force_next_flush()") + .execute(pool) + .await + .expect("force a statistics flush"); + sqlx::query("SELECT 1") + .execute(pool) + .await + .expect("flush statistics"); + sqlx::query_scalar( + "SELECT COALESCE(toast.idx_scan, 0)::bigint + FROM pg_class AS batches + JOIN pg_namespace AS namespace ON namespace.oid = batches.relnamespace + JOIN pg_stat_all_tables AS toast ON toast.relid = batches.reltoastrelid + WHERE namespace.nspname = $1 AND batches.relname = $2", + ) + .bind(schema) + .bind(format!("lease_claim_batches_{slot}")) + .fetch_one(pool) + .await + .expect("claim batch array reads") +} + +/// A deadline backlog in a few compact batches, each larger than the scan +/// window: every rescue transaction examines the batch at the cursor, not +/// every expired batch in the slot, and the backlog still drains with each +/// claim closed once. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_rescue_deadline_batch_scan_examines_a_bounded_number_of_members() { + const BATCHES: i64 = 4; + const BATCH_MEMBERS: i64 = 1_500; + const SCAN_LIMIT: i64 = 1_000; + const RESCUE_LIMIT: i64 = 500; + let (_db_guard, pool) = setup_pool(10).await; + let schema = "awa_rescue_deadline_batch_bound"; + let queue = "rescue_deadline_batch_bound"; + 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; + let params: Vec<_> = (0..BATCHES * BATCH_MEMBERS) + .map(|id| { + insert::params_with( + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue the backlog"); + let mut slot = None; + for _ in 0..BATCHES { + let claims = store + .claim_runtime_batch(&pool, queue, BATCH_MEMBERS, Duration::from_secs(30)) + .await + .expect("claim one large batch"); + assert_eq!(claims.len() as i64, BATCH_MEMBERS); + slot = Some(claims[0].claim.claim_slot); + } + let slot = slot.expect("claim slot"); + sqlx::query(audited_sql(format!( + "UPDATE {schema}.lease_claim_batches SET deadline_at = now() - interval '1 minute'" + ))) + .execute(&pool) + .await + .expect("expire the batches"); + let shape: (i64, i64, i64, i64) = sqlx::query_as(audited_sql(format!( + "SELECT count(*), min(claimed_count)::bigint, max(claimed_count)::bigint, + count(DISTINCT claim_slot) + FROM {schema}.lease_claim_batches" + ))) + .fetch_one(&pool) + .await + .expect("backlog shape"); + assert_eq!( + shape, + (BATCHES, BATCH_MEMBERS, BATCH_MEMBERS, 1), + "a few large compact batches in one claim slot" + ); + + // One connection, so the statistics read sees the rescue's own reads. + let single = PgPoolOptions::new() + .max_connections(1) + .connect_with((*pool.connect_options()).clone()) + .await + .expect("single-connection pool"); + let options = awa::model::RescueBatchOptions { + cursor_scan_limit: Some(SCAN_LIMIT as usize), + ..Default::default() + }; + // The window holds the batches after the cursor up to the one that + // crosses the scan limit. Each of the two rescue statements expands a + // window batch's six member arrays once. + let window_batches = (SCAN_LIMIT + BATCH_MEMBERS - 1) / BATCH_MEMBERS; + let read_bound = 2 * 6 * window_batches; + assert!( + window_batches < BATCHES, + "the window must be smaller than the backlog for the bound to mean anything" + ); + let mut rescued = HashSet::new(); + let mut transactions = 0; + while (rescued.len() as i64) < BATCHES * BATCH_MEMBERS { + let before = claim_batch_array_reads(&single, schema, slot).await; + let batch = store + .rescue_expired_deadlines_batch(&single, &|_| false, &options) + .await + .expect("rescue batch"); + let reads = claim_batch_array_reads(&single, schema, slot).await - before; + transactions += 1; + assert!( + reads > 0 && reads <= read_bound, + "rescue transaction {transactions} read {reads} member arrays, more than the \ + {read_bound} its window of {window_batches} of the {BATCHES} batches allows" + ); + assert_eq!(batch.rescued.len() as i64, RESCUE_LIMIT); + assert!(batch.saturated); + for job in &batch.rescued { + assert_eq!(job.state, JobState::Retryable); + assert!(rescued.insert(job.id), "job {} was rescued twice", job.id); + } + assert!(transactions <= BATCHES * BATCH_MEMBERS / RESCUE_LIMIT); + } + + // With nothing left, the walk wraps to the first batch and stays bounded. + for _ in 0..2 { + let before = claim_batch_array_reads(&single, schema, slot).await; + let batch = store + .rescue_expired_deadlines_batch(&single, &|_| false, &options) + .await + .expect("rescue batch over closed history"); + let reads = claim_batch_array_reads(&single, schema, slot).await - before; + assert!(batch.rescued.is_empty()); + assert!( + reads <= read_bound, + "a transaction over closed batches read {reads} member arrays" + ); + } + single.close().await; + + let closed: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT COALESCE(sum(closed_count), 0)::bigint + FROM {schema}.lease_claim_closure_batches WHERE outcome = 'deadline_expired'" + ))) + .fetch_one(&pool) + .await + .expect("closed claims"); + assert_eq!(closed, BATCHES * BATCH_MEMBERS, "each claim is closed once"); + let retryable: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.deferred_jobs WHERE state = 'retryable'" + ))) + .fetch_one(&pool) + .await + .expect("retryable jobs"); + assert_eq!(retryable, BATCHES * BATCH_MEMBERS); +} + +#[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] + ); + } +} + +/// A stale-heartbeat batch given `heartbeated_before` leaves an attempt +/// heartbeated after that instant alone, however long the batch waits for a +/// connection before it reads the clock. The attempt is stale by the +/// staleness alone, so a batch without the cutoff takes it. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_rescue_stale_heartbeat_cutoff_holds_while_the_transaction_waits() { + let (_db_guard, pool) = setup_pool(10).await; + let schema = "awa_rescue_heartbeat_cutoff"; + let queue = "rescue_heartbeat_cutoff"; + 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; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await; + let claims = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim"); + assert_eq!(claims.len(), 1); + age_receipt_claim( + &pool, + &store, + job_id, + claims[0].job.run_lease, + Duration::from_secs(3600), + ) + .await; + let heartbeat_at: DateTime = sqlx::query_scalar(audited_sql(format!( + "INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + VALUES ($1, $2, clock_timestamp() - interval '10 seconds', clock_timestamp()) + RETURNING heartbeat_at" + ))) + .bind(job_id) + .bind(claims[0].job.run_lease) + .fetch_one(&pool) + .await + .expect("last heartbeat"); + let wait = Duration::from_secs(2); + let cutoff = heartbeat_at - chrono::Duration::milliseconds(500); + let staleness = Duration::from_secs(1); + + let single = PgPoolOptions::new() + .max_connections(1) + .connect_with((*pool.connect_options()).clone()) + .await + .expect("single-connection pool"); + let held = single.acquire().await.expect("hold the only connection"); + let release = async { + tokio::time::sleep(wait).await; + drop(held); + }; + let with_cutoff = awa::model::RescueBatchOptions { + heartbeated_before: Some(cutoff), + ..Default::default() + }; + let started = Instant::now(); + let (batch, ()) = tokio::join!( + store.rescue_stale_heartbeats_batch(&single, staleness, &|_| false, &with_cutoff), + release + ); + assert!( + started.elapsed() >= wait, + "the batch waited for its connection" + ); + assert!( + batch.expect("batch with cutoff").rescued.is_empty(), + "an attempt heartbeated after the cutoff was rescued once the batch had waited {wait:?}" + ); + + let batch = store + .rescue_stale_heartbeats_batch(&single, staleness, &|_| false, &Default::default()) + .await + .expect("batch without cutoff"); + assert_eq!( + batch.rescued.iter().map(|job| job.id).collect::>(), + vec![job_id] + ); + single.close().await; +} + #[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; @@ -9206,23 +9829,114 @@ async fn test_queue_storage_receipt_rescue_cursor_sweeps_past_fresh_claims_and_w "second rescue should wrap and close the formerly fresh claim" ); - let cursor: (DateTime, i64, i64) = sqlx::query_as(audited_sql(format!( - r#" - SELECT rescue_cursor_claimed_at, rescue_cursor_job_id, rescue_cursor_run_lease - FROM {schema}.claim_ring_slots - WHERE slot = $1 - "# + let cursor: (DateTime, i64, i64) = sqlx::query_as(audited_sql(format!( + r#" + SELECT rescue_cursor_claimed_at, rescue_cursor_job_id, rescue_cursor_run_lease + FROM {schema}.claim_ring_slots + WHERE slot = $1 + "# + ))) + .bind(claim_a.claim.claim_slot) + .fetch_one(&pool) + .await + .expect("read rescue cursor after second rescue"); + assert_eq!( + (cursor.1, cursor.2), + (job_c, claim_c.job.run_lease), + "cursor should advance through the rescued blocker and already-closed tail" + ); + assert_eq!(open_receipt_claim_count(&pool, &store).await, 0); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_flush_progress_waits_for_a_locked_receipt_attempt() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_flush_progress_locked"; + let schema = "awa_qs_flush_progress_locked"; + let store = Arc::new( + create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await, + ); + + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 78 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claim = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim receipt-backed job") + .into_iter() + .next() + .expect("claimed job"); + assert!(claim.claim.lease_claim_receipt); + let run_lease = claim.job.run_lease; + + let mut attempt_lock = pool.begin().await.expect("begin attempt lock tx"); + sqlx::query( + r#" + SELECT pg_catalog.pg_advisory_xact_lock( + pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), + 0 + ) + ) + "#, + ) + .bind(job_id) + .bind(run_lease) + .execute(attempt_lock.as_mut()) + .await + .expect("hold the attempt lock"); + + let flush = { + let store = store.clone(); + let pool = pool.clone(); + tokio::spawn(async move { + store + .flush_progress(&pool, job_id, run_lease, serde_json::json!({"step": 3})) + .await + }) + }; + tokio::time::sleep(Duration::from_millis(200)).await; + assert!( + !flush.is_finished(), + "an explicit flush must wait for the attempt lock rather than skip the attempt" + ); + attempt_lock + .commit() + .await + .expect("release the attempt lock"); + flush + .await + .expect("flush task") + .expect("flush progress after the lock is released"); + + let progress: Option = sqlx::query_scalar(audited_sql(format!( + "SELECT progress FROM {schema}.attempt_state WHERE job_id = $1 AND run_lease = $2" ))) - .bind(claim_a.claim.claim_slot) - .fetch_one(&pool) + .bind(job_id) + .bind(run_lease) + .fetch_optional(&pool) .await - .expect("read rescue cursor after second rescue"); - assert_eq!( - (cursor.1, cursor.2), - (job_c, claim_c.job.run_lease), - "cursor should advance through the rescued blocker and already-closed tail" - ); - assert_eq!(open_receipt_claim_count(&pool, &store).await, 0); + .expect("read attempt state"); + assert_eq!(progress, Some(serde_json::json!({"step": 3}))); } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] @@ -14492,6 +15206,119 @@ async fn test_queue_storage_retry_jobs_by_ids_dedupes_duplicate_ids() { assert_eq!(counts.terminal, 0); } +/// Waits until another backend in this database is blocked on a lock while +/// running a statement that contains `fragment`. +async fn wait_for_lock_waiter(pool: &sqlx::PgPool, fragment: &str) { + let started = Instant::now(); + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS ( + SELECT 1 FROM pg_stat_activity + WHERE datname = current_database() + AND pid <> pg_backend_pid() + AND wait_event_type = 'Lock' + AND query LIKE '%' || $1 || '%' + )", + ) + .bind(fragment) + .fetch_one(pool) + .await + .expect("poll pg_stat_activity for a lock waiter"); + if waiting { + return; + } + assert!( + started.elapsed() < ci_timing::scaled_timeout(Duration::from_secs(10)), + "no backend waited on a lock while running {fragment:?}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +/// `TRUNCATE` is not MVCC-safe: a read whose snapshot predates a prune and +/// reaches the pruned slot after the prune commits would see the slot empty, +/// and could not see the failed row prune re-homed into the live slot either. +/// A point read locks the ring before it takes its snapshot, so one that +/// starts while prune holds the slot waits for the prune and must find the +/// failed job in the live slot. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_load_job_finds_failed_row_carried_by_concurrent_prune() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_load_job_prune_carry"; + let schema = "awa_qs_load_job_prune_carry"; + let store = create_store(&pool, schema) + .await + .with_prune_lock_timeout(Duration::from_secs(30)) + .expect("prune lock timeout"); + + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim job") + .into_iter() + .next() + .expect("missing claimed job"); + store + .fail_terminal(&pool, job_id, claimed.job.run_lease, "boom", None) + .await + .expect("fail job") + .expect("running job should fail"); + assert!(matches!( + store.rotate(&pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + + // Hold prune after it has locked slot 0 and before it carries the failed + // row into slot 1, then start a read so it queues behind prune's slot 0 + // lock. + let mut holder = pool.begin().await.expect("begin carry holder"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.done_entries_1 IN SHARE MODE" + ))) + .execute(holder.as_mut()) + .await + .expect("hold the live done slot"); + + let prune = store.prune_oldest(&pool, Duration::from_secs(3600)); + let read = async { + wait_for_lock_waiter(&pool, "WITH carried AS").await; + let release = async { + wait_for_lock_waiter(&pool, "IN ACCESS SHARE MODE").await; + holder.rollback().await.expect("release carry holder"); + }; + let (job, ()) = tokio::join!(store.load_job(&pool, job_id), release); + job + }; + let (pruned, job) = tokio::join!(prune, read); + + assert!( + matches!( + pruned.expect("prune_oldest"), + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } + ), + "prune should carry the failed row out of slot 0" + ); + let job = job + .expect("load_job") + .expect("a failed job carried by a concurrent prune must stay visible"); + assert_eq!(job.id, job_id); + assert_eq!(job.state, JobState::Failed); +} + /// #337: a mix of completed and failed terminal rows under a generous /// retention floor. The completed rows fold into `pruned_completed_count`; /// the failed rows (including a NARROW ready-backed one) are carried into @@ -17543,6 +18370,72 @@ async fn test_queue_prune_releases_segment_lock_on_every_tail_conflict() { assert_eq!(store.queue_counts(&pool, "q").await.unwrap().terminal, 1); } +struct CallbackProgressWorker; + +#[async_trait::async_trait] +impl Worker for CallbackProgressWorker { + fn kind(&self) -> &'static str { + "callback_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let callback = ctx + .register_callback(Duration::from_secs(30)) + .await + .map_err(JobError::retryable)?; + ctx.set_progress(60, "awaiting callback"); + Ok(JobResult::WaitForCallback(callback)) + } +} + +/// Progress a handler sets after registering its callback, with no heartbeat +/// in between, is stored when the attempt is parked. +async fn callback_wait_stores_the_final_progress(receipts: bool) { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_callback_progress"; + let store_config = QueueStorageConfig { + schema: "awa_qs_callback_progress".to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: receipts, + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + let job_id = enqueue_job( + &pool, + &store, + &CallbackJob { id: 5 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let client = queue_storage_client(&pool, queue, store_config, CallbackProgressWorker); + client.start().await.expect("start callback client"); + let waiting = wait_for_callback_job(&store, &pool, job_id, Duration::from_secs(10)).await; + client.shutdown(Duration::from_secs(5)).await; + assert_eq!( + waiting + .progress + .as_ref() + .and_then(|progress| progress.get("percent")), + Some(&serde_json::json!(60)), + "parked attempt's progress: {:?}", + waiting.progress + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_callback_wait_stores_the_final_progress() { + callback_wait_stores_the_final_progress(true).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_legacy_callback_wait_stores_the_final_progress() { + callback_wait_stores_the_final_progress(false).await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_runtime_complete_external() { let (_db_guard, pool) = setup_pool(10).await; @@ -24024,6 +24917,208 @@ async fn test_point_read_snapshot_blocks_lease_prune() { point_read_snapshot_survives_ring_prune(false).await; } +/// 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]); +} /// A lane move leaves immutable source rows behind. Wrap the ring so the /// destination is in an earlier partition: point reads must use live evidence, /// rather than depending on which equal-timestamp ready row the plan emits last. diff --git a/docs/configuration.md b/docs/configuration.md index 6f83f9aaa..4304c9f80 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -485,6 +485,29 @@ 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. | + +In practice the database sets the rate, not the cap. The figures in this paragraph were measured on 0.6.16 with PostgreSQL 17, where each claim with a deadline has its own claim row; 0.7 writes those claims as compact batches, and the same drain has not been measured on it. A batch of 500 such claims spent about 0.1–0.2 s in its rescue statements and wrote about 0.3 MB of WAL for its closures, so a default pass took several batches before its second of rest. With default intervals, 1.6 million abandoned claims past their deadlines drained in about 17 minutes (1,600 a second), against about 33 a second (500 per kind per interval) without passes. Rescue kinds run one after another on the leader, not in parallel. Raising `rescue_pass_budget` towards its cap and `rescue_pass_max_rows` keeps rescue running for more of each second, but each pass holds the leader for its whole budget, so promotion, rotation and prune wait as long. After a large drain the rings take a minute or two to move past the slots the backlog pinned, while promotion admits the rescued retries. + +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 f1d031b28..cf0929a7a 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -258,6 +258,8 @@ In Prometheus they appear under the OpenTelemetry [OTLP-to-Prometheus translatio | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | | `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` | +| `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 | ### Heartbeat diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index e00ddfae0..0b9fb35d7 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -312,6 +312,32 @@ 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 +- queue prune skips the slot those claims reference with reason `queue.unclosed_claim_refs`, and the queue ring stops wrapping too +- 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_budget` (at most half the rescue interval, so 15 s by default) and `rescue_pass_max_rows`, 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. The [measured rates](configuration.md#rescue-backlog) show what to expect. 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 1b2639f8d..ea8f935c8 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -153,6 +153,28 @@ 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