diff --git a/CHANGELOG.md b/CHANGELOG.md index 41880eab4..e2ad89da3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,10 +12,13 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Added +- `ClientBuilder::rescue_pass_budget` and `rescue_pass_max_rows` (Python `rescue_pass_budget_ms` and `rescue_pass_max_rows`) control how long and how far one queue-storage rescue pass drains. `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. - Durable operator halt/resume for `maintenance.prune`, in drain mode. Compatible leaders observe it on their existing session; dispatch, cron, promotion, rescue and enqueue continue. Status lists live runtimes without supported, recent reports and cannot report the fleet halted while any remain. Older binaries ignore the control. See [runtime controls](docs/runtime-controls.md). ### Fixed +- 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 (`queue.unclosed_claim_refs`) and claim prune skips their claim slots, so neither ring can wrap. Rescue closed at most 500 attempts per kind per 30-second interval: about 33 a second, so 200,000 abandoned row-local claims took about 100 minutes. A rescue pass now 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. Rescue's cursor scans also seek to their cursor through the index instead of rereading every claim already swept, so a batch costs the same at the end of a drain as at the start, and they run without JIT, whose compilation took seconds per statement. With the defaults, 1.6 million abandoned row-local claims past their deadlines drained in about 17 minutes (1,600 a second), and both rings moved past the pinned slots about two minutes later. 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, about 3 to 3.5 minutes after a runtime stops reporting with default settings. Retried attempts keep their backoffs, so their due times can cluster. Zero for either control keeps one batch per interval. See [draining a backlog of abandoned claims](docs/configuration.md#rescue-backlog). +- Heartbeats no longer queue behind a completion. Compact receipt claims are heartbeated by their claim batch and item index, an attempt whose closer holds its lock is skipped rather than waited for, and lease and progress writes skip rows a finalizer has locked and retry them on the next pass. One stalled completion could otherwise hold up a whole heartbeat pass until other attempts went stale and were rescued. A finished attempt stops flushing progress once its handler returns. - `insert_many_copy` and `insert_many_skip_unique_conflicts` no longer deadlock with per-row producers (generic `insert_many`, single-job inserts, 0.6.14 and earlier binaries) that share their unique keys or lanes. Since 0.6.15, when a per-row transaction held a lane and then inserted a key a waiting batch had claimed, PostgreSQL aborted one of the two with `deadlock detected` (`40P01`). The batch function now waits at most 500 ms (or half of its session's `deadlock_timeout`, if lower) while it holds claims or lanes: on a longer wait it releases them and tries again, so the per-row transaction proceeds and the batch then treats the keys it committed as already claimed. The same applies to a transaction that enqueues lanes in separate statements in the opposite order to a batch. This covers per-row sessions whose `deadlock_timeout` is above 500 ms, as PostgreSQL's default of 1 s is, and batches whose unique claims insert within that budget (10,000 to 20,000 keys on a fast server, fewer on a slow or busy one); larger batches behave as in 0.6.15. The protection lasts for 2 s of retries; then the batch makes one last attempt with ordinary lock waits, as in 0.6.15. A per-row transaction that holds a lane for longer than that before it inserts one of the batch's keys can therefore still deadlock with it, as in 0.6.15. A cycle through a lock the batch's own transaction took in an earlier statement also ends in `deadlock detected`, after about 3 to 4 s instead of 1 to 2 s, and usually aborts the batch's transaction. The caller's `lock_timeout` applies to each lock wait, as before. Results and the lane lock window are unchanged. - Claim prune no longer treats a partial GIN index on `lease_claim_batches.job_ids`, or a session with `quote_all_identifiers` on, as the claim batch index. It recognises only a valid, non-partial GIN index on `job_ids` alone. When another index or relation already has the index's name, `ensure_claim_batch_job_index` leaves it in place, returns false and raises a warning naming it, instead of reporting the slot indexed; the slot keeps scanning until an operator drops or renames it. - Refreshing `ensure_claim_batch_job_index` as a migration role that owns it but may only `SET ROLE` to the table owner (PostgreSQL 16+ membership `WITH INHERIT FALSE`) no longer leaves `EXECUTE` with `PUBLIC`. The refresh used to hand the function to the table owner first, after which that role could no longer revoke or grant on it, and PostgreSQL only warned. It now resets `EXECUTE` after the transfer, acting as the new owner when it cannot as itself, so exactly the roles holding `TRUNCATE` on `lease_claim_batches` can execute it, the former owner included. Migrations run with the owner as the effective role, as documented, were not affected. diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 2751baf16..1cb640642 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -61,6 +61,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"; @@ -174,6 +176,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. @@ -433,6 +443,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") @@ -751,6 +774,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( @@ -1418,7 +1463,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, @@ -1433,6 +1478,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 ac1dfaedf..5f2a46ab5 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -51,7 +51,8 @@ pub use partitioned_queue::{ }; pub use queue_storage::{ ClaimedEntry, ClaimedRuntimeJob, PromoteOutcome, PromotionLane, PruneOutcome, QueueCounts, - QueueStorage, QueueStorageConfig, RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, + QueueStorage, QueueStorageConfig, RescueBatch, RescueBatchOptions, RescueScanCursor, + RescueScanProgress, RescueSlotOrder, RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, }; pub use storage::StorageStatus; diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index f1c64fcd9..9b3646f16 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -25,11 +25,15 @@ const COPY_NULL_SENTINEL: &str = "__AWA_NULL__"; const COPY_CHUNK_TARGET_BYTES: usize = 256 * 1024; 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); @@ -154,6 +158,39 @@ pub struct ClaimedRuntimeJob { pub unique_states: Option, } +/// Identity of a compact receipt claim for heartbeat purposes: addresses +/// `lease_claim_batches` by primary key and the attempt by item index, so a +/// heartbeat does not have to scan every open claim batch. Built with +/// [`ClaimHeartbeatRef::from_claim`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub struct ClaimHeartbeatRef { + pub job_id: i64, + pub run_lease: i64, + pub claim_slot: i32, + pub claim_batch_id: i64, + pub claim_batch_index: i32, + pub receipt_id: i64, +} + +impl ClaimHeartbeatRef { + /// The heartbeat identity of a compact receipt claim, or `None` for + /// legacy or row-local claims. + pub fn from_claim(job_id: i64, run_lease: i64, claim: &ClaimedEntry) -> Option { + if !claim.lease_claim_receipt { + return None; + } + Some(Self { + job_id, + run_lease, + claim_slot: claim.claim_slot, + claim_batch_id: claim.claim_batch_id?, + claim_batch_index: claim.claim_batch_index?, + receipt_id: claim.receipt_id?, + }) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, sqlx::FromRow)] pub struct QueueClaimerLease { pub claimer_slot: i16, @@ -343,6 +380,112 @@ pub enum PruneOutcome { }, } +/// 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, +} + +/// 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, @@ -2469,7 +2612,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, @@ -10295,6 +10438,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>, @@ -10591,7 +10775,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 @@ -10640,6 +10823,19 @@ impl QueueStorage { .iter() .map(|(job_id, run_lease, _)| (*job_id, *run_lease)) .collect(); + let locked: HashSet<_> = self + .try_lock_receipt_attempt_keys_tx(tx, &receipt_jobs) + .await? + .into_iter() + .collect(); + let jobs: Vec<_> = jobs + .iter() + .filter(|(job_id, run_lease, _)| locked.contains(&(*job_id, *run_lease))) + .collect(); + if jobs.is_empty() { + return Ok(0); + } + let receipt_jobs: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect(); self.lock_receipt_attempts_tx(tx, &receipt_jobs).await?; let schema = self.schema(); @@ -10709,7 +10905,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 @@ -11735,57 +11930,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 @@ -11807,8 +12035,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, @@ -11825,8 +12077,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, @@ -11846,6 +12108,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, @@ -11855,119 +12118,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.* @@ -12008,7 +12346,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) @@ -12020,7 +12358,7 @@ impl QueueStorage { .map_err(map_sqlx_error)? .into_iter() .unzip(); - let rescued: Vec = sqlx::query_as(&format!( + let (rescued, scan_full, cursor_at, cursor_job_id, cursor_run_lease): ReceiptRescueScanRow = sqlx::query_as(&format!( r#" WITH {candidate_ctes}, stale_row_locked AS ( @@ -12216,11 +12554,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, @@ -12239,12 +12575,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) @@ -12253,10 +12595,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 @@ -12279,54 +12629,56 @@ 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?; - - 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_expired_receipt_deadlines_for_slot_tx( - tx, slot, remaining, excluded, selected, + tx, + slot, + options + .cursor_scan_limit + .map_or(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT, |limit| { + limit.max(1) as i64 + }), + remaining, + excluded, + selected, ) .await?; - remaining = remaining.saturating_sub(slot_rescued.len() as i64); - rescued.append(&mut slot_rescued); - if remaining == 0 { - return Ok(rescued); - } - } - - for slot in 0..self.claim_slot_count() { - let slot = slot as i32; - if Some(slot) == preferred_slot { - continue; + batch.saturated |= scan_full; + if !scan_full && slot_rescued.is_empty() { + batch.idle_slots.push(slot); } - - let mut slot_rescued = self - .rescue_expired_receipt_deadlines_for_slot_tx( - tx, slot, remaining, excluded, selected, - ) - .await?; + 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 { + 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)>(&format!( + let ring_cursor = sqlx::query_as::<_, (i32, i64, i32)>(&format!( r#" SELECT current_slot, generation, slot_count FROM {schema}.claim_ring_state @@ -12335,26 +12687,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(), false, None)); } let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = @@ -12388,15 +12756,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 @@ -12413,7 +12780,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 @@ -12421,9 +12787,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) @@ -12497,7 +12863,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) @@ -12509,7 +12875,7 @@ impl QueueStorage { .map_err(map_sqlx_error)? .into_iter() .unzip(); - let rescued: Vec = sqlx::query_as(&format!( + let (rescued, scan_full, cursor_at, cursor_job_id, cursor_run_lease): ReceiptRescueScanRow = sqlx::query_as(&format!( r#" WITH {candidate_ctes}, expired_locked AS ( @@ -12587,11 +12953,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, @@ -12610,11 +12974,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) @@ -12623,10 +12993,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)), + )) } async fn active_job_read_row_tx( @@ -14302,48 +14680,187 @@ impl QueueStorage { Ok(()) } - pub async fn heartbeat_batch( + /// 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). Attempts that + /// a closer holds right now are skipped rather than waited for. + pub async fn heartbeat_claims_batch( &self, pool: &PgPool, - jobs: &[(i64, i64)], + claims: &[ClaimHeartbeatRef], ) -> Result { - if jobs.is_empty() { + if claims.is_empty() { return Ok(0); } - + let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let mut updated = 0_usize; - if self.lease_claim_receipts() { - // #169 B1: in receipts mode, attempt_state is the - // authoritative heartbeat home. Skip the `UPDATE leases SET - // heartbeat_at = ...` entirely — that write was the - // dominant per-heartbeat non-HOT update on the partitioned - // `leases` table (every state-indexed partial index pays 2 - // dead entries per write under sustained churn). Compat - // reads in `LeaseJobRow` SELECTs and the `awa.jobs` view - // COALESCE attempt_state.heartbeat_at first. - updated += self - .upsert_attempt_state_from_receipts_tx(&mut tx, jobs) - .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 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(&format!( - r#" - WITH inflight AS ( - SELECT * FROM unnest($1::bigint[], $2::bigint[]) AS v(job_id, run_lease) - ) - UPDATE {table} - SET heartbeat_at = clock_timestamp() - FROM inflight - WHERE {table}.job_id = inflight.job_id - AND {table}.run_lease = inflight.run_lease - AND {table}.state = 'running' - "#, + let pairs: Vec<(i64, i64)> = claims + .iter() + .map(|claim| (claim.job_id, claim.run_lease)) + .collect(); + 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(); + let batch_ids: Vec = claims.iter().map(|claim| claim.claim_batch_id).collect(); + let batch_indices: Vec = claims.iter().map(|claim| claim.claim_batch_index).collect(); + let receipt_ids: Vec = claims.iter().map(|claim| claim.receipt_id).collect(); + let updated: i64 = sqlx::query_scalar(&format!( + r#" + WITH inflight(job_id, run_lease, claim_slot, batch_id, batch_index, receipt_id) AS ( + SELECT * FROM unnest($1::bigint[], $2::bigint[], $3::int[], $4::bigint[], $5::int[], $6::bigint[]) + ), + open_claims AS ( + SELECT inflight.job_id, inflight.run_lease + FROM inflight + JOIN {schema}.lease_claim_batches AS claim_batches + ON claim_batches.claim_slot = inflight.claim_slot + AND claim_batches.batch_id = inflight.batch_id + WHERE inflight.batch_index BETWEEN 1 AND claim_batches.claimed_count + AND claim_batches.job_ids[inflight.batch_index] = inflight.job_id + AND claim_batches.run_leases[inflight.batch_index] = inflight.run_lease + AND claim_batches.receipt_ids[inflight.batch_index] = inflight.receipt_id + AND NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closures AS closures + WHERE closures.claim_slot = inflight.claim_slot + AND closures.job_id = inflight.job_id + AND closures.run_lease = inflight.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches + WHERE closure_batches.claim_slot = inflight.claim_slot + AND closure_batches.receipt_ranges @> inflight.receipt_id + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.done_entries AS done + WHERE done.job_id = inflight.job_id AND done.run_lease = inflight.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.deferred_jobs AS deferred + WHERE deferred.job_id = inflight.job_id AND deferred.run_lease = inflight.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.dlq_entries AS dlq + WHERE dlq.job_id = inflight.job_id AND dlq.run_lease = inflight.run_lease + ) + ), + upserted AS ( + INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + SELECT open_claims.job_id, open_claims.run_lease, clock_timestamp(), clock_timestamp() + FROM open_claims + ON CONFLICT (job_id, run_lease) + DO UPDATE SET + heartbeat_at = clock_timestamp(), + updated_at = clock_timestamp() + RETURNING job_id + ) + SELECT count(*)::bigint FROM upserted + "# + )) + .bind(&job_ids) + .bind(&run_leases) + .bind(&claim_slots) + .bind(&batch_ids) + .bind(&batch_indices) + .bind(&receipt_ids) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + 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)>(&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, + jobs: &[(i64, i64)], + ) -> Result { + if jobs.is_empty() { + return Ok(0); + } + + let mut tx = pool.begin().await.map_err(map_sqlx_error)?; + let mut updated = 0_usize; + if self.lease_claim_receipts() { + // #169 B1: in receipts mode, attempt_state is the + // authoritative heartbeat home. Skip the `UPDATE leases SET + // heartbeat_at = ...` entirely — that write was the + // dominant per-heartbeat non-HOT update on the partitioned + // `leases` table (every state-indexed partial index pays 2 + // 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, &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(&format!( + r#" + WITH inflight AS ( + SELECT * FROM unnest($1::bigint[], $2::bigint[]) AS v(job_id, run_lease) + ) + UPDATE {table} + SET heartbeat_at = clock_timestamp() + FROM inflight + WHERE {table}.job_id = inflight.job_id + AND {table}.run_lease = inflight.run_lease + AND {table}.state = 'running' + "#, table = self.leases_table(), )) .bind(&job_ids) @@ -14380,9 +14897,19 @@ impl QueueStorage { self.upsert_attempt_state_progress_from_receipts_tx(&mut tx, jobs) .await? } else { - // Legacy non-receipts mode keeps the old CTE shape that - // updates leases.heartbeat_at + leases.progress - // (via attempt_state upsert) in a single round-trip. + // Prelock independently so a finalizer cannot stall the entire + // heartbeat pass. Skipped progress snapshots remain unacknowledged + // in the worker and are retried on the next pass. + let keys: Vec<_> = jobs.iter().map(|(id, lease, _)| (*id, *lease)).collect(); + let locked: HashSet<_> = self + .lock_live_leases_for_heartbeat_tx(&mut tx, &keys) + .await? + .into_iter() + .collect(); + let jobs: Vec<_> = jobs + .iter() + .filter(|(id, lease, _)| locked.contains(&(*id, *lease))) + .collect(); let schema = self.schema(); let job_ids: Vec = jobs.iter().map(|(job_id, _, _)| *job_id).collect(); let run_leases: Vec = jobs.iter().map(|(_, run_lease, _)| *run_lease).collect(); @@ -14617,59 +15144,196 @@ impl QueueStorage { done.into_job_row() } - /// 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) } /// Rescue-path fallback for unique-claim conflicts: insert the rescued @@ -15429,18 +16093,41 @@ 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}")))?; @@ -15480,73 +16167,83 @@ impl QueueStorage { .unwrap_or_default(); let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(&format!( - r#" - DELETE FROM {schema}.leases AS target - WHERE (target.job_id, target.run_lease) IN ( - 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)?; + let mut saturated = 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, + idle_slots, + scan_progress, + .. + } = receipts; selected.get_or_insert_with(|| { deleted @@ -15559,7 +16256,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 @@ -15585,34 +16287,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, + }); } } @@ -15627,17 +16332,34 @@ 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; loop { @@ -15649,53 +16371,55 @@ impl QueueStorage { .unwrap_or_default(); let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(&format!( - r#" - DELETE FROM {schema}.leases - WHERE (job_id, 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)?; + let mut saturated = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; // Receipts-mode short-path claims hold their deadline on // `lease_claims.deadline_at` rather than on a `leases` row, so @@ -15707,16 +16431,24 @@ 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, + idle_slots, + scan_progress, + .. + } = receipts; selected.get_or_insert_with(|| { deleted @@ -15729,7 +16461,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 @@ -15743,36 +16480,37 @@ impl QueueStorage { continue; } }; - moved.extend( - match self - .try_hydrate_deleted_leases_tx(&mut attempt, receipt_deleted) - .await? - { - Ok(moved) => moved, - Err(busy) => { - attempt.rollback().await.map_err(map_sqlx_error)?; - busy.exclude(&mut excluded, &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) + .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, + }); } } @@ -15787,17 +16525,35 @@ 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; loop { @@ -15809,53 +16565,55 @@ impl QueueStorage { .unwrap_or_default(); let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; let deleted: Vec = sqlx::query_as(&format!( - r#" - DELETE FROM {schema}.leases - WHERE (job_id, 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)?; + let saturated = deleted.len() as i64 >= RECEIPT_RESCUE_BATCH_LIMIT; selected.get_or_insert_with(|| { deleted @@ -15867,7 +16625,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 @@ -15882,22 +16640,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 177df916f..00f506efa 100644 --- a/awa-python/python/awa/_awa.pyi +++ b/awa-python/python/awa/_awa.pyi @@ -574,6 +574,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 dba2a6424..6f6ca6d85 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -858,6 +858,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, @@ -901,6 +903,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 @@ -932,6 +939,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 6109868e0..811ee2d02 100644 --- a/awa-python/src/client.rs +++ b/awa-python/src/client.rs @@ -1666,7 +1666,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, 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, storage_transition_role=None))] #[allow(clippy::too_many_arguments)] fn start<'py>( &self, @@ -1685,6 +1685,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, @@ -1804,6 +1806,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 ce90ba39f..4e11c3348 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -144,6 +144,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, global_max_workers: Option, leader_election_interval: Option, @@ -203,6 +205,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(), global_max_workers: None, leader_election_interval: None, @@ -763,6 +767,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 @@ -1119,6 +1143,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), dispatch_cancel: CancellationToken::new(), service_cancel: CancellationToken::new(), @@ -1216,6 +1242,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>, /// Cancellation token for dispatchers only — stops claiming new jobs. dispatch_cancel: CancellationToken, @@ -1712,6 +1740,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 59ccd581a..e95072848 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -370,6 +370,14 @@ impl JobExecutor { let in_flight_state = InFlightState { cancel: cancel.clone(), progress: progress_state.clone(), + claim: queue_storage_claim.as_ref().and_then(|claim| { + awa_model::queue_storage::ClaimHeartbeatRef::from_claim( + job_id, + job_run_lease, + claim, + ) + }), + finalizing: false, }; in_flight.insert((job_id, job_run_lease), in_flight_state); let heartbeat_guard = storage.queue_storage().map(|_| FinalizationHeartbeatGuard { @@ -453,6 +461,9 @@ impl JobExecutor { guard.clone_latest() }; + // 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 capacity and heartbeats until a bounded finalizer slot // is available. Otherwise returning handlers can create an unbounded // set of tasks retaining job payloads while completion is stalled. @@ -3644,6 +3655,8 @@ mod tests { InFlightState { cancel: Arc::new(AtomicBool::new(false)), progress: Arc::new(std::sync::Mutex::new(ProgressState::new(None))), + claim: None, + finalizing: false, }, ); } diff --git a/awa-worker/src/heartbeat.rs b/awa-worker/src/heartbeat.rs index 462b33eb7..655e02729 100644 --- a/awa-worker/src/heartbeat.rs +++ b/awa-worker/src/heartbeat.rs @@ -1,5 +1,6 @@ use crate::runtime::{InFlightMap, RunLease}; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; +use awa_model::queue_storage::ClaimHeartbeatRef; use sqlx::PgPool; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; @@ -76,8 +77,9 @@ impl HeartbeatService { fields(interval_ms = self.interval.as_millis() as u64) )] async fn heartbeat_once(&self) { - let mut all_keys: Vec<(i64, RunLease)> = self.in_flight.keys(); - all_keys.sort_unstable(); + let mut in_flight = self.in_flight.claim_refs(); + in_flight.sort_unstable_by_key(|(key, _)| *key); + let all_keys: Vec<(i64, RunLease)> = in_flight.iter().map(|(key, _)| *key).collect(); if all_keys.is_empty() { return; @@ -104,7 +106,14 @@ impl HeartbeatService { .await; } RuntimeStorage::QueueStorage(runtime) => { - self.heartbeat_once_queue_storage(runtime, &heartbeat_only, &pending_progress) + let (claimed, plain): (Vec<_>, Vec<_>) = in_flight + .iter() + .filter(|(key, _)| !progress_keys.contains(key)) + .partition(|(_, claim)| claim.is_some()); + let claimed: Vec = + claimed.into_iter().filter_map(|(_, claim)| claim).collect(); + let plain: Vec<(i64, RunLease)> = plain.into_iter().map(|(key, _)| key).collect(); + self.heartbeat_once_queue_storage(runtime, &claimed, &plain, &pending_progress) .await; } } @@ -131,7 +140,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() @@ -178,7 +187,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(), @@ -204,7 +213,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"); @@ -222,9 +237,26 @@ impl HeartbeatService { async fn heartbeat_once_queue_storage( &self, runtime: &QueueStorageRuntime, + claimed: &[ClaimHeartbeatRef], heartbeat_only: &[(i64, RunLease)], pending_progress: &[(i64, i64, u64, serde_json::Value)], ) { + for chunk in claimed.chunks(self.batch_size) { + match runtime + .store + .heartbeat_claims_batch(&self.pool, chunk) + .await + { + Ok(updated) => { + self.metrics.heartbeat_batches.add(1, &[]); + debug!(count = chunk.len(), updated, "Heartbeat claim batch sent"); + } + Err(err) => { + warn!(error = %err, "Failed to send heartbeat claim batch"); + } + } + } + for chunk in heartbeat_only.chunks(self.batch_size) { match runtime.store.heartbeat_batch(&self.pool, chunk).await { Ok(updated) => { @@ -258,7 +290,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 7de7f7bb8..ed5cad751 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}; @@ -231,6 +232,271 @@ 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, + batches_started: usize, +} + +#[derive(Debug, Default)] +struct RescueSweepCursor { + first: Option<(DateTime, i64, i64)>, + previous: Option<(DateTime, i64, i64)>, + wrapped: bool, + done: bool, +} + +impl RescueSweep { + /// A long drain must revisit slots whose claims have become eligible + /// since their last scan. Keep unfinished cursors so a large history + /// sweep still converges across these ordinary-cadence revisits. + fn revisit(&mut self, now: Instant, interval: Duration, origin: RescueTickOrigin) { + if origin == RescueTickOrigin::Ordinary + || self + .last_revisit + .is_none_or(|last| now.saturating_duration_since(last) >= interval) + { + self.idle_slots.clear(); + for cursor in self.cursors.values_mut().flatten() { + if cursor.done { + *cursor = RescueSweepCursor::default(); + } + } + self.last_revisit = Some(now); + } + } + + fn observe(&mut self, progress: RescueScanProgress) { + let cursors = self.cursors.entry(progress.slot).or_default(); + cursors.resize_with(progress.cursors.len(), RescueSweepCursor::default); + for (cursor, progress) in cursors.iter_mut().zip(progress.cursors) { + if !progress.saturated { + cursor.done = true; + } + if let Some(position) = progress.position { + if let Some(first) = cursor.first { + let crossed = cursor.previous.is_some_and(|previous| position <= previous); + cursor.done |= (cursor.wrapped && crossed) + || ((cursor.wrapped || crossed) && position >= first); + cursor.wrapped |= crossed; + } else { + cursor.first = Some(position); + } + cursor.previous = Some(position); + } + } + if !cursors.is_empty() && cursors.iter().all(|cursor| cursor.done) { + self.idle_slots.push(progress.slot); + } + } +} + +/// An early re-arm may run after owner evidence disappears. Keep ordinary +/// first batches on their original cadence; every accelerated batch needs +/// the same owner cutoff as follow-up batches. +#[derive(Debug, Default)] +struct RescueHeartbeatCadence { + last_baseline: Option, +} + +impl RescueHeartbeatCadence { + fn select( + &mut self, + now: Instant, + interval: Duration, + cutoff: Option>, + origin: RescueTickOrigin, + ) -> Option>> { + if origin == RescueTickOrigin::Ordinary + || self + .last_baseline + .is_none_or(|last| now.saturating_duration_since(last) >= interval) + { + self.last_baseline = Some(now); + Some(None) + } else { + cutoff.map(Some) + } + } +} + /// Per-queue or global retention policy for completed and failed/cancelled jobs. #[derive(Debug, Clone)] pub struct RetentionPolicy { @@ -698,6 +964,12 @@ pub struct MaintenanceService { heartbeat_rescue_interval: Duration, deadline_rescue_interval: Duration, callback_rescue_interval: Duration, + rescue_pass_budget: Duration, + rescue_pass_max_rows: usize, + rescue_owner_evidence: std::sync::Mutex, + rescue_heartbeat_cadence: std::sync::Mutex, + rescue_sweeps: std::sync::Mutex>, + rescue_others_first_offset: AtomicUsize, promote_interval: Duration, cleanup_interval: Duration, cron_sync_interval: Duration, @@ -998,6 +1270,12 @@ impl MaintenanceService { heartbeat_rescue_interval: Duration::from_secs(30), deadline_rescue_interval: DEFAULT_DEADLINE_RESCUE_INTERVAL, callback_rescue_interval: Duration::from_secs(30), + rescue_pass_budget: DEFAULT_RESCUE_PASS_BUDGET, + rescue_pass_max_rows: DEFAULT_RESCUE_PASS_MAX_ROWS, + rescue_owner_evidence: std::sync::Mutex::new(RescueOwnerEvidence::default()), + rescue_heartbeat_cadence: std::sync::Mutex::new(RescueHeartbeatCadence::default()), + rescue_sweeps: std::sync::Mutex::new(HashMap::new()), + rescue_others_first_offset: AtomicUsize::new(0), promote_interval: Duration::from_millis(250), cleanup_interval: Duration::from_secs(60), cron_sync_interval: Duration::from_secs(60), @@ -1114,6 +1392,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 @@ -1242,10 +1536,9 @@ impl MaintenanceService { return; } // 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); @@ -1289,6 +1582,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 @@ -1370,37 +1677,43 @@ impl MaintenanceService { // Dropping the observer closes the session holding leadership. return; } - _ = heartbeat_rescue_timer.tick() => { + origin = heartbeat_rescue_timer.tick() => { if leadership_lost.is_cancelled() || self.cancel.is_cancelled() { return; } if let Some(timer) = branch_tracker.try_begin(MaintenanceBranch::RescueHeartbeat.classification().0, self.heartbeat_rescue_interval, &self.metrics) { - if let Some(admission) = self.control.admit(MaintenanceBranch::RescueHeartbeat) { - admission.run(async { - self.rescue_stale_heartbeats().await; - }).await; - } + let backlog = match self.control.admit(MaintenanceBranch::RescueHeartbeat) { + Some(admission) => admission.run(self.rescue_stale_heartbeats(origin)).await, + None => false, + }; 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 leadership_lost.is_cancelled() || self.cancel.is_cancelled() { return; } if let Some(timer) = branch_tracker.try_begin(MaintenanceBranch::RescueDeadline.classification().0, self.deadline_rescue_interval, &self.metrics) { - if let Some(admission) = self.control.admit(MaintenanceBranch::RescueDeadline) { - admission.run(async { - self.rescue_expired_deadlines().await; - }).await; - } + let backlog = match self.control.admit(MaintenanceBranch::RescueDeadline) { + Some(admission) => admission.run(self.rescue_expired_deadlines(origin)).await, + None => false, + }; 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 leadership_lost.is_cancelled() || self.cancel.is_cancelled() { return; } if let Some(timer) = branch_tracker.try_begin(MaintenanceBranch::RescueCallback.classification().0, self.callback_rescue_interval, &self.metrics) { - if let Some(admission) = self.control.admit(MaintenanceBranch::RescueCallback) { - admission.run(async { - self.rescue_expired_callbacks().await; - }).await; - } + let backlog = match self.control.admit(MaintenanceBranch::RescueCallback) { + Some(admission) => admission.run(self.rescue_expired_callbacks(origin)).await, + None => false, + }; timer.finish(); + if backlog { + callback_rescue_timer.reset_after(RESCUE_BACKLOG_INTERVAL.min(self.callback_rescue_interval)); + } } } _ = promote_timer.tick() => { @@ -1779,42 +2092,303 @@ impl MaintenanceService { /// `canonical_live_backlog` above zero and wedge `finalize --wait`, /// and stuck queue-storage rows would wait out a drain-only leader. /// Both cross-plane sweeps are cheap no-ops once their tables are empty. + /// + /// Returns whether the queue-storage pass left a backlog it may keep + /// draining, so the next pass should run before the full interval. #[tracing::instrument(skip(self), name = "maintenance.rescue_stale")] - async fn rescue_stale_heartbeats(&self) { + async fn rescue_stale_heartbeats(&self, origin: RescueTickOrigin) -> bool { match &self.storage { RuntimeStorage::Canonical => { let outcome = self.rescue_canonical_stale_heartbeats().await; self.process_heartbeat_rescues(outcome, true, None).await; if let Some(runtime) = self.active_standby_queue_storage().await { - let outcome = runtime - .store - .rescue_stale_heartbeats_with_dlq( - &self.pool, - self.heartbeat_staleness, - &|queue| self.dlq_policy.enabled_for(queue), - ) - .await; - self.process_heartbeat_rescues(outcome, false, Some(runtime)) + return self + .rescue_queue_storage_pass(RescueKind::Heartbeat, runtime, false, origin) .await; } + false } RuntimeStorage::QueueStorage(runtime) => { - let outcome = runtime - .store - .rescue_stale_heartbeats_with_dlq( - &self.pool, - self.heartbeat_staleness, - &|queue| self.dlq_policy.enabled_for(queue), - ) - .await; - self.process_heartbeat_rescues(outcome, true, Some(runtime)) + let backlog = self + .rescue_queue_storage_pass(RescueKind::Heartbeat, runtime, true, origin) .await; let outcome = self.rescue_canonical_stale_heartbeats().await; self.process_heartbeat_rescues(outcome, false, None).await; + backlog } } } + /// One queue-storage rescue pass: batches of up to 500 attempts per + /// source, each in its own transaction, until a batch comes back short, + /// the pass budget is spent, or the pass reaches `rescue_pass_max_rows`. Most batches visit the oldest claim slot + /// first, since the ring cannot rotate back into it while it holds open + /// claims; every [`RESCUE_OTHERS_FIRST_EVERY`]th visits the others first. + /// A slot a batch found nothing in is left out for the rest of the pass. + /// + /// A stale heartbeat cannot tell a dead runtime from a stalled one, so a + /// stale-heartbeat pass takes more than one batch only when + /// [`RescueOwnerEvidence`] holds, and its later batches take only claims + /// last heartbeated before the most recently gone runtime stopped + /// reporting. This reduces the chance of draining a lagging live runtime, + /// but snapshot freshness and sampling cannot prove an owner is dead. + /// Deadline and callback rescues act on a deadline that has passed whoever + /// owns the attempt. + /// + /// Returns whether the pass stopped with a backlog it may keep draining, + /// so the next pass should run before the full interval. + async fn rescue_queue_storage_pass( + &self, + kind: RescueKind, + runtime: &QueueStorageRuntime, + signal_local: bool, + origin: RescueTickOrigin, + ) -> bool { + let started = Instant::now(); + let interval = match kind { + RescueKind::Heartbeat => self.heartbeat_rescue_interval, + RescueKind::Deadline => self.deadline_rescue_interval, + RescueKind::Callback => self.callback_rescue_interval, + }; + let budget = self.rescue_pass_budget.min(interval / 2); + let dlq_enabled = |queue: &str| self.dlq_policy.enabled_for(queue); + let bulk_heartbeat_cutoff = if kind == RescueKind::Heartbeat { + self.stale_claim_owners_gone().await + } else { + None + }; + let first_heartbeat_cutoff = if kind == RescueKind::Heartbeat { + let selected = self + .rescue_heartbeat_cadence + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .select(started, interval, bulk_heartbeat_cutoff, origin); + let Some(cutoff) = selected else { + self.rescue_sweeps + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&kind); + return false; + }; + cutoff + } else { + None + }; + let mut rescued_total = 0usize; + let mut oldest_attempt: Option> = None; + let mut batches = 0usize; + let mut sweep = self + .rescue_sweeps + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&kind) + .unwrap_or_default(); + sweep.revisit(started, interval, origin); + let pass_end = loop { + let slot_order = if sweep.batches_started % RESCUE_OTHERS_FIRST_EVERY + == RESCUE_OTHERS_FIRST_EVERY - 1 + { + RescueSlotOrder::OthersFirst( + self.rescue_others_first_offset + .fetch_add(1, Ordering::Relaxed), + ) + } else { + RescueSlotOrder::OldestFirst + }; + let options = RescueBatchOptions { + slot_order, + skip_slots: sweep.idle_slots.clone(), + 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), + }; + let outcome = match kind { + RescueKind::Heartbeat => { + // Follow-up batches take only claims last heartbeated + // before a runtime known to be gone stopped reporting. + let cutoff = if batches == 0 { + first_heartbeat_cutoff + } else { + bulk_heartbeat_cutoff + }; + let staleness = cutoff + .and_then(|cutoff| (Utc::now() - cutoff).to_std().ok()) + .map_or(self.heartbeat_staleness, |since_cutoff| { + since_cutoff.max(self.heartbeat_staleness) + }); + runtime + .store + .rescue_stale_heartbeats_batch( + &self.pool, + staleness, + &dlq_enabled, + &options, + ) + .await + } + RescueKind::Deadline => { + runtime + .store + .rescue_expired_deadlines_batch(&self.pool, &dlq_enabled, &options) + .await + } + RescueKind::Callback => { + runtime + .store + .rescue_expired_callbacks_batch(&self.pool, &dlq_enabled, &options) + .await + } + }; + batches += 1; + sweep.batches_started = sweep.batches_started.wrapping_add(1); + let RescueBatch { + rescued, + saturated, + idle_slots: batch_idle_slots, + scan_progress, + } = match outcome { + Ok(batch) => batch, + Err(err) => { + self.process_queue_storage_rescues(kind, Err(err), signal_local, runtime) + .await; + break RescuePassEnd::Error; + } + }; + sweep.idle_slots.extend(batch_idle_slots); + for progress in scan_progress { + sweep.observe(progress); + } + rescued_total += rescued.len(); + if let Some(attempted_at) = rescued.iter().filter_map(|job| job.attempted_at).min() { + oldest_attempt = + Some(oldest_attempt.map_or(attempted_at, |oldest| oldest.min(attempted_at))); + } + self.process_queue_storage_rescues(kind, Ok(rescued), signal_local, runtime) + .await; + + if !saturated { + break RescuePassEnd::Drained; + } + if budget.is_zero() || self.rescue_pass_max_rows == 0 { + break RescuePassEnd::TimeBudget; + } + if kind == RescueKind::Heartbeat && bulk_heartbeat_cutoff.is_none() { + break RescuePassEnd::OwnersUnconfirmed; + } + if rescued_total >= self.rescue_pass_max_rows { + break RescuePassEnd::RowLimit; + } + if started.elapsed() >= budget { + break RescuePassEnd::TimeBudget; + } + }; + + let oldest_attempt_age = oldest_attempt + .and_then(|attempted_at| (Utc::now() - attempted_at).to_std().ok()) + .unwrap_or_default(); + self.metrics.record_rescue_pass( + kind.as_str(), + pass_end.as_str(), + rescued_total as u64, + oldest_attempt_age, + ); + if batches > 1 { + info!( + rescue_kind = kind.as_str(), + rescued = rescued_total, + batches, + pass_end = pass_end.as_str(), + elapsed_ms = started.elapsed().as_millis() as u64, + oldest_attempt_age_s = oldest_attempt_age.as_secs(), + "Queue-storage rescue pass drained a backlog" + ); + } + let rearm = pass_end.backlog_remains() + && !budget.is_zero() + && self.rescue_pass_max_rows > 0 + && (kind != RescueKind::Heartbeat || bulk_heartbeat_cutoff.is_some()); + if rearm { + self.rescue_sweeps + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(kind, sweep); + } + rearm + } + + async fn process_queue_storage_rescues( + &self, + kind: RescueKind, + outcome: Result, awa_model::AwaError>, + signal_local: bool, + runtime: &QueueStorageRuntime, + ) { + match kind { + RescueKind::Heartbeat => { + self.process_heartbeat_rescues(outcome, signal_local, Some(runtime)) + .await + } + RescueKind::Deadline => { + self.process_deadline_rescues(outcome, signal_local, Some(runtime)) + .await + } + RescueKind::Callback => { + self.process_callback_rescues(outcome, signal_local, Some(runtime)) + .await + } + } + } + + /// Sample the runtime snapshots and update [`RescueOwnerEvidence`]. + /// Returns the heartbeat cutoff for accelerated stale-heartbeat batches + /// when the evidence holds and some runtime is known to be gone; `None` + /// keeps the pass to one batch. A failed read counts as no evidence. + async fn stale_claim_owners_gone(&self) -> Option> { + let staleness_ms = i64::try_from(self.heartbeat_staleness.as_millis()).unwrap_or(i64::MAX); + let (settled, last_gone_heartbeat) = match sqlx::query_as::< + _, + (i64, i64, Option>), + >(RESCUE_RUNTIME_EVIDENCE_SQL) + .bind(staleness_ms) + .fetch_one(&self.pool) + .await + { + Ok((fresh_without_heartbeat, unconfirmed, last_gone_heartbeat)) => { + if fresh_without_heartbeat > 0 || unconfirmed > 0 { + debug!( + fresh_without_heartbeat, + unconfirmed, + "Stale-heartbeat rescue kept to one batch: runtime snapshots do not show the claim owners gone" + ); + } + ( + fresh_without_heartbeat == 0 && unconfirmed == 0, + last_gone_heartbeat, + ) + } + Err(err) => { + warn!(error = %err, "Failed to read runtime snapshots for stale-heartbeat rescue"); + (false, None) + } + }; + let steady = self + .rescue_owner_evidence + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .observe(Instant::now(), settled, self.heartbeat_staleness); + if steady { + last_gone_heartbeat + } else { + None + } + } + /// Batched canonical stale-heartbeat rescue with the #388 row-at-a-time /// fallback on unique-claim conflict. Exhausted attempts fail in place, /// as callback timeouts do, so a job that keeps killing or hanging its @@ -1978,33 +2552,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 } } } @@ -2151,33 +2717,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 } } } @@ -4304,8 +4862,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 @@ -4558,6 +5120,645 @@ mod tests { ); } + #[tokio::test(start_paused = true)] + async fn ordinary_rescue_ticks_survive_jitter_and_revisit_finished_slots() { + let interval = Duration::from_secs(30); + let mut timer = RescueTimer::new(interval); + timer.tick().await; + let mut cadence = RescueHeartbeatCadence::default(); + let mut sweep = RescueSweep::default(); + tokio::time::advance(interval + Duration::from_millis(2)).await; + let first_origin = timer.tick().await; + assert_eq!(first_origin, RescueTickOrigin::Ordinary); + let first_start = tokio::time::Instant::now().into_std(); + assert_eq!( + cadence.select(first_start, interval, None, first_origin), + Some(None) + ); + sweep.revisit(first_start, interval, first_origin); + let stamp = Utc::now(); + for (slot, saturated) in [(0, false), (1, true)] { + sweep.observe(RescueScanProgress { + slot, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 100, 1)), + saturated, + }], + }); + } + assert_eq!(sweep.idle_slots, vec![0]); + tokio::time::advance(interval - Duration::from_millis(2)).await; + let next_origin = timer.tick().await; + let next_start = tokio::time::Instant::now().into_std(); + assert_eq!( + next_start.duration_since(first_start), + interval - Duration::from_millis(2) + ); + assert_eq!(next_origin, RescueTickOrigin::Ordinary); + assert_eq!( + cadence.select(next_start, interval, None, next_origin), + Some(None), + "an ordinary tick must not disappear because the prior tick started 2ms late" + ); + sweep.revisit(next_start, interval, next_origin); + assert!( + sweep.idle_slots.is_empty(), + "newly eligible finished slots must be revisited" + ); + assert!(!sweep.cursors[&0][0].done); + assert_eq!( + sweep.cursors[&1][0].first, + Some((stamp, 100, 1)), + "ordinary revisits must retain unfinished sweep progress" + ); + } + + #[test] + fn ordinary_rescue_revisit_preserves_unfinished_cursor_under_jitter() { + let interval = Duration::from_secs(30); + let now = Instant::now(); + let first_start = now + interval + Duration::from_millis(2); + let second_start = now + interval * 2; + let stamp = Utc::now(); + let mut sweep = RescueSweep::default(); + sweep.revisit(first_start, interval, RescueTickOrigin::Ordinary); + for (slot, saturated) in [(0, false), (1, true)] { + sweep.observe(RescueScanProgress { + slot, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 100, 1)), + saturated, + }], + }); + } + assert_eq!(sweep.idle_slots, vec![0]); + sweep.revisit(second_start, interval, RescueTickOrigin::Ordinary); + assert!( + sweep.idle_slots.is_empty(), + "finished slot must be revisited on the ordinary tick" + ); + assert!(!sweep.cursors[&0][0].done); + assert_eq!( + sweep.cursors[&1][0].first, + Some((stamp, 100, 1)), + "unfinished sweep progress must survive the same ordinary revisit" + ); + } + + #[tokio::test(start_paused = true)] + async fn delayed_rescue_baseline_does_not_turn_early_rearm_into_baseline() { + let interval = Duration::from_secs(30); + let mut timer = RescueTimer::new(interval); + timer.tick().await; + tokio::time::advance(Duration::from_secs(52)).await; + let first_origin = timer.tick().await; + let first_start = tokio::time::Instant::now().into_std(); + let mut cadence = RescueHeartbeatCadence::default(); + assert_eq!( + cadence.select(first_start, interval, None, first_origin), + Some(None) + ); + timer.reset_after(Duration::from_secs(8)); + let origin = timer.tick().await; + let now = tokio::time::Instant::now().into_std(); + assert_eq!(now.duration_since(first_start), Duration::from_secs(8)); + assert_eq!(origin, RescueTickOrigin::Backlog); + assert_eq!(cadence.select(now, interval, None, origin), None, + "a tick scheduled 30s but serviced at52s must not make backlog60s an unguarded baseline"); + let cutoff = Utc::now(); + assert_eq!( + cadence.select(now, interval, Some(cutoff), origin), + Some(Some(cutoff)), + "an early re-arm must retain the owner cutoff" + ); + let origin = timer.tick().await; + let now = tokio::time::Instant::now().into_std(); + assert_eq!(origin, RescueTickOrigin::Ordinary); + assert_eq!(now.duration_since(first_start), Duration::from_secs(38)); + assert_eq!(cadence.select(now, interval, None, origin), Some(None)); + } + + #[tokio::test(start_paused = true)] + async fn rescue_backlog_tick_is_consumed_even_when_branch_cooldown_skips_it() { + let (tracker, clock, cooldown_until) = tracker_flipped_by_250ms_bodies(); + let metrics = metrics_for_test(); + let mut timer = RescueTimer::new(Duration::from_secs(30)); + timer.tick().await; + timer.reset_after(TICK_100MS); + let origin = timer.tick().await; + assert_eq!(origin, RescueTickOrigin::Backlog); + let skipped_at = clock + Duration::from_nanos(1); + assert!(skipped_at < cooldown_until); + assert!(tracker + .try_begin_at("cleanup", TICK_100MS, &metrics, true, skipped_at) + .is_none()); + // tick() consumed the marker before try_begin could skip the work. + assert_eq!(timer.tick().await, RescueTickOrigin::Ordinary); + } + + #[test] + fn rescue_accelerated_first_batch_requires_owner_evidence() { + let mut cadence = RescueHeartbeatCadence::default(); + let now = Instant::now(); + let interval = Duration::from_secs(30); + let cutoff = Utc::now(); + assert_eq!( + cadence.select(now, interval, None, RescueTickOrigin::Backlog), + Some(None) + ); + assert_eq!( + cadence.select( + now + Duration::from_secs(1), + interval, + Some(cutoff), + RescueTickOrigin::Backlog + ), + Some(Some(cutoff)) + ); + // The preceding accelerated pass re-armed the timer, but owners + // became unconfirmed before it fired: no ordinary batch may slip in. + assert_eq!( + cadence.select( + now + Duration::from_secs(2), + interval, + None, + RescueTickOrigin::Backlog + ), + None + ); + assert_eq!( + cadence.select(now + interval, interval, None, RescueTickOrigin::Backlog), + Some(None) + ); + } + + #[test] + fn rescue_closed_history_sweep_stops_after_wrap() { + let mut sweep = RescueSweep::default(); + let stamp = Utc::now(); + // First window ends midway around a slot. Sweep state survives + // the pass's time budget and the next early timer fire. + for id in [20, 40, 60, 10] { + sweep.observe(RescueScanProgress { + slot: 1, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, id, 1)), + saturated: true, + }], + }); + assert!(sweep.idle_slots.is_empty()); + } + sweep.observe(RescueScanProgress { + slot: 1, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 20, 1)), + saturated: true, + }], + }); + assert_eq!(sweep.idle_slots, vec![1]); + + let mut one = RescueSweep::default(); + for _ in 0..2 { + one.observe(RescueScanProgress { + slot: 0, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 1, 1)), + saturated: true, + }], + }); + } + assert_eq!(one.idle_slots, vec![0]); + } + + #[test] + fn rescue_sweep_revisits_newly_eligible_slots_at_the_ordinary_interval() { + let mut sweep = RescueSweep::default(); + let now = Instant::now(); + let interval = Duration::from_secs(30); + let stamp = Utc::now(); + sweep.revisit(now, interval, RescueTickOrigin::Backlog); + // Slot A's deadlines are in the future; B keeps a large pass open. + sweep.observe(RescueScanProgress { + slot: 0, + cursors: vec![awa_model::RescueScanCursor { + position: None, + saturated: false, + }], + }); + sweep.observe(RescueScanProgress { + slot: 1, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, 100, 1)), + saturated: true, + }], + }); + sweep.revisit(now + interval / 2, interval, RescueTickOrigin::Backlog); + assert_eq!(sweep.idle_slots, vec![0]); + // A can now be expired (or contain claims from a reused generation). + sweep.revisit(now + interval, interval, RescueTickOrigin::Backlog); + assert!(sweep.idle_slots.is_empty()); + assert!(!sweep.cursors[&0][0].done); + assert_eq!( + sweep.cursors[&1][0].first, + Some((stamp, 100, 1)), + "revisiting an idle slot must not restart an unfinished sweep" + ); + } + + #[test] + fn rescue_sweep_detects_a_window_crossing_its_origin() { + let mut sweep = RescueSweep::default(); + let stamp = Utc::now(); + // A window can cross the origin and wrap again before its last key. + // Its end alone need never equal or exceed the original last key. + for id in [10_000, 1_999, 3_999, 5_999, 7_999, 9_999, 1_998] { + sweep.observe(RescueScanProgress { + slot: 0, + cursors: vec![awa_model::RescueScanCursor { + position: Some((stamp, id, 1)), + saturated: true, + }], + }); + } + assert_eq!(sweep.idle_slots, vec![0]); + } + + struct SparseHeartbeatFixture { + pool: PgPool, + runtime: QueueStorageRuntime, + service: MaintenanceService, + stale_job: i64, + fresh_job: i64, + fresh_runtime: Uuid, + } + + impl SparseHeartbeatFixture { + async fn new(budget: Duration, max_rows: usize, confirmed: bool) -> Self { + let pool = setup_pool(4).await; + reset_schema(&pool).await; + migrations::run(&pool).await.expect("migrations"); + let config = QueueStorageConfig::default(); + let store = QueueStorage::new(config.clone()).expect("queue storage"); + store.install(&pool).await.expect("install"); + let jobs: Vec<_> = (0..2) + .map(|id| awa_model::InsertParams { + kind: "sparse_heartbeat".into(), + args: serde_json::json!({ "id": id }), + opts: awa_model::InsertOpts { + queue: "sparse_heartbeat".into(), + max_attempts: 3, + ..Default::default() + }, + }) + .collect(); + store + .enqueue_params_batch(&pool, &jobs) + .await + .expect("enqueue"); + let claims = store + .claim_runtime_batch(&pool, "sparse_heartbeat", 2, Duration::from_secs(3600)) + .await + .expect("claim real jobs"); + assert_eq!(claims.len(), 2); + let slot = claims[0].claim.claim_slot; + assert_eq!(claims[1].claim.claim_slot, slot); + let stamp = Utc::now() - chrono::Duration::hours(2); + + // Keep the old row-local representation valid on both versions. + // Closed history needs no payload; the two open claims retain their + // real ready rows so rescue exercises the actual disposition path. + sqlx::query("DELETE FROM awa.lease_claim_batches WHERE queue = 'sparse_heartbeat'") + .execute(&pool) + .await + .expect("replace compact claims with legacy rows"); + for (index, claim) in claims.iter().enumerate() { + sqlx::query( + "INSERT INTO awa.lease_claims + (claim_slot, receipt_id, job_id, run_lease, ready_slot, ready_generation, + queue, priority, attempt, max_attempts, lane_seq, enqueue_shard, claimed_at, deadline_at) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,clock_timestamp()+interval '1 hour') + ON CONFLICT (claim_slot,job_id,run_lease) DO UPDATE + SET claimed_at=EXCLUDED.claimed_at, deadline_at=EXCLUDED.deadline_at", + ) + .bind(slot).bind(claim.claim.receipt_id).bind(claim.job.id).bind(claim.job.run_lease) + .bind(claim.claim.ready_slot).bind(claim.claim.ready_generation).bind(&claim.claim.queue) + .bind(claim.claim.priority).bind(claim.job.attempt).bind(claim.job.max_attempts) + .bind(claim.claim.lane_seq).bind(claim.claim.enqueue_shard) + .bind(stamp + chrono::Duration::microseconds(if index == 0 { 3500 } else { 4500 })) + .execute(&pool).await.expect("legacy open receipt"); + } + sqlx::query( + "INSERT INTO awa.lease_claims + (claim_slot,job_id,run_lease,ready_slot,ready_generation,queue,priority, + attempt,max_attempts,lane_seq,claimed_at,closed_at) + SELECT $1,1000000+id,1,0,0,'sparse_history',0,1,3,id, + $2::timestamptz+id*interval '1 microsecond',clock_timestamp() + FROM generate_series(1,6000) AS id", + ) + .bind(slot) + .bind(stamp) + .execute(&pool) + .await + .expect("sparse closed prefix"); + let stale_job = claims[0].job.id; + let fresh_job = claims[1].job.id; + sqlx::query( + "INSERT INTO awa.attempt_state(job_id,run_lease,heartbeat_at,updated_at) + VALUES($1,$2,clock_timestamp(),clock_timestamp())", + ) + .bind(fresh_job) + .bind(claims[1].job.run_lease) + .execute(&pool) + .await + .expect("protect fresh claim"); + let fresh_runtime = Uuid::new_v4(); + for (id, seen, heartbeat_alive) in [ + (Uuid::new_v4(), stamp + chrono::Duration::hours(1), false), + (fresh_runtime, Utc::now(), true), + ] { + sqlx::query( + "INSERT INTO awa.runtime_instances + (instance_id,hostname,pid,version,started_at,last_seen_at,snapshot_interval_ms, + healthy,postgres_connected,poll_loop_alive,heartbeat_alive,shutting_down,leader) + VALUES($1,'sparse_test',1,'0.6.13',$2,$2,1000,true,true,true,$3,false,false)", + ) + .bind(id).bind(seen).bind(heartbeat_alive) + .execute(&pool).await.expect("runtime evidence"); + } + let runtime = + QueueStorageRuntime::new(config, Duration::from_secs(1), Duration::from_secs(1)) + .expect("runtime"); + let service = MaintenanceService::new( + pool.clone(), + metrics_for_test(), + Arc::new(AtomicBool::new(true)), + Arc::new(AtomicBool::new(true)), + CancellationToken::new(), + Arc::new(Vec::new()), + InFlightMap::default(), + RuntimeStorage::QueueStorage(runtime.clone()), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + ) + .rescue_pass_budget(budget) + .rescue_pass_max_rows(max_rows); + if confirmed { + let now = Instant::now(); + *service + .rescue_owner_evidence + .lock() + .expect("owner evidence") = RescueOwnerEvidence { + steady_since: Some(now - service.heartbeat_staleness * 2), + last_sample: Some(now), + }; + } + Self { + pool, + runtime, + service, + stale_job, + fresh_job, + fresh_runtime, + } + } + + async fn pass(&self) -> bool { + sqlx::query("UPDATE awa.runtime_instances SET last_seen_at=clock_timestamp() WHERE instance_id=$1") + .bind(self.fresh_runtime).execute(&self.pool).await.expect("fresh owner snapshot"); + self.service + .rescue_queue_storage_pass( + RescueKind::Heartbeat, + &self.runtime, + false, + RescueTickOrigin::Backlog, + ) + .await + } + + async fn assert_dispositions(&self) { + let stale = self + .runtime + .store + .load_job(&self.pool, self.stale_job) + .await + .expect("stale job lookup") + .expect("stale job"); + assert_eq!(stale.state, JobState::Retryable); + let fresh = self + .runtime + .store + .load_job(&self.pool, self.fresh_job) + .await + .expect("fresh job lookup") + .expect("fresh job"); + assert_eq!( + fresh.state, + JobState::Running, + "fresh heartbeat stays protected" + ); + let closures: i64 = + sqlx::query_scalar("SELECT count(*) FROM awa.lease_claim_closures WHERE job_id=$1") + .bind(self.stale_job) + .fetch_one(&self.pool) + .await + .expect("closure count"); + assert_eq!(closures, 1, "sparse stale attempt is rescued exactly once"); + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_confirmed_heartbeat_sparse_sweep_progresses_and_stops() { + let _guard = db_test_mutex().lock().await; + // A positive budget below any database round trip makes each pass + // finish its first batch and rearm; no elapsed-time assertion is used. + let run = SparseHeartbeatFixture::new(Duration::from_nanos(1), 5000, true).await; + assert!(run.pass().await, "closed prefix leaves a resumable sweep"); + let cursor: i64 = sqlx::query_scalar( + "SELECT rescue_cursor_job_id FROM awa.claim_ring_slots WHERE slot=0", + ) + .fetch_one(&run.pool) + .await + .expect("first persisted cursor"); + assert_eq!( + cursor, 1002000, + "confirmed bulk starts with one small window" + ); + let stale = run + .runtime + .store + .load_job(&run.pool, run.stale_job) + .await + .expect("unvisited job lookup") + .expect("unvisited job"); + assert_eq!( + stale.state, + JobState::Running, + "sparse job lies beyond the first window" + ); + let mut drained = false; + for _ in 0..12 { + if !run.pass().await { + drained = true; + break; + } + } + assert!( + drained, + "persisted sparse sweep converges instead of rearming forever" + ); + assert!(run.service.rescue_sweeps.lock().expect("sweeps").is_empty()); + run.assert_dispositions().await; + run.pool.close().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_disabled_acceleration_keeps_wide_first_heartbeat_window() { + let _guard = db_test_mutex().lock().await; + for (budget, max_rows) in [(Duration::ZERO, 5000), (Duration::from_secs(1), 0)] { + let run = SparseHeartbeatFixture::new(budget, max_rows, true).await; + assert!(!run.pass().await, "disabled acceleration never rearms"); + run.assert_dispositions().await; + run.pool.close().await; + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_unconfirmed_heartbeat_keeps_wide_first_window() { + let _guard = db_test_mutex().lock().await; + let run = SparseHeartbeatFixture::new(Duration::from_nanos(1), 5000, false).await; + assert!( + !run.pass().await, + "unconfirmed owners keep the ordinary one-batch cadence" + ); + run.assert_dispositions().await; + run.pool.close().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn rescue_closed_history_stops_rearming() { + let _guard = db_test_mutex().lock().await; + let pool = setup_pool(4).await; + reset_schema(&pool).await; + migrations::run(&pool).await.expect("migrations"); + let config = QueueStorageConfig::default(); + let store = QueueStorage::new(config.clone()).expect("queue storage"); + store.install(&pool).await.expect("install"); + sqlx::query( + "INSERT INTO awa.lease_claims + (claim_slot, job_id, run_lease, ready_slot, ready_generation, + queue, priority, attempt, max_attempts, lane_seq, claimed_at, closed_at, deadline_at) + SELECT 0, id, 1, 0, 0, 'closed_history', 0, 1, 3, id, + now() - interval '1 hour', now(), now() - interval '1 minute' + FROM generate_series(1, 10001) AS id", + ) + .execute(&pool) + .await + .expect("seed more closed history than a scan window"); + let runtime = + QueueStorageRuntime::new(config, Duration::from_secs(1), Duration::from_secs(1)) + .expect("runtime"); + let service = MaintenanceService::new( + pool.clone(), + metrics_for_test(), + Arc::new(AtomicBool::new(true)), + Arc::new(AtomicBool::new(true)), + CancellationToken::new(), + Arc::new(Vec::new()), + InFlightMap::default(), + RuntimeStorage::QueueStorage(runtime.clone()), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + ) + .rescue_pass_budget(Duration::from_millis(200)); + let mut drained = false; + for _ in 0..20 { + if !service + .rescue_queue_storage_pass( + RescueKind::Deadline, + &runtime, + false, + RescueTickOrigin::Backlog, + ) + .await + { + drained = true; + break; + } + } + assert!( + drained, + "closed history must finish a bounded sweep instead of rearming forever" + ); + assert!(service.rescue_sweeps.lock().expect("sweeps").is_empty()); + // Zero is a documented way to retain the ordinary one-batch cadence. + let service = service.rescue_pass_max_rows(0); + assert!( + !service + .rescue_queue_storage_pass( + RescueKind::Deadline, + &runtime, + false, + RescueTickOrigin::Backlog + ) + .await + ); + assert!(service.rescue_sweeps.lock().expect("sweeps").is_empty()); + pool.close().await; + } + + #[test] + fn rescue_owner_evidence_needs_a_full_settled_staleness_window() { + let staleness = Duration::from_secs(90); + let sample_every = Duration::from_secs(30); + let mut evidence = RescueOwnerEvidence::default(); + let mut clock = Instant::now(); + + for _ in 0..3 { + assert!(!evidence.observe(clock, true, staleness)); + clock += sample_every; + } + assert!( + evidence.observe(clock, true, staleness), + "settled at every sample across the staleness window" + ); + + clock += sample_every; + assert!( + !evidence.observe(clock, false, staleness), + "an unsettled sample closes the gate at once" + ); + for _ in 0..3 { + clock += sample_every; + assert!( + !evidence.observe(clock, true, staleness), + "the window restarts after an unsettled sample" + ); + } + clock += sample_every; + assert!(evidence.observe(clock, true, staleness)); + } + + #[test] + fn rescue_owner_evidence_restarts_after_an_unsampled_gap() { + let staleness = Duration::from_secs(90); + let mut evidence = RescueOwnerEvidence::default(); + let mut clock = Instant::now(); + for _ in 0..4 { + evidence.observe(clock, true, staleness); + clock += Duration::from_secs(30); + } + assert!(evidence.observe(clock, true, staleness)); + + // The leader did not sample for longer than half the staleness: a + // stall in that gap would have been invisible, so start over. + clock += Duration::from_secs(46); + assert!(!evidence.observe(clock, true, staleness)); + clock += Duration::from_secs(45); + assert!(!evidence.observe(clock, true, staleness)); + clock += Duration::from_secs(45); + assert!(evidence.observe(clock, true, staleness)); + } + #[test] fn branch_tracker_without_cooldown_tracks_overrun_but_keeps_running() { let tracker = MaintenanceBranchTracker::new(); diff --git a/awa-worker/src/runtime.rs b/awa-worker/src/runtime.rs index 4de28684a..5a87d4428 100644 --- a/awa-worker/src/runtime.rs +++ b/awa-worker/src/runtime.rs @@ -5,6 +5,8 @@ use std::sync::atomic::AtomicBool; use std::sync::Arc; use std::sync::RwLock; +use awa_model::queue_storage::ClaimHeartbeatRef; + const IN_FLIGHT_SHARDS: usize = 64; pub(crate) type RunLease = i64; @@ -17,6 +19,13 @@ type InFlightShard = RwLock>; pub(crate) struct InFlightState { pub cancel: CancelFlag, pub progress: Arc>, + /// Compact receipt claim identity, when the attempt was dispatched from + /// 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. @@ -170,13 +179,26 @@ impl InFlightRegistry { guard.get(&key).map(|s| s.cancel.clone()) } - pub fn keys(&self) -> Vec { - let mut keys = Vec::new(); + /// 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 { let guard = shard.read().expect("in_flight shard poisoned"); - keys.extend(guard.keys().copied()); + refs.extend(guard.iter().map(|(key, state)| (*key, state.claim))); + } + 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; } - keys } pub fn flags(&self) -> Vec { @@ -198,6 +220,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 = lock_progress(&state.progress); if progress.has_pending() && progress.in_flight.is_none() { if let Some(ref value) = progress.latest { @@ -283,6 +308,8 @@ mod tests { InFlightState { cancel: Arc::new(AtomicBool::new(false)), progress: progress.clone(), + claim: None, + finalizing: false, }, ); diff --git a/awa-worker/src/runtime_control.rs b/awa-worker/src/runtime_control.rs index 9004dae2d..7fd7560ff 100644 --- a/awa-worker/src/runtime_control.rs +++ b/awa-worker/src/runtime_control.rs @@ -970,7 +970,10 @@ mod tests { .next() .unwrap(); // Every select arm except cancellation must register and gate its work. - let selected = leader.matches("\n _ = ").count(); + let selected = leader.matches("\n _ = ").count() + + leader + .matches("\n origin = ") + .count(); assert_eq!( selected, leader.matches("branch_tracker.try_begin").count() + 2 diff --git a/awa/tests/bulk_rescue_test.rs b/awa/tests/bulk_rescue_test.rs new file mode 100644 index 000000000..2fc103c02 --- /dev/null +++ b/awa/tests/bulk_rescue_test.rs @@ -0,0 +1,1053 @@ +//! 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::model::{migrations, QueueStorage, QueueStorageConfig}; +use awa::{Client, InsertOpts, JobArgs, JobEvent, JobResult, QueueConfig}; +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::{PgConnection, PgPoolOptions}; +use sqlx::{Connection, PgPool}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +const ABANDONED: i64 = 50_000; +const ABANDONED_CHUNK: i64 = 10_000; +/// Every tenth abandoned job has one attempt, so its rescue exhausts it. +const EXHAUSTED_EVERY: i64 = 10; +const ABANDONED_QUEUE: &str = "bulk_rescue_abandoned"; +const LIVE_QUEUE: &str = "bulk_rescue_live"; + +/// What each run of a job saw, by job id. +type RunsByJob = HashMap>; + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct AbandonedJob { + n: i64, +} + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct LiveJob { + n: i64, + hold_ms: u64, +} + +fn queue_storage_builder( + pool: &PgPool, + staleness: Duration, + heartbeat_interval: Duration, +) -> awa::ClientBuilder { + Client::builder(pool.clone()) + .queue_storage( + QueueStorageConfig::default(), + Duration::from_secs(1), + Duration::from_secs(1), + ) + // Rotate the claim ring often so the abandoned claims spread across + // slots and the ring wraps onto the oldest of them. + .claim_rotate_interval(Duration::from_millis(200)) + .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(&sql) + .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 = BulkRescueDatabase::create().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_state WHERE singleton".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") +} + +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 +} + +fn database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) +} + +fn with_database(url: &str, database: &str) -> String { + let (base, _) = url + .rsplit_once('/') + .expect("database URL has a database name"); + format!("{base}/{database}") +} + +/// A database of its own with queue storage active, dropped afterwards. +struct BulkRescueDatabase { + name: String, + pool: PgPool, +} + +impl BulkRescueDatabase { + async fn create() -> Self { + static NEXT_DATABASE: AtomicU64 = AtomicU64::new(0); + let name = format!( + "awa_bulk_rescue_{}_{}", + std::process::id(), + NEXT_DATABASE.fetch_add(1, Ordering::Relaxed) + ); + let mut admin = PgConnection::connect(&with_database(&database_url(), "postgres")) + .await + .expect("admin connection"); + sqlx::raw_sql(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&mut admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(&format!("CREATE DATABASE {name}")) + .execute(&mut admin) + .await + .expect("create database"); + let pool = PgPoolOptions::new() + .max_connections(5) + .connect(&with_database(&database_url(), &name)) + .await + .expect("test pool"); + migrations::run(&pool).await.expect("migrate"); + QueueStorage::new(QueueStorageConfig::default()) + .expect("queue storage") + .install(&pool) + .await + .expect("install queue storage"); + awa_testing::setup::activate_queue_storage(&pool).await; + Self { name, pool } + } + + fn pool(&self) -> &PgPool { + &self.pool + } + + fn url(&self) -> String { + with_database(&database_url(), &self.name) + } +} + +impl Drop for BulkRescueDatabase { + fn drop(&mut self) { + let name = self.name.clone(); + let admin_url = with_database(&database_url(), "postgres"); + let cleanup = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("cleanup runtime"); + runtime.block_on(async { + if let Ok(mut admin) = PgConnection::connect(&admin_url).await { + let _ = sqlx::raw_sql(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&mut admin) + .await; + } + }); + }); + let _ = cleanup.join(); + } +} + +/// A client on its own Tokio runtime, so it can be killed the way SIGKILL +/// would: every task is dropped mid-flight, no shutdown runs, and its +/// connections close with their transactions rolled back. +struct DoomedClient { + kill: Option>, + thread: Option>, +} + +impl DoomedClient { + fn start(url: String, staleness: Duration, heartbeat_interval: Duration) -> Self { + let (kill, killed) = tokio::sync::oneshot::channel::<()>(); + let (started_tx, started_rx) = std::sync::mpsc::channel::<()>(); + let thread = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .expect("doomed client runtime"); + let client = runtime.block_on(async { + let pool = PgPoolOptions::new() + .max_connections(20) + .connect(&url) + .await + .expect("doomed client pool"); + let client = queue_storage_builder(&pool, staleness, heartbeat_interval) + .queue( + ABANDONED_QUEUE, + QueueConfig { + max_workers: ABANDONED as u32 + 1_000, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claimers: 4, + claim_batch_size: 1_000, + ..QueueConfig::default() + }, + ) + .register::(|_, _| async { + std::future::pending::>().await + }) + .build() + .expect("doomed client"); + client.start().await.expect("start doomed client"); + client + }); + started_tx.send(()).expect("report doomed client start"); + runtime.block_on(async { + let _ = killed.await; + }); + runtime.shutdown_background(); + drop(client); + }); + started_rx.recv().expect("doomed client started"); + Self { + kill: Some(kill), + thread: Some(thread), + } + } + + fn kill(mut self) { + if let Some(kill) = self.kill.take() { + let _ = kill.send(()); + } + if let Some(thread) = self.thread.take() { + thread.join().expect("doomed client thread"); + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn rescue_drains_an_abandoned_backlog_without_rescuing_live_claims() { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(); + let exporter = InMemoryMetricExporter::default(); + let meter_provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + opentelemetry::global::set_meter_provider(meter_provider.clone()); + + let db = BulkRescueDatabase::create().await; + let pool = db.pool().clone(); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let claim_slots = store.claim_slot_count() as i64; + + // Once the doomed runtime has claimed its backlog, live jobs run for + // three staleness windows, so a rescue that ignored their heartbeats + // would take them. Until then they finish at once, so their claims do + // not hold up the claim ring while the doomed runtime fills it. + let staleness = scaled_staleness(Duration::from_secs(6)); + let heartbeat_interval = staleness / 8; + let live_job_duration = staleness * 3; + + let live_pool = PgPoolOptions::new() + .max_connections(30) + .connect(&db.url()) + .await + .expect("live pool"); + let live_runs: Arc>> = Arc::default(); + // The live runtime leads maintenance, so it sees every rescue. + let rescues: Arc>> = Arc::default(); + let live = { + let live_runs = live_runs.clone(); + let live_rescues = rescues.clone(); + let abandoned_rescues = rescues.clone(); + queue_storage_builder(&live_pool, staleness, heartbeat_interval) + // Under half the staleness, so the leader samples runtime + // snapshots often enough to confirm the dead owner. + .heartbeat_rescue_interval(Duration::from_secs(2)) + .deadline_rescue_interval(Duration::from_secs(2)) + .leader_election_interval(Duration::from_millis(100)) + .queue( + LIVE_QUEUE, + QueueConfig { + max_workers: 1_000, + poll_interval: Duration::from_millis(25), + deadline_duration: Duration::ZERO, + ..QueueConfig::default() + }, + ) + .on_event::(move |event| { + if let JobEvent::Rescued { job, .. } = event { + live_rescues + .lock() + .expect("rescues") + .entry(job.id) + .or_default() + .push((LIVE_QUEUE, job.attempt)); + } + async {} + }) + .on_event::(move |event| { + if let JobEvent::Rescued { job, .. } = event { + abandoned_rescues + .lock() + .expect("rescues") + .entry(job.id) + .or_default() + .push((ABANDONED_QUEUE, job.attempt)); + } + async {} + }) + .register::(move |job, ctx| { + live_runs + .lock() + .expect("live runs") + .entry(ctx.job.id) + .or_default() + .push(ctx.job.attempt); + async move { + tokio::time::sleep(Duration::from_millis(job.hold_ms)).await; + Ok(JobResult::Completed) + } + }) + .build() + .expect("live client") + }; + live.start().await.expect("start live client"); + wait_until( + scaled_timeout(Duration::from_secs(10)), + "live leader", + || async { live.health_check().await.leader }, + ) + .await; + + // The live runtime keeps claiming new work for the whole test. + let producing = Arc::new(AtomicBool::new(true)); + let live_hold_ms = Arc::new(AtomicU64::new(0)); + let producer = { + let pool = pool.clone(); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let producing = producing.clone(); + let live_hold_ms = live_hold_ms.clone(); + tokio::spawn(async move { + let mut n = 0; + let mut produced = 0usize; + while producing.load(Ordering::SeqCst) { + let params: Vec<_> = (0..5) + .map(|i| { + awa::model::insert::params_with( + &LiveJob { + n: n + i, + hold_ms: live_hold_ms.load(Ordering::SeqCst), + }, + InsertOpts { + queue: LIVE_QUEUE.into(), + ..Default::default() + }, + ) + .expect("live params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue live jobs"); + n += 5; + produced += params.len(); + tokio::time::sleep(Duration::from_millis(500)).await; + } + produced + }) + }; + + // The doomed runtime claims the backlog in chunks, each after the claim + // ring has rotated, so the abandoned claims span several slots. + let generation_before = claim_ring_generation(&pool).await; + let doomed = DoomedClient::start(db.url(), staleness, heartbeat_interval); + let claim_started = Instant::now(); + for chunk_start in (0..ABANDONED).step_by(ABANDONED_CHUNK as usize) { + let chunk_end = (chunk_start + ABANDONED_CHUNK).min(ABANDONED); + let params: Vec<_> = (chunk_start..chunk_end) + .map(|n| { + awa::model::insert::params_with( + &AbandonedJob { n }, + InsertOpts { + queue: ABANDONED_QUEUE.into(), + max_attempts: if n % EXHAUSTED_EVERY == 0 { 1 } else { 3 }, + ..Default::default() + }, + ) + .expect("abandoned params") + }) + .collect(); + let generation = claim_ring_generation(&pool).await; + for batch in params.chunks(5_000) { + store + .enqueue_params_batch(&pool, batch) + .await + .expect("enqueue abandoned jobs"); + } + wait_until( + scaled_timeout(Duration::from_secs(60)), + "the doomed runtime to claim the chunk", + || async { running(&pool, &store, ABANDONED_QUEUE).await >= chunk_end }, + ) + .await; + // Give the ring a chance to move on before the next chunk; once it + // wraps onto a slot of abandoned claims it stays there. + let rotate_deadline = Instant::now() + scaled_timeout(Duration::from_secs(3)); + while claim_ring_generation(&pool).await <= generation && Instant::now() < rotate_deadline { + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + let claimed_in = claim_started.elapsed(); + assert_eq!( + rescued_rows(&pool, ABANDONED_QUEUE).await, + 0, + "rescue took a claim while its runtime was alive and heartbeating" + ); + let generation_at_kill = claim_ring_generation(&pool).await; + live_hold_ms.store(live_job_duration.as_millis() as u64, Ordering::SeqCst); + doomed.kill(); + let slots = abandoned_claim_slots(&pool).await; + assert!( + slots >= 2, + "the abandoned claims should span several claim slots: {slots} \ + (ring generation {generation_before} -> {generation_at_kill})" + ); + + // Before the drain: the claims become stale one staleness after the + // kill, and the owner evidence needs one more staleness window. After + // that the backlog must drain at the multi-batch rate, not 500 claims + // per heartbeat-rescue interval. + let drained_in = wait_until( + staleness * 2 + scaled_timeout(Duration::from_secs(90)), + "rescue to drain the abandoned claims", + || async { running(&pool, &store, ABANDONED_QUEUE).await == 0 }, + ) + .await; + let generation_drained = claim_ring_generation(&pool).await; + eprintln!( + "claimed {ABANDONED} in {claimed_in:?}; drained {drained_in:?} after the kill \ + (claims stale after {staleness:?})" + ); + + wait_until( + scaled_timeout(Duration::from_secs(10)), + "a rescue event for every abandoned job", + || async { rescues.lock().expect("rescues").len() as i64 >= ABANDONED }, + ) + .await; + type Rescued = (&'static str, i16); + let rescued_twice = |rescues: &RunsByJob| -> Vec<(i64, Vec)> { + rescues + .iter() + .filter(|(_, rescued)| rescued.as_slice() != [(ABANDONED_QUEUE, 1)]) + .take(5) + .map(|(id, rescued)| (*id, rescued.clone())) + .collect() + }; + { + let rescues = rescues.lock().expect("rescues"); + assert_eq!(rescues.len() as i64, ABANDONED); + let unexpected = rescued_twice(&rescues); + assert!( + unexpected.is_empty(), + "each abandoned job is rescued once and nothing else is: {unexpected:?}" + ); + } + + let exhausted = ABANDONED / EXHAUSTED_EVERY; + assert_eq!( + count( + &pool, + format!( + "SELECT count(*) FROM awa.dlq_entries + WHERE queue = '{ABANDONED_QUEUE}' AND dlq_reason = 'heartbeat_stale' AND attempt = 1" + ), + ) + .await, + exhausted, + "every one-attempt job is dead-lettered by its rescue" + ); + let (retryable, promoted) = waiting_retry_counts(&pool, ABANDONED_QUEUE).await; + assert_eq!( + retryable + promoted, + ABANDONED - exhausted, + "every other abandoned job is waiting to retry" + ); + meter_provider.force_flush().expect("flush metrics"); + let largest_pass = largest_heartbeat_rescue_pass(&exporter); + assert!( + largest_pass > 500, + "a stale-heartbeat pass should take more than one 500-claim batch once the \ + owner is confirmed gone; the largest took {largest_pass}" + ); + + // Live claims can hold a slot for a whole live job. + let rotated_in = wait_until( + live_job_duration * 2 + scaled_timeout(Duration::from_secs(30)), + "the claim ring to rotate through every slot after the drain", + || async { claim_ring_generation(&pool).await >= generation_drained + claim_slots }, + ) + .await; + eprintln!("claim ring rotated a full revolution {rotated_in:?} after the drain"); + + // A fresh runtime runs every retried job exactly once, on its second attempt. + let recovered: Arc>> = Arc::default(); + let recovery_pool = PgPoolOptions::new() + .max_connections(20) + .connect(&db.url()) + .await + .expect("recovery pool"); + let recovery = { + let recovered = recovered.clone(); + queue_storage_builder(&recovery_pool, staleness, heartbeat_interval) + .queue( + ABANDONED_QUEUE, + // Small claims, so finished attempts never wait on a long + // completion backlog while the recovery runtime is under test. + QueueConfig { + max_workers: 100, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: 100, + ..QueueConfig::default() + }, + ) + .register::(move |_, ctx| { + let rescue_recorded = ctx + .job + .errors + .as_ref() + .and_then(|errors| errors.last()) + .and_then(|error| error.get("error")) + .and_then(|error| error.as_str()) + .is_some_and(|error| error.contains("presumed dead")); + recovered + .lock() + .expect("recovered") + .entry(ctx.job.id) + .or_default() + .push((ctx.job.attempt, rescue_recorded)); + async { Ok(JobResult::Completed) } + }) + .build() + .expect("recovery client") + }; + recovery.start().await.expect("start recovery client"); + wait_until( + scaled_timeout(Duration::from_secs(120)), + "the retried jobs to complete", + || async { + count(&pool, format!( + "SELECT count(*) FROM awa.terminal_jobs WHERE queue = '{ABANDONED_QUEUE}' AND state = 'completed'" + )).await == ABANDONED - exhausted + }, + ) + .await; + recovery.shutdown(Duration::from_secs(10)).await; + { + let recovered = recovered.lock().expect("recovered"); + assert_eq!(recovered.len() as i64, ABANDONED - exhausted); + let unexpected: Vec<_> = recovered + .iter() + .filter(|(_, runs)| runs.as_slice() != [(2, true)]) + .take(5) + .collect(); + assert!( + unexpected.is_empty(), + "each retried job runs 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 ROW_LOCAL_QUEUE: &str = "bulk_rescue_row_local"; + +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_state WHERE singleton"), + ) + .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_claims + WHERE queue = '{ROW_LOCAL_QUEUE}' AND run_lease = 1 AND closed_at IS NULL" + ), + ) + .await +} + +/// A killed runtime left hundreds of thousands of row-local claims: claims +/// with a deadline, so each has its own `lease_claims` row rather than a +/// compact batch. 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_ROW_LOCAL_RESCUE_CLAIMS` sets the backlog (default 20,000); +/// `AWA_ROW_LOCAL_RESCUE_PASS_BUDGET_MS` and `AWA_ROW_LOCAL_RESCUE_PASS_MAX_ROWS` +/// override the pass controls, and `AWA_ROW_LOCAL_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_row_local_claims_and_unpins_both_rings() { + let _ = tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(); + let abandoned = env_u64("AWA_ROW_LOCAL_RESCUE_CLAIMS").unwrap_or(20_000) as i64; + let observe_for = env_u64("AWA_ROW_LOCAL_RESCUE_OBSERVE_SECS").map(Duration::from_secs); + + let db = BulkRescueDatabase::create().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: ROW_LOCAL_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, ROW_LOCAL_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(&format!( + "UPDATE awa.lease_claims + SET claimed_at = claimed_at - interval '20 hours', + deadline_at = deadline_at - interval '20 hours' + WHERE queue = '{ROW_LOCAL_QUEUE}'; + INSERT INTO awa.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + SELECT job_id, run_lease, claimed_at + interval '1 minute', claimed_at + interval '1 minute' + FROM awa.lease_claims WHERE queue = '{ROW_LOCAL_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_claims; + ANALYZE awa.attempt_state;" + )) + .execute(&pool) + .await + .expect("age the abandoned claims"); + let shape: (i64, i64, i64, i64) = sqlx::query_as(&format!( + "SELECT count(*), count(DISTINCT claim_slot), count(DISTINCT ready_slot), + (SELECT count(*) FROM awa.lease_claim_batches) + FROM awa.lease_claims WHERE queue = '{ROW_LOCAL_QUEUE}'" + )) + .fetch_one(&pool) + .await + .expect("backlog shape"); + assert_eq!( + shape, + (abandoned, 2, 1, 0), + "row-local claims 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 row-local 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( + ROW_LOCAL_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_ROW_LOCAL_RESCUE_PASS_BUDGET_MS") { + builder = builder.rescue_pass_budget(Duration::from_millis(budget)); + } + if let Some(max_rows) = env_u64("AWA_ROW_LOCAL_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..489f29677 --- /dev/null +++ b/awa/tests/heartbeat_pass_benchmark_test.rs @@ -0,0 +1,187 @@ +//! 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::{InsertOpts, JobArgs}; +use serde::{Deserialize, Serialize}; +use sqlx::postgres::{PgConnection, PgPoolOptions}; +use sqlx::{Connection, PgPool}; +use std::time::{Duration, Instant}; + +const HEARTBEAT_CHUNK: usize = 500; +const PASSES: usize = 6; + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct HeartbeatBenchJob { + n: i64, +} + +fn database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) +} + +fn with_database(url: &str, database: &str) -> String { + let (base, _) = url + .rsplit_once('/') + .expect("database URL has a database name"); + format!("{base}/{database}") +} + +async fn fresh_database(name: &str) -> PgPool { + let mut admin = PgConnection::connect(&with_database(&database_url(), "postgres")) + .await + .expect("admin connection"); + sqlx::raw_sql(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&mut admin) + .await + .expect("drop leftover database"); + sqlx::raw_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) { + let mut admin = PgConnection::connect(&with_database(&database_url(), "postgres")) + .await + .expect("admin connection"); + let _ = sqlx::raw_sql(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&mut admin) + .await; +} + +#[tokio::test] +#[ignore = "benchmark; creates and drops its own database"] +async fn heartbeat_pass_time_by_in_flight_claims() { + let sizes: Vec = std::env::var("AWA_HEARTBEAT_BENCH_SIZES") + .ok() + .map(|sizes| { + sizes + .split(',') + .filter_map(|size| size.trim().parse().ok()) + .collect() + }) + .unwrap_or_else(|| vec![1_000, 10_000, 25_000]); + let claim_batch: i64 = std::env::var("AWA_HEARTBEAT_BENCH_CLAIM_BATCH") + .ok() + .and_then(|size| size.parse().ok()) + .unwrap_or(16); + let server: String = { + let pool = fresh_database(&format!("awa_hb_bench_{}", std::process::id())).await; + let version = sqlx::query_scalar("SHOW server_version") + .fetch_one(&pool) + .await + .expect("server version"); + pool.close().await; + version + }; + for size in sizes { + let name = format!("awa_hb_bench_{}", std::process::id()); + let pool = fresh_database(&name).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + store.install(&pool).await.expect("install queue storage"); + awa_testing::setup::activate_queue_storage(&pool).await; + let queue = "heartbeat_bench"; + for first in (0..size as i64).step_by(1_000) { + let params: Vec<_> = (first..(first + 1_000).min(size as i64)) + .map(|n| { + insert::params_with( + &HeartbeatBenchJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue"); + } + let mut claimed = Vec::with_capacity(size); + while claimed.len() < size { + let batch = store + .claim_runtime_batch(&pool, queue, claim_batch, Duration::ZERO) + .await + .expect("claim"); + assert!(!batch.is_empty(), "claimed {} of {size}", claimed.len()); + claimed.extend(batch); + } + sqlx::query("ANALYZE") + .execute(&pool) + .await + .expect("analyze"); + + // The first pass creates every attempt's heartbeat row; later passes + // refresh them. + let mut passes = Vec::with_capacity(PASSES); + let mut refreshed = 0; + for _ in 0..PASSES { + let started = Instant::now(); + refreshed = heartbeat_pass(&store, &pool, &claimed).await; + passes.push(started.elapsed()); + } + let first = passes[0]; + let mut later = passes[1..].to_vec(); + later.sort_unstable(); + eprintln!( + "heartbeat pass: server={server} in_flight={size} claim_batch={claim_batch} \ + refreshed={refreshed} first_ms={} median_ms={} max_ms={}", + first.as_millis(), + later[later.len() / 2].as_millis(), + later[later.len() - 1].as_millis() + ); + assert_eq!(refreshed, size, "every in-flight attempt is refreshed"); + pool.close().await; + drop_database(&name).await; + } +} + +async fn heartbeat_pass( + store: &QueueStorage, + pool: &PgPool, + claimed: &[awa::model::ClaimedRuntimeJob], +) -> usize { + // As the heartbeat service splits them: compact receipt claims by their + // batch identity, anything else by attempt key. + let mut claims = Vec::new(); + let mut plain = Vec::new(); + for job in claimed { + match ClaimHeartbeatRef::from_claim(job.job.id, job.job.run_lease, &job.claim) { + Some(claim) => claims.push(claim), + None => plain.push((job.job.id, job.job.run_lease)), + } + } + let mut refreshed = 0; + for chunk in claims.chunks(HEARTBEAT_CHUNK) { + refreshed += store + .heartbeat_claims_batch(pool, chunk) + .await + .expect("heartbeat claims"); + } + for chunk in plain.chunks(HEARTBEAT_CHUNK) { + refreshed += store.heartbeat_batch(pool, chunk).await.expect("heartbeat"); + } + refreshed +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index c1b712520..cc913eccd 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -2523,6 +2523,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(&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 busy-check, and `prune_oldest_claims` on an /// empty schema: rotate cycles through every slot, prune is a noop @@ -6085,6 +6250,188 @@ async fn test_queue_storage_receipt_deadline_rescue_cursor_advances_over_termina ); } +#[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(&format!("UPDATE {schema}.{table} SET claimed_at = now() - interval '1 hour', deadline_at = now() - interval '1 minute'")) + .execute(&pool).await.expect("expire attempts"); + } + let batch = store + .rescue_expired_deadlines_batch(&pool, &|_| false, &Default::default()) + .await + .expect("first rescue batch"); + assert_eq!(batch.rescued.len(), 500); + assert!(batch.saturated); + let slot = claims[0].claim.claim_slot; + assert!(!batch.idle_slots.contains(&slot)); + assert!( + batch + .scan_progress + .iter() + .filter(|scan| scan.slot == slot) + .flat_map(|scan| &scan.cursors) + .any(|cursor| cursor.saturated), + "a partially rescued candidate window must not complete that slot's sweep" + ); + let tail = store + .rescue_expired_deadlines_batch(&pool, &|_| false, &Default::default()) + .await + .expect("rescue the remaining attempt"); + assert_eq!(tail.rescued.len(), 1); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_rescue_batch_continues_past_closed_scan_windows() { + for (deadline, fresh_history) in [(false, false), (true, false), (false, true)] { + let (_db_guard, pool) = setup_pool(10).await; + let suffix = if deadline { + "deadline" + } else if fresh_history { + "fresh" + } else { + "heartbeat" + }; + let schema = format!("awa_rescue_sparse_{suffix}"); + let queue = format!("rescue_sparse_{suffix}"); + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.clone(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await; + for id in 0..3 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.clone(), + ..Default::default() + }, + ) + .await; + } + let claims = store + .claim_runtime_batch(&pool, &queue, 3, 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(&format!( + "UPDATE {schema}.lease_claims SET deadline_at = now() - interval '1 second'" + )) + .execute(&pool) + .await + .expect("expire deadlines"); + } + if fresh_history { + for claim in &claims[..2] { + sqlx::query(&format!("INSERT INTO {schema}.attempt_state (job_id, run_lease, heartbeat_at, updated_at) VALUES ($1, $2, now(), now())")) + .bind(claim.job.id).bind(claim.job.run_lease).execute(&pool).await.expect("fresh heartbeat history"); + } + } else { + store + .complete_runtime_batch(&pool, &claims[..2]) + .await + .expect("complete history"); + } + let options = awa::model::RescueBatchOptions { + cursor_scan_limit: Some(1), + ..Default::default() + }; + let batch = if deadline { + store + .rescue_expired_deadlines_batch(&pool, &|_| false, &options) + .await + } else { + store + .rescue_stale_heartbeats_batch(&pool, Duration::from_secs(90), &|_| false, &options) + .await + } + .expect("walk first closed window"); + assert!(batch.rescued.is_empty()); + assert!( + batch.saturated, + "a full closed scan window must continue toward unseen abandoned claims ({suffix})" + ); + assert!( + !batch.idle_slots.contains(&claims[0].claim.claim_slot), + "cursor progress is not an idle slot" + ); + let mut rescued = Vec::new(); + for _ in 0..3 { + let batch = if deadline { + store + .rescue_expired_deadlines_batch(&pool, &|_| false, &options) + .await + } else { + store + .rescue_stale_heartbeats_batch( + &pool, + Duration::from_secs(90), + &|_| false, + &options, + ) + .await + } + .expect("continue sparse scan"); + rescued.extend(batch.rescued); + } + assert_eq!( + rescued.iter().map(|job| job.id).collect::>(), + vec![claims[2].job.id] + ); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_receipt_rescue_cursor_sweeps_past_fresh_claims_and_wraps() { let (_db_guard, pool) = setup_pool(10).await; @@ -17370,6 +17717,209 @@ async fn test_aged_materialized_lease_is_visible_to_admin_reads() { assert_eq!(cancelled.id, job_id); } +/// 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(&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(&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(&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 2f2dd6cad..b6e8938a6 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -474,6 +474,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. On PostgreSQL 17, a batch of 500 row-local claims (claims with a deadline) spends about 0.1–0.2 s in its rescue statements and writes about 0.3 MB of WAL for its closures, so a default pass takes several batches before its second of rest. Rescue kinds run one after another on the leader, not in parallel. Measured with default intervals, 1.6 million abandoned row-local 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. 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 81dcb0f19..59f3d4d07 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -158,6 +158,8 @@ All metrics use the `awa` OTel meter name and are exported via OTLP to your conf | `awa.maintenance.promote_duration` | Histogram (s) | state | Promotion time | | `awa.maintenance.promote.lane_contention` | Counter | state, queue, outcome | Promotion lanes whose lock was busy: `handed_off` to a dedicated waiter, `waiter_limit` when every waiter was busy, `waiter_timed_out` when a waiter's 30 s wait ran out | | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | +| `awa.maintenance.rescue.pass_size` | Histogram | rescue_kind, rescue_pass_end | Attempts rescued per queue-storage rescue pass, and why the pass stopped (`drained`, `time_budget`, `row_limit`, `owners_unconfirmed`, `error`) | +| `awa.maintenance.rescue.oldest_attempt_age` | Gauge (s) | rescue_kind | Age of the oldest attempt the latest queue-storage rescue pass on this process rescued; `0` when it rescued none. A former leader retains its last value until it runs another pass; filter to the current leader | ### Heartbeat diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 7f3bb3e67..1ceaf20b7 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -320,6 +320,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/docs/upgrade-0.6.15-to-0.6.16.md b/docs/upgrade-0.6.15-to-0.6.16.md index b0a60c7b4..8902d6a5b 100644 --- a/docs/upgrade-0.6.15-to-0.6.16.md +++ b/docs/upgrade-0.6.15-to-0.6.16.md @@ -92,6 +92,19 @@ Potentially unsettled reclamation or admin metadata refresh sets `settlement_unk A `halted` status is not proof that PostgreSQL released every lock or that stale hosts are fenced. Runtime credentials can forge observations or delete runtime registrations, so confirmation also depends on trusting those credentials. The operator boundary protects the requested halt, not the reported proof. Inspect backend activity and cancel or fence lingering sessions if lock relief is required. See [runtime controls](runtime-controls.md#read-status) for the status fields and states and [runtime control privileges](security.md#runtime-control-privileges) for the trust boundary. +## Rescue drains abandoned claims in passes + +This part needs no migration and applies once the maintenance leader runs 0.6.16. Queue-storage rescue no longer stops at 500 attempts per kind per interval. Each pass takes 500-attempt batches, each in its own short transaction, for up to `rescue_pass_budget` (default 1 s, at most half the rescue interval) or `rescue_pass_max_rows` (default 5,000) rescued attempts, and while a backlog remains the next pass runs a second later. A runtime killed while holding claims therefore stops pinning the queue and claim rings for hours: with default settings, 1.6 million abandoned row-local claims drained in about 17 minutes, where 0.6.15 needed about 13 hours. + +What changes on a fleet with a backlog: + +- Claims past their deadline drain from the first deadline pass after the leader starts. Stale-heartbeat passes, the only route for claims without a deadline (`deadline_duration` zero), take more than one batch only once runtime snapshots have shown the dead runtime gone for a full staleness window, about 3–3.5 minutes with default settings. +- Rescued jobs return to `retryable` with their ordinary backoff, so they become due in bursts and the workers that run them see the backlog arrive within minutes rather than over hours. On DLQ-enabled queues, exhausted attempts move to the DLQ just as quickly. +- Closures, retry rows and their WAL arrive at the drain rate, about 0.6 KB of closure WAL per claim plus the retry rows, which replicas and CDC consumers see as a burst. +- Each pass holds the maintenance leader for up to its budget plus one batch, so promotion, rotation and prune run between passes. + +`rescue_pass_budget` and `rescue_pass_max_rows` (Python `rescue_pass_budget_ms` and `rescue_pass_max_rows`) tune the drain; zero for either keeps the 0.6.15 cadence of one batch per interval. Watch it with `awa.maintenance.rescue.pass_size` and `awa.maintenance.rescue.oldest_attempt_age`. See [draining a backlog of abandoned claims](configuration.md#rescue-backlog). + ## Roll back Before rolling back to 0.6.15, fence processes that could prune or explicitly accept losing the halt guarantee: 0.6.15 ignores the control, even when a halt remains stored. Leave the patches installed; the `prune_controls` tables and functions are safe for older binaries, and replay does not resume the fleet. 0.6.15 binaries call the new `ensure_claim_batch_job_index` unchanged, and a 0.6.15 `awa migrate` leaves it in place.