diff --git a/awa/tests/persona_soak/abandon_oracle.rs b/awa/tests/persona_soak/abandon_oracle.rs new file mode 100644 index 000000000..ba35cb24f --- /dev/null +++ b/awa/tests/persona_soak/abandon_oracle.rs @@ -0,0 +1,1222 @@ +//! Checks for the claim-and-abandon persona: a process holding thousands of claims is +//! SIGKILLed again and again. After each kill the rescue backlog has to clear within a +//! bound, the claim ring has to keep rotating, no job whose owner is alive may be +//! rescued, and attempts and the dead-letter queue have to follow `max_attempts`. + +use super::config::Persona; +use super::shared::Shared; +use serde::Serialize; +use sqlx::PgPool; +use std::collections::HashMap; +use std::io::Write as _; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; +use std::sync::Mutex; +use std::time::{SystemTime, UNIX_EPOCH}; + +const MAX_ROWS: usize = 25; +/// The ring rotates about once a second; this long without a rotation is a stall. +const RING_STALL_AFTER_MS: i64 = 6_000; +/// Owner of attempts run by the surviving in-process runtime, which is never killed. +pub const FINISHER: u32 = u32::MAX; + +pub fn epoch_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_millis() as i64) + .unwrap_or(0) +} + +#[derive(Debug, Clone, Copy)] +pub struct Attempt { + pub attempt: i16, + pub owner: u32, + pub at_ms: i64, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Kill { + pub number: u32, + pub generation: u32, + /// Milliseconds since the run started, and since the epoch, when SIGKILL was sent. + pub at_run_ms: i64, + pub signalled_epoch_ms: i64, + /// Handlers the process had started, which is what it held when it died. + pub started_by_process: i64, + /// Open claims on the persona's queues right after the kill. + pub open_after_kill: i64, + /// Claims of the process that were rescued before it was killed. + pub rescued_while_alive: i64, + pub expired_peak: i64, + pub oldest_expired_peak_s: f64, + pub cleared_after_s: Option, + pub reported_slow: bool, + pub load_ended: bool, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RingEpisode { + pub start_ms: i64, + pub end_ms: Option, + pub generation: i64, + pub slot: i32, + /// The last time the slot the ring has to enter next held a claim whose owner is + /// alive, which explains a stall. Starts at the beginning of the episode. + pub last_live_ms: i64, + pub expired_in_next_slot: i64, + pub reported: bool, +} + +#[derive(Default)] +struct RingMonitor { + last: Option<(i32, i64)>, + last_change_ms: i64, + episode: Option, + episodes: Vec, +} + +#[derive(Debug, Clone, Copy, Serialize)] +pub struct Sample { + pub at_ms: i64, + pub open: i64, + pub expired: i64, + pub oldest_expired_s: f64, + pub ring_generation: i64, + pub probe_ms: i64, +} + +#[derive(Default)] +pub struct AbandonState { + /// Open claims on the persona's queues at the last probe. + pub open_claims: AtomicI64, + /// A child process is running. + pub child_alive: AtomicBool, + /// Handlers running in the child process at its last report. + pub child_active: AtomicI64, + started_by_process: Mutex>, + /// The running child process and the epoch milliseconds at which it was started. + live: Mutex>, + /// Most claims of a process seen rescued while it was alive, per process. + live_lost: Mutex>, + attempts: Mutex>>, + max_attempts: Mutex>, + /// Epoch milliseconds at which each process generation was killed. + killed_at: Mutex>, + kills: Mutex>, + samples: Mutex>, + ring: Mutex, + csv: Mutex>, + passes_csv: Mutex>, + heartbeat_last: Mutex>, +} + +pub struct Probe { + pub open: i64, + pub expired: i64, + pub oldest_expired_s: f64, + /// Open claims taken at or after the `since` the probe was given. + pub open_since: i64, +} + +/// The process that is alive, and how many handlers it had started before the probe. +#[derive(Debug, Clone, Copy)] +pub struct LiveCheck { + pub generation: u32, + pub started: i64, +} + +impl AbandonState { + pub fn set_csv(&self, path: PathBuf) { + if let Ok(mut file) = std::fs::File::create(&path) { + let _ = writeln!( + file, + "run_s,open_claims,expired_claims,oldest_expired_s,claim_ring_generation,probe_ms" + ); + } + let passes = path.with_file_name("abandon-heartbeat-passes.csv"); + if let Ok(mut file) = std::fs::File::create(&passes) { + let _ = writeln!(file, "run_s,handlers_running,claims_sampled,pass_ms"); + } + *self.passes_csv.lock().expect("passes csv") = Some(passes); + *self.csv.lock().expect("csv") = Some(path); + } + + /// Folds in one reading of the oldest and newest heartbeat (epoch seconds) among the live + /// process's claims. While a pass runs, the oldest heartbeat is still one the previous pass + /// wrote, so it cannot be newer than that pass's newest. The first reading whose oldest + /// heartbeat is newer was taken just after a pass finished, and the spread from oldest to + /// newest is its duration (to within the last chunk). The samples are bucketed by handler + /// count so the report shows how pass time grows with what a process holds. `now` is the + /// database clock at the reading. + pub fn note_heartbeat_sample( + &self, + shared: &Shared, + rows: i64, + oldest: f64, + newest: f64, + now: f64, + active: i64, + ) { + let mut last = self.heartbeat_last.lock().expect("heartbeat last"); + let Some(settled_newest) = *last else { + *last = Some(newest); + return; + }; + if oldest <= settled_newest || active <= 0 { + return; + } + *last = Some(newest); + drop(last); + // A reading taken long after the pass finished may include the next pass's first + // chunks, so only a fresh one measures the pass. + if now - newest > 2.0 { + shared + .stats(Persona::ClaimAbandon) + .add("heartbeat_passes_unmeasured", 1); + return; + } + let ms = ((newest - oldest) * 1000.0).round() as i64; + let stats = shared.stats(Persona::ClaimAbandon); + stats.observe("heartbeat_pass_ms", ms); + let bucket = match active { + ..=4_999 => "under 5k", + 5_000..=9_999 => "5k-10k", + 10_000..=14_999 => "10k-15k", + 15_000..=19_999 => "15k-20k", + _ => "20k and over", + }; + stats.observe(&format!("heartbeat_pass_ms (handlers {bucket})"), ms); + if let Some(path) = self.passes_csv.lock().expect("passes csv").as_ref() { + if let Ok(mut file) = std::fs::OpenOptions::new().append(true).open(path) { + let _ = writeln!( + file, + "{:.1},{active},{rows},{ms}", + shared.now_ms() as f64 / 1000.0 + ); + } + } + } + + pub fn register_job(&self, job_id: i64, max_attempts: i16) { + self.max_attempts + .lock() + .expect("max attempts") + .insert(job_id, max_attempts); + } + + pub fn note_attempt(&self, job_id: i64, attempt: Attempt) { + if attempt.owner != FINISHER { + *self + .started_by_process + .lock() + .expect("started") + .entry(attempt.owner) + .or_insert(0) += 1; + } + self.attempts + .lock() + .expect("attempts") + .entry(job_id) + .or_default() + .push(attempt); + } + + pub fn note_generation_killed(&self, generation: u32, at_epoch_ms: i64) { + self.killed_at + .lock() + .expect("killed at") + .insert(generation, at_epoch_ms); + } + + /// Signal and publish the death before a waiting probe can classify its result. + /// The live-state lock also fences the probe's post-query check; no await occurs + /// between a successful signal and publishing its generation and timestamp. + pub fn signal_generation_killed( + &self, + generation: u32, + signal: impl FnOnce() -> std::io::Result<()>, + ) -> std::io::Result { + let mut live = self.live.lock().expect("live"); + signal()?; + let signalled_at = epoch_ms(); + self.note_generation_killed(generation, signalled_at); + if live.is_some_and(|(current, _)| current == generation) { + *live = None; + } + Ok(signalled_at) + } + + pub fn set_live(&self, live: Option<(u32, i64)>) { + *self.heartbeat_last.lock().expect("heartbeat last") = None; + *self.live.lock().expect("live") = live; + } + + pub fn live(&self) -> Option<(u32, i64)> { + *self.live.lock().expect("live") + } + + pub fn live_lost(&self, generation: u32) -> i64 { + self.live_lost + .lock() + .expect("live lost") + .get(&generation) + .copied() + .unwrap_or(0) + } + + pub fn record_kill(&self, kill: Kill) { + self.kills.lock().expect("kills").push(kill); + } + + pub fn kills(&self) -> Vec { + self.kills.lock().expect("kills").clone() + } + + pub fn started_by(&self, generation: u32) -> i64 { + self.started_by_process + .lock() + .expect("started") + .get(&generation) + .copied() + .unwrap_or(0) + } + + /// Whether the ring has been still for long enough to count as stalled. + pub fn ring_stalled(&self) -> bool { + self.ring.lock().expect("ring").episode.is_some() + } + + pub fn current_ring_episode(&self) -> Option { + self.ring.lock().expect("ring").episode.clone() + } + + pub fn episodes(&self) -> Vec { + let ring = self.ring.lock().expect("ring"); + let mut all = ring.episodes.clone(); + all.extend(ring.episode.clone()); + all + } + + /// Records one probe and advances the per-kill clocks. + pub fn observe( + &self, + shared: &Shared, + probe: &Probe, + live: Option, + ring_generation: i64, + probe_ms: i64, + ) { + let now = shared.now_ms(); + self.open_claims.store(probe.open, Ordering::Relaxed); + let sample = Sample { + at_ms: now, + open: probe.open, + expired: probe.expired, + oldest_expired_s: probe.oldest_expired_s, + ring_generation, + probe_ms, + }; + self.samples.lock().expect("samples").push(sample); + if let Some(path) = self.csv.lock().expect("csv").as_ref() { + if let Ok(mut file) = std::fs::OpenOptions::new().append(true).open(path) { + let _ = writeln!( + file, + "{:.1},{},{},{:.1},{},{}", + now as f64 / 1000.0, + probe.open, + probe.expired, + probe.oldest_expired_s, + ring_generation, + probe_ms + ); + } + } + let stats = shared.stats(Persona::ClaimAbandon); + stats.set_max("peak_open_claims", probe.open); + stats.set_max("peak_expired_claims", probe.expired); + stats.set_max( + "peak_oldest_expired_s", + probe.oldest_expired_s.round() as i64, + ); + stats.observe("probe_ms", probe_ms); + + // A probe can begin while the child is alive and finish after its kill. Only + // a generation still alive after the query proves missing claims were taken + // from a live owner. Keep this check serialized with signaling the kill. + if let Some(live) = live { + let lost = (live.started - probe.open_since).max(0); + let current = self.live.lock().expect("live"); + let still_live = current.is_some_and(|(generation, _)| generation == live.generation) + && !self + .killed_at + .lock() + .expect("killed at") + .contains_key(&live.generation); + if lost > 0 && !still_live { + shared.incomplete( + format!( + "claim probe for process {} crossed its kill or replacement; {lost} missing claim(s) cannot be classified as a live-owner rescue", + live.generation + ), + vec![format!( + "{{\"process\":{},\"started\":{},\"claims_still_open\":{}}}", + live.generation, live.started, probe.open_since + )], + ); + } else { + let mut peaks = self.live_lost.lock().expect("live lost"); + let peak = peaks.entry(live.generation).or_insert(0); + let first = lost > 0 && *peak == 0; + *peak = (*peak).max(lost); + drop(peaks); + stats.set_max("live_claims_rescued_peak", lost); + if first { + shared.violation( + "live_owner_rescued", + format!( + "claims of process {} were rescued while it was alive: it had started {} handler(s) and only {} of its claims were still open ({} claim(s) were eligible for rescue)", + live.generation, live.started, probe.open_since, probe.expired + ), + vec![format!( + "{{\"process\":{},\"started\":{},\"claims_still_open\":{},\"eligible_for_rescue\":{}}}", + live.generation, live.started, probe.open_since, probe.expired + )], + ); + } + } + } + let staleness_ms = shared.config.heartbeat_staleness.as_millis() as i64 + 3_000; + let epoch = epoch_ms(); + let bound = shared.config.abandon_bound; + let mut slow: Vec = Vec::new(); + { + let mut kills = self.kills.lock().expect("kills"); + for kill in kills + .iter_mut() + .filter(|kill| kill.cleared_after_s.is_none()) + { + kill.expired_peak = kill.expired_peak.max(probe.expired); + kill.oldest_expired_peak_s = kill.oldest_expired_peak_s.max(probe.oldest_expired_s); + let since_kill_ms = epoch - kill.signalled_epoch_ms; + // Claims of the dead process only count as stale after the staleness + // window, so a probe sooner than that proves nothing. + if since_kill_ms >= staleness_ms && probe.expired == 0 { + kill.cleared_after_s = Some(since_kill_ms as f64 / 1000.0); + continue; + } + let mut allowed = bound; + if shared.chaos_active() { + allowed += shared.config.chaos_slack; + } + if !kill.reported_slow && since_kill_ms as u64 > allowed.as_millis() as u64 { + kill.reported_slow = true; + slow.push(format!( + "{{\"kill\":{},\"seconds_since_kill\":{:.0},\"open_claims\":{},\"expired_claims\":{},\"oldest_expired_s\":{:.0},\"bound_s\":{}}}", + kill.number, + since_kill_ms as f64 / 1000.0, + probe.open, + probe.expired, + probe.oldest_expired_s, + bound.as_secs() + )); + } + } + } + for row in slow { + shared.violation( + "rescue_backlog_slow", + format!( + "the rescue backlog left by a killed process was not cleared within {}s of the kill ({} expired claim(s) remain, the oldest eligible for {:.0}s)", + bound.as_secs(), + probe.expired, + probe.oldest_expired_s + ), + vec![row], + ); + } + } + + /// Watches the claim ring; called about once a second. + /// + /// `next_slot` is the number of live and expired open claims in the slot the ring + /// enters next, known only while the ring is stalled. + pub fn observe_ring( + &self, + shared: &Shared, + current_slot: i32, + generation: i64, + next_slot: Option<(i64, i64)>, + ) { + let now = shared.now_ms(); + let mut report: Option = None; + { + let mut ring = self.ring.lock().expect("ring"); + let moved = ring.last != Some((current_slot, generation)); + if moved { + ring.last = Some((current_slot, generation)); + ring.last_change_ms = now; + if let Some(mut episode) = ring.episode.take() { + episode.end_ms = Some(now); + ring.episodes.push(episode); + } + return; + } + let mut threshold = crate::ci_timing::scaled_timeout(std::time::Duration::from_millis( + RING_STALL_AFTER_MS as u64, + )) + .as_millis() as i64; + if shared.chaos_active() { + threshold += shared.config.chaos_slack.as_millis() as i64; + } + if now - ring.last_change_ms >= threshold && ring.episode.is_none() { + ring.episode = Some(RingEpisode { + start_ms: ring.last_change_ms, + end_ms: None, + generation, + slot: current_slot, + last_live_ms: now, + expired_in_next_slot: 0, + reported: false, + }); + } + let bound_ms = shared.config.abandon_bound.as_millis() as i64; + if let Some(episode) = ring.episode.as_mut() { + if let Some((live, expired)) = next_slot { + if live > 0 { + episode.last_live_ms = now; + } + episode.expired_in_next_slot = expired; + } + // A stall is expected while a live process holds the next slot; it is a + // violation once nothing alive holds it and the ring still does not move. + if !episode.reported && now - episode.last_live_ms > bound_ms { + episode.reported = true; + report = Some(episode.clone()); + } + } + } + if let Some(episode) = report { + shared.violation( + "claim_ring_stalled", + format!( + "the claim ring has not rotated for more than {}s with no live claim in the slot it has to enter next ({} expired claim(s) there); stuck at slot {}, generation {}", + shared.config.abandon_bound.as_secs(), + episode.expired_in_next_slot, + episode.slot, + episode.generation + ), + vec![format!( + "{{\"stalled_since_s\":{:.1},\"slot\":{},\"generation\":{},\"expired_in_next_slot\":{}}}", + episode.start_ms as f64 / 1000.0, + episode.slot, + episode.generation, + episode.expired_in_next_slot + )], + ); + } + } + + /// Closes the ring episode left open when the run ends. + pub fn finish_ring(&self, shared: &Shared) { + let now = shared.now_ms(); + let mut ring = self.ring.lock().expect("ring"); + if let Some(mut episode) = ring.episode.take() { + episode.end_ms = Some(now); + ring.episodes.push(episode); + } + } + + /// Stall time and resume time of the ring in each kill's window. + fn ring_by_kill(&self) -> Vec<(f64, Option)> { + let kills = self.kills(); + let episodes = self.episodes(); + kills + .iter() + .enumerate() + .map(|(index, kill)| { + let from = kill.at_run_ms; + let until = kills.get(index + 1).map_or(i64::MAX, |next| next.at_run_ms); + let mut stalled = 0i64; + let mut resumed: Option = None; + for episode in &episodes { + let end = episode.end_ms.unwrap_or(i64::MAX); + if end <= from || episode.start_ms >= until { + continue; + } + stalled += end.min(until) - episode.start_ms.max(from); + if let Some(end) = episode.end_ms { + resumed = Some(resumed.map_or(end, |seen| seen.max(end))); + } + } + ( + stalled as f64 / 1000.0, + resumed.map(|end| (end - from) as f64 / 1000.0), + ) + }) + .collect() + } + + /// The per-kill table for the report. + pub fn render(&self) -> String { + let kills = self.kills(); + if kills.is_empty() { + return String::new(); + } + let ring = self.ring_by_kill(); + let samples = self.samples.lock().expect("samples"); + let mut out = String::new(); + out.push_str("### Kills\n\n"); + out.push_str("| kill | held by the process | rescued while alive | open claims after kill | peak expired | oldest expired (s) | time to clear (s) | ring stalled (s) | ring resumed after (s) |\n"); + out.push_str("| -- | -- | -- | -- | -- | -- | -- | -- | -- |\n"); + for (kill, (stalled, resumed)) in kills.iter().zip(ring) { + out.push_str(&format!( + "| {}{} | {} | {} | {} | {} | {:.0} | {} | {:.0} | {} |\n", + kill.number, + if kill.load_ended { " (load end)" } else { "" }, + kill.started_by_process, + kill.rescued_while_alive, + kill.open_after_kill, + kill.expired_peak, + kill.oldest_expired_peak_s, + kill.cleared_after_s + .map_or("not cleared".to_string(), |seconds| format!("{seconds:.0}")), + stalled, + resumed.map_or("-".to_string(), |seconds| format!("{seconds:.0}")), + )); + } + let ring_episodes = self.episodes(); + out.push_str(&format!( + "\nSamples: {}. Ring stall episodes: {}.\n", + samples.len(), + ring_episodes.len() + )); + out + } + + pub fn to_json(&self) -> serde_json::Value { + serde_json::json!({ + "kills": self.kills(), + "ring_episodes": self.episodes(), + "samples": self.samples.lock().expect("samples").len(), + }) + } + + /// Final checks over everything the child processes and the survivors reported. + pub async fn analyse(&self, shared: &Shared, pool: &PgPool, queues: &[String]) { + let killed_at = self.killed_at.lock().expect("killed at").clone(); + let max_attempts = self.max_attempts.lock().expect("max attempts").clone(); + let attempts = self.attempts.lock().expect("attempts").clone(); + let stats = shared.stats(Persona::ClaimAbandon); + + let mut live_rescued: Vec = Vec::new(); + let mut live_rescued_count = 0usize; + let mut duplicates_from_dead = 0i64; + let mut over_max: Vec = Vec::new(); + let mut gaps: Vec = Vec::new(); + let mut exhausted_by_attempts: Vec = Vec::new(); + for (job_id, list) in &attempts { + let mut list = list.clone(); + list.sort_by_key(|attempt| attempt.at_ms); + for pair in list.windows(2) { + let (previous, next) = (pair[0], pair[1]); + let previous_died_at = if previous.owner == FINISHER { + None + } else { + killed_at.get(&previous.owner).copied() + }; + match previous_died_at { + Some(died_at) if next.at_ms >= died_at => duplicates_from_dead += 1, + _ => { + live_rescued_count += 1; + if live_rescued.len() < MAX_ROWS { + live_rescued.push(format!( + "{{\"job_id\":{job_id},\"attempt\":{},\"owner\":\"{}\",\"owner_alive_until_ms\":{},\"next_attempt\":{},\"next_owner\":\"{}\",\"next_started_ms\":{}}}", + previous.attempt, + owner_name(previous.owner), + previous_died_at.map_or("null".to_string(), |at| at.to_string()), + next.attempt, + owner_name(next.owner), + next.at_ms + )); + } + } + } + if next.attempt != previous.attempt + 1 && gaps.len() < MAX_ROWS { + gaps.push(format!( + "{{\"job_id\":{job_id},\"attempt\":{},\"next_attempt\":{}}}", + previous.attempt, next.attempt + )); + } + } + let highest = list + .iter() + .map(|attempt| attempt.attempt) + .max() + .unwrap_or(0); + if let Some(max) = max_attempts.get(job_id) { + if highest > *max && over_max.len() < MAX_ROWS { + over_max.push(format!( + "{{\"job_id\":{job_id},\"attempts_seen\":{highest},\"max_attempts\":{max}}}" + )); + } + let last = list.last().copied(); + if highest == *max + && last.is_some_and(|last| last.owner != FINISHER && last.attempt == *max) + { + exhausted_by_attempts.push(*job_id); + } + } + } + stats.add("duplicate_executions_after_kill", duplicates_from_dead); + stats.add("live_owner_rescues", live_rescued_count as i64); + if live_rescued_count > 0 { + shared.violation( + "live_owner_rescued", + format!("{live_rescued_count} job(s) were run again while the process that held them was still alive"), + live_rescued, + ); + } + if !over_max.is_empty() { + shared.violation( + "attempts_exceeded_max", + format!( + "{} job(s) ran more attempts than max_attempts", + over_max.len() + ), + over_max, + ); + } + if !gaps.is_empty() { + shared.finding( + "abandon_attempt_gap", + format!( + "{} job(s) skipped an attempt number between executions", + gaps.len() + ), + gaps, + ); + } + + // Jobs whose last attempt was abandoned at max_attempts have to be dead-lettered. + let dlq: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, attempt, max_attempts FROM {}.dlq_entries WHERE queue = ANY($1)", + schema(shared) + )) + .bind(queues) + .fetch_all(pool) + .await; + match dlq { + Ok(rows) => { + stats.add("dead_lettered", rows.len() as i64); + let early: Vec = rows + .iter() + .filter(|(_, attempt, max)| attempt < max) + .take(MAX_ROWS) + .map(|(job, attempt, max)| { + format!("{{\"job_id\":{job},\"attempt\":{attempt},\"max_attempts\":{max}}}") + }) + .collect(); + if !early.is_empty() { + shared.violation( + "dead_lettered_early", + format!( + "{} job(s) were dead-lettered before reaching max_attempts", + early.len() + ), + early, + ); + } + let in_dlq: std::collections::HashSet = + rows.iter().map(|(job, _, _)| *job).collect(); + let missing: Vec = exhausted_by_attempts + .iter() + .filter(|job| !in_dlq.contains(job)) + .take(MAX_ROWS) + .map(|job| { + format!( + "{{\"job_id\":{job},\"max_attempts\":{}}}", + max_attempts.get(job).copied().unwrap_or(0) + ) + }) + .collect(); + stats.add("exhausted_after_kill", exhausted_by_attempts.len() as i64); + if !missing.is_empty() && !shared.quiesced.load(Ordering::Relaxed) { + // The rescue was still working when the run ended. + shared.finding( + "abandon_rescue_unfinished", + format!( + "{} job(s) abandoned on their last attempt were not yet dead-lettered when the drain grace ended", + missing.len() + ), + missing, + ); + } else if !missing.is_empty() { + shared.violation( + "exhausted_not_dead_lettered", + format!( + "{} job(s) abandoned on their last attempt are not in the dead-letter queue", + missing.len() + ), + missing, + ); + } + } + Err(error) => shared.incomplete(format!("dead-letter check failed: {error}"), vec![]), + } + } +} + +fn owner_name(owner: u32) -> String { + if owner == FINISHER { + "survivor".to_string() + } else { + format!("process-{owner}") + } +} + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Which claims a backlog probe looks at. +pub enum Scope<'a> { + /// Open receipt claims and running leases of these queues. + Queues(&'a [String]), + /// Open receipt identities in one claim-ring slot, including attempts + /// materialized into running or callback-waiting leases. + Slot(i32), +} + +/// Open claims in `scope`, and the ones the rescue is already entitled to take: claims +/// whose last heartbeat (or claim time, when there never was one) is older than +/// `staleness_secs`, or whose deadline has passed. +/// +/// `since_epoch_ms` splits out the open claims taken at or after that time. +pub async fn probe_backlog( + pool: &PgPool, + schema: &str, + scope: Scope<'_>, + staleness_secs: f64, + since_epoch_ms: i64, +) -> Result { + let (batch_filter, row_filter, receipt_lease_filter, lease_leg) = match &scope { + Scope::Queues(_) => ( + "b.queue = ANY($1)", + "c.queue = ANY($1)", + "l.job_id IS NULL", + format!( + "UNION ALL + SELECT COALESCE(a.heartbeat_at, l.heartbeat_at, l.attempted_at), l.deadline_at, + COALESCE(l.attempted_at, l.heartbeat_at), NULL::timestamptz + FROM {schema}.leases l + LEFT JOIN {schema}.attempt_state a ON a.job_id = l.job_id AND a.run_lease = l.run_lease + WHERE l.queue = ANY($1) AND l.state = 'running'" + ), + ), + Scope::Slot(_) => ( + "b.claim_slot = $4", + "c.claim_slot = $4", + "l.job_id IS NULL OR l.state IN ('running', 'waiting_external')", + String::new(), + ), + }; + let sql = format!( + r#" + WITH bounds AS ( + SELECT clock_timestamp() AS now_ts, + clock_timestamp() - make_interval(secs => $2) AS stale_before + ), + batch_claims AS ( + SELECT b.claim_slot, i.receipt_id, i.job_id, i.run_lease, b.queue, b.claimed_at, b.deadline_at + FROM {schema}.lease_claim_batches b + CROSS JOIN LATERAL unnest(b.job_ids, b.run_leases, b.receipt_ids) AS i(job_id, run_lease, receipt_id) + WHERE {batch_filter} + ), + row_claims AS ( + SELECT c.claim_slot, c.receipt_id, c.job_id, c.run_lease, c.queue, c.claimed_at, c.deadline_at + FROM {schema}.lease_claims c + WHERE {row_filter} AND c.closed_at IS NULL + ), + open_receipts AS ( + SELECT x.* FROM (SELECT * FROM batch_claims UNION ALL SELECT * FROM row_claims) x + WHERE NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closures cl + WHERE cl.claim_slot = x.claim_slot AND cl.job_id = x.job_id AND cl.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closure_batches cb + WHERE cb.claim_slot = x.claim_slot AND cb.receipt_ranges @> x.receipt_id + -- Keep membership correlated: joining by slot can multiply + -- receipt members by retained closure batches. + OFFSET 0) + AND NOT EXISTS (SELECT 1 FROM {schema}.done_entries d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.deferred_jobs d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.dlq_entries d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + ), + holders AS ( + SELECT CASE WHEN l.state = 'waiting_external' THEN NULL + ELSE COALESCE(a.heartbeat_at, l.heartbeat_at, l.attempted_at, x.claimed_at) + END AS beat_at, + CASE WHEN l.job_id IS NULL THEN x.deadline_at + WHEN l.state = 'running' THEN l.deadline_at + END AS deadline_at, + x.claimed_at, + CASE WHEN l.state = 'waiting_external' THEN l.callback_timeout_at END + AS callback_timeout_at + FROM open_receipts x + LEFT JOIN {schema}.leases l ON l.job_id = x.job_id AND l.run_lease = x.run_lease + LEFT JOIN {schema}.attempt_state a ON a.job_id = x.job_id AND a.run_lease = x.run_lease + WHERE {receipt_lease_filter} + {lease_leg} + ), + eligible AS ( + SELECT LEAST( + CASE WHEN beat_at < stale_before THEN beat_at + make_interval(secs => $2) END, + CASE WHEN deadline_at < now_ts THEN deadline_at END, + CASE WHEN callback_timeout_at < now_ts THEN callback_timeout_at END + ) AS eligible_at, + claimed_at >= to_timestamp($3::float8 / 1000.0) AS since_start + FROM holders, bounds + ) + SELECT count(*)::bigint, + count(eligible_at)::bigint, + COALESCE(EXTRACT(EPOCH FROM (clock_timestamp() - min(eligible_at))), 0)::float8, + count(*) FILTER (WHERE since_start)::bigint + FROM eligible + "# + ); + let query = sqlx::query_as::<_, (i64, i64, f64, i64)>(&sql); + let query = match scope { + Scope::Queues(queues) => query + .bind(queues.to_vec()) + .bind(staleness_secs) + .bind(since_epoch_ms as f64) + .bind(0_i32), + Scope::Slot(slot) => query + .bind(Vec::::new()) + .bind(staleness_secs) + .bind(since_epoch_ms as f64) + .bind(slot), + }; + let (open, expired, oldest, open_since) = query.fetch_one(pool).await?; + Ok(Probe { + open, + expired, + oldest_expired_s: oldest.max(0.0), + open_since, + }) +} + +#[cfg(test)] +mod tests { + use super::{epoch_ms, probe_backlog, LiveCheck, Probe, Scope}; + use crate::persona_soak::config::{replace_database_name, SoakConfig}; + use crate::persona_soak::db::{self, Budget, ServerInfo}; + use crate::persona_soak::shared::Shared; + use awa::model::{ClaimedRuntimeJob, QueueStorage, QueueStorageConfig}; + use sqlx::PgPool; + use std::sync::atomic::Ordering; + use std::sync::Arc; + use std::time::Duration; + + fn live_probe_fixture(generation: u32) -> (Arc, LiveCheck) { + let config = SoakConfig::for_test( + "postgres://postgres:test@localhost/awa_soak_live_probe".to_string(), + ); + let info = ServerInfo::default(); + let budget = Budget::new(&info, 0); + let shared = Shared::new(config, info, budget); + shared.abandon.set_live(Some((generation, epoch_ms()))); + let live = LiveCheck { + generation: shared.abandon.live().expect("live generation").0, + started: 1_000, + }; + (shared, live) + } + + fn observe_live_probe(shared: &Shared, live: LiveCheck, open: i64) { + shared.abandon.observe( + shared, + &Probe { + open, + expired: 0, + oldest_expired_s: 0.0, + open_since: open, + }, + Some(live), + 0, + 1, + ); + } + + fn live_rescue_reports(shared: &Shared) -> usize { + shared + .violations + .lock() + .expect("violations") + .iter() + .filter(|violation| violation.kind == "live_owner_rescued") + .count() + } + + #[test] + fn live_probe_crossing_its_generation_kill_is_incomplete() { + let (shared, captured) = live_probe_fixture(1); + // The query was dispatched with this owner alive, but its snapshot + // boundary relative to the signal is unknown when the result arrives. + shared.abandon.note_generation_killed(1, epoch_ms()); + shared.abandon.set_live(None); + observe_live_probe(&shared, captured, 500); + + assert!(shared.oracle_incomplete.load(Ordering::Relaxed)); + assert_eq!(live_rescue_reports(&shared), 0); + assert_eq!(shared.abandon.live_lost(1), 0); + // The usable backlog evidence is retained even though liveness is unknown. + assert_eq!(shared.abandon.open_claims.load(Ordering::Relaxed), 500); + } + + #[test] + fn live_probe_still_alive_missing_claims_remain_a_violation() { + let (shared, captured) = live_probe_fixture(1); + observe_live_probe(&shared, captured, 500); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(1), 500); + assert!(!shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_other_generation_kill_does_not_hide_live_loss() { + let (shared, captured) = live_probe_fixture(2); + shared.abandon.note_generation_killed(1, epoch_ms()); + observe_live_probe(&shared, captured, 500); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(2), 500); + assert!(!shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_kill_overlap_preserves_prior_proven_loss() { + let (shared, captured) = live_probe_fixture(1); + observe_live_probe(&shared, captured, 500); + assert_eq!(shared.abandon.live_lost(1), 500); + + shared.abandon.note_generation_killed(1, epoch_ms()); + shared.abandon.set_live(None); + observe_live_probe(&shared, captured, 100); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(1), 500); + assert!(shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_successful_signal_publishes_death_before_returning() { + let (shared, _) = live_probe_fixture(1); + let at = shared + .abandon + .signal_generation_killed(1, || Ok(())) + .expect("successful signal"); + + assert_eq!(shared.abandon.live(), None); + assert_eq!( + shared.abandon.killed_at.lock().expect("killed at").get(&1), + Some(&at) + ); + } + + #[test] + fn live_probe_failed_signal_does_not_publish_a_death() { + let (shared, _) = live_probe_fixture(1); + let before = shared.abandon.live(); + let result = shared + .abandon + .signal_generation_killed(1, || Err(std::io::Error::other("signal failed"))); + + assert!(result.is_err()); + assert_eq!(shared.abandon.live(), before); + assert!(shared + .abandon + .killed_at + .lock() + .expect("killed at") + .is_empty()); + } + + struct SlotProbe { + config: SoakConfig, + pool: PgPool, + store: QueueStorage, + } + + impl SlotProbe { + async fn create() -> Self { + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_slot_probe_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + db::recreate_database(&config) + .await + .expect("isolated database"); + let pool = db::connect_pool(&config.database_url, "slot_probe_regression", 4) + .await + .expect("connect"); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("store"); + Self { + config, + pool, + store, + } + } + + async fn claim(&self, deadline: Duration) -> ClaimedRuntimeJob { + self.store + .enqueue_batch(&self.pool, "slot_probe", 2, 1) + .await + .expect("enqueue"); + self.store + .claim_runtime_batch(&self.pool, "slot_probe", 1, deadline) + .await + .expect("claim") + .pop() + .expect("one claim") + } + + async fn assert_counts(&self, slot: i32, open: i64, expired: i64) { + let probe = probe_backlog(&self.pool, "awa", Scope::Slot(slot), 30.0, 0) + .await + .expect("slot probe"); + assert_eq!((probe.open, probe.expired), (open, expired)); + } + + async fn age_heartbeat(&self, claim: &ClaimedRuntimeJob) { + sqlx::query( + "UPDATE awa.leases + SET heartbeat_at = clock_timestamp() - interval '1 minute' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&self.pool) + .await + .expect("age lease heartbeat"); + sqlx::query( + "INSERT INTO awa.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + VALUES ($1, $2, clock_timestamp() - interval '1 minute', clock_timestamp()) + ON CONFLICT (job_id, run_lease) DO UPDATE + SET heartbeat_at = EXCLUDED.heartbeat_at", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&self.pool) + .await + .expect("age attempt heartbeat"); + } + + async fn drop(self) { + self.pool.close().await; + let admin = db::connect_pool(&self.config.admin_url(), "slot_probe_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(&format!( + "DROP DATABASE {} WITH (FORCE)", + self.config.database_name() + )) + .execute(&admin) + .await + .expect("drop isolated database"); + admin.close().await; + } + } + + #[tokio::test] + async fn slot_probe_keeps_materialized_running_claims_until_durable_completion() { + let run = SlotProbe::create().await; + // Zero-deadline receipts are compact; deadline receipts are row-local. + for deadline in [Duration::ZERO, Duration::from_secs(3600)] { + let claim = run.claim(deadline).await; + run.store + .register_callback( + &run.pool, + claim.job.id, + claim.job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("materialize the running receipt"); + let slot = claim.claim.claim_slot; + run.assert_counts(slot, 1, 0).await; + run.assert_counts((slot + 1) % 8, 0, 0).await; + run.age_heartbeat(&claim).await; + run.assert_counts(slot, 1, 1).await; + run.store + .heartbeat_batch(&run.pool, &[(claim.job.id, claim.job.run_lease)]) + .await + .expect("refresh authoritative attempt heartbeat"); + // The stale leases.heartbeat_at must not override attempt_state. + run.assert_counts(slot, 1, 0).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&claim)) + .await + .expect("durably complete"); + run.assert_counts(slot, 0, 0).await; + } + run.drop().await; + } + + #[tokio::test] + async fn slot_probe_uses_callback_timeout_for_waiting_materialized_claims() { + let run = SlotProbe::create().await; + for deadline in [Duration::ZERO, Duration::from_secs(3600)] { + let claim = run.claim(deadline).await; + let callback = run + .store + .register_callback( + &run.pool, + claim.job.id, + claim.job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("materialize the receipt"); + assert!(run + .store + .enter_callback_wait(&run.pool, claim.job.id, claim.job.run_lease, callback) + .await + .expect("wait for callback")); + run.age_heartbeat(&claim).await; + sqlx::query( + "UPDATE awa.leases SET deadline_at = clock_timestamp() - interval '1 second' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&run.pool) + .await + .expect("expire the original job deadline"); + let slot = claim.claim.claim_slot; + // Waiting ignores the old heartbeat and deadline until its callback times out. + run.assert_counts(slot, 1, 0).await; + run.assert_counts((slot + 1) % 8, 0, 0).await; + sqlx::query( + "UPDATE awa.leases SET callback_timeout_at = clock_timestamp() - interval '1 second' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&run.pool) + .await + .expect("expire callback"); + run.assert_counts(slot, 1, 1).await; + run.store + .complete_external(&run.pool, callback, None, Some(claim.job.run_lease), false) + .await + .expect("durably complete callback"); + run.assert_counts(slot, 0, 0).await; + } + run.drop().await; + } +} diff --git a/awa/tests/persona_soak/adminserver.rs b/awa/tests/persona_soak/adminserver.rs new file mode 100644 index 000000000..79e6b0946 --- /dev/null +++ b/awa/tests/persona_soak/adminserver.rs @@ -0,0 +1,132 @@ +//! The real admin HTTP server (`awa serve`), run as a child process against the soak +//! database so the UI personas exercise the same code the web UI talks to. + +use super::config::SoakConfig; +use std::net::TcpListener; +use std::path::PathBuf; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +pub const APPLICATION_NAME: &str = "soak_ui_server"; + +pub struct AdminServer { + child: Child, + pub base_url: String, + pub statement_timeout: Duration, + pub pool_max: u32, +} + +/// Looks for the `awa` binary next to the test executable, where +/// `cargo build --release -p awa-cli` puts it, unless `AWA_SOAK_AWA_BIN` names one. +pub fn locate_binary(config: &SoakConfig) -> Result { + if let Some(path) = &config.awa_bin { + return path + .exists() + .then(|| path.clone()) + .ok_or_else(|| format!("AWA_SOAK_AWA_BIN {} does not exist", path.display())); + } + let exe = std::env::current_exe().map_err(|error| error.to_string())?; + let candidate = exe + .parent() + .and_then(|deps| deps.parent()) + .map(|profile| profile.join("awa")); + match candidate { + Some(path) if path.exists() => Ok(path), + _ => Err( + "the admin UI personas need the awa binary: run `cargo build --release -p awa-cli` \ + or set AWA_SOAK_AWA_BIN" + .to_string(), + ), + } +} + +impl AdminServer { + pub async fn start(config: &SoakConfig) -> Result { + let binary = locate_binary(config)?; + let port = TcpListener::bind("127.0.0.1:0") + .and_then(|listener| listener.local_addr()) + .map_err(|error| format!("pick a port: {error}"))? + .port(); + let separator = if config.database_url.contains('?') { + '&' + } else { + '?' + }; + let url = format!( + "{}{separator}application_name={APPLICATION_NAME}", + config.database_url + ); + let statement_timeout = config.ui_statement_timeout; + let log = std::fs::File::create(config.report_dir.join("awa-serve.log")) + .or_else(|_| { + std::fs::create_dir_all(&config.report_dir)?; + std::fs::File::create(config.report_dir.join("awa-serve.log")) + }) + .map_err(|error| format!("create awa-serve.log: {error}"))?; + let child = Command::new(&binary) + .args([ + "--database-url", + &url, + "serve", + "--host", + "127.0.0.1", + "--port", + &port.to_string(), + ]) + .args([ + "--statement-timeout", + &statement_timeout.as_secs().to_string(), + ]) + .args([ + "--pool-max", + &config.ui_pool_max.to_string(), + "--pool-min", + "2", + ]) + .stdout(Stdio::from( + log.try_clone().map_err(|error| error.to_string())?, + )) + .stderr(Stdio::from(log)) + .spawn() + .map_err(|error| format!("spawn {}: {error}", binary.display()))?; + let mut server = Self { + child, + base_url: format!("http://127.0.0.1:{port}"), + statement_timeout, + pool_max: config.ui_pool_max, + }; + let client = reqwest::Client::new(); + let deadline = Instant::now() + Duration::from_secs(60); + loop { + if let Ok(Some(status)) = server.child.try_wait() { + return Err(format!("awa serve exited during startup: {status}")); + } + if let Ok(response) = client + .get(format!("{}/api/capabilities", server.base_url)) + .send() + .await + { + if response.status().is_success() { + return Ok(server); + } + } + if Instant::now() >= deadline { + return Err("awa serve did not start listening within 60s".to_string()); + } + tokio::time::sleep(Duration::from_millis(300)).await; + } + } + + pub fn is_running(&mut self) -> bool { + matches!(self.child.try_wait(), Ok(None)) + } + + pub fn stop(self) {} +} + +impl Drop for AdminServer { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} diff --git a/awa/tests/persona_soak/config.rs b/awa/tests/persona_soak/config.rs new file mode 100644 index 000000000..aaac98c38 --- /dev/null +++ b/awa/tests/persona_soak/config.rs @@ -0,0 +1,327 @@ +//! Run configuration, read from the environment. + +use crate::ci_timing::{scaled_staleness, scaled_timeout}; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::Duration; + +/// The harness's fast heartbeat cadence, so that rescues happen many times in a short run. +pub const DEFAULT_HEARTBEAT_SECONDS: u64 = 2; +pub const DEFAULT_STALENESS_SECONDS: u64 = 12; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum Persona { + Mia, + Tom, + Priya, + Sam, + Alex, + Chaos, + /// Admin UI users, through the real HTTP server. + Watchers, + Support, + UiOps, + Impatient, + /// A separate Python process: producers, async workers and the migrate path. + Python, + /// Handlers that batch-insert unique child jobs into their own queue. + Fanout, + /// An operator changing `enqueue_shards` on hot queues under live load, with the + /// actors running as different non-superuser roles. + ShardOps, + /// A separate process that claims as much as it can and is SIGKILLed again and again. + ClaimAbandon, +} + +impl Persona { + pub const ALL: [Persona; 14] = [ + Persona::Mia, + Persona::Tom, + Persona::Priya, + Persona::Sam, + Persona::Alex, + Persona::Chaos, + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + Persona::Python, + Persona::Fanout, + Persona::ShardOps, + Persona::ClaimAbandon, + ]; + + pub fn name(self) -> &'static str { + match self { + Persona::Mia => "mia", + Persona::Tom => "tom", + Persona::Priya => "priya", + Persona::Sam => "sam", + Persona::Alex => "alex", + Persona::Chaos => "chaos", + Persona::Watchers => "watchers", + Persona::Support => "support", + Persona::UiOps => "uiops", + Persona::Impatient => "impatient", + Persona::Python => "python", + Persona::Fanout => "fanout", + Persona::ShardOps => "shardops", + Persona::ClaimAbandon => "abandon", + } + } + + pub fn index(self) -> u8 { + self as u8 + } + + pub fn from_index(index: u8) -> Persona { + Persona::ALL[(index as usize).min(Persona::ALL.len() - 1)] + } +} + +impl FromStr for Persona { + type Err = String; + + fn from_str(value: &str) -> Result { + Persona::ALL + .into_iter() + .find(|persona| persona.name() == value.trim().to_ascii_lowercase()) + .ok_or_else(|| { + format!("unknown persona {value:?}; expected one of mia,tom,priya,sam,alex,chaos,watchers,support,uiops,impatient,python,fanout,shardops,abandon") + }) + } +} + +#[derive(Debug, Clone)] +pub struct SoakConfig { + /// URL of the disposable soak database. Its database name must start with `awa_soak`. + pub database_url: String, + pub personas: Vec, + /// Length of the load phase. + pub load: Duration, + /// How long the oracle waits for the system to quiesce after load stops. + pub drain_grace: Duration, + pub seed: u64, + /// Multiplies each persona's rates and queue counts. + pub scale: f64, + pub report_dir: PathBuf, + /// Postgres server log file readable from this host. + pub pg_log_file: Option, + /// Docker container running the server, used to read its log. + pub pg_container: Option, + /// Log file path inside `pg_container`; when absent `docker logs` is used. + pub pg_container_log: Option, + /// Seconds a scheduled job may be overdue before the promotion check fails. + pub promotion_overdue: Duration, + /// Seconds a queue may hold available jobs with idle workers before the stall check fails. + pub stall_window: Duration, + /// The `awa` binary that serves the admin UI; located next to the test binary when unset. + pub awa_bin: Option, + /// Python interpreter of a virtualenv with the `awa` extension installed. + pub python_bin: Option, + /// Server-side statement timeout of the admin server, and the matching client timeout. + pub ui_statement_timeout: Duration, + pub ui_client_timeout: Duration, + /// Connections of the admin server's pool, and seconds between the operator's admin reads. + pub ui_pool_max: u32, + pub admin_read_every: Duration, + /// The load phase stops early when the report directory's filesystem has less free space. + pub min_free_disk_gb: i64, + /// Extra tolerance added to the stall and promotion thresholds while chaos is active. + pub chaos_slack: Duration, + /// Multiplies the shard-ops persona's rates independently of `scale`, which is + /// usually lowered to keep the other personas light. + pub shardops_scale: f64, + /// Whether the application roles may create objects in the schema. Without it the + /// operator provisions every lane before the shard count is raised, as the + /// restricted-runtime procedure in `docs/security.md` requires. + pub shardops_schema_create: bool, + /// Claim-and-abandon persona: claims held by each killed process, seconds between + /// kills, how many kills, and the bound within which the rescue backlog has to clear + /// and the claim ring has to resume after a kill. Independent of `scale`. + pub abandon_claims: u64, + pub abandon_kill_every: Duration, + pub abandon_kills: u32, + pub abandon_bound: Duration, + /// Pause between a kill and the replacement process. + pub abandon_restart_delay: Duration, + /// Overrides the heartbeat and deadline rescue intervals of every runtime, to run the + /// soak with production's cadence instead of the harness's fast one. + pub rescue_interval: Option, + /// Heartbeat interval and the heartbeat staleness after which a claim may be rescued, + /// for every runtime. The library's own defaults are 30 s and 90 s. + pub heartbeat_interval: Duration, + pub heartbeat_staleness: Duration, +} + +fn env_string(name: &str) -> Option { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) +} + +fn env_parse(name: &str) -> Option { + env_string(name).and_then(|value| value.trim().parse().ok()) +} + +impl SoakConfig { + pub fn from_env() -> Result { + let database_url = env_string("AWA_SOAK_DATABASE_URL").ok_or( + "AWA_SOAK_DATABASE_URL is required (postgres://user:pass@host:port/awa_soak_)", + )?; + Self::from_database_url(database_url) + } + + #[cfg(test)] + pub fn for_test(database_url: String) -> Self { + let mut config = Self::from_database_url(database_url).expect("test database URL"); + config.personas.clear(); + config + } + + fn from_database_url(database_url: String) -> Result { + let database_name = database_name(&database_url); + if !database_name.starts_with("awa_soak") + || !database_name + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || ch == '_') + { + return Err(format!( + "refusing to run: the soak drops and recreates its database, so the name must \ + match awa_soak[A-Za-z0-9_]*, got {database_name:?}" + )); + } + + let personas = match env_string("AWA_SOAK_PERSONAS") { + Some(list) => list + .split(',') + .map(str::parse) + .collect::, _>>()?, + None => Persona::ALL.to_vec(), + }; + let seed = env_parse("AWA_SOAK_SEED").unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|elapsed| elapsed.as_nanos() as u64) + .unwrap_or(1) + }); + let report_dir = env_string("AWA_SOAK_REPORT_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| std::env::temp_dir().join(format!("awa-soak-{seed}"))); + + Ok(Self { + database_url, + personas, + load: Duration::from_secs(env_parse("AWA_SOAK_SECONDS").unwrap_or(900)), + drain_grace: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_DRAIN_GRACE_SECONDS").unwrap_or(120), + )), + seed, + scale: env_parse("AWA_SOAK_SCALE").unwrap_or(1.0), + report_dir, + pg_log_file: env_string("AWA_SOAK_PG_LOG").map(PathBuf::from), + pg_container: env_string("AWA_SOAK_PG_CONTAINER"), + pg_container_log: env_string("AWA_SOAK_PG_CONTAINER_LOG"), + awa_bin: env_string("AWA_SOAK_AWA_BIN").map(PathBuf::from), + python_bin: env_string("AWA_SOAK_PYTHON").map(PathBuf::from), + ui_statement_timeout: Duration::from_secs( + env_parse("AWA_SOAK_UI_STATEMENT_TIMEOUT_SECONDS").unwrap_or(30), + ), + ui_client_timeout: Duration::from_secs( + env_parse("AWA_SOAK_UI_CLIENT_TIMEOUT_SECONDS").unwrap_or(45), + ), + ui_pool_max: env_parse("AWA_SOAK_UI_POOL_MAX").unwrap_or(10), + admin_read_every: Duration::from_secs( + env_parse("AWA_SOAK_ADMIN_READ_SECONDS").unwrap_or(2), + ), + min_free_disk_gb: env_parse("AWA_SOAK_MIN_FREE_DISK_GB").unwrap_or(8), + promotion_overdue: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_PROMOTION_OVERDUE_SECONDS").unwrap_or(10), + )), + stall_window: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_STALL_SECONDS").unwrap_or(30), + )), + chaos_slack: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_CHAOS_SLACK_SECONDS").unwrap_or(45), + )), + shardops_scale: env_parse("AWA_SOAK_SHARDOPS_SCALE").unwrap_or(1.0), + shardops_schema_create: env_parse::("AWA_SOAK_SHARDOPS_SCHEMA_CREATE") + .map(|value| value != 0) + .unwrap_or(true), + abandon_claims: env_parse("AWA_SOAK_ABANDON_CLAIMS").unwrap_or(4_000), + abandon_kill_every: Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_KILL_SECONDS").unwrap_or(180), + ), + abandon_kills: env_parse("AWA_SOAK_ABANDON_KILLS").unwrap_or(4), + abandon_bound: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_BOUND_SECONDS").unwrap_or(120), + )), + abandon_restart_delay: Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_RESTART_SECONDS").unwrap_or(3), + ), + rescue_interval: env_parse::("AWA_SOAK_RESCUE_INTERVAL_SECONDS") + .map(Duration::from_secs), + heartbeat_interval: Duration::from_secs( + env_parse("AWA_SOAK_HEARTBEAT_SECONDS").unwrap_or(DEFAULT_HEARTBEAT_SECONDS), + ), + heartbeat_staleness: scaled_staleness(Duration::from_secs( + env_parse("AWA_SOAK_STALENESS_SECONDS").unwrap_or(DEFAULT_STALENESS_SECONDS), + )), + }) + } + + pub fn has(&self, persona: Persona) -> bool { + self.personas.contains(&persona) + } + + /// Scales a count, never below one. + pub fn scaled(&self, base: u64) -> u64 { + ((base as f64 * self.scale).round() as u64).max(1) + } + + pub fn database_name(&self) -> String { + database_name(&self.database_url) + } + + /// The same server, connected to the maintenance database. + pub fn admin_url(&self) -> String { + replace_database_name(&self.database_url, "postgres") + } +} + +pub fn database_name(url: &str) -> String { + let without_query = url.split_once('?').map_or(url, |(prefix, _)| prefix); + without_query + .rsplit_once('/') + .map(|(_, name)| name.to_string()) + .unwrap_or_default() +} + +pub fn replace_database_name(url: &str, database_name: &str) -> String { + let (without_query, query) = match url.split_once('?') { + Some((prefix, query)) => (prefix, Some(query)), + None => (url, None), + }; + let base = without_query + .rsplit_once('/') + .map_or(without_query, |(base, _)| base); + match query { + Some(query) => format!("{base}/{database_name}?{query}"), + None => format!("{base}/{database_name}"), + } +} + +/// `url` with its user and password replaced. +pub fn replace_credentials(url: &str, user: &str, password: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return url.to_string(); + }; + let authority_end = rest.find('/').unwrap_or(rest.len()); + let host = rest[..authority_end] + .rsplit_once('@') + .map_or(&rest[..authority_end], |(_, host)| host); + format!( + "{scheme}://{user}:{password}@{host}{}", + &rest[authority_end..] + ) +} diff --git a/awa/tests/persona_soak/db.rs b/awa/tests/persona_soak/db.rs new file mode 100644 index 000000000..d195a6c1a --- /dev/null +++ b/awa/tests/persona_soak/db.rs @@ -0,0 +1,293 @@ +//! Database lifecycle, connection budgeting and the per-client runtime host. + +use super::config::{replace_database_name, Persona, SoakConfig}; +use super::roles; +use awa::Client; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::collections::BTreeMap; +use std::str::FromStr; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +/// Opens a pool whose connections carry `application_name`, so the oracle and the +/// chaos persona can tell the runtimes apart in `pg_stat_activity`. +pub async fn connect_pool( + url: &str, + application_name: &str, + max_connections: u32, +) -> Result { + connect_pool_with_timeout(url, application_name, max_connections, None).await +} + +/// As [`connect_pool`], with a server-side `statement_timeout` on every connection. +pub async fn connect_pool_with_timeout( + url: &str, + application_name: &str, + max_connections: u32, + statement_timeout: Option, +) -> Result { + let mut options = PgConnectOptions::from_str(url)?.application_name(application_name); + if let Some(timeout) = statement_timeout { + options = options.options([("statement_timeout", timeout.as_millis().to_string())]); + } + PgPoolOptions::new() + .max_connections(max_connections) + .acquire_timeout(Duration::from_secs(60)) + .connect_with(options) + .await +} + +/// Drops and recreates the soak database, applies migrations and activates queue +/// storage, the production engine. +pub async fn recreate_database(config: &SoakConfig) -> Result<(), String> { + let name = config.database_name(); + assert!(name.starts_with("awa_soak"), "database name guard"); + let admin = connect_pool(&config.admin_url(), "awa_soak_admin", 2) + .await + .map_err(|error| format!("connect to maintenance database: {error}"))?; + sqlx::query(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&admin) + .await + .map_err(|error| format!("drop database: {error}"))?; + let restricted_roles = config.has(Persona::ShardOps); + if restricted_roles { + roles::create_roles(&admin).await?; + sqlx::query(&format!("CREATE DATABASE {name} OWNER {}", roles::OWNER)) + .execute(&admin) + .await + .map_err(|error| format!("create database: {error}"))?; + } else { + sqlx::query(&format!("CREATE DATABASE {name}")) + .execute(&admin) + .await + .map_err(|error| format!("create database: {error}"))?; + } + admin.close().await; + + // With the shard-ops persona the schema belongs to a non-superuser owner role, as in + // a production install, and the application roles get the documented grants. + let setup_url = if restricted_roles { + roles::url_for_role(config, roles::OWNER) + } else { + config.database_url.clone() + }; + let pool = connect_pool(&setup_url, "awa_soak_setup", 2) + .await + .map_err(|error| format!("connect to soak database: {error}"))?; + awa::model::migrations::run(&pool) + .await + .map_err(|error| format!("migrate: {error}"))?; + awa_testing::setup::activate_queue_storage(&pool).await; + pool.close().await; + if restricted_roles { + let superuser = connect_pool(&config.database_url, "awa_soak_setup_grants", 1) + .await + .map_err(|error| format!("connect to soak database: {error}"))?; + let mixed_personas = config + .personas + .iter() + .any(|persona| *persona != Persona::ShardOps); + roles::grant_runtime( + &superuser, + &name, + config.shardops_schema_create, + mixed_personas, + ) + .await?; + superuser.close().await; + } + Ok(()) +} + +#[derive(Debug, Clone, Default)] +pub struct ServerInfo { + pub version: String, + pub max_connections: u32, + pub superuser_reserved: u32, + pub connections_at_start: i64, + pub postmaster_start: String, + pub engine: String, + pub queue_storage_schema: Option, + pub awa_version: String, +} + +pub async fn server_info(pool: &PgPool) -> Result { + let version: String = sqlx::query_scalar("SHOW server_version") + .fetch_one(pool) + .await?; + let max_connections: String = sqlx::query_scalar("SHOW max_connections") + .fetch_one(pool) + .await?; + let reserved: String = sqlx::query_scalar("SHOW superuser_reserved_connections") + .fetch_one(pool) + .await?; + let connections: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE datname IS NOT NULL", + ) + .fetch_one(pool) + .await?; + let postmaster_start: String = sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await?; + let schema: Option = sqlx::query_scalar("SELECT awa.active_queue_storage_schema()") + .fetch_one(pool) + .await?; + Ok(ServerInfo { + version, + max_connections: max_connections.parse().unwrap_or(100), + superuser_reserved: reserved.parse().unwrap_or(3), + connections_at_start: connections, + postmaster_start, + engine: if schema.is_some() { + "queue_storage" + } else { + "canonical" + } + .to_string(), + queue_storage_schema: schema, + awa_version: env!("CARGO_PKG_VERSION").to_string(), + }) +} + +/// Hands out connections to the personas so that all of them together stay inside +/// the server's `max_connections`. +pub struct Budget { + usable: u32, + allocations: Mutex>, +} + +impl Budget { + /// `headroom` connections are kept free for the oracle, operator probes and chaos + /// sessions on top of the server's superuser reservation. + pub fn new(info: &ServerInfo, headroom: u32) -> Self { + let usable = info + .max_connections + .saturating_sub(info.superuser_reserved) + .saturating_sub(headroom) + .saturating_sub(info.connections_at_start.max(0) as u32); + Self { + usable, + allocations: Mutex::new(BTreeMap::new()), + } + } + + /// Reserves up to `want` connections under `label`, never fewer than `min`. + /// Returns the amount granted; panics when even `min` does not fit, because the + /// run would exhaust the server before it tested anything. + pub fn reserve(&self, label: &str, want: u32, min: u32) -> u32 { + let mut allocations = self.allocations.lock().expect("budget lock"); + let used: u32 = allocations.values().sum(); + let available = self.usable.saturating_sub(used); + assert!( + available >= min, + "connection budget exhausted: {label} needs at least {min}, {available} of {} left; \ + raise max_connections or lower AWA_SOAK_SCALE / persona count. Allocations: {allocations:?}", + self.usable + ); + let granted = want.min(available).max(min); + allocations.insert(label.to_string(), granted); + granted + } + + pub fn release(&self, label: &str) { + self.allocations.lock().expect("budget lock").remove(label); + } + + pub fn usable(&self) -> u32 { + self.usable + } + + pub fn table(&self) -> Vec<(String, u32)> { + self.allocations + .lock() + .expect("budget lock") + .iter() + .map(|(label, count)| (label.clone(), *count)) + .collect() + } +} + +/// A function that builds one runtime's client over its own pool. +pub type ClientFactory = Arc Client + Send + Sync>; + +/// One awa runtime living on its own tokio runtime, so it can be stopped gracefully +/// or killed the way a crashed process dies: tasks vanish mid-job without any +/// shutdown handshake and the connections drop. +pub struct ClientHost { + pub label: String, + runtime: tokio::runtime::Runtime, + client: Arc, + pool: PgPool, +} + +impl ClientHost { + pub async fn spawn( + url: String, + label: String, + pool_size: u32, + worker_threads: usize, + factory: ClientFactory, + ) -> Result { + let thread_label = label.clone(); + tokio::task::spawn_blocking(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .thread_name(format!("soak-{thread_label}")) + .enable_all() + .build() + .map_err(|error| format!("build runtime: {error}"))?; + let (pool, client) = runtime.block_on(async { + let pool = connect_pool(&url, &format!("soak_{thread_label}"), pool_size) + .await + .map_err(|error| format!("connect pool: {error}"))?; + let client = factory(pool.clone()); + client + .start() + .await + .map_err(|error| format!("start client: {error}"))?; + Ok::<_, String>((pool, client)) + })?; + Ok(Self { + label: thread_label, + runtime, + client: Arc::new(client), + pool, + }) + }) + .await + .map_err(|error| format!("spawn host task: {error}"))? + } + + pub fn client(&self) -> &Client { + &self.client + } + + pub fn pool(&self) -> &PgPool { + &self.pool + } + + /// Graceful stop: stop claiming, drain in-flight jobs, then close the pool. + pub async fn shutdown(self, timeout: Duration) { + let client = self.client.clone(); + let pool = self.pool.clone(); + let _ = self + .runtime + .handle() + .spawn(async move { + client.shutdown(timeout).await; + pool.close().await; + }) + .await; + self.runtime.shutdown_background(); + } + + /// Abrupt stop: no shutdown handshake, as when a process is killed. + pub fn kill(self) { + self.runtime.shutdown_background(); + } +} + +pub fn url_for_database(config: &SoakConfig, name: &str) -> String { + replace_database_name(&config.database_url, name) +} diff --git a/awa/tests/persona_soak/fleet.rs b/awa/tests/persona_soak/fleet.rs new file mode 100644 index 000000000..b7e709ca1 --- /dev/null +++ b/awa/tests/persona_soak/fleet.rs @@ -0,0 +1,192 @@ +//! A persona's group of runtimes, with the capacity bookkeeping the stall check +//! needs and the lifecycle operations the operator and chaos personas use. + +use super::config::Persona; +use super::db::{ClientFactory, ClientHost}; +use super::registry::QueueInfo; +use super::shared::Shared; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Mutex; + +pub struct Fleet { + pub name: String, + pub persona: Persona, + /// Every runtime in the fleet serves these queues with this many workers. + served: Vec<(Arc, i64)>, + factory: ClientFactory, + pool_size: u32, + worker_threads: usize, + url: String, + hosts: Mutex>>, + budget_label: String, +} + +impl Fleet { + #[allow(clippy::too_many_arguments)] + pub async fn start( + shared: &Arc, + name: &str, + persona: Persona, + replicas: usize, + pool_size: u32, + worker_threads: usize, + served: Vec<(Arc, i64)>, + factory: ClientFactory, + ) -> Result, String> { + let url = shared.config.database_url.clone(); + Self::start_with_url( + shared, + url, + name, + persona, + replicas, + pool_size, + worker_threads, + served, + factory, + ) + .await + } + + /// As [`Fleet::start`], with every runtime connecting through `url`, for a fleet that + /// runs as a role other than the one in the run's database URL. + #[allow(clippy::too_many_arguments)] + pub async fn start_with_url( + shared: &Arc, + url: String, + name: &str, + persona: Persona, + replicas: usize, + pool_size: u32, + worker_threads: usize, + served: Vec<(Arc, i64)>, + factory: ClientFactory, + ) -> Result, String> { + let budget_label = name.to_string(); + let granted = shared.budget.reserve( + &budget_label, + pool_size * replicas as u32, + 2 * replicas as u32, + ); + let pool_size = (granted / replicas as u32).max(2); + let fleet = Arc::new(Self { + name: name.to_string(), + persona, + served, + factory, + pool_size, + worker_threads, + url, + hosts: Mutex::new(Vec::new()), + budget_label, + }); + for index in 0..replicas { + let host = fleet.spawn_host(index).await?; + fleet.hosts.lock().await.push(Some(host)); + } + shared + .fleets + .lock() + .expect("fleets") + .insert(name.to_string(), fleet.clone()); + Ok(fleet) + } + + async fn spawn_host(&self, index: usize) -> Result { + let host = ClientHost::spawn( + self.url.clone(), + format!("{}_{index}", self.name), + self.pool_size, + self.worker_threads, + self.factory.clone(), + ) + .await?; + for (queue, workers) in &self.served { + queue.capacity.fetch_add(*workers, Ordering::Relaxed); + } + Ok(host) + } + + fn release_capacity(&self) { + for (queue, workers) in &self.served { + queue.capacity.fetch_sub(*workers, Ordering::Relaxed); + } + } + + pub async fn replicas(&self) -> usize { + self.hosts.lock().await.len() + } + + pub async fn live_replicas(&self) -> Vec { + self.hosts + .lock() + .await + .iter() + .enumerate() + .filter_map(|(index, host)| host.is_some().then_some(index)) + .collect() + } + + pub async fn pool_of(&self, index: usize) -> Option { + self.hosts + .lock() + .await + .get(index)? + .as_ref() + .map(|host| host.pool().clone()) + } + + /// Stops runtime `index` gracefully and starts a fresh one in its place. + pub async fn restart(&self, index: usize, timeout: Duration) -> Result { + let old = self + .hosts + .lock() + .await + .get_mut(index) + .and_then(Option::take); + let Some(old) = old else { return Ok(false) }; + self.release_capacity(); + old.shutdown(timeout).await; + let fresh = self.spawn_host(index).await?; + self.hosts.lock().await[index] = Some(fresh); + Ok(true) + } + + /// Kills runtime `index` with no shutdown handshake, then (optionally) starts a + /// replacement after `downtime`. + pub async fn kill(&self, index: usize, downtime: Option) -> Result { + let old = self + .hosts + .lock() + .await + .get_mut(index) + .and_then(Option::take); + let Some(old) = old else { return Ok(false) }; + self.release_capacity(); + old.kill(); + if let Some(downtime) = downtime { + tokio::time::sleep(downtime).await; + let fresh = self.spawn_host(index).await?; + self.hosts.lock().await[index] = Some(fresh); + } + Ok(true) + } + + pub async fn shutdown_all(&self, timeout: Duration, shared: &Shared) { + let hosts: Vec = self + .hosts + .lock() + .await + .iter_mut() + .filter_map(Option::take) + .collect(); + for host in hosts { + self.release_capacity(); + host.shutdown(timeout).await; + } + shared.budget.release(&self.budget_label); + } +} diff --git a/awa/tests/persona_soak/ledger.rs b/awa/tests/persona_soak/ledger.rs new file mode 100644 index 000000000..a1c8ca236 --- /dev/null +++ b/awa/tests/persona_soak/ledger.rs @@ -0,0 +1,266 @@ +//! The client-side ledger: every job a persona enqueued, what the workers saw, and +//! what the system reported back. The database may prune terminal rows long before +//! the run ends, so the ledger is the oracle's memory of outcomes. + +use super::config::Persona; +use std::collections::HashMap; +use std::sync::Mutex; + +const SHARDS: usize = 64; + +/// NO_EVENTS workloads keep their fast path; durable loss detection covers +/// this deterministic sample, chosen from the immutable assigned job ID. +pub const TERMINAL_PROOF_STRIDE: i64 = 64; + +pub fn samples_terminal_proof(job_id: i64) -> bool { + job_id.rem_euclid(TERMINAL_PROOF_STRIDE) == 0 +} + +/// What a job is allowed to end as. +pub mod expect { + pub const COMPLETED: u8 = 1; + /// Failed, which with the DLQ enabled means a DLQ row. + pub const FAILED: u8 = 2; + pub const CANCELLED: u8 = 4; + pub const ANY_TERMINAL: u8 = COMPLETED | FAILED | CANCELLED; +} + +/// Facts about a job, accumulated as the run proceeds. +pub mod flag { + /// An operator cancel, drain or purge may legitimately have removed the job. + pub const OPERATOR_REMOVED: u32 = 1 << 0; + pub const OPERATOR_CANCEL_OK: u32 = 1 << 1; + pub const OPERATOR_MOVED: u32 = 1 << 2; + pub const DLQ_RETRIED: u32 = 1 << 3; + pub const HANDLER_CANCELLED: u32 = 1 << 4; + pub const RESCUED: u32 = 1 << 5; + pub const EXHAUSTED: u32 = 1 << 6; + pub const COMPLETED_EVENT: u32 = 1 << 7; + pub const RETURNED_COMPLETED: u32 = 1 << 8; + /// The job is served by a firehose queue where lifecycle events are not collected. + pub const NO_EVENTS: u32 = 1 << 9; + /// First seen at execution time (cron-created); not enqueued by a producer. + pub const UNREGISTERED: u32 = 1 << 10; + /// A DLQ row for this job was purged by an operator. + pub const DLQ_PURGED: u32 = 1 << 11; + pub const CHAOS_VICTIM: u32 = 1 << 12; + /// Terminal rows read from the store for jobs whose runtime reports no lifecycle events. + pub const DB_SEEN_COMPLETED: u32 = 1 << 13; + pub const DB_SEEN_FAILED: u32 = 1 << 14; + pub const DB_SEEN_CANCELLED: u32 = 1 << 15; + /// How a firehose job was enqueued. + pub const SRC_BURST: u32 = 1 << 16; + pub const SRC_COPY: u32 = 1 << 17; + pub const SRC_STEADY: u32 = 1 << 18; +} + +#[derive(Debug, Clone)] +pub struct JobRecord { + pub persona: Persona, + pub queue: u32, + pub enqueued_ms: u32, + pub expect: u8, + pub flags: u32, + pub executions: u16, + pub in_flight: u8, + pub overlapped: bool, + pub max_attempt: i16, + pub latest_run_lease: i64, + pub terminal_run_lease: Option, + pub first_start_ms: u32, + pub last_start_ms: u32, + pub completed_events: u8, + pub exhausted_events: u8, + pub cancel_events: u8, + pub completed_event_ms: u32, + /// When an operator cancel on this job committed, 0 if none. + pub operator_cancel_ms: u32, +} + +impl JobRecord { + pub fn new(persona: Persona, queue: u32, enqueued_ms: u32, expect: u8, flags: u32) -> Self { + Self { + persona, + queue, + enqueued_ms, + expect, + flags, + executions: 0, + in_flight: 0, + overlapped: false, + max_attempt: 0, + latest_run_lease: 0, + terminal_run_lease: None, + first_start_ms: 0, + last_start_ms: 0, + completed_events: 0, + exhausted_events: 0, + cancel_events: 0, + completed_event_ms: 0, + operator_cancel_ms: 0, + } + } + + pub fn note_terminal(&mut self, run_lease: i64) { + if self.terminal_run_lease.is_none_or(|old| run_lease > old) { + // DB outcome bits belong to the newest proven attempt. A later + // lifecycle proof also supersedes an older observed DB state. + self.flags &= + !(flag::DB_SEEN_COMPLETED | flag::DB_SEEN_FAILED | flag::DB_SEEN_CANCELLED); + self.terminal_run_lease = Some(run_lease); + } + } + + pub fn note_observed_terminal(&mut self, run_lease: i64, state: &str) { + self.note_terminal(run_lease); + if self.terminal_run_lease == Some(run_lease) { + // The query is unordered: an older row can follow a newer proof. + // Different states at the same lease remain visible as a conflict. + self.flags |= match state { + "completed" => flag::DB_SEEN_COMPLETED, + "failed" => flag::DB_SEEN_FAILED, + "cancelled" => flag::DB_SEEN_CANCELLED, + _ => 0, + }; + } + } + + pub fn has_durable_terminal(&self) -> bool { + self.terminal_run_lease + .is_some_and(|lease| lease >= self.latest_run_lease) + } + + pub fn has(&self, flag: u32) -> bool { + self.flags & flag != 0 + } +} + +pub struct Ledger { + shards: Vec>>, + /// `(cron name, fire time)` -> job ids executed for it; more than one is a fencing failure. + cron_fires: Mutex>>, +} + +impl Default for Ledger { + fn default() -> Self { + Self::new() + } +} + +impl Ledger { + pub fn new() -> Self { + Self { + shards: (0..SHARDS).map(|_| Mutex::new(HashMap::new())).collect(), + cron_fires: Mutex::new(HashMap::new()), + } + } + + fn shard(&self, id: i64) -> &Mutex> { + &self.shards[(id as u64 % SHARDS as u64) as usize] + } + + /// Records an enqueue. A fast worker can run the job before its producer gets here, so + /// an existing record keeps what the workers already reported. + pub fn register(&self, id: i64, record: JobRecord) { + let mut shard = self.shard(id).lock().expect("ledger"); + match shard.get_mut(&id) { + Some(existing) => { + existing.persona = record.persona; + existing.queue = record.queue; + existing.enqueued_ms = record.enqueued_ms; + existing.expect = record.expect; + existing.flags = (existing.flags | record.flags) & !flag::UNREGISTERED; + } + None => { + shard.insert(id, record); + } + } + } + + /// Applies `update` to the job's record, creating an unregistered record when the + /// job was not enqueued by a producer (cron fires, operator moves). + pub fn with( + &self, + id: i64, + default: impl FnOnce() -> JobRecord, + update: impl FnOnce(&mut JobRecord) -> R, + ) -> R { + let mut shard = self.shard(id).lock().expect("ledger"); + let record = shard.entry(id).or_insert_with(|| { + let mut record = default(); + record.flags |= flag::UNREGISTERED; + record + }); + update(record) + } + + pub fn mark(&self, id: i64, flags: u32) { + if let Some(record) = self.shard(id).lock().expect("ledger").get_mut(&id) { + record.flags |= flags; + } + } + + pub fn mark_many(&self, ids: &[i64], flags: u32) { + for id in ids { + self.mark(*id, flags); + } + } + + /// Adds `flags` to every record for which `predicate` holds. + pub fn mark_where(&self, predicate: impl Fn(&JobRecord) -> bool, flags: u32) -> usize { + let mut marked = 0; + for shard in &self.shards { + for record in shard.lock().expect("ledger").values_mut() { + if predicate(record) { + record.flags |= flags; + marked += 1; + } + } + } + marked + } + + pub fn get(&self, id: i64) -> Option { + self.shard(id).lock().expect("ledger").get(&id).cloned() + } + + pub fn len(&self) -> usize { + self.shards + .iter() + .map(|shard| shard.lock().expect("ledger").len()) + .sum() + } + + /// Calls `visit` for every record. Shards are locked one at a time. + pub fn for_each(&self, mut visit: impl FnMut(i64, &JobRecord)) { + for shard in &self.shards { + let shard = shard.lock().expect("ledger"); + for (id, record) in shard.iter() { + visit(*id, record); + } + } + } + + pub fn note_cron_fire(&self, name: String, fire_time: String, job_id: i64) -> bool { + let mut fires = self.cron_fires.lock().expect("cron fires"); + let ids = fires.entry((name, fire_time)).or_default(); + if !ids.contains(&job_id) { + ids.push(job_id); + } + ids.len() == 1 + } + + pub fn cron_duplicates(&self) -> Vec<((String, String), Vec)> { + self.cron_fires + .lock() + .expect("cron fires") + .iter() + .filter(|(_, ids)| ids.len() > 1) + .map(|(key, ids)| (key.clone(), ids.clone())) + .collect() + } + + pub fn cron_fire_count(&self) -> usize { + self.cron_fires.lock().expect("cron fires").len() + } +} diff --git a/awa/tests/persona_soak/mod.rs b/awa/tests/persona_soak/mod.rs new file mode 100644 index 000000000..e45c4b7b1 --- /dev/null +++ b/awa/tests/persona_soak/mod.rs @@ -0,0 +1,23 @@ +//! Persona-driven soak harness with an invariant oracle. See the test binary +//! `persona_soak_test.rs` for how to run it. + +pub mod abandon_oracle; +pub mod adminserver; +pub mod config; +pub mod db; +pub mod fleet; +pub mod ledger; +pub mod oracle; +pub mod personas; +pub mod pgserver; +pub mod python_probe; +pub mod registry; +pub mod report; +pub mod ring_locks; +pub mod rng; +pub mod roles; +pub mod run; +pub mod shard_oracle; +pub mod shared; +pub mod stop; +pub mod terminal_proof; diff --git a/awa/tests/persona_soak/oracle.rs b/awa/tests/persona_soak/oracle.rs new file mode 100644 index 000000000..aa70b2e5b --- /dev/null +++ b/awa/tests/persona_soak/oracle.rs @@ -0,0 +1,1392 @@ +//! The invariant oracle: periodic checks while the load runs and a final +//! reconciliation of the client ledger against the database. + +use super::ledger::{expect, flag, samples_terminal_proof, JobRecord}; +use super::pgserver::PgWatch; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +const MAX_ROWS: usize = 25; + +/// count, min and max job id, min and max lane sequence. +type ForensicRow = (i64, Option, Option, Option, Option); + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Every terminal outcome the store records, one row per `(job, attempt)`. +fn terminal_cte(schema: &str) -> String { + format!( + "terminal AS ( + SELECT job_id, run_lease, 'done:' || state::text AS src FROM {schema}.done_entries + UNION ALL SELECT job_id, run_lease, 'dlq' FROM {schema}.dlq_entries + UNION ALL SELECT b.job_id, b.run_lease, 'receipt_completed' + FROM {schema}.receipt_completion_batches c, + LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + )" + ) +} + +/// Jobs the lane cursors still consider claimable: at or after the claim cursor and +/// not tombstoned. +fn available_cte(schema: &str) -> String { + format!( + "heads AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS claim_seq + FROM {schema}.queue_claim_heads + ), + available AS MATERIALIZED ( + SELECT r.job_id, r.run_lease, r.queue + FROM {schema}.ready_entries r + JOIN heads h ON h.queue = r.queue AND h.priority = r.priority AND h.enqueue_shard = r.enqueue_shard + WHERE r.lane_seq >= h.claim_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones t + WHERE t.ready_slot = r.ready_slot AND t.ready_generation = r.ready_generation + AND t.queue = r.queue AND t.priority = r.priority + AND t.enqueue_shard = r.enqueue_shard AND t.lane_seq = r.lane_seq + ) + )" + ) +} + +#[derive(Default)] +struct StallTracker { + /// queue -> (claim cursor, millis when the cursor last moved, already reported) + cursors: HashMap, +} + +pub struct Oracle { + shared: Arc, + pool: PgPool, + watch: Arc, + stall: StallTracker, + last_cancel_scan_ms: i64, + reported_duplicates: std::collections::HashSet<(i64, i64)>, + reported_cancel: std::collections::HashSet, + initial_stats_reset: std::sync::Mutex>>, +} + +impl Oracle { + pub fn new(shared: Arc, pool: PgPool, watch: Arc) -> Self { + Self { + shared, + pool, + watch, + stall: StallTracker::default(), + last_cancel_scan_ms: 0, + reported_duplicates: Default::default(), + reported_cancel: Default::default(), + initial_stats_reset: Default::default(), + } + } + + /// Runs the periodic checks until `stop`. + pub async fn run(mut self, stop: Stop) -> Self { + let mut tick = 0u64; + loop { + if stop.sleep(Duration::from_secs(5)).await { + return self; + } + tick += 1; + self.probe_server().await; + self.check_stall().await; + self.check_promotion().await; + if tick.is_multiple_of(2) { + self.check_terminal_uniqueness().await; + self.check_unexplained_cancels().await; + } + if tick.is_multiple_of(3) { + self.sample_resources().await; + } + } + } + + /// Crash recovery discards the cumulative statistics, so a changed reset time means the + /// server restarted its processes even though the postmaster, and its start time, survived. + async fn probe_stats_reset(&self) { + let reset: Result, _> = + sqlx::query_scalar("SELECT stats_reset::text FROM pg_stat_bgwriter") + .fetch_one(&self.pool) + .await; + let reset = match reset { + Ok(reset) => reset, + Err(error) => { + self.shared + .incomplete(format!("crash-recovery probe failed: {error}"), vec![]); + return; + } + }; + let mut initial = self.initial_stats_reset.lock().expect("stats reset"); + match initial.as_ref() { + None => *initial = Some(reset), + Some(first) if *first != reset => { + let message = format!("pg_stat_bgwriter.stats_reset changed from {first:?} to {reset:?}: the server went through crash recovery"); + *initial = Some(reset); + drop(initial); + self.shared + .violation("server_crash_recovery", message, vec![]); + } + Some(_) => {} + } + } + + async fn probe_server(&self) { + self.probe_stats_reset().await; + let mut last_error = String::new(); + for attempt in 0..4 { + match PgWatch::postmaster_start(&self.pool).await { + Ok(start) if start == self.watch.started_at => return, + Ok(start) => { + self.shared.violation( + "server_restart", + format!( + "pg_postmaster_start_time changed from {} to {start}", + self.watch.started_at + ), + vec![], + ); + return; + } + Err(error) => { + last_error = error.to_string(); + tokio::time::sleep(Duration::from_secs(1 + attempt)).await; + } + } + } + self.shared.run.add("server_probe_failures", 1); + self.shared.incomplete( + format!("server probe failed repeatedly: {last_error}"), + vec![], + ); + } + + async fn check_terminal_uniqueness(&mut self) { + let schema = schema(&self.shared); + let terminal = terminal_cte(&schema); + let started = std::time::Instant::now(); + + let duplicates: Result)>, _> = sqlx::query_as(&format!( + "WITH {terminal} + SELECT job_id, run_lease, array_agg(src ORDER BY src) + FROM terminal GROUP BY job_id, run_lease HAVING count(*) > 1 + ORDER BY job_id LIMIT {MAX_ROWS}" + )) + .fetch_all(&self.pool) + .await; + match duplicates { + Ok(rows) => { + let fresh: Vec<_> = rows + .iter() + .filter(|(job_id, run_lease, _)| { + self.reported_duplicates.insert((*job_id, *run_lease)) + }) + .collect(); + if !fresh.is_empty() { + self.shared.violation( + "duplicate_terminal_rows", + format!("{} job attempt(s) have more than one terminal row", fresh.len()), + fresh + .iter() + .map(|(job_id, run_lease, sources)| { + format!("{{\"job_id\":{job_id},\"run_lease\":{run_lease},\"terminal_rows\":{sources:?}}}") + }) + .collect(), + ); + } + } + Err(error) => self + .shared + .incomplete(format!("terminal uniqueness query failed: {error}"), vec![]), + } + + let available = available_cte(&schema); + let overlap: Result, _> = sqlx::query_as(&format!( + "WITH {terminal}, {available} + SELECT t.job_id, 'available'::text, t.run_lease, a.run_lease + FROM terminal t JOIN available a ON a.job_id = t.job_id AND a.run_lease <= t.run_lease + UNION ALL + SELECT t.job_id, 'deferred:' || d.state::text, t.run_lease, d.run_lease + FROM terminal t JOIN {schema}.deferred_jobs d ON d.job_id = t.job_id AND d.run_lease <= t.run_lease + UNION ALL + SELECT t.job_id, 'leased:' || l.state::text, t.run_lease, l.run_lease + FROM terminal t JOIN {schema}.leases l ON l.job_id = t.job_id AND l.run_lease <= t.run_lease + LIMIT {MAX_ROWS}" + )) + .fetch_all(&self.pool) + .await; + match overlap { + Ok(rows) if !rows.is_empty() => { + // Confirm on a second read: a single statement is one snapshot, so a + // repeat that still shows the overlap is not a race with this probe. + tokio::time::sleep(Duration::from_millis(500)).await; + let ids: Vec = rows.iter().map(|row| row.0).collect(); + let still: Result, _> = sqlx::query_as(&format!( + "WITH {terminal}, {available} + SELECT t.job_id, 'available'::text FROM terminal t JOIN available a ON a.job_id = t.job_id AND a.run_lease <= t.run_lease WHERE t.job_id = ANY($1) + UNION ALL SELECT t.job_id, 'deferred' FROM terminal t JOIN {schema}.deferred_jobs d ON d.job_id = t.job_id AND d.run_lease <= t.run_lease WHERE t.job_id = ANY($1) + UNION ALL SELECT t.job_id, 'leased' FROM terminal t JOIN {schema}.leases l ON l.job_id = t.job_id AND l.run_lease <= t.run_lease WHERE t.job_id = ANY($1)" + )) + .bind(&ids) + .fetch_all(&self.pool) + .await; + let confirmed: std::collections::HashSet = match still { + Ok(rows) => rows.into_iter().map(|row| row.0).collect(), + Err(error) => { + self.shared.incomplete( + format!("terminal/live confirmation failed: {error}"), + vec![], + ); + Default::default() + } + }; + let confirmed_rows: Vec = rows + .iter() + .filter(|row| confirmed.contains(&row.0)) + .filter(|row| self.reported_duplicates.insert((row.0, -row.2 - 1))) + .map(|(job_id, place, terminal_lease, live_lease)| { + format!("{{\"job_id\":{job_id},\"also_in\":\"{place}\",\"terminal_run_lease\":{terminal_lease},\"live_run_lease\":{live_lease}}}") + }) + .collect(); + if !confirmed_rows.is_empty() { + self.shared.violation( + "terminal_and_live", + format!( + "{} job(s) are terminal and still available, deferred or leased", + confirmed_rows.len() + ), + confirmed_rows, + ); + } + } + Ok(_) => {} + Err(error) => self + .shared + .incomplete(format!("terminal/live query failed: {error}"), vec![]), + } + self.shared.run.observe( + "oracle_terminal_check_ms", + started.elapsed().as_millis() as i64, + ); + } + + /// A `cancelled` terminal row nobody asked for: no operator cancel, drain or purge + /// touched the job and its handler did not cancel it. + async fn check_unexplained_cancels(&mut self) { + let schema = schema(&self.shared); + let since_ms = self.last_cancel_scan_ms; + let scan_started_ms = self.shared.now_ms(); + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, queue, run_lease FROM {schema}.done_entries + WHERE state = 'cancelled' AND finalized_at > now() - make_interval(secs => $1) + LIMIT 2000" + )) + .bind(((scan_started_ms - since_ms) as f64 / 1000.0 + 15.0).max(20.0)) + .fetch_all(&self.pool) + .await; + self.last_cancel_scan_ms = scan_started_ms; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + let mut unexplained = Vec::new(); + for (job_id, queue_name, run_lease) in rows { + let Some(record) = self.shared.ledger.get(job_id) else { + continue; + }; + let queue = self.shared.queues.get(&queue_name); + let drained = queue.as_ref().is_some_and(|queue| { + let drained_at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + drained_at > 0 && i64::from(record.enqueued_ms) <= drained_at + 1000 + }); + let explained = record.has(flag::OPERATOR_REMOVED) + || record.has(flag::HANDLER_CANCELLED) + || record.has(flag::OPERATOR_MOVED) + || drained; + if !explained { + // An operator flag is set just after its cancel commits; look again before accusing. + let lag = if self + .shared + .ledger + .get(job_id) + .is_some_and(|record| record.persona == super::config::Persona::Python) + { + 4500 + } else { + 1500 + }; + tokio::time::sleep(Duration::from_millis(lag)).await; + } + let explained = explained + || self.shared.ledger.get(job_id).is_some_and(|record| { + record.has(flag::OPERATOR_REMOVED) || record.has(flag::HANDLER_CANCELLED) + }); + if !explained && self.reported_cancel.insert(job_id) { + unexplained.push(format!( + "{{\"job_id\":{job_id},\"queue\":\"{queue_name}\",\"run_lease\":{run_lease},\"executions\":{},\"running_now\":{},\"completed_events\":{}}}", + record.executions, record.in_flight, record.completed_events + )); + } + } + if !unexplained.is_empty() { + self.shared.violation( + "unexplained_cancel", + format!( + "{} job(s) recorded as cancelled with no operator or handler cancel", + unexplained.len() + ), + unexplained.into_iter().take(MAX_ROWS).collect(), + ); + } + } + + async fn check_promotion(&self) { + let schema = schema(&self.shared); + let mut threshold = self.shared.config.promotion_overdue; + if self.shared.chaos_active() { + threshold += self.shared.config.chaos_slack; + } + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, queue, extract(epoch FROM now() - run_at)::float8 + FROM {schema}.deferred_jobs + WHERE state IN ('scheduled', 'retryable') AND run_at < now() - make_interval(secs => $1) + ORDER BY run_at LIMIT {MAX_ROWS}" + )) + .bind(threshold.as_secs_f64()) + .fetch_all(&self.pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + if rows.is_empty() { + return; + } + // Re-read so that a job promoted a moment ago is not reported. + tokio::time::sleep(Duration::from_secs(2)).await; + let ids: Vec = rows.iter().map(|row| row.0).collect(); + let still: Result, _> = sqlx::query_scalar(&format!( + "SELECT job_id FROM {schema}.deferred_jobs WHERE job_id = ANY($1) AND state IN ('scheduled','retryable')" + )) + .bind(&ids) + .fetch_all(&self.pool) + .await; + let still = match still { + Ok(still) => still, + Err(error) => { + self.shared + .incomplete(format!("promotion confirmation failed: {error}"), vec![]); + return; + } + }; + let stuck: Vec<_> = rows.iter().filter(|row| still.contains(&row.0)).collect(); + if stuck.is_empty() { + return; + } + let worst = stuck.iter().map(|row| row.2).fold(0.0, f64::max); + self.shared + .run + .set_max("promotion_overdue_max_s", worst as i64); + self.shared.violation( + "promotion_stall", + format!( + "{} scheduled job(s) overdue beyond {}s, worst {worst:.1}s", + stuck.len(), + threshold.as_secs() + ), + stuck + .iter() + .map(|(job_id, queue, overdue)| { + format!( + "{{\"job_id\":{job_id},\"queue\":\"{queue}\",\"overdue_s\":{overdue:.1}}}" + ) + }) + .collect(), + ); + } + + async fn check_stall(&mut self) { + let schema = schema(&self.shared); + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT c.queue, + COALESCE(sum({schema}.sequence_next_value(e.seq_name) - {schema}.sequence_next_value(c.seq_name)), 0)::bigint, + COALESCE(sum({schema}.sequence_next_value(c.seq_name)), 0)::bigint + FROM {schema}.queue_claim_heads c + JOIN {schema}.queue_enqueue_heads e + ON e.queue = c.queue AND e.priority = c.priority AND e.enqueue_shard = c.enqueue_shard + GROUP BY c.queue" + )) + .fetch_all(&self.pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + let now = self.shared.now_ms(); + let mut window = self.shared.config.stall_window; + if self.shared.chaos_active() { + window += self.shared.config.chaos_slack; + } + let mut candidates = Vec::new(); + for (queue_name, backlog, cursor) in rows { + let entry = self + .stall + .cursors + .entry(queue_name.clone()) + .or_insert((cursor, now, false)); + if entry.0 != cursor { + *entry = (cursor, now, false); + continue; + } + let Some(queue) = self.shared.queues.get(&queue_name) else { + continue; + }; + let capacity = queue.capacity.load(std::sync::atomic::Ordering::Relaxed); + let in_flight = queue.in_flight.load(std::sync::atomic::Ordering::Relaxed); + let held = queue.is_held() || queue.rate_limited; + let idle_workers = capacity > 0 && in_flight < capacity; + if !entry.2 + && backlog > 0 + && idle_workers + && !held + && (now - entry.1) as u64 >= window.as_millis() as u64 + { + candidates.push((queue_name, backlog, cursor, entry.1, capacity, in_flight)); + } else if !idle_workers || held || backlog <= 0 { + // Busy or deliberately held: the stall clock restarts. + entry.1 = now; + } + } + for (queue_name, backlog, cursor, since, capacity, in_flight) in candidates { + let exact: Result = sqlx::query_scalar(&format!( + "WITH {} SELECT count(*)::bigint FROM available WHERE queue = $1", + available_cte(&schema) + )) + .bind(&queue_name) + .fetch_one(&self.pool) + .await; + let available = match exact { + Ok(available) => available, + Err(error) => { + self.shared + .incomplete(format!("exact claimable-row check failed: {error}"), vec![]); + continue; + } + }; + if available > 0 { + if let Some(entry) = self.stall.cursors.get_mut(&queue_name) { + entry.2 = true; + } + self.shared.run.add("stalls", 1); + self.shared.violation( + "stall", + format!( + "queue {queue_name}: {available} available job(s), claim cursor frozen at {cursor} for {}s with {}/{} workers busy", + (now - since) / 1000, + in_flight, + capacity + ), + vec![format!( + "{{\"queue\":\"{queue_name}\",\"available\":{available},\"lane_backlog\":{backlog},\"claim_cursor\":{cursor},\"workers_busy\":{in_flight},\"worker_capacity\":{capacity}}}" + )], + ); + } else if let Some(entry) = self.stall.cursors.get_mut(&queue_name) { + entry.1 = now; + } + } + } + + /// Stops the run before it can fill the disk it shares with other work. + fn check_disk(&self) { + let output = std::process::Command::new("df") + .args(["--output=avail", "-BM"]) + .arg(&self.shared.config.report_dir) + .output(); + let Ok(output) = output else { return }; + let text = String::from_utf8_lossy(&output.stdout); + let Some(free_mb) = text + .lines() + .nth(1) + .and_then(|line| line.trim().trim_end_matches('M').parse::().ok()) + else { + return; + }; + self.shared.run.observe("free_disk_gb", free_mb / 1024); + if free_mb < self.shared.config.min_free_disk_gb * 1024 + && !self.shared.load_stop.is_stopped() + { + self.shared.violation( + "disk_low", + format!( + "only {} GB free on the report directory's filesystem; stopping the load phase", + free_mb / 1024 + ), + vec![], + ); + self.shared.load_stop.trigger(); + } + } + + async fn sample_resources(&self) { + self.check_disk(); + let connections: Result = sqlx::query_scalar( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE backend_type = 'client backend'", + ) + .fetch_one(&self.pool) + .await; + if let Ok(connections) = connections { + self.shared.run.set_max("peak_connections", connections); + self.shared.run.observe("connections", connections); + } + let size: Result = + sqlx::query_scalar("SELECT pg_database_size(current_database())") + .fetch_one(&self.pool) + .await; + if let Ok(size) = size { + self.shared.run.set_max("peak_db_mb", size / (1024 * 1024)); + self.shared.run.observe("db_mb", size / (1024 * 1024)); + } + if let Ok(status) = std::fs::read_to_string("/proc/self/status") { + if let Some(kb) = status + .lines() + .find_map(|line| line.strip_prefix("VmRSS:")) + .and_then(|rest| { + rest.trim() + .trim_end_matches("kB") + .trim() + .parse::() + .ok() + }) + { + self.shared.run.set_max("peak_harness_rss_mb", kb / 1024); + } + } + if let Ok(load) = std::fs::read_to_string("/proc/loadavg") { + if let Some(one_minute) = load + .split_whitespace() + .next() + .and_then(|v| v.parse::().ok()) + { + self.shared + .run + .observe("loadavg_x100", (one_minute * 100.0) as i64); + } + } + } + + /// Waits for the system to quiesce after load stops. Returns true when no live + /// work remains. + pub async fn wait_quiescent(&mut self, timeout: Duration) -> bool { + let schema = schema(&self.shared); + let deadline = std::time::Instant::now() + timeout; + let mut quiet_polls = 0; + loop { + let live: Result = sqlx::query_scalar(&format!( + "WITH {} SELECT (SELECT count(*) FROM available) + (SELECT count(*) FROM {schema}.leases) + + (SELECT count(*) FROM {schema}.deferred_jobs WHERE run_at < now() + interval '1 second')", + available_cte(&schema) + )) + .fetch_one(&self.pool) + .await; + if let Err(error) = &live { + self.shared + .incomplete(format!("quiescence query failed: {error}"), vec![]); + } + let mut pending_terminal_proof = 0usize; + self.shared.ledger.for_each(|id, record| { + if record.has(flag::RETURNED_COMPLETED) + && !record.has_durable_terminal() + && !record.has(flag::OPERATOR_REMOVED) + && !record.has(flag::DLQ_PURGED) + && (!record.has(flag::NO_EVENTS) || samples_terminal_proof(id)) + { + pending_terminal_proof += 1; + } + }); + let in_flight: i64 = self + .shared + .queues + .all() + .iter() + .map(|queue| queue.in_flight.load(std::sync::atomic::Ordering::Relaxed)) + .sum(); + let python_active = self + .shared + .python_active + .load(std::sync::atomic::Ordering::Relaxed); + // Claims abandoned by a killed process sit in no table the live count reads. + let abandoned_open = self + .shared + .abandon + .open_claims + .load(std::sync::atomic::Ordering::Relaxed); + match live { + Ok(count) + if in_flight == 0 + && python_active == 0 + && abandoned_open == 0 + && pending_terminal_proof == 0 + && count_is_idle(count) => + { + quiet_polls += 1 + } + _ => quiet_polls = 0, + } + if quiet_polls >= 3 { + return true; + } + if std::time::Instant::now() >= deadline { + self.shared.note( + "oracle", + format!( + "not quiescent after {}s: handlers running={in_flight}, live rows={live:?}", + timeout.as_secs() + ), + ); + return false; + } + self.probe_server().await; + self.check_stall().await; + tokio::time::sleep(Duration::from_secs(2)).await; + } + } + + /// Final reconciliation of the ledger against the database. + pub async fn final_check(&mut self, quiesced: bool) { + let schema = schema(&self.shared); + self.check_terminal_uniqueness().await; + self.check_unexplained_cancels().await; + self.check_promotion().await; + self.probe_server().await; + + self.shared.run.add( + "cron_fires_executed", + self.shared.ledger.cron_fire_count() as i64, + ); + for ((name, fire), ids) in self.shared.ledger.cron_duplicates() { + self.shared.violation( + "cron_double_fire", + format!("cron {name} fire {fire} produced jobs {ids:?}"), + vec![], + ); + } + + // Jobs with no terminal evidence on the client side have to be found in the database. + let mut unresolved: Vec = Vec::new(); + let mut unexpected: Vec = Vec::new(); + let mut sampled = 0i64; + let mut confirmed = 0i64; + let mut excluded = 0i64; + self.shared.ledger.for_each(|id, record| { + if record.has(flag::OPERATOR_REMOVED) || record.has(flag::DLQ_PURGED) { + return; + } + let sampled_no_events = record.has(flag::NO_EVENTS) && samples_terminal_proof(id); + sampled += i64::from(sampled_no_events); + if record.has_durable_terminal() { + confirmed += i64::from(sampled_no_events); + for outcome in unexpected_terminal_outcomes(record) { + if unexpected.len() >= MAX_ROWS { + break; + } + let queue = self.shared.queues.by_index(record.queue); + unexpected.push(serde_json::json!({ + "job_id": id, + "outcome": outcome, + "allowed": record.expect, + "persona": record.persona.name(), + "queue": queue.as_ref().map(|queue| queue.name.chars().take(256).collect::()), + "flags": record.flags, + "latest_run_lease": record.latest_run_lease, + "terminal_run_lease": record.terminal_run_lease, + "executions": record.executions, + "in_flight": record.in_flight, + "max_attempt": record.max_attempt, + "completed_events": record.completed_events, + "exhausted_events": record.exhausted_events, + "cancel_events": record.cancel_events, + "enqueued_ms": record.enqueued_ms, + "first_start_ms": record.first_start_ms, + "last_start_ms": record.last_start_ms, + "completed_event_ms": record.completed_event_ms, + "operator_cancel_ms": record.operator_cancel_ms, + "queue_drained_at_ms": queue.map(|queue| queue.drained_at_ms.load(std::sync::atomic::Ordering::Relaxed)), + }).to_string()); + } + } else if record.has(flag::NO_EVENTS) + && record.has(flag::RETURNED_COMPLETED) + && !sampled_no_events + { + // This optimized workload is sampled, not proven per job. + // A return is execution evidence only, never commit evidence. + excluded += 1; + } else { + unresolved.push(id); + } + }); + if !unexpected.is_empty() { + self.shared.violation( + "unexpected_outcome", + format!( + "{} job(s) ended in an outcome their persona did not allow", + unexpected.len() + ), + unexpected, + ); + } + + let mut live_stuck = Vec::new(); + let mut forensic_ids: Vec = Vec::new(); + let mut lost = Vec::new(); + let mut lost_after_run = Vec::new(); + let mut terminal_in_db = 0usize; + for chunk in unresolved.chunks(5000) { + let found: Result)>, _> = sqlx::query_as(&format!( + "SELECT job_id, string_agg(place, ','), max(terminal_run_lease) FROM ( + SELECT job_id, 'done:' || state::text AS place, run_lease AS terminal_run_lease FROM {schema}.done_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'dlq', run_lease FROM {schema}.dlq_entries WHERE job_id = ANY($1) + UNION ALL SELECT b.job_id, 'receipt_completed', b.run_lease + FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + WHERE c.job_ids && $1 AND b.job_id = ANY($1) + UNION ALL SELECT job_id, 'deferred:' || state::text, NULL::bigint FROM {schema}.deferred_jobs WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'ready', NULL::bigint FROM {schema}.ready_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'leased:' || state::text, NULL::bigint FROM {schema}.leases WHERE job_id = ANY($1) + ) places GROUP BY job_id" + )) + .bind(chunk) + .fetch_all(&self.pool) + .await; + let found: HashMap)> = match found { + Ok(rows) => rows + .into_iter() + .map(|(id, place, lease)| (id, (place, lease))) + .collect(), + Err(error) => { + self.shared.run.add("final_lookup_errors", 1); + self.shared.incomplete( + format!( + "final lookup failed for {} unresolved jobs: {error}", + chunk.len() + ), + vec![], + ); + continue; + } + }; + for id in chunk { + let record = self.shared.ledger.get(*id); + let describe = |place: &str| { + let record = record.clone(); + format!( + "{{\"job_id\":{id},\"found\":\"{place}\",\"executions\":{},\"persona\":\"{}\",\"flags\":{},\"enqueued_at_s\":{:.1},\"queue\":\"{}\"}}", + record.as_ref().map_or(0, |r| r.executions), + record.as_ref().map_or("?", |r| r.persona.name()), + record.as_ref().map_or(0, |r| r.flags), + record.as_ref().map_or(0.0, |r| f64::from(r.enqueued_ms) / 1000.0), + record + .as_ref() + .and_then(|r| self.shared.queues.by_index(r.queue)) + .map_or("?".to_string(), |queue| queue.name.clone()) + ) + }; + if record + .as_ref() + .is_some_and(|record| record.has_durable_terminal()) + { + confirmed += i64::from( + record + .as_ref() + .is_some_and(|record| record.has(flag::NO_EVENTS)) + && samples_terminal_proof(*id), + ); + continue; + } + match found.get(id) { + Some((place, terminal_lease)) + if terminal_lease.is_some_and(|lease| { + lease >= record.as_ref().map_or(0, |record| record.latest_run_lease) + }) => + { + terminal_in_db += 1; + confirmed += i64::from( + record + .as_ref() + .is_some_and(|record| record.has(flag::NO_EVENTS)) + && samples_terminal_proof(*id), + ); + if place.contains("done:cancelled") { + if let Some(record) = &record { + let queue = self.shared.queues.by_index(record.queue); + let drained = queue.is_some_and(|queue| { + let at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + at > 0 && i64::from(record.enqueued_ms) <= at + 1000 + }); + if !(record.has(flag::OPERATOR_REMOVED) || drained) + && record.expect & expect::CANCELLED == 0 + { + lost_after_run.push(describe(place)); + forensic_ids.push(*id); + } + } + } + } + Some((place, _)) => live_stuck.push(describe(place)), + None => { + let executed = record.as_ref().is_some_and(|r| r.executions > 0); + let returned = record + .as_ref() + .is_some_and(|r| r.has(flag::RETURNED_COMPLETED)); + let drained = record.as_ref().is_some_and(|r| { + self.shared.queues.by_index(r.queue).is_some_and(|queue| { + let at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + at > 0 && i64::from(r.enqueued_ms) <= at + 1000 + }) + }); + if drained { + continue; + } + if !executed { + lost.push(describe("absent")); + forensic_ids.push(*id); + } else if !returned { + lost_after_run.push(describe("absent")); + forensic_ids.push(*id); + } else { + self.shared.incomplete( + "a returned handler has no durable terminal proof; a missed pruned outcome cannot be distinguished from a lost finalize", + vec![describe("absent without durable terminal evidence")], + ); + } + } + } + } + } + if !forensic_ids.is_empty() { + self.record_forensics(&schema, &forensic_ids).await; + } + self.shared + .run + .add("terminal_proof_no_events_sampled", sampled); + self.shared + .run + .add("terminal_proof_no_events_confirmed", confirmed); + self.shared + .run + .add("terminal_proof_no_events_excluded", excluded); + self.shared + .run + .add("final_unresolved_ids", unresolved.len() as i64); + self.shared + .run + .add("final_terminal_in_db", terminal_in_db as i64); + let report = |kind: &str, what: &str, rows: Vec| { + if !rows.is_empty() { + let count = rows.len(); + self.shared.violation( + kind, + format!("{count} job(s) {what}"), + rows.into_iter().take(MAX_ROWS).collect(), + ); + } + }; + if quiesced { + report( + "stuck_job", + "never reached a terminal state although the system went quiet (available/deferred/leased)", + live_stuck, + ); + } else if !live_stuck.is_empty() { + // Work still queued when the grace ended is backlog, not loss: the run was + // sized beyond what the server could clear. + self.shared.finding( + "backlog_at_end", + format!( + "{} job(s) were still queued when the drain grace ended", + live_stuck.len() + ), + live_stuck.into_iter().take(MAX_ROWS).collect(), + ); + } + report( + "lost_job", + "were enqueued, never ran and are in no table", + lost, + ); + report("lost_job_after_run", "ran without a recorded terminal outcome and are in no table, or ended cancelled without cause", lost_after_run); + } + + /// For jobs that look lost: everywhere the store could still mention them, written to + /// the timeline so the reproducer carries it. + async fn record_forensics(&self, schema: &str, ids: &[i64]) { + let sample: Vec = ids.iter().copied().take(200).collect(); + let tables: [(&str, String); 8] = [ + ("ready_entries", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.ready_entries WHERE job_id = ANY($1)")), + ("ready_tombstones", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.ready_tombstones WHERE job_id = ANY($1)")), + ("lease_claims", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.lease_claims WHERE job_id = ANY($1)")), + ("lease_claim_batches", format!("SELECT count(*), min(j), max(j), 0::bigint, 0::bigint FROM {schema}.lease_claim_batches b, LATERAL unnest(b.job_ids) AS j WHERE b.job_ids && $1 AND j = ANY($1)")), + ("lease_claim_closures", format!("SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM {schema}.lease_claim_closures WHERE job_id = ANY($1)")), + ("attempt_state", format!("SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM {schema}.attempt_state WHERE job_id = ANY($1)")), + ("receipt_completion_tombstones", format!("SELECT count(*), 0::bigint, 0::bigint, 0::bigint, 0::bigint FROM {schema}.receipt_completion_tombstones")), + ("job_unique_claims", "SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM awa.job_unique_claims WHERE job_id = ANY($1)".to_string()), + ]; + let mut lines = Vec::new(); + for (name, sql) in tables { + let row: Result = sqlx::query_as(&sql) + .bind(&sample) + .fetch_one(&self.pool) + .await; + match row { + Ok((count, min_id, max_id, min_seq, max_seq)) => lines.push(format!("{name}: {count} row(s) for the sampled ids (job ids {min_id:?}..{max_id:?}, lane_seq {min_seq:?}..{max_seq:?})")), + Err(error) => lines.push(format!("{name}: query failed: {error}")), + } + } + self.shared.note( + "oracle", + format!( + "forensics for {} lost-looking job(s), first {:?}: {}", + ids.len(), + &ids[..ids.len().min(5)], + lines.join("; ") + ), + ); + } + + pub fn crash_scan(&self) -> Vec { + match self.watch.log_since_start(&self.shared.config) { + Some(log) => super::pgserver::crash_excerpts(&log), + None => Vec::new(), + } + } +} + +/// The final verdict's outcome predicates, shared by deterministic lease-order controls. +fn unexpected_terminal_outcomes(record: &JobRecord) -> Vec<&'static str> { + if !record.has_durable_terminal() { + return Vec::new(); + } + let completed = record.completed_events > 0 + || record.has(flag::RETURNED_COMPLETED) + || record.has(flag::DB_SEEN_COMPLETED); + [ + ( + "cancelled", + record.has(flag::DB_SEEN_CANCELLED) + && record.expect & expect::CANCELLED == 0 + && !record.has(flag::HANDLER_CANCELLED), + ), + ( + "completed", + completed && record.expect & expect::COMPLETED == 0, + ), + ( + "failed", + (record.exhausted_events > 0 || record.has(flag::DB_SEEN_FAILED)) + && record.expect & expect::FAILED == 0, + ), + ] + .into_iter() + .filter_map(|(outcome, unexpected)| unexpected.then_some(outcome)) + .collect() +} + +fn count_is_idle(count: i64) -> bool { + count <= 0 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::persona_soak::config::{replace_database_name, Persona, SoakConfig}; + use crate::persona_soak::db::{self, Budget}; + use awa::model::{ClaimedRuntimeJob, QueueStorage, QueueStorageConfig}; + use std::sync::atomic::Ordering; + + fn current_terminal_record() -> JobRecord { + let mut record = JobRecord::new(Persona::Tom, 0, 0, expect::COMPLETED, 0); + record.latest_run_lease = 1; + record + } + + #[test] + fn terminal_state_old_cancel_then_new_completion_is_completed() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "cancelled"); + record.note_observed_terminal(1, "completed"); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert_eq!(record.terminal_run_lease, Some(1)); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + } + + #[test] + fn terminal_state_new_completion_then_old_cancel_is_completed() { + let mut record = current_terminal_record(); + record.note_observed_terminal(1, "completed"); + record.note_observed_terminal(0, "cancelled"); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + } + + #[test] + fn terminal_state_lifecycle_completion_supersedes_old_db_cancel() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "cancelled"); + record.note_terminal(1); + record.completed_events += 1; + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + record.note_observed_terminal(0, "cancelled"); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + } + + #[test] + fn terminal_state_new_cancel_remains_an_unexpected_outcome() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "completed"); + record.note_observed_terminal(1, "cancelled"); + assert!(record.has(flag::DB_SEEN_CANCELLED)); + assert_eq!(unexpected_terminal_outcomes(&record), vec!["cancelled"]); + } + + #[test] + fn terminal_state_same_lease_conflicting_states_remain_visible() { + for states in [["cancelled", "completed"], ["completed", "cancelled"]] { + let mut record = current_terminal_record(); + for state in states { + record.note_observed_terminal(1, state); + } + assert!(record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + assert_eq!(unexpected_terminal_outcomes(&record), vec!["cancelled"]); + } + } + + struct TestRun { + shared: Arc, + pool: PgPool, + store: QueueStorage, + claim: ClaimedRuntimeJob, + oracle: Oracle, + } + + impl TestRun { + async fn create() -> Self { + Self::create_with(true, true).await + } + + async fn create_with(no_events: bool, sampled: bool) -> Self { + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_oracle_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + db::recreate_database(&config) + .await + .expect("create isolated oracle database"); + let pool = db::connect_pool(&config.database_url, "soak_oracle_regression", 4) + .await + .expect("connect oracle"); + let info = db::server_info(&pool).await.expect("server info"); + let watch = Arc::new(PgWatch::begin(&config, &pool).await.expect("watch server")); + let budget = Budget::new(&info, 0); + let shared = Shared::new(config, info, budget); + let queue = shared + .queues + .register("oracle_proof", Persona::Priya, false); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("store"); + store + .enqueue_batch(&pool, "oracle_proof", 2, 64) + .await + .expect("enqueue sample candidates"); + let claimed = store + .claim_runtime_batch(&pool, "oracle_proof", 64, Duration::ZERO) + .await + .expect("claim sample"); + let claim = claimed + .into_iter() + .find(|claim| samples_terminal_proof(claim.job.id) == sampled) + .expect("immutable job ID selects the requested sample status"); + shared.ledger.register( + claim.job.id, + JobRecord::new( + Persona::Priya, + queue.index, + 0, + expect::COMPLETED, + if no_events { flag::NO_EVENTS } else { 0 }, + ), + ); + // The handler returned, but its caller has not finalized the claim. + let execution = shared.begin_execution(Persona::Priya, &claim.job, no_events); + shared.ledger.mark(claim.job.id, flag::RETURNED_COMPLETED); + drop(execution); + let oracle = Oracle::new(shared.clone(), pool.clone(), watch); + Self { + shared, + pool, + store, + claim, + oracle, + } + } + + async fn drop(self) { + self.pool.close().await; + let admin = db::connect_pool(&self.shared.config.admin_url(), "oracle_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(&format!( + "DROP DATABASE {} WITH (FORCE)", + self.shared.config.database_name() + )) + .execute(&admin) + .await + .expect("drop oracle database"); + admin.close().await; + } + } + + #[tokio::test] + async fn returned_handler_without_finalization_cannot_pass_the_oracle() { + let mut run = TestRun::create().await; + run.oracle.final_check(true).await; + assert!( + run.shared + .violations + .lock() + .expect("violations") + .iter() + .any(|violation| violation.kind == "stuck_job"), + "a handler return must not stand in for a terminal commit" + ); + run.drop().await; + } + + #[tokio::test] + async fn sampled_terminal_pruned_before_observation_is_inconclusive() { + let mut run = TestRun::create().await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + run.store + .reset(&run.pool) + .await + .expect("remove retained evidence"); + run.oracle.final_check(true).await; + assert!( + run.shared.oracle_incomplete.load(Ordering::Relaxed), + "an unobserved pruned result cannot establish a pass or a lost-job claim" + ); + run.drop().await; + } + + #[tokio::test] + async fn observed_terminal_survives_pruning_in_the_client_ledger() { + let mut run = TestRun::create().await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(run + .shared + .ledger + .get(run.claim.job.id) + .expect("sample record") + .has_durable_terminal()); + run.store + .reset(&run.pool) + .await + .expect("prune retained evidence"); + run.oracle.final_check(true).await; + assert!( + run.shared.violations.lock().expect("violations").is_empty(), + "captured durable proof must survive pruning" + ); + run.drop().await; + } + + #[tokio::test] + async fn eventful_observed_terminal_without_callback_survives_pruning() { + let mut run = TestRun::create_with(false, false).await; + assert!(!samples_terminal_proof(run.claim.job.id)); + let completed = run + .store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion without a lifecycle callback"); + assert_eq!(completed, vec![(run.claim.job.id, run.claim.job.run_lease)]); + let durable: bool = sqlx::query_scalar( + "SELECT EXISTS ( + SELECT 1 FROM awa.receipt_completion_batches AS batch, + LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + WHERE item.job_id = $1 AND item.run_lease = $2 + )", + ) + .bind(run.claim.job.id) + .bind(run.claim.job.run_lease) + .fetch_one(&run.pool) + .await + .expect("read committed completion independently of the observer"); + assert!(durable); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + let record = run + .shared + .ledger + .get(run.claim.job.id) + .expect("eventful record"); + assert!(record.has_durable_terminal()); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + assert_eq!(record.completed_events, 0); + run.store + .reset(&run.pool) + .await + .expect("prune retained evidence"); + run.oracle.final_check(true).await; + assert!(run.shared.violations.lock().expect("violations").is_empty()); + run.drop().await; + } + + #[tokio::test] + async fn eventful_return_without_finalization_cannot_gain_terminal_proof() { + let mut run = TestRun::create_with(false, false).await; + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.oracle.final_check(true).await; + assert!(run + .shared + .violations + .lock() + .expect("violations") + .iter() + .any(|violation| violation.kind == "stuck_job")); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_pruned_before_observation_stays_inconclusive() { + let mut run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + run.store + .reset(&run.pool) + .await + .expect("remove evidence before observation"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + run.oracle.final_check(true).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_observer_query_failure_stays_inconclusive() { + let run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("commit completion"); + sqlx::query("ALTER TABLE awa.receipt_completion_batches RENAME TO receipt_completion_batches_unavailable") + .execute(&run.pool).await.expect("inject observer query failure"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_observer_requires_the_exact_attempt() { + let run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("complete old attempt"); + run.shared.ledger.with( + run.claim.job.id, + || unreachable!("registered"), + |record| { + record.latest_run_lease += 1; + }, + ); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + let record = run.shared.ledger.get(run.claim.job.id).expect("record"); + assert!(!record.has_durable_terminal()); + assert!(!record.has(flag::DB_SEEN_COMPLETED)); + run.drop().await; + } + + #[tokio::test] + async fn unsampled_no_events_job_does_not_gain_periodic_terminal_proof() { + let run = TestRun::create_with(true, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("commit completion"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.drop().await; + } + + #[tokio::test] + async fn failed_final_lookup_cannot_omit_unchecked_jobs_and_pass() { + let mut run = TestRun::create().await; + sqlx::query("ALTER TABLE awa.ready_entries RENAME TO ready_entries_unavailable") + .execute(&run.pool) + .await + .expect("inject a lookup failure"); + run.oracle.final_check(true).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + assert_eq!( + run.shared.run.get("final_lookup_errors"), + 1, + "the failed lookup chunk must be accounted for" + ); + run.drop().await; + } +} diff --git a/awa/tests/persona_soak/personas/abandon.rs b/awa/tests/persona_soak/personas/abandon.rs new file mode 100644 index 000000000..7cc48f023 --- /dev/null +++ b/awa/tests/persona_soak/personas/abandon.rs @@ -0,0 +1,1020 @@ +//! Claim and abandon: a runtime that claims as much as it can, holds thousands of jobs +//! with deadlines and heartbeats, and is SIGKILLed again and again. +//! +//! The runtime is a separate process: the test binary started again with only the +//! `claim_abandon_child` test selected, so that a kill leaves its claims genuinely +//! abandoned (no shutdown handshake, connections dropped by the kernel). Between kills a +//! surviving in-process runtime and the other personas keep working on the same +//! database. After the last kill an in-process runtime takes over the abandoned queues +//! and finishes whatever the rescue returns to them. +//! +//! The checks are in `abandon_oracle.rs`: how long the rescue backlog takes to clear +//! after each kill, whether the claim ring keeps rotating, whether any job was rescued +//! while its owner was alive, and whether attempts and the dead-letter queue follow +//! `max_attempts`. + +use super::super::abandon_oracle::{ + epoch_ms, probe_backlog, Attempt, Kill, LiveCheck, Scope, FINISHER, +}; +use super::super::config::{Persona, DEFAULT_HEARTBEAT_SECONDS, DEFAULT_STALENESS_SECONDS}; +use super::super::db::{connect_pool, connect_pool_with_timeout, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{backoff, queue, register_rows, timings, tune}; +use awa::model::{insert_many, insert_many_copy_from_pool}; +use awa::{ + Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, QueueConfig, Worker, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::{Child, Command}; + +/// Zero-deadline claims: written as compact claim batches. +pub const QUEUE_PLAIN: &str = "awa_soak_abandon_plain"; +/// Claims with a deadline: written as one claim row per job. +pub const QUEUE_DEADLINE: &str = "awa_soak_abandon_deadline"; +/// Short jobs served by the survivors, to see claim latency while the backlog drains. +pub const QUEUE_CANARY: &str = "awa_soak_abandon_canary"; + +/// The name of the test the child process runs. +pub const CHILD_TEST: &str = "claim_abandon_child"; +const ENV_URL: &str = "AWA_SOAK_ABANDON_URL"; +const ENV_GENERATION: &str = "AWA_SOAK_ABANDON_GENERATION"; +const ENV_LEDGER: &str = "AWA_SOAK_ABANDON_LEDGER"; +const ENV_WORKERS: &str = "AWA_SOAK_ABANDON_WORKERS"; +const ENV_POOL: &str = "AWA_SOAK_ABANDON_POOL"; +const ENV_HEARTBEAT: &str = "AWA_SOAK_HEARTBEAT_SECONDS"; +const ENV_STALENESS: &str = "AWA_SOAK_STALENESS_SECONDS"; + +/// Longer than any process lives, so a job is only ever finished by the survivors. +const HOLD: Duration = Duration::from_secs(20 * 60); +const DEADLINE: Duration = Duration::from_secs(15 * 60); +const CHILD_POOL: u32 = 16; +const FINISHER_REPLICAS: usize = 2; +const FINISHER_POOL: u32 = 16; +/// Percentage of jobs that may only be attempted twice, so that repeated kills +/// dead-letter them. +const SHORT_RETRY_PERCENT: u64 = 25; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AbandonJob { + pub seq: u64, + pub cycle: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AbandonCanaryJob { + pub seq: u64, +} + +/// The worker in the child process: records every start in the child ledger and then +/// holds the job until the process is killed. +struct ChildWorker { + generation: u32, + ledger: Mutex, + active: Arc, + peak: Arc, +} + +#[async_trait::async_trait] +impl Worker for ChildWorker { + fn kind(&self) -> &'static str { + "abandon_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let line = format!( + "{{\"e\":\"start\",\"id\":{},\"attempt\":{},\"gen\":{},\"q\":\"{}\",\"t\":{}}}\n", + ctx.job.id, + ctx.job.attempt, + self.generation, + ctx.job.queue, + epoch_ms() + ); + // One write per start, straight to the file: nothing is buffered that a kill could + // lose, and the parent never sees a half line it cannot skip. + if let Ok(mut file) = self.ledger.lock() { + let _ = file.write_all(line.as_bytes()); + } + let now = self.active.fetch_add(1, Ordering::Relaxed) + 1; + self.peak.fetch_max(now, Ordering::Relaxed); + tokio::time::sleep(HOLD).await; + self.active.fetch_sub(1, Ordering::Relaxed); + Ok(JobResult::Completed) + } +} + +/// Entry point of the child process; see [`CHILD_TEST`]. Returns at once when the +/// environment is not set up by the parent, so a plain `--ignored` run is unaffected. +pub async fn run_child() { + let (Ok(url), Ok(generation), Ok(ledger), Ok(workers)) = ( + std::env::var(ENV_URL), + std::env::var(ENV_GENERATION), + std::env::var(ENV_LEDGER), + std::env::var(ENV_WORKERS), + ) else { + return; + }; + let generation: u32 = generation.parse().expect("generation"); + let workers: u32 = workers.parse().expect("workers"); + let pool_size: u32 = std::env::var(ENV_POOL) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(CHILD_POOL); + + // Exit with the harness, so a harness that died leaves no process behind. + let parent = std::os::unix::process::parent_id(); + std::thread::spawn(move || loop { + std::thread::sleep(Duration::from_secs(1)); + if std::os::unix::process::parent_id() != parent { + std::process::exit(0); + } + }); + + let pool = connect_pool(&url, &format!("soak_abandon_{generation}"), pool_size) + .await + .expect("child connect"); + let file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&ledger) + .expect("child ledger"); + let active = Arc::new(AtomicI64::new(0)); + let peak = Arc::new(AtomicI64::new(0)); + let claiming = |deadline_duration: Duration| QueueConfig { + max_workers: workers, + poll_interval: Duration::from_millis(50), + deadline_duration, + claimers: 2, + claim_batch_size: 1_000, + ..QueueConfig::default() + }; + // The child keeps away from maintenance leadership: the survivors do the rescuing, as + // the other replicas of an application would. + let seconds = |name: &str, default: u64| { + Duration::from_secs( + std::env::var(name) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(default), + ) + }; + let client = timings( + Client::builder(pool) + .queue(QUEUE_PLAIN, claiming(Duration::ZERO)) + .queue(QUEUE_DEADLINE, claiming(DEADLINE)) + .register_worker(ChildWorker { + generation, + ledger: Mutex::new(file), + active: active.clone(), + peak: peak.clone(), + }), + seconds(ENV_HEARTBEAT, DEFAULT_HEARTBEAT_SECONDS), + seconds(ENV_STALENESS, DEFAULT_STALENESS_SECONDS), + ) + .leader_election_interval(Duration::from_secs(3600)) + .build() + .expect("build child client"); + client.start().await.expect("start child client"); + println!("READY gen={generation}"); + loop { + tokio::time::sleep(Duration::from_secs(2)).await; + println!( + "STAT gen={generation} active={} peak={}", + active.load(Ordering::Relaxed), + peak.load(Ordering::Relaxed) + ); + } +} + +/// Finishes the abandoned jobs once the child processes are gone. +struct FinisherWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for FinisherWorker { + fn kind(&self) -> &'static str { + "abandon_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ClaimAbandon, &ctx.job, true); + self.shared.abandon.note_attempt( + ctx.job.id, + Attempt { + attempt: ctx.job.attempt, + owner: FINISHER, + at_ms: epoch_ms(), + }, + ); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +struct CanaryWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for CanaryWorker { + fn kind(&self) -> &'static str { + "abandon_canary_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ClaimAbandon, &ctx.job, true); + let waited = (chrono::Utc::now() - ctx.job.created_at).num_milliseconds(); + self.shared + .stats(Persona::ClaimAbandon) + .observe("canary_wait_ms", waited); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +/// Copies the running test binary into the report directory. Every process is started +/// from the copy, so a rebuild of the binary during a long run cannot change what is +/// started, or leave nothing to start. +fn copy_executable(directory: &std::path::Path) -> Result { + let copy = directory.join("claim-abandon-child"); + let _ = std::fs::remove_file(©); + let source = std::path::Path::new("/proc/self/exe"); + let source = if source.exists() { + source.to_path_buf() + } else { + std::env::current_exe().map_err(|error| format!("current_exe: {error}"))? + }; + std::fs::copy(&source, ©) + .map_err(|error| format!("copy {} to {}: {error}", source.display(), copy.display()))?; + Ok(copy) +} + +fn spawn_child( + shared: &Arc, + executable: &PathBuf, + generation: u32, + ledger: &PathBuf, + workers_per_queue: u32, + pool_size: u32, +) -> Result { + let url = &shared.config.database_url; + let mut child = Command::new(executable) + .args([ + "--ignored", + "--exact", + CHILD_TEST, + "--nocapture", + "--test-threads=1", + ]) + .env(ENV_URL, url) + .env(ENV_GENERATION, generation.to_string()) + .env(ENV_LEDGER, ledger) + .env(ENV_WORKERS, workers_per_queue.to_string()) + .env(ENV_POOL, pool_size.to_string()) + .env( + ENV_HEARTBEAT, + shared.config.heartbeat_interval.as_secs().to_string(), + ) + .env( + ENV_STALENESS, + shared.config.heartbeat_staleness.as_secs().to_string(), + ) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|error| format!("spawn the claiming process: {error}"))?; + for stream in [ + child + .stdout + .take() + .map(|s| Box::new(s) as Box), + child + .stderr + .take() + .map(|s| Box::new(s) as Box), + ] + .into_iter() + .flatten() + { + let shared = shared.clone(); + tokio::spawn(async move { + let mut lines = BufReader::new(stream).lines(); + while let Ok(Some(line)) = lines.next_line().await { + absorb_output(&shared, &line); + } + }); + } + Ok(child) +} + +fn absorb_output(shared: &Shared, line: &str) { + let stats = shared.stats(Persona::ClaimAbandon); + let field = |name: &str| { + line.split_whitespace() + .find_map(|part| part.strip_prefix(&format!("{name}="))) + .and_then(|value| value.parse::().ok()) + }; + if line.starts_with("STAT ") { + if let Some(active) = field("active") { + shared.abandon.child_active.store(active, Ordering::Relaxed); + } + if let Some(peak) = field("peak") { + stats.set_max("process_peak_held", peak); + } + } else if line.contains("panicked") || line.contains("Error") { + if stats.get("process_error_lines") < 5 { + shared.note("abandon", line.chars().take(200).collect::()); + } + stats.add("process_error_lines", 1); + } +} + +/// Folds the new lines of the child ledger into the harness ledger. +struct Importer { + path: PathBuf, + offset: u64, +} + +impl Importer { + fn import(&mut self, shared: &Shared) { + let Ok(mut file) = std::fs::File::open(&self.path) else { + return; + }; + if file.seek(SeekFrom::Start(self.offset)).is_err() { + return; + } + let mut text = String::new(); + if file.read_to_string(&mut text).is_err() { + return; + } + for line in text.split_inclusive('\n') { + if !line.ends_with('\n') { + break; + } + self.offset += line.len() as u64; + if let Ok(event) = serde_json::from_str::(line) { + self.apply(shared, &event); + } + } + } + + fn apply(&self, shared: &Shared, event: &Value) { + if event.get("e").and_then(Value::as_str) != Some("start") { + return; + } + let (Some(id), Some(attempt), Some(generation), Some(at_ms)) = ( + event.get("id").and_then(Value::as_i64), + event.get("attempt").and_then(Value::as_i64), + event.get("gen").and_then(Value::as_u64), + event.get("t").and_then(Value::as_i64), + ) else { + return; + }; + let queue = event + .get("q") + .and_then(Value::as_str) + .and_then(|name| shared.queues.get(name)); + let queue_index = queue.as_ref().map_or(u32::MAX, |queue| queue.index); + if let Some(queue) = &queue { + queue.started.fetch_add(1, Ordering::Relaxed); + } + let now = shared.now_ms() as u32; + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::ClaimAbandon, + queue_index, + now, + expect::COMPLETED | expect::FAILED, + flag::NO_EVENTS, + ) + }, + |record| { + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(attempt as i16); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + }, + ); + shared.abandon.note_attempt( + id, + Attempt { + attempt: attempt as i16, + owner: generation as u32, + at_ms, + }, + ); + shared + .stats(Persona::ClaimAbandon) + .add("process_executions", 1); + } +} + +fn params(rng: &mut Rng, seq: &mut u64, cycle: u32, count: usize) -> Vec { + (0..count) + .map(|index| { + *seq += 1; + awa::insert::params_with( + &AbandonJob { seq: *seq, cycle }, + InsertOpts { + queue: if index % 2 == 0 { + QUEUE_PLAIN + } else { + QUEUE_DEADLINE + } + .to_string(), + max_attempts: if rng.chance(SHORT_RETRY_PERCENT) { + 2 + } else { + 25 + }, + ..Default::default() + }, + ) + .expect("abandon params") + }) + .collect() +} + +/// Enqueues `count` jobs for the next process to claim, in chunks that COPY handles +/// comfortably. +async fn enqueue_cycle( + shared: &Arc, + pool: &sqlx::PgPool, + rng: &mut Rng, + seq: &mut u64, + cycle: u32, + count: usize, +) -> Result { + let mut enqueued = 0; + while enqueued < count { + let chunk = (count - enqueued).min(10_000); + let batch = params(rng, seq, cycle, chunk); + let rows = insert_many_copy_from_pool(pool, &batch) + .await + .map_err(|error| format!("enqueue: {error}"))?; + for row in &rows { + shared.abandon.register_job(row.id, row.max_attempts); + if let Some(queue) = shared.queues.get(&row.queue) { + queue.enqueued.fetch_add(1, Ordering::Relaxed); + } + } + register_rows( + shared, + Persona::ClaimAbandon, + &rows, + |_| expect::COMPLETED | expect::FAILED, + flag::NO_EVENTS, + ); + shared + .stats(Persona::ClaimAbandon) + .add("enqueued", rows.len() as i64); + enqueued += rows.len(); + } + Ok(enqueued) +} + +async fn canary_producer(shared: Arc, pool: sqlx::PgPool) { + let mut seq = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(200)).await { + let batch: Vec = (0..2) + .map(|_| { + seq += 1; + awa::insert::params_with( + &AbandonCanaryJob { seq }, + InsertOpts { + queue: QUEUE_CANARY.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .expect("canary params") + }) + .collect(); + match insert_many(&pool, &batch).await { + Ok(rows) => { + shared + .stats(Persona::ClaimAbandon) + .add("canary_enqueued", rows.len() as i64); + register_rows( + &shared, + Persona::ClaimAbandon, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS, + ); + } + Err(error) => backoff(&shared, error, Persona::ClaimAbandon).await, + } + } +} + +/// Measures the live process's heartbeat passes from the timestamps its heartbeats write. +/// A pass updates every in-flight claim with the time of its own chunk, so once a pass has +/// finished the oldest heartbeat is newer than anything the previous pass wrote, and the +/// spread from oldest to newest is the pass's duration (to within the last chunk). Only +/// claims taken since the process started are read, so the dead processes' open claims, +/// whose heartbeats stopped, do not count. They are the deadline queue's, which are +/// interleaved with the other queue's in the pass's job-id order and so span the whole pass. +async fn sample_heartbeats( + shared: &Arc, + pool: &sqlx::PgPool, + schema: &str, + since_epoch_ms: i64, +) { + let row: Result<(i64, Option, Option, f64), _> = sqlx::query_as(&format!( + "SELECT count(*)::bigint, + EXTRACT(EPOCH FROM min(a.heartbeat_at))::float8, + EXTRACT(EPOCH FROM max(a.heartbeat_at))::float8, + EXTRACT(EPOCH FROM clock_timestamp())::float8 + FROM {schema}.lease_claims c + JOIN {schema}.attempt_state a ON a.job_id = c.job_id AND a.run_lease = c.run_lease + WHERE c.queue = $1 AND c.closed_at IS NULL + AND c.claimed_at >= to_timestamp($2::float8 / 1000.0)" + )) + .bind(QUEUE_DEADLINE) + .bind(since_epoch_ms as f64) + .fetch_one(pool) + .await; + let Ok((rows, Some(oldest), Some(newest), now)) = row else { + return; + }; + let active = shared.abandon.child_active.load(Ordering::Relaxed); + shared + .abandon + .note_heartbeat_sample(shared, rows, oldest, newest, now, active); +} + +/// Probes the backlog every few seconds and the claim ring every second, until the +/// run is over. +async fn monitor(shared: Arc, queues: Vec) { + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()); + let Ok(pool) = connect_pool_with_timeout( + &shared.config.database_url, + "soak_abandon_monitor", + 2, + Some(Duration::from_secs(60)), + ) + .await + else { + shared.violation( + "persona_failed", + "abandon monitor could not connect", + vec![], + ); + return; + }; + let staleness_secs = shared.config.heartbeat_staleness.as_secs_f64(); + let mut tick = 0u64; + let mut lock_captures = super::super::ring_locks::CaptureBudget::default(); + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + tick += 1; + let ring: Result<(i32, i64, i32), _> = sqlx::query_as(&format!( + "SELECT current_slot, generation, slot_count FROM {schema}.claim_ring_state WHERE singleton" + )) + .fetch_one(&pool) + .await; + let (slot, ring_generation, slot_count) = match ring { + Ok(ring) => ring, + Err(error) => { + shared.incomplete(format!("claim-ring monitor query failed: {error}"), vec![]); + continue; + } + }; + if let Some((_, spawned)) = shared.abandon.live() { + sample_heartbeats(&shared, &pool, &schema, spawned).await; + } + // While the ring is still, look at the slot it has to enter next: live claims + // there explain the stall. + let next_slot = if shared.abandon.ring_stalled() && tick.is_multiple_of(2) { + probe_backlog( + &pool, + &schema, + Scope::Slot((slot + 1).rem_euclid(slot_count)), + staleness_secs, + 0, + ) + .await + .map(|probe| (probe.open - probe.expired, probe.expired)) + .map_err(|error| { + shared.incomplete(format!("claim-slot monitor query failed: {error}"), vec![]) + }) + .ok() + } else { + None + }; + shared + .abandon + .observe_ring(&shared, slot, ring_generation, next_slot); + super::super::ring_locks::observe( + &shared, + &pool, + &schema, + slot_count, + next_slot, + &mut lock_captures, + ) + .await; + if tick.is_multiple_of(3) { + // Read what the live process has started before the probe, so a claim started + // meanwhile cannot make a claim look missing. + let live = shared.abandon.live(); + let live_check = live.map(|(generation, _)| LiveCheck { + generation, + started: shared.abandon.started_by(generation), + }); + let started = Instant::now(); + match probe_backlog( + &pool, + &schema, + Scope::Queues(&queues), + staleness_secs, + live.map_or(0, |(_, spawned)| spawned), + ) + .await + { + Ok(probe) => shared.abandon.observe( + &shared, + &probe, + live_check, + ring_generation, + started.elapsed().as_millis() as i64, + ), + Err(error) => { + shared.incomplete( + format!("abandoned-claim backlog query failed: {error}"), + vec![], + ); + shared.stats(Persona::ClaimAbandon).add("probe_errors", 1); + if shared.stats(Persona::ClaimAbandon).get("probe_errors") <= 3 { + shared.note("abandon", format!("backlog probe failed: {error}")); + } + } + } + } + } + shared.abandon.finish_ring(&shared); + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::ClaimAbandon; + let config = shared.config.clone(); + let staleness_secs = config.heartbeat_staleness.as_secs_f64(); + let plain = shared.queues.register(QUEUE_PLAIN, persona, false); + let deadline = shared.queues.register(QUEUE_DEADLINE, persona, false); + let canary = shared.queues.register(QUEUE_CANARY, persona, false); + let queues = vec![QUEUE_PLAIN.to_string(), QUEUE_DEADLINE.to_string()]; + if config.abandon_kill_every >= DEADLINE { + return Err(format!( + "AWA_SOAK_ABANDON_KILL_SECONDS must stay below the {}s job deadline, or the owner's own deadline rescue would count as a false rescue", + DEADLINE.as_secs() + )); + } + let target = config.abandon_claims.max(2); + let workers_per_queue = (target / 2).max(1) as u32; + + std::fs::create_dir_all(&config.report_dir).map_err(|error| error.to_string())?; + let ledger_path = config.report_dir.join("abandon-child-ledger.jsonl"); + let _ = std::fs::remove_file(&ledger_path); + let executable = copy_executable(&config.report_dir)?; + shared + .abandon + .set_csv(config.report_dir.join("abandon-backlog.csv")); + + let child_pool = shared.budget.reserve("abandon-child", CHILD_POOL, 4); + // Held for the runtime that takes over at the end; Fleet::start replaces it. + shared.budget.reserve( + "abandon-finisher", + FINISHER_POOL * FINISHER_REPLICAS as u32, + FINISHER_REPLICAS as u32 * 2, + ); + + // The survivors: other replicas of the application. They lead maintenance, so they do + // the rescuing, and they serve a canary queue. + let canary_shared = shared.clone(); + let survivor_factory: ClientFactory = Arc::new(move |pool| { + tune( + Client::builder(pool) + .queue(QUEUE_CANARY, queue(4)) + .register_worker(CanaryWorker { + shared: canary_shared.clone(), + }), + &canary_shared, + ) + .build() + .expect("build abandon survivor client") + }); + let survivors = Fleet::start( + &shared, + "abandon-survivor", + persona, + 2, + 8, + 2, + vec![(canary, 4)], + survivor_factory, + ) + .await?; + + let producer_pool = connect_pool( + &config.database_url, + "soak_abandon_producer", + shared.budget.reserve("abandon-producer", 4, 2), + ) + .await + .map_err(|error| error.to_string())?; + let canary_task = tokio::spawn(canary_producer(shared.clone(), producer_pool.clone())); + let monitor_task = tokio::spawn(monitor(shared.clone(), queues.clone())); + let backlog_pool = connect_pool_with_timeout( + &config.database_url, + "soak_abandon_control", + 2, + Some(Duration::from_secs(60)), + ) + .await + .map_err(|error| error.to_string())?; + + let mut rng = Rng::for_task(config.seed, "abandon"); + let mut importer = Importer { + path: ledger_path.clone(), + offset: 0, + }; + let mut seq = 0u64; + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()); + shared.note( + "abandon", + format!( + "armed: {} kill(s) every {}s, each leaving about {target} open claims ({} plain, {} with a deadline); rescue must clear within {}s", + config.abandon_kills, + config.abandon_kill_every.as_secs(), + target / 2, + target - target / 2, + config.abandon_bound.as_secs() + ), + ); + + let mut generation = 0u32; + let mut load_ended = false; + for number in 1..=config.abandon_kills { + if shared.load_stop.is_stopped() { + break; + } + // The first process finds a full queue. Later ones find the rescued jobs of the + // earlier kills and only need fresh jobs to top their capacity up, so that jobs + // are run, and abandoned, more than once. + let fresh = if number == 1 { target } else { target / 2 }; + match enqueue_cycle( + &shared, + &producer_pool, + &mut rng, + &mut seq, + number, + fresh as usize, + ) + .await + { + Ok(count) => shared.note( + "abandon", + format!("enqueued {count} fresh job(s) for process {number}"), + ), + Err(error) => { + backoff(&shared, error, persona).await; + continue; + } + } + generation += 1; + let mut child = spawn_child( + &shared, + &executable, + generation, + &ledger_path, + workers_per_queue, + child_pool, + )?; + shared.abandon.child_alive.store(true, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + shared + .abandon + .set_live(Some((generation, epoch_ms() - 1_000))); + let spawned = Instant::now(); + let mut ramped = false; + shared.note( + "abandon", + format!("process {generation} started with {workers_per_queue} workers on each queue"), + ); + let mut unexpected_exit = false; + loop { + importer.import(&shared); + if shared.load_stop.sleep(Duration::from_secs(1)).await { + load_ended = true; + break; + } + if let Ok(Some(status)) = child.try_wait() { + shared.violation( + "abandon_process_exited", + format!("the claiming process {generation} exited on its own: {status}"), + vec![], + ); + unexpected_exit = true; + break; + } + if !ramped + && shared.abandon.child_active.load(Ordering::Relaxed) >= (target as i64) * 95 / 100 + { + ramped = true; + shared.note( + "abandon", + format!( + "process {generation} holds {} jobs after {:.0}s", + shared.abandon.child_active.load(Ordering::Relaxed), + spawned.elapsed().as_secs_f64() + ), + ); + shared + .stats(persona) + .observe("ramp_seconds", spawned.elapsed().as_secs() as i64); + } + if spawned.elapsed() >= config.abandon_kill_every { + break; + } + } + + let held = shared.abandon.child_active.load(Ordering::Relaxed); + let signalled_epoch_ms = if unexpected_exit { + let at = epoch_ms(); + shared.abandon.note_generation_killed(generation, at); + shared.abandon.set_live(None); + at + } else { + // SIGKILL: no shutdown handshake, no chance to release a claim. + match shared + .abandon + .signal_generation_killed(generation, || child.start_kill()) + { + Ok(at) => at, + Err(error) => { + shared.incomplete( + format!("could not signal claiming process {generation}: {error}"), + vec![], + ); + // No confirmed signal means no valid death boundary. Stop this + // persona and clean up its owned tasks instead of inventing one. + shared.abandon.set_live(None); + shared.abandon.child_alive.store(false, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + drop(child); // kill_on_drop remains a best-effort cleanup. + canary_task.abort(); + monitor_task.abort(); + let _ = canary_task.await; + let _ = monitor_task.await; + survivors + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + backlog_pool.close().await; + return Ok(()); + } + } + }; + let signalled_run_ms = shared.now_ms(); + shared.abandon.child_alive.store(false, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + if !unexpected_exit { + if let Err(error) = child.wait().await { + shared.incomplete( + format!( + "could not reap claiming process {generation} after signaling it: {error}" + ), + vec![], + ); + } + } + importer.import(&shared); + let open_after_kill = probe_backlog( + &backlog_pool, + &schema, + Scope::Queues(&queues), + staleness_secs, + 0, + ) + .await + .map(|probe| probe.open) + .unwrap_or_else(|error| { + shared.incomplete(format!("post-kill backlog query failed: {error}"), vec![]); + -1 + }); + let started = shared.abandon.started_by(generation); + shared.abandon.record_kill(Kill { + number, + generation, + at_run_ms: signalled_run_ms, + signalled_epoch_ms, + started_by_process: started, + rescued_while_alive: shared.abandon.live_lost(generation), + open_after_kill, + expired_peak: 0, + oldest_expired_peak_s: 0.0, + cleared_after_s: None, + reported_slow: false, + load_ended, + }); + shared.stats(persona).add("kills", 1); + shared.stats(persona).add("claims_abandoned", started); + shared.note( + "abandon", + format!( + "SIGKILL process {generation} (kill {number}): {started} handlers started, {held} running at the last report, {open_after_kill} open claims on its queues" + ), + ); + if load_ended { + break; + } + if shared.load_stop.sleep(config.abandon_restart_delay).await { + break; + } + } + + // Everything the kills left behind is finished by an in-process runtime that is + // never killed. + let finisher_shared = shared.clone(); + let finisher_factory: ClientFactory = Arc::new(move |pool| { + tune( + Client::builder(pool) + .queue(QUEUE_PLAIN, queue(256)) + .queue(QUEUE_DEADLINE, queue(256)) + .register_worker(FinisherWorker { + shared: finisher_shared.clone(), + }), + &finisher_shared, + ) + .build() + .expect("build abandon finisher client") + }); + shared.note( + "abandon", + "kills over; starting the runtime that finishes the rescued jobs", + ); + let finisher = Fleet::start( + &shared, + "abandon-finisher", + persona, + FINISHER_REPLICAS, + FINISHER_POOL, + 4, + vec![(plain, 256), (deadline, 256)], + finisher_factory, + ) + .await?; + + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + importer.import(&shared); + } + importer.import(&shared); + canary_task.abort(); + let _ = monitor_task.await; + shared + .abandon + .analyse(&shared, &backlog_pool, &queues) + .await; + finisher + .shutdown_all(Duration::from_secs(20), &shared) + .await; + survivors + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + backlog_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/alex.rs b/awa/tests/persona_soak/personas/alex.rs new file mode 100644 index 000000000..88498cf04 --- /dev/null +++ b/awa/tests/persona_soak/personas/alex.rs @@ -0,0 +1,756 @@ +//! Alex, the operator: under live load, pauses and resumes queues, drains, cancels, +//! moves and reprioritises jobs, retries and purges DLQ rows, runs the admin read +//! queries the UI issues, and rolls a client restart through another persona's fleet. + +use super::super::config::Persona; +use super::super::db::{connect_pool, connect_pool_with_timeout, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::registry::QueueInfo; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::admin::{self, ListJobsFilter}; +use awa::model::batch_operations::{ + get_batch_operation, submit_batch_operation, BatchOperationFilter, BatchOperationSpec, + BatchOperationState, SubmitBatchOperation, +}; +use awa::model::dlq::{self, ListDlqFilter}; +use awa::model::{insert_many, JobState, RetryFromDlqOpts}; +use awa::{Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const QUEUE_A: &str = "awa_soak_alex_a"; +pub const QUEUE_B: &str = "awa_soak_alex_b"; +pub const QUEUE_DRAIN: &str = "awa_soak_alex_drain"; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AlexJob { + pub seq: u64, + pub work_ms: u64, +} + +struct AlexWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for AlexWorker { + fn kind(&self) -> &'static str { + "alex_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Alex, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(100); + tokio::time::sleep(Duration::from_millis(work_ms)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE_A, queue(4)) + .queue(QUEUE_B, queue(4)) + .queue(QUEUE_DRAIN, queue(4)) + .register_worker(AlexWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["alex_job"]) + .build() + .expect("build alex client") + }) +} + +struct Context { + shared: Arc, + pool: PgPool, + diagnostics: PgPool, + stalled_rows: std::sync::Mutex>, + rng: Rng, +} + +/// pid, application, state, wait event, seconds running, blocking pids, query. +type BackendRow = (i32, String, String, String, i32, String, String); + +/// An admin call this slow under load is a finding in its own right. +const SLOW_CALL: Duration = Duration::from_secs(30); +/// Server-side limit on the operator's statements; hitting it is reported as a finding. +const ADMIN_STATEMENT_TIMEOUT: Duration = Duration::from_secs(150); +const STALLED_CALL: Duration = Duration::from_secs(120); + +impl Context { + fn stats(&self) -> &super::super::shared::Stats { + self.shared.stats(Persona::Alex) + } + + async fn timed( + &self, + name: &str, + call: impl std::future::Future>, + ) -> Option { + let started = Instant::now(); + let mut call = std::pin::pin!(call); + let result = tokio::select! { + result = &mut call => result, + _ = tokio::time::sleep(SLOW_CALL) => { + self.diagnose_slow_call(name).await; + call.await + } + }; + let ms = started.elapsed().as_millis() as i64; + if started.elapsed() > STALLED_CALL { + self.shared.finding( + "admin_call_stalled", + format!("admin call {name} took {ms} ms under load"), + self.stalled_rows.lock().expect("stalled rows").clone(), + ); + } + self.stats().observe(name, ms); + if ms > 5_000 { + self.stats().add("slow_admin_calls_over_5s", 1); + self.shared.note("alex", format!("{name} took {ms} ms")); + } + match result { + Ok(value) => Some(value), + Err(error) => { + if error.to_string().contains("statement timeout") { + self.shared.finding( + "admin_call_timeout", + format!( + "admin call {name} was cancelled by the {}s statement timeout", + ADMIN_STATEMENT_TIMEOUT.as_secs() + ), + self.stalled_rows.lock().expect("stalled rows").clone(), + ); + } + self.stats().add(&format!("{name}_errors"), 1); + if self.stats().get(&format!("{name}_errors")) <= 3 { + self.shared.note("alex", format!("{name} failed: {error}")); + } + None + } + } + } + + /// Records what the operator's backend is waiting on and who blocks it. + async fn diagnose_slow_call(&self, name: &str) { + let rows: Result, _> = sqlx::query_as( + "SELECT a.pid, a.application_name, COALESCE(a.state, ''), + COALESCE(a.wait_event_type, '') || ':' || COALESCE(a.wait_event, ''), + COALESCE(extract(epoch FROM now() - a.query_start)::int, 0), + pg_blocking_pids(a.pid)::text, left(a.query, 200) + FROM pg_stat_activity a + WHERE a.pid <> pg_backend_pid() + AND (a.application_name = 'soak_alex_operator' + OR a.pid IN (SELECT unnest(pg_blocking_pids(op.pid)) FROM pg_stat_activity op + WHERE op.application_name = 'soak_alex_operator'))", + ) + .fetch_all(&self.diagnostics) + .await; + let lines: Vec = match rows { + Ok(rows) => rows + .into_iter() + .map(|(pid, app, state, wait, secs, blockers, query)| { + format!("{{\"pid\":{pid},\"application\":\"{app}\",\"state\":\"{state}\",\"wait\":\"{wait}\",\"running_s\":{secs},\"blocked_by\":\"{blockers}\",\"query\":{query:?}}}") + }) + .collect(), + Err(error) => vec![format!("diagnostic query failed: {error}")], + }; + self.shared.note( + "alex", + format!( + "{name} still running after {}s; backend state: {}", + SLOW_CALL.as_secs(), + lines.join(" | ") + ), + ); + *self.stalled_rows.lock().expect("stalled rows") = lines; + } + + /// A random queue of a persona, other than Alex's own. + fn queue_of(&mut self, persona: Persona) -> Option> { + let queues = self.shared.queues.of(persona); + self.rng.pick(&queues).cloned() + } + + async fn admin_reads(&mut self) { + let pool = self.pool.clone(); + self.timed("admin_queue_overviews", admin::queue_overviews(&pool)) + .await; + self.timed("admin_state_counts", admin::state_counts(&pool)) + .await; + self.timed("admin_job_kind_overviews", admin::job_kind_overviews(&pool)) + .await; + self.timed("admin_runtime_overview", admin::runtime_overview(&pool)) + .await; + self.timed( + "admin_queue_runtime_summary", + admin::queue_runtime_summary(&pool), + ) + .await; + self.timed("admin_dlq_depth_by_queue", dlq::dlq_depth_by_queue(&pool)) + .await; + if let Some(queue) = self.queue_of(Persona::Tom) { + let filter = ListJobsFilter { + queue: Some(queue.name.clone()), + limit: Some(50), + ..Default::default() + }; + self.timed("admin_list_jobs", admin::list_jobs(&pool, &filter)) + .await; + } + } + + /// Resumes a queue, retrying: a queue left paused would look like a stall. + async fn resume_with_retry(&self, queue: &str) { + for attempt in 0..8 { + if self + .timed("resume_queue", admin::resume_queue(&self.pool, queue)) + .await + .is_some() + { + return; + } + tokio::time::sleep(Duration::from_secs(2 + attempt)).await; + } + self.shared.violation( + "resume_failed", + format!("queue {queue} could not be resumed after eight attempts"), + vec![], + ); + } + + async fn pause_resume(&mut self) { + let target = match self.rng.below(3) { + 0 => self.shared.queues.get(QUEUE_B), + 1 => self.queue_of(Persona::Tom), + _ => self + .shared + .queues + .get(super::priya::HOT_QUEUE) + .or_else(|| self.queue_of(Persona::Mia)), + }; + let Some(queue) = target else { return }; + queue.hold(); + let pool = self.pool.clone(); + let paused = self + .timed( + "pause_queue", + admin::pause_queue(&pool, &queue.name, Some("soak")), + ) + .await + .is_some(); + if paused { + self.stats().add("pauses", 1); + self.shared + .stop_sleep(Duration::from_secs(self.rng.range(3, 8))) + .await; + } + // A pause whose answer was lost may still have applied, so always resume. + self.resume_with_retry(&queue.name).await; + // Give the claimers a moment to pick the queue back up before stall checks resume. + self.shared.stop_sleep(Duration::from_secs(3)).await; + queue.release(); + } + + async fn drain(&mut self) { + let target = if self.rng.chance(50) { + self.shared.queues.get(QUEUE_DRAIN) + } else { + self.queue_of(Persona::Tom) + }; + let Some(queue) = target else { return }; + // Until the drain returns, any job enqueued so far may legitimately be cancelled. + queue.drained_at_ms.store(i64::MAX / 4, Ordering::Relaxed); + let pool = self.pool.clone(); + if let Some(count) = self + .timed("drain_queue", admin::drain_queue(&pool, &queue.name)) + .await + { + self.stats().add("drains", 1); + self.shared + .note("alex", format!("drained {} ({count} job(s))", queue.name)); + self.stats().add("drained_jobs", count as i64); + } + queue + .drained_at_ms + .store(self.shared.now_ms(), Ordering::Relaxed); + } + + /// Cancels live jobs of other personas, some of them already running. + async fn cancel(&mut self) { + let persona = *self + .rng + .pick(&[Persona::Mia, Persona::Tom, Persona::Sam, Persona::Priya]) + .expect("personas"); + let Some(queue) = self.queue_of(persona) else { + return; + }; + let states = [ + JobState::Available, + JobState::Running, + JobState::Scheduled, + JobState::Retryable, + ]; + let state = *self.rng.pick(&states).expect("states"); + let pool = self.pool.clone(); + let filter = ListJobsFilter { + queue: Some(queue.name.clone()), + state: Some(state), + limit: Some(30), + ..Default::default() + }; + let Some(rows) = self + .timed("admin_list_jobs", admin::list_jobs(&pool, &filter)) + .await + else { + return; + }; + let ids: Vec = rows.iter().map(|row| row.id).take(8).collect(); + if ids.is_empty() { + return; + } + self.shared.note( + "alex", + format!("cancelling jobs {ids:?} ({state:?}) in {}", queue.name), + ); + if self.rng.chance(50) { + let started_ms = self.shared.now_ms() as u32; + let first = ids[0]; + let completed_before = self + .shared + .ledger + .get(first) + .is_some_and(|record| record.completed_events > 0); + if let Some(Some(job)) = self.timed("cancel", admin::cancel(&pool, first)).await { + self.after_cancel(job.id, completed_before, started_ms); + } + } else if let Some(cancelled) = self + .timed("bulk_cancel", admin::bulk_cancel(&pool, &ids)) + .await + { + let now = self.shared.now_ms() as u32; + for job in &cancelled { + self.after_cancel(job.id, false, now); + } + } + } + + fn after_cancel(&self, job_id: i64, completed_before: bool, call_started_ms: u32) { + self.shared.record_operator_cancel( + Persona::Alex, + job_id, + completed_before, + call_started_ms, + ); + } + + async fn dlq_operations(&mut self) { + let pool = self.pool.clone(); + let filter = ListDlqFilter { + queue: Some(super::mia::INFER_QUEUE.to_string()), + limit: Some(20), + ..Default::default() + }; + let Some(rows) = self.timed("dlq_list", dlq::list_dlq(&pool, &filter)).await else { + return; + }; + let ids: Vec = rows.iter().map(|row| row.job.id).collect(); + self.stats().set_max("dlq_listed_max", ids.len() as i64); + match self.rng.below(4) { + 0 => { + if let Some(&id) = ids.first() { + self.shared.ledger.mark(id, flag::DLQ_RETRIED); + let opts = RetryFromDlqOpts::default(); + if self + .timed("dlq_retry", dlq::retry_from_dlq(&pool, id, &opts)) + .await + .is_some() + { + self.stats().add("dlq_retried", 1); + } + } + } + 1 => { + for id in ids.iter().take(3) { + self.shared.ledger.mark(*id, flag::DLQ_PURGED); + if self + .timed("dlq_purge_job", dlq::purge_dlq_job(&pool, *id)) + .await + == Some(true) + { + self.stats().add("dlq_purged", 1); + } + } + } + 2 => { + // Bulk retry: every exhausted Mia job may come back. + self.shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_RETRIED, + ); + let bulk = ListDlqFilter { + queue: Some(super::mia::INFER_QUEUE.to_string()), + ..Default::default() + }; + if let Some(count) = self + .timed( + "dlq_bulk_retry", + dlq::bulk_retry_from_dlq(&pool, &bulk, false), + ) + .await + { + self.stats().add("dlq_bulk_retried", count as i64); + } + } + _ => { + self.shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_PURGED, + ); + let bulk = ListDlqFilter { + queue: Some(super::mia::TRAIN_QUEUE.to_string()), + ..Default::default() + }; + if let Some(count) = self + .timed("dlq_bulk_purge", dlq::purge_dlq(&pool, &bulk, false)) + .await + { + self.stats().add("dlq_bulk_purged", count as i64); + } + } + } + } + + async fn batch_operations(&mut self) { + let (spec, filter, label, source) = if self.rng.chance(50) { + ( + BatchOperationSpec::MoveQueue { + queue: QUEUE_B.to_string(), + priority: None, + }, + BatchOperationFilter { + queue: Some(QUEUE_A.to_string()), + state: Some(JobState::Available), + ..Default::default() + }, + "move_queue", + QUEUE_A, + ) + } else { + ( + BatchOperationSpec::SetPriority { priority: 1 }, + BatchOperationFilter { + queue: Some(QUEUE_B.to_string()), + state: Some(JobState::Available), + ..Default::default() + }, + "set_priority", + QUEUE_B, + ) + }; + // Hold the source queue so a backlog builds for the operation to act on. + let pool = self.pool.clone(); + let held = self.shared.queues.get(source); + if let Some(queue) = &held { + queue.hold(); + let _ = self + .timed( + "pause_queue", + admin::pause_queue(&pool, &queue.name, Some("soak")), + ) + .await; + self.shared.stop_sleep(Duration::from_secs(6)).await; + } + let request = SubmitBatchOperation { + spec, + filter, + submitted_by: Some("soak".into()), + allow_all: false, + }; + self.run_batch(request, label).await; + if let Some(queue) = &held { + self.resume_with_retry(&queue.name).await; + self.shared.stop_sleep(Duration::from_secs(3)).await; + queue.release(); + } + } + + async fn run_batch(&mut self, request: SubmitBatchOperation, label: &str) { + let pool = self.pool.clone(); + let Some(operation) = self + .timed("batch_submit", submit_batch_operation(&pool, request)) + .await + else { + return; + }; + let started = Instant::now(); + loop { + let Some(current) = self + .timed("batch_get", get_batch_operation(&pool, operation.id)) + .await + else { + return; + }; + match current.state { + BatchOperationState::Completed | BatchOperationState::Cancelled => { + self.stats().add(&format!("batch_{label}"), 1); + self.stats() + .add(&format!("batch_{label}_processed"), current.processed); + self.stats() + .observe("batch_operation_ms", started.elapsed().as_millis() as i64); + return; + } + BatchOperationState::Failed => { + self.shared.violation( + "batch_operation_failed", + format!( + "{label} batch operation {} failed: {:?}", + operation.id, current.last_error + ), + vec![], + ); + return; + } + _ => {} + } + if started.elapsed() > Duration::from_secs(90) { + self.stats().add("batch_operations_incomplete", 1); + self.shared.note( + "alex", + format!( + "{label} batch operation {} still {:?} after 90s", + operation.id, current.state + ), + ); + return; + } + if self.shared.stop_sleep(Duration::from_secs(1)).await { + return; + } + } + } + + async fn rolling_restart(&mut self) { + let names = ["mia", "tom", "priya", "sam", "sam-stripe"]; + let name = *self.rng.pick(&names).expect("fleet names"); + let Some(fleet) = self.shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + let Some(&index) = self.rng.pick(&live) else { + return; + }; + self.shared + .note("alex", format!("rolling restart: {name} replica {index}")); + self.shared.chaos_window(Duration::from_secs(20)); + let started = Instant::now(); + let restarted = fleet.restart(index, Duration::from_secs(30)).await; + // If the restarted runtime led, the next election can take a full interval. + self.shared.chaos_window(Duration::from_secs(45)); + match restarted { + Ok(_) => { + self.stats().add("rolling_restarts", 1); + self.stats() + .observe("rolling_restart_ms", started.elapsed().as_millis() as i64); + } + Err(error) => { + self.shared.violation( + "restart_failed", + format!("restarting {name} replica {index}: {error}"), + vec![], + ); + } + } + } +} + +impl Shared { + /// Sleeps unless the load phase ends first; true when it ended. + pub async fn stop_sleep(&self, duration: Duration) -> bool { + self.load_stop.sleep(duration).await + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Alex; + let queue_a = shared.queues.register(QUEUE_A, persona, false); + let queue_b = shared.queues.register(QUEUE_B, persona, false); + let queue_drain = shared.queues.register(QUEUE_DRAIN, persona, false); + let fleet = Fleet::start( + &shared, + "alex", + persona, + 2, + 10, + 2, + vec![(queue_a, 4), (queue_b, 4), (queue_drain, 4)], + client_factory(shared.clone()), + ) + .await?; + + let pool = connect_pool_with_timeout( + &shared.config.database_url, + "soak_alex_operator", + shared.budget.reserve("alex-operator", 4, 2), + Some(ADMIN_STATEMENT_TIMEOUT), + ) + .await + .map_err(|error| error.to_string())?; + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_alex_producer", + shared.budget.reserve("alex-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + + // Alex's own traffic: a stream the move and reprioritise operations act on, and a + // stream that gets drained while it is still being produced. + let producer = { + let shared = shared.clone(); + tokio::spawn(async move { + let mut rng = Rng::for_task(shared.config.seed, "alex-producer"); + let mut seq = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(200)).await { + let mut params = Vec::new(); + for queue in [ + QUEUE_A, + QUEUE_A, + QUEUE_A, + QUEUE_A, + QUEUE_B, + QUEUE_DRAIN, + QUEUE_DRAIN, + QUEUE_DRAIN, + ] { + seq += 1; + params.push( + awa::insert::params_with( + &AlexJob { + seq, + work_ms: rng.range(50, 300), + }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .expect("alex params"), + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(Persona::Alex) + .add("enqueued", rows.len() as i64); + register_rows( + &shared, + Persona::Alex, + &rows, + |_| expect::ANY_TERMINAL, + flag::OPERATOR_MOVED, + ); + } + Err(error) => backoff(&shared, error, Persona::Alex).await, + } + } + producer_pool.close().await; + }) + }; + + let diagnostics = connect_pool( + &shared.config.database_url, + "soak_alex_diagnostics", + shared.budget.reserve("alex-diagnostics", 1, 1), + ) + .await + .map_err(|error| error.to_string())?; + let mut context = Context { + shared: shared.clone(), + pool, + diagnostics, + stalled_rows: Default::default(), + rng: Rng::for_task(shared.config.seed, "alex-operator"), + }; + let load = shared.config.load.as_secs_f64(); + // Scale the cadence down for short runs so each action still happens a few times. + let pace = (load / 900.0).clamp(0.1, 1.0); + let mut schedule: Vec<(&str, f64, Instant)> = vec![ + ( + "reads", + shared.config.admin_read_every.as_secs_f64(), + Instant::now(), + ), + ( + "pause", + 60.0 * pace, + Instant::now() + Duration::from_secs_f64(20.0 * pace), + ), + ( + "drain", + 75.0 * pace, + Instant::now() + Duration::from_secs_f64(35.0 * pace), + ), + ( + "cancel", + 10.0 * pace, + Instant::now() + Duration::from_secs_f64(15.0 * pace), + ), + ( + "dlq", + 45.0 * pace, + Instant::now() + Duration::from_secs_f64(60.0 * pace), + ), + ( + "batch", + 40.0 * pace, + Instant::now() + Duration::from_secs_f64(25.0 * pace), + ), + ( + "restart", + 150.0 * pace, + Instant::now() + Duration::from_secs_f64(90.0 * pace), + ), + ]; + shared.note("alex", "operating: reads, pause/resume, drain, cancel, move/reprioritise, DLQ retry/purge, rolling restarts"); + while !shared.load_stop.is_stopped() { + let now = Instant::now(); + for (name, interval, due) in schedule.iter_mut() { + if *due > now || shared.load_stop.is_stopped() { + continue; + } + *due = Instant::now() + Duration::from_secs_f64(*interval); + match *name { + "reads" => context.admin_reads().await, + "pause" => context.pause_resume().await, + "drain" => context.drain().await, + "cancel" => context.cancel().await, + "dlq" => context.dlq_operations().await, + "batch" => context.batch_operations().await, + "restart" => context.rolling_restart().await, + _ => {} + } + } + shared.load_stop.sleep(Duration::from_millis(250)).await; + } + + let _ = producer.await; + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + context.pool.close().await; + context.diagnostics.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/chaos.rs b/awa/tests/persona_soak/personas/chaos.rs new file mode 100644 index 000000000..fee1c63eb --- /dev/null +++ b/awa/tests/persona_soak/personas/chaos.rs @@ -0,0 +1,242 @@ +//! Chaos: kills a runtime without a shutdown handshake, terminates the maintenance +//! leader's connection, terminates random runtime connections, exhausts a runtime's +//! pool and holds a long idle-in-transaction session. + +use super::super::config::Persona; +use super::super::db::connect_pool; +use super::super::rng::Rng; +use super::super::shared::Shared; +use sqlx::PgPool; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// The advisory lock the maintenance leader holds for as long as it leads. +const MAINTENANCE_LEADER_LOCK_KEY: i64 = 0x4157_415f_4d41_494e; +const RUNTIME_FLEETS: [&str; 6] = ["mia", "tom", "priya", "sam", "sam-stripe", "fanout"]; + +async fn kill_runtime(shared: &Arc, rng: &mut Rng) { + let Some(&name) = rng.pick(&RUNTIME_FLEETS) else { + return; + }; + let Some(fleet) = shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + // Never take the last runtime of a fleet down; the point is a survivor that rescues. + if live.len() < 2 { + return; + } + let Some(&index) = rng.pick(&live) else { + return; + }; + shared.note( + "chaos", + format!("killing {name} replica {index} without shutdown"), + ); + shared.chaos_window(Duration::from_secs(45)); + match fleet.kill(index, Some(Duration::from_secs(4))).await { + Ok(_) => shared.stats(Persona::Chaos).add("runtime_kills", 1), + Err(error) => shared.violation( + "restart_failed", + format!("restarting killed {name} replica {index}: {error}"), + vec![], + ), + } +} + +async fn kill_python(shared: &Arc) { + let handle = shared.python.lock().expect("python handle").clone(); + let Some(handle) = handle else { return }; + shared.note("chaos", "killing the Python process without shutdown"); + shared.chaos_window(Duration::from_secs(45)); + if let Err(error) = handle + .kill_and_restart(shared, Duration::from_secs(4)) + .await + { + shared.violation( + "restart_failed", + format!("restarting the Python process: {error}"), + vec![], + ); + } +} + +async fn terminate_leader(shared: &Arc, pool: &PgPool) { + // Maintenance leadership is an advisory lock held by one runtime connection. + let holders: Result, _> = sqlx::query_as( + "SELECT l.pid, COALESCE(a.application_name, '') + FROM pg_locks l JOIN pg_stat_activity a ON a.pid = l.pid + WHERE l.locktype = 'advisory' AND l.granted AND a.application_name LIKE 'soak\\_%' + AND ((l.classid::bigint << 32) | l.objid::bigint) = $1 + ORDER BY l.pid", + ) + .bind(MAINTENANCE_LEADER_LOCK_KEY) + .fetch_all(pool) + .await; + let Ok(holders) = holders else { return }; + let Some((pid, application)) = holders.first().cloned() else { + shared.note("chaos", "no advisory-lock holder found to terminate"); + return; + }; + shared.note( + "chaos", + format!("terminating maintenance leader backend {pid} ({application})"), + ); + shared.chaos_window(Duration::from_secs(40)); + let _ = sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(pool) + .await; + shared.stats(Persona::Chaos).add("leader_terminations", 1); +} + +async fn terminate_runtime_connections(shared: &Arc, pool: &PgPool, rng: &mut Rng) { + let count = rng.range(3, 12) as i64; + let terminated: Result, _> = sqlx::query_as( + "SELECT pid FROM ( + SELECT pid FROM pg_stat_activity + WHERE application_name LIKE 'soak\\_%' AND application_name NOT LIKE '%producer%' + AND application_name NOT LIKE '%operator%' AND application_name NOT LIKE '%oracle%' + AND application_name NOT LIKE '%chaos%' AND application_name NOT LIKE 'soak\\_abandon\\_%' + AND backend_type = 'client backend' + AND pid <> pg_backend_pid() + ORDER BY random() LIMIT $1 + ) victims WHERE pg_terminate_backend(pid)", + ) + .bind(count) + .fetch_all(pool) + .await; + let killed = terminated.map(|rows| rows.len()).unwrap_or(0); + shared.note( + "chaos", + format!("pg_terminate_backend on {killed} runtime connection(s)"), + ); + shared.chaos_window(Duration::from_secs(30)); + shared + .stats(Persona::Chaos) + .add("connections_terminated", killed as i64); +} + +/// Takes every connection of one runtime's pool for `hold`, starving its heartbeats +/// and claims the way a connection-pool exhaustion bug would. +async fn exhaust_pool(shared: &Arc, rng: &mut Rng, hold: Duration) { + let Some(&name) = rng.pick(&RUNTIME_FLEETS) else { + return; + }; + let Some(fleet) = shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + let Some(&index) = rng.pick(&live) else { + return; + }; + let Some(pool) = fleet.pool_of(index).await else { + return; + }; + let size = pool.options().get_max_connections() as usize; + shared.note( + "chaos", + format!( + "exhausting the {size}-connection pool of {name} replica {index} for {}s", + hold.as_secs() + ), + ); + shared.chaos_window(hold + Duration::from_secs(40)); + let mut held = Vec::new(); + let deadline = Instant::now() + Duration::from_secs(20); + while held.len() < size && Instant::now() < deadline { + if let Ok(Ok(connection)) = + tokio::time::timeout(Duration::from_secs(2), pool.acquire()).await + { + held.push(connection); + } + } + shared.stats(Persona::Chaos).add("pool_exhaustions", 1); + shared + .stats(Persona::Chaos) + .add("pool_connections_held", held.len() as i64); + shared.load_stop.sleep(hold).await; + drop(held); +} + +/// One session that begins a transaction, takes a snapshot and sits idle, pinning the +/// xmin horizon so vacuum and prune cannot reclaim anything newer. +async fn idle_in_transaction(shared: Arc, hold: Duration) { + let Ok(pool) = connect_pool(&shared.config.database_url, "soak_chaos_idle", 1).await else { + return; + }; + let Ok(mut transaction) = pool.begin().await else { + return; + }; + if sqlx::query("SELECT txid_current()") + .execute(&mut *transaction) + .await + .is_err() + { + return; + } + shared.note( + "chaos", + format!( + "holding an idle-in-transaction session for {}s", + hold.as_secs() + ), + ); + shared + .stats(Persona::Chaos) + .add("idle_in_transaction_holds", 1); + shared.load_stop.sleep(hold).await; + let _ = transaction.rollback().await; + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let pool = connect_pool( + &shared.config.database_url, + "soak_chaos_control", + shared.budget.reserve("chaos", 3, 2), + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "chaos"); + let load = shared.config.load.as_secs_f64(); + let pace = (load / 900.0).clamp(0.15, 1.0); + let second = |seconds: f64| Duration::from_secs_f64(seconds * pace); + let mut schedule: Vec<(&str, Duration, Instant)> = vec![ + ("kill", second(120.0), Instant::now() + second(60.0)), + ("python", second(200.0), Instant::now() + second(110.0)), + ("leader", second(150.0), Instant::now() + second(90.0)), + ("terminate", second(75.0), Instant::now() + second(45.0)), + ("pool", second(240.0), Instant::now() + second(150.0)), + ("idle", second(300.0), Instant::now() + second(30.0)), + ]; + shared.note("chaos", "armed: runtime kill, leader termination, connection termination, pool exhaustion, idle transaction"); + while !shared.load_stop.is_stopped() { + let now = Instant::now(); + for (name, interval, due) in schedule.iter_mut() { + if *due > now || shared.load_stop.is_stopped() { + continue; + } + *due = Instant::now() + *interval; + match *name { + "kill" => kill_runtime(&shared, &mut rng).await, + "python" => kill_python(&shared).await, + "leader" => terminate_leader(&shared, &pool).await, + "terminate" => terminate_runtime_connections(&shared, &pool, &mut rng).await, + "pool" => { + exhaust_pool(&shared, &mut rng, Duration::from_secs(20).max(second(20.0))).await + } + "idle" => { + tokio::spawn(idle_in_transaction( + shared.clone(), + second(90.0).max(Duration::from_secs(30)), + )); + } + _ => {} + } + } + shared.load_stop.sleep(Duration::from_millis(500)).await; + } + pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/fanout.rs b/awa/tests/persona_soak/personas/fanout.rs new file mode 100644 index 000000000..8b3146300 --- /dev/null +++ b/awa/tests/persona_soak/personas/fanout.rs @@ -0,0 +1,278 @@ +//! Fanout, replay-style fan-out: each parent job batch-inserts a hundred or so unique +//! child jobs into the queue it runs on, from many concurrent handlers. Some parents +//! then replay a batch that reuses an already-claimed unique key, fail with a unique +//! violation and fall back to inserting that batch row by row. Stresses concurrent +//! batch enqueue into one lane, unique claims and the insert-from-handler path. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert_many, insert_with}; +use awa::{ + AwaError, Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, JobRow, + UniqueOpts, Worker, +}; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const QUEUE: &str = "awa_soak_fanout"; +const MIN_CHILDREN: u64 = 100; +const MAX_CHILDREN: u64 = 185; +const REPLAY_PERCENT: u64 = 10; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct FanoutParent { + pub seq: u64, + pub children: u64, + pub replay: bool, +} + +/// Distinct `(parent_seq, index)` per child, so the unique key (queue + args) is +/// shared only by a deliberate replay. +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct FanoutChild { + pub parent_seq: u64, + pub index: u64, +} + +fn child_opts() -> InsertOpts { + InsertOpts { + queue: QUEUE.to_string(), + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + ..UniqueOpts::default() + }), + ..Default::default() + } +} + +fn is_unique_violation(error: &AwaError) -> bool { + match error { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(database)) => { + database.code().as_deref() == Some("23505") + } + _ => false, + } +} + +struct ParentWorker { + shared: Arc, +} + +impl ParentWorker { + /// Records rows the handler enqueued. A child can run before this is reached; the + /// ledger merges the two. + fn record(&self, rows: &[JobRow]) { + let persona = Persona::Fanout; + if let Some(queue) = self.shared.queues.get(QUEUE) { + queue + .enqueued + .fetch_add(rows.len() as i64, Ordering::Relaxed); + } + self.shared + .stats(persona) + .add("enqueued_children", rows.len() as i64); + register_rows( + &self.shared, + persona, + rows, + |_| expect::COMPLETED, + flag::NO_EVENTS, + ); + } + + async fn insert_batch( + &self, + pool: &PgPool, + parent_seq: u64, + indexes: impl Iterator + Clone, + ) -> Result<(), JobError> { + let stats = self.shared.stats(Persona::Fanout); + let params: Vec = indexes + .clone() + .map(|index| { + awa::insert::params_with(&FanoutChild { parent_seq, index }, child_opts()) + .map_err(JobError::retryable) + }) + .collect::>()?; + let started = Instant::now(); + match insert_many(pool, ¶ms).await { + Ok(rows) => { + stats.observe("batch_insert_ms", started.elapsed().as_millis() as i64); + self.record(&rows); + Ok(()) + } + Err(error) if is_unique_violation(&error) => { + stats.add("batch_unique_conflicts", 1); + for index in indexes { + match insert_with(pool, &FanoutChild { parent_seq, index }, child_opts()).await + { + Ok(row) => { + stats.add("fallback_rows_inserted", 1); + self.record(&[row]); + } + Err(error) if is_unique_violation(&error) => { + stats.add("fallback_rows_skipped", 1) + } + Err(error) => return Err(JobError::retryable(error)), + } + } + Ok(()) + } + Err(error) => { + stats.add("batch_insert_errors", 1); + Err(JobError::retryable(error)) + } + } + } +} + +#[async_trait::async_trait] +impl Worker for ParentWorker { + fn kind(&self) -> &'static str { + "fanout_parent" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::Fanout, &ctx.job, false); + let args: FanoutParent = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + self.insert_batch(ctx.pool(), args.seq, 0..args.children) + .await?; + if args.replay { + // Index 0 is already claimed, so the batch fails as a whole. + let fresh = args.children..args.children * 2 - 1; + self.insert_batch(ctx.pool(), args.seq, std::iter::once(0).chain(fresh)) + .await?; + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +struct ChildWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for ChildWorker { + fn kind(&self) -> &'static str { + "fanout_child" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + // No lifecycle hooks: a hook per child would add a read per job. + let _guard = self.shared.begin_execution(Persona::Fanout, &ctx.job, true); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE, queue(16)) + .register_worker(ParentWorker { + shared: shared.clone(), + }) + .register_worker(ChildWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["fanout_parent"]) + .build() + .expect("build fanout client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Fanout; + let queue_info = shared.queues.register(QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "fanout", + persona, + 2, + 12, + 2, + vec![(queue_info.clone(), 16)], + client_factory(shared.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_fanout_producer", + shared.budget.reserve("fanout-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "fanout-producer"); + let tick = Duration::from_millis(500); + let parents_per_tick = shared.config.scaled(2) as usize; + let mut seq = 0u64; + shared.note( + "fanout", + format!( + "{parents_per_tick} parents per {} ms tick, {MIN_CHILDREN}-{MAX_CHILDREN} unique children each, {REPLAY_PERCENT}% replay a used key", + tick.as_millis() + ), + ); + + while !shared.load_stop.sleep(tick).await { + let mut params = Vec::with_capacity(parents_per_tick); + for _ in 0..parents_per_tick { + seq += 1; + params.push( + awa::insert::params_with( + &FanoutParent { + seq, + children: rng.range(MIN_CHILDREN, MAX_CHILDREN), + replay: rng.chance(REPLAY_PERCENT), + }, + InsertOpts { + queue: QUEUE.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + queue_info + .enqueued + .fetch_add(rows.len() as i64, Ordering::Relaxed); + shared + .stats(persona) + .add("enqueued_parents", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note( + "fanout", + format!("load stopped after {seq} parents; draining"), + ); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/mia.rs b/awa/tests/persona_soak/personas/mia.rs new file mode 100644 index 000000000..07d77f276 --- /dev/null +++ b/awa/tests/persona_soak/personas/mia.rs @@ -0,0 +1,260 @@ +//! Mia, an ML pipeline: jobs from seconds to about a minute (compressed), deadlines +//! that sometimes expire, heartbeats and progress, retries into the DLQ and latency +//! spikes. Stresses deadline and stale rescues racing heartbeats, completions and +//! the DLQ. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert, insert_many}; +use awa::{Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::Duration; + +pub const TRAIN_QUEUE: &str = "awa_soak_mia_train"; +pub const INFER_QUEUE: &str = "awa_soak_mia_infer"; +pub const DEADLINE_QUEUE: &str = "awa_soak_mia_deadline"; +/// Deadline given to jobs in the deadline queue; they run longer than this. +const DEADLINE: chrono::Duration = chrono::Duration::seconds(6); + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct MiaJob { + pub seq: u64, + pub behavior: String, + pub work_ms: u64, +} + +struct MiaWorker { + shared: Arc, +} + +/// Sleeps in slices, reporting progress and stopping early when the job is cancelled +/// or the runtime is shutting down. +async fn work(ctx: &JobContext, total_ms: u64) -> bool { + let slice = 500; + let mut elapsed = 0; + while elapsed < total_ms { + if ctx.is_cancelled() { + return false; + } + let step = slice.min(total_ms - elapsed); + tokio::time::sleep(Duration::from_millis(step)).await; + elapsed += step; + if total_ms >= 4_000 && elapsed % 2_000 < slice { + ctx.set_progress(((elapsed * 100) / total_ms).min(99) as u8, "training"); + let _ = ctx.flush_progress().await; + } + } + true +} + +#[async_trait::async_trait] +impl Worker for MiaWorker { + fn kind(&self) -> &'static str { + "mia_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Mia, &ctx.job, false); + let args: MiaJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + let attempt = ctx.job.attempt; + let finished = match args.behavior.as_str() { + "flaky" if attempt < 3 => { + work(ctx, args.work_ms.min(1_500)).await; + return Err(JobError::retryable_msg("transient model server error")); + } + "doomed" => { + work(ctx, args.work_ms.min(1_500)).await; + return Err(JobError::retryable_msg("model artefact missing")); + } + "cancel" => { + work(ctx, args.work_ms.min(1_000)).await; + return Ok(JobResult::Cancel("rejected by validation".into())); + } + // The first attempt overruns its deadline; later attempts often recover. + "overrun" if attempt >= 2 && !args.seq.is_multiple_of(3) => work(ctx, 1_000).await, + _ => work(ctx, args.work_ms).await, + }; + if !finished { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(TRAIN_QUEUE, queue(24)) + .queue(INFER_QUEUE, queue(16)) + .queue(DEADLINE_QUEUE, queue(12)) + .register_worker(MiaWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["mia_job"]) + .build() + .expect("build mia client") + }) +} + +/// One enqueued job: queue, behaviour, how long it works, deadline, max attempts and +/// the outcomes it may end in. +struct Plan { + queue: &'static str, + behavior: &'static str, + work_ms: u64, + max_attempts: i16, + expect: u8, +} + +fn plan(rng: &mut Rng) -> Plan { + let roll = rng.below(100); + if roll < 45 { + Plan { + queue: INFER_QUEUE, + behavior: "fast", + work_ms: rng.range(100, 2_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 70 { + Plan { + queue: TRAIN_QUEUE, + behavior: "train", + work_ms: rng.range(2_000, 20_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 78 { + // Latency spike: a long job that heartbeats and reports progress throughout. + Plan { + queue: TRAIN_QUEUE, + behavior: "spike", + work_ms: rng.range(30_000, 60_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 88 { + Plan { + queue: DEADLINE_QUEUE, + behavior: "overrun", + work_ms: rng.range(10_000, 25_000), + max_attempts: 3, + expect: expect::COMPLETED | expect::FAILED, + } + } else if roll < 94 { + Plan { + queue: INFER_QUEUE, + behavior: "flaky", + work_ms: 500, + max_attempts: 5, + expect: expect::COMPLETED | expect::FAILED, + } + } else if roll < 98 { + Plan { + queue: INFER_QUEUE, + behavior: "doomed", + work_ms: 500, + max_attempts: 3, + expect: expect::FAILED, + } + } else { + Plan { + queue: INFER_QUEUE, + behavior: "cancel", + work_ms: 500, + max_attempts: 3, + expect: expect::CANCELLED, + } + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Mia; + let train = shared.queues.register(TRAIN_QUEUE, persona, false); + let infer = shared.queues.register(INFER_QUEUE, persona, false); + let deadline = shared.queues.register(DEADLINE_QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "mia", + persona, + 2, + 14, + 2, + vec![(train, 24), (infer, 16), (deadline, 12)], + client_factory(shared.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_mia_producer", + shared.budget.reserve("mia-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + let rate_per_second = (6.0 * shared.config.scale).max(0.5); + let tick = Duration::from_millis(250); + let mut rng = Rng::for_task(shared.config.seed, "mia-producer"); + let mut seq = 0u64; + let mut carry = 0.0; + shared.note("mia", format!("producing about {rate_per_second:.1} jobs/s across {TRAIN_QUEUE}, {INFER_QUEUE} and {DEADLINE_QUEUE}")); + + while !shared.load_stop.sleep(tick).await { + carry += rate_per_second * tick.as_secs_f64(); + let count = carry as usize; + carry -= count as f64; + if count == 0 { + continue; + } + let mut params = Vec::with_capacity(count); + let mut expects = Vec::with_capacity(count); + for _ in 0..count { + seq += 1; + let plan = plan(&mut rng); + let args = MiaJob { + seq, + behavior: plan.behavior.to_string(), + work_ms: plan.work_ms, + }; + params.push( + insert::params_with( + &args, + InsertOpts { + queue: plan.queue.to_string(), + max_attempts: plan.max_attempts, + deadline_duration: (plan.queue == DEADLINE_QUEUE).then_some(DEADLINE), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + expects.push(plan.expect); + shared + .stats(persona) + .add(&format!("planned_{}", plan.behavior), 1); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows(&shared, persona, &rows, |index| expects[index], 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note("mia", format!("load stopped after {seq} jobs; draining")); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(30), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/mod.rs b/awa/tests/persona_soak/personas/mod.rs new file mode 100644 index 000000000..bfebd3462 --- /dev/null +++ b/awa/tests/persona_soak/personas/mod.rs @@ -0,0 +1,117 @@ +//! The personas. Each one owns its queues, job kinds, runtimes and producers, and +//! reports what it did through [`Shared::stats`]. + +pub mod abandon; +pub mod alex; +pub mod chaos; +pub mod fanout; +pub mod mia; +pub mod priya; +pub mod python; +pub mod sam; +pub mod shardops; +pub mod tom; +pub mod ui; + +use super::config::Persona; +use super::ledger::{expect, JobRecord}; +use super::shared::Shared; +use awa::{ClientBuilder, JobRow, QueueConfig}; +use std::sync::Arc; +use std::time::Duration; + +/// Registers the lifecycle event hook for each kind so the ledger learns about +/// terminal outcomes as they commit. +pub fn attach_events( + mut builder: ClientBuilder, + shared: &Arc, + kinds: &[&str], +) -> ClientBuilder { + for kind in kinds { + let shared = shared.clone(); + builder = builder.on_event_kind(*kind, move |event| { + let shared = shared.clone(); + async move { shared.on_event(&event) } + }); + } + builder +} + +/// Runtime timings shared by every fleet: fast enough that rescues, promotion and +/// rotation happen many times in a short run, but with realistic ratios between them. +pub fn tune(builder: ClientBuilder, shared: &Shared) -> ClientBuilder { + // The maintenance leader evaluates only the schedules its own runtime declares, and + // any runtime can lead, so every fleet declares the same set. + let builder = if shared.config.has(Persona::Tom) { + tom::declare_cron(builder) + } else { + builder + }; + let builder = timings( + builder, + shared.config.heartbeat_interval, + shared.config.heartbeat_staleness, + ); + match shared.config.rescue_interval { + Some(interval) => builder + .heartbeat_rescue_interval(interval) + .deadline_rescue_interval(interval), + None => builder, + } +} + +/// The timings of [`tune`] that do not depend on the run, for runtimes that live in a +/// process of their own. +pub fn timings( + builder: ClientBuilder, + heartbeat_interval: Duration, + heartbeat_staleness: Duration, +) -> ClientBuilder { + builder + .heartbeat_interval(heartbeat_interval) + .heartbeat_staleness(heartbeat_staleness) + .heartbeat_rescue_interval(Duration::from_secs(3)) + .deadline_rescue_interval(Duration::from_secs(2)) + .promote_interval(Duration::from_millis(250)) + .runtime_snapshot_interval(Duration::from_secs(5)) + .dlq_enabled_by_default(true) +} + +pub fn queue(max_workers: u32) -> QueueConfig { + QueueConfig { + max_workers, + ..QueueConfig::default() + } +} + +/// Records freshly enqueued rows in the ledger. +pub fn register_rows( + shared: &Shared, + persona: Persona, + rows: &[JobRow], + expect_for: impl Fn(usize) -> u8, + flags: u32, +) { + let now = shared.now_ms() as u32; + for (index, row) in rows.iter().enumerate() { + let queue = shared + .queues + .get(&row.queue) + .map_or(u32::MAX, |queue| queue.index); + shared.ledger.register( + row.id, + JobRecord::new(persona, queue, now, expect_for(index), flags), + ); + } +} + +pub const ANY: u8 = expect::ANY_TERMINAL; + +/// Short pause used by producers after an error, so a dead connection does not spin. +pub async fn backoff(shared: &Shared, error: impl std::fmt::Display, persona: Persona) { + shared.stats(persona).add("enqueue_errors", 1); + if shared.stats(persona).get("enqueue_errors") <= 5 { + shared.note(persona.name(), format!("enqueue error: {error}")); + } + shared.load_stop.sleep(Duration::from_millis(500)).await; +} diff --git a/awa/tests/persona_soak/personas/priya.rs b/awa/tests/persona_soak/personas/priya.rs new file mode 100644 index 000000000..e674d9a31 --- /dev/null +++ b/awa/tests/persona_soak/personas/priya.rs @@ -0,0 +1,203 @@ +//! Priya, firehose: one hot queue, large batch inserts (above the bind-parameter +//! limit), COPY inserts, tiny jobs and several clients. Stresses bind limits, ring +//! rotation and prune under load and claim contention. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{backoff, queue, register_rows, tune}; +use awa::model::{insert_many, insert_many_copy_from_pool}; +use awa::{Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const HOT_QUEUE: &str = "awa_soak_priya_hot"; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct PriyaJob { + pub seq: u64, + pub source: String, +} + +struct PriyaWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for PriyaWorker { + fn kind(&self) -> &'static str { + "priya_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Priya, &ctx.job, true); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + // No lifecycle hooks: a hook per completion would add a read per job and + // change the load this persona exists to apply. + tune(Client::builder(pool).queue(HOT_QUEUE, queue(48)), &shared) + .register_worker(PriyaWorker { + shared: shared.clone(), + }) + .build() + .expect("build priya client") + }) +} + +fn batch(seq: &mut u64, count: usize, source: &str) -> Vec { + (0..count) + .map(|_| { + *seq += 1; + awa::insert::params_with( + &PriyaJob { + seq: *seq, + source: source.to_string(), + }, + InsertOpts { + queue: HOT_QUEUE.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .expect("priya params") + }) + .collect() +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Priya; + let hot = shared.queues.register(HOT_QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "priya", + persona, + 3, + 12, + 3, + vec![(hot, 48)], + client_factory(shared.clone()), + ) + .await?; + + let pool_size = shared.budget.reserve("priya-producer", 6, 2); + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_priya_producer", + pool_size, + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "priya-producer"); + let burst_size = shared.config.scaled(100_000) as usize; + let copy_size = shared.config.scaled(10_000) as usize; + let steady_per_tick = shared.config.scaled(30) as usize; + let burst_every = Duration::from_secs(180); + let copy_every = Duration::from_secs(45); + let mut next_burst = Instant::now() + Duration::from_secs(10); + let mut next_copy = Instant::now() + Duration::from_secs(25); + let mut seq = 0u64; + shared.note( + "priya", + format!("steady {} jobs/s, {burst_size}-row insert_many burst every {}s, {copy_size}-row COPY every {}s", steady_per_tick * 10, burst_every.as_secs(), copy_every.as_secs()), + ); + + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + let now = Instant::now(); + if now >= next_burst { + next_burst = now + burst_every; + // Vary the size around the target so chunk boundaries move between bursts. + let size = burst_size + rng.below(5_000) as usize; + let params = batch(&mut seq, size, "burst"); + let started = Instant::now(); + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + shared.stats(persona).add("bursts", 1); + shared + .stats(persona) + .observe("burst_insert_ms", started.elapsed().as_millis() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_BURST, + ); + shared.note( + "priya", + format!( + "burst of {} rows inserted in {:?}", + rows.len(), + started.elapsed() + ), + ); + } + Err(error) if shared.chaos_active() => { + backoff(&shared, format!("burst of {size}: {error}"), persona).await + } + Err(error) => { + shared.violation( + "insert_many_failed", + format!("insert_many of {size} rows failed: {error}"), + vec![], + ); + } + } + } + if now >= next_copy { + next_copy = now + copy_every; + let params = batch(&mut seq, copy_size, "copy"); + let started = Instant::now(); + match insert_many_copy_from_pool(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + shared.stats(persona).add("copy_batches", 1); + shared + .stats(persona) + .observe("copy_insert_ms", started.elapsed().as_millis() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_COPY, + ); + } + Err(error) => { + backoff(&shared, format!("copy of {copy_size}: {error}"), persona).await + } + } + } + let params = batch(&mut seq, steady_per_tick, "steady"); + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_STEADY, + ); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note("priya", format!("load stopped after {seq} jobs; draining")); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/python.rs b/awa/tests/persona_soak/personas/python.rs new file mode 100644 index 000000000..7ece76691 --- /dev/null +++ b/awa/tests/persona_soak/personas/python.rs @@ -0,0 +1,887 @@ +//! Python users: a separate process built on the `awa-python` bindings, started and +//! stopped by the harness. It produces (single inserts, `insert_many` above the old +//! bind-parameter limit, COPY, and enqueues inside transactions that sometimes roll +//! back), runs async workers (completing, retryable, `RetryAfter`, `Snooze`, progress, +//! deadline overruns and a handler that ignores its deadline) and exercises the +//! migrate-then-query startup path. One queue is shared with a Rust runtime so Rust and +//! Python workers claim from the same lanes. The process writes its own per-job ledger, +//! which this module folds into the harness ledger for the oracle. + +use super::super::config::Persona; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::python_probe::{self, Evidence, Phase}; +use super::super::shared::Shared; +use super::{attach_events, queue, tune}; +use awa::{Client, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashSet; +use std::io::{BufRead, Seek, SeekFrom}; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::{Child, Command}; + +pub const QUEUE_WORK: &str = "awa_soak_py_work"; +pub const QUEUE_DEADLINE: &str = "awa_soak_py_deadline"; +pub const QUEUE_SHARED: &str = "awa_soak_py_shared"; +const SCRIPT: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/persona_soak/python/soak_python.py" +); + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct PySharedJob { + pub seq: u64, + pub behavior: String, + pub work_ms: u64, +} + +struct SharedQueueWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for SharedQueueWorker { + fn kind(&self) -> &'static str { + "py_shared_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::Python, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(50); + match ctx.job.args["behavior"].as_str().unwrap_or("ok") { + "retry" if ctx.job.attempt < 3 => { + tokio::time::sleep(Duration::from_millis(50)).await; + return Err(JobError::retryable_msg("transient failure")); + } + "progress" => { + for step in 1..=4u8 { + tokio::time::sleep(Duration::from_millis(work_ms / 4)).await; + ctx.set_progress(step * 25, "step"); + let _ = ctx.flush_progress().await; + } + } + _ => tokio::time::sleep(Duration::from_millis(work_ms)).await, + } + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +/// Handle to the running Python process, shared with the chaos persona. +pub struct PythonProcess { + child: Option, + ledger_path: PathBuf, + python: PathBuf, + restarts: u32, +} + +fn spawn( + shared: &Arc, + python: &PathBuf, + ledger_path: &PathBuf, + role: &str, +) -> Result { + spawn_process(shared, python, ledger_path, role, "soak_python", None) +} + +fn spawn_process( + shared: &Arc, + python: &PathBuf, + ledger_path: &PathBuf, + role: &str, + application_name: &str, + probe_phase: Option>, +) -> Result { + let mut child = Command::new(python) + .arg(SCRIPT) + .env( + "SOAK_DATABASE_URL", + format!( + "{}{}application_name={}", + shared.config.database_url, + if shared.config.database_url.contains('?') { + '&' + } else { + '?' + }, + application_name + ), + ) + .env("SOAK_LEDGER", ledger_path) + .env("SOAK_SEED", shared.config.seed.to_string()) + .env("SOAK_SCALE", shared.config.scale.to_string()) + .env( + "SOAK_PACE", + (shared.config.load.as_secs_f64() / 900.0) + .clamp(0.1, 1.0) + .to_string(), + ) + .env("SOAK_ROLE", role) + .env( + "SOAK_DRAIN", + if shared.load_stop.is_stopped() { + "1" + } else { + "0" + }, + ) + .env("SOAK_QUEUE_WORK", QUEUE_WORK) + .env("SOAK_QUEUE_DEADLINE", QUEUE_DEADLINE) + .env("SOAK_QUEUE_SHARED", QUEUE_SHARED) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|error| format!("spawn {}: {error}", python.display()))?; + for stream in [ + child + .stdout + .take() + .map(|s| Box::new(s) as Box), + child + .stderr + .take() + .map(|s| Box::new(s) as Box), + ] + .into_iter() + .flatten() + { + let shared = shared.clone(); + let probe_phase = probe_phase.clone(); + let application_name = application_name.to_string(); + tokio::spawn(async move { + let mut lines = BufReader::new(stream).lines(); + while let Ok(Some(line)) = lines.next_line().await { + if let (Some(phase), Some(marker)) = (&probe_phase, Phase::marker(&line)) { + // A close marker alone does not prove migrate-and-query + // reached readiness. IDLE also follows that completed call. + if marker != Phase::Closed || phase.load(Ordering::Acquire) > 0 { + phase.fetch_max(marker as u8, Ordering::Release); + } + shared.note("python", format!("{application_name}: {line}")); + } + absorb_output(&shared, &line); + } + }); + } + Ok(child) +} + +/// Turns the interesting stdout lines of the Python process into report counters. +fn absorb_output(shared: &Shared, line: &str) { + let stats = shared.stats(Persona::Python); + let field = |name: &str| { + line.split_whitespace() + .find_map(|part| part.strip_prefix(&format!("{name}="))) + .and_then(|value| value.parse::().ok()) + }; + if line.starts_with("STAT ") { + if let Some(active) = field("active") { + shared + .python_active + .store(active as i64, std::sync::atomic::Ordering::Relaxed); + } + if let Some(max_active) = field("max_active") { + stats.set_max("python_max_concurrent_handlers", max_active as i64); + } + if let Some(errors) = field("errors") { + stats.set_max("python_errors", errors as i64); + } + } else if line.starts_with("BATCH ") { + if let (Some(rows), Some(seconds)) = (field("rows"), field("seconds")) { + stats.add("python_batches", 1); + stats.add("python_batch_rows", rows as i64); + stats.observe("python_batch_insert_ms", (seconds * 1000.0) as i64); + } + } else if line.starts_with("MIGRATED") { + stats.add("python_migrations", 1); + } else if line.starts_with("ERROR") || line.contains("Traceback") { + if stats.get("python_error_lines") < 5 { + shared.note("python", line.chars().take(200).collect::()); + } + stats.add("python_error_lines", 1); + } +} + +pub struct PythonHandle { + process: tokio::sync::Mutex, +} + +// Used in both phases: after load stops a dead process still needs a worker +// replacement, but spawn_process disables its producers through SOAK_DRAIN. +fn restart_if_needed( + child: &mut Option, + start: impl FnOnce() -> Result, +) -> Result, String> { + let reason = match child.as_mut() { + Some(child) => match child.try_wait().map_err(|error| error.to_string())? { + Some(status) => format!("process exited with {status}"), + None => return Ok(None), + }, + None => "process was absent".to_string(), + }; + child.take(); + *child = Some(start()?); + Ok(Some(reason)) +} + +fn finished_probe_errors(probes: &mut tokio::task::JoinSet<()>) -> Vec { + let mut errors = Vec::new(); + while let Some(result) = probes.try_join_next() { + if let Err(error) = result { + errors.push(error.to_string()); + } + } + errors +} + +async fn begin_drain( + stop_producers: impl std::future::Future, + probes: &mut tokio::task::JoinSet<()>, +) -> Vec { + // Startup probes may still be blocked. Stop production before waiting for + // them, and leave their observation tasks running beside worker recovery. + stop_producers.await; + finished_probe_errors(probes) +} + +impl PythonHandle { + async fn ensure_running(&self, shared: &Arc) { + let mut process = self.process.lock().await; + let python = process.python.clone(); + let ledger_path = process.ledger_path.clone(); + match restart_if_needed(&mut process.child, || { + spawn(shared, &python, &ledger_path, "full") + }) { + Ok(Some(reason)) => { + process.restarts += 1; + shared + .stats(Persona::Python) + .add("python_unexpected_exits", 1); + shared.note( + "python", + format!( + "{reason}; restarted (draining={})", + shared.load_stop.is_stopped() + ), + ); + } + Ok(None) => {} + Err(error) => shared.incomplete( + format!("Python worker liveness/restart check failed: {error}"), + vec![], + ), + } + } + /// Kills the process without a shutdown handshake and starts a new one after `downtime`. + pub async fn kill_and_restart( + &self, + shared: &Arc, + downtime: Duration, + ) -> Result<(), String> { + let mut process = self.process.lock().await; + if let Some(mut child) = process.child.take() { + let _ = child.kill().await; + } + tokio::time::sleep(downtime).await; + shared + .python_active + .store(0, std::sync::atomic::Ordering::Relaxed); + process.restarts += 1; + let child = spawn( + shared, + &process.python.clone(), + &process.ledger_path.clone(), + "full", + )?; + process.child = Some(child); + shared.stats(Persona::Python).add("python_kills", 1); + Ok(()) + } +} + +/// Folds new lines of the Python ledger into the harness ledger. +struct Importer { + path: PathBuf, + offset: u64, + rolled_back: HashSet, +} + +impl Importer { + fn import(&mut self, shared: &Shared) { + let Ok(mut file) = std::fs::File::open(&self.path) else { + return; + }; + if file.seek(SeekFrom::Start(self.offset)).is_err() { + return; + } + let mut reader = std::io::BufReader::new(file); + let mut line = String::new(); + loop { + line.clear(); + match reader.read_line(&mut line) { + Ok(0) | Err(_) => break, + Ok(read) => { + if !line.ends_with('\n') { + break; + } + self.offset += read as u64; + if let Ok(event) = serde_json::from_str::(&line) { + self.apply(shared, &event); + } + } + } + } + } + + fn apply(&mut self, shared: &Shared, event: &Value) { + let now = shared.now_ms() as u32; + let id = event.get("id").and_then(Value::as_i64); + let queue_of = |name: Option<&str>| { + name.and_then(|name| shared.queues.get(name)) + .map_or(u32::MAX, |queue| queue.index) + }; + match event.get("e").and_then(Value::as_str) { + Some("enq") => { + let Some(id) = id else { return }; + let allowed = event.get("x").and_then(Value::as_u64).unwrap_or(7) as u8; + let queue = shared + .queues + .get(event.get("q").and_then(Value::as_str).unwrap_or("")); + if let Some(queue) = &queue { + queue + .enqueued + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + let index = queue.map_or(u32::MAX, |queue| queue.index); + let mut record = + JobRecord::new(Persona::Python, index, now, allowed, flag::NO_EVENTS); + record.max_attempt = 0; + // Keep what the Rust side may have seen already (shared queue). + shared.ledger.with( + id, + || record.clone(), + |existing| { + existing.expect = allowed; + existing.flags &= !flag::UNREGISTERED; + existing.persona = Persona::Python; + }, + ); + shared.stats(Persona::Python).add("python_enqueued", 1); + } + Some("start") => { + let Some(id) = id else { return }; + let queue = queue_of(event.get("q").and_then(Value::as_str)); + let attempt = event.get("attempt").and_then(Value::as_i64).unwrap_or(0) as i16; + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + queue, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| { + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(attempt); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + }, + ); + if let Some(queue) = event + .get("q") + .and_then(Value::as_str) + .and_then(|name| shared.queues.get(name)) + { + queue + .started + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + shared.stats(Persona::Python).add("python_executions", 1); + } + Some("end") => { + let Some(id) = id else { return }; + let result = event.get("r").and_then(Value::as_str).unwrap_or(""); + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + u32::MAX, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| match result { + // A completed return is evidence only for jobs that must complete; for + // the rest the outcome is read from the database, since a late return + // after a deadline rescue is discarded by the store. + "completed" if record.expect == expect::COMPLETED => { + record.flags |= flag::RETURNED_COMPLETED + } + "cancel" => { + record.flags |= flag::HANDLER_CANCELLED; + record.cancel_events = record.cancel_events.saturating_add(1); + } + _ => {} + }, + ); + } + Some("rolled_back") => { + if let Some(ids) = event.get("ids").and_then(Value::as_array) { + for id in ids.iter().filter_map(Value::as_i64) { + self.rolled_back.insert(id); + } + shared + .stats(Persona::Python) + .add("python_rolled_back_jobs", ids.len() as i64); + } + } + Some("error") => shared.stats(Persona::Python).add("python_logged_errors", 1), + _ => {} + } + } +} + +/// A migrate-then-query run in a throwaway process; its database sessions must not be +/// left stranded while it idles, and must be gone once it exits. +async fn migrate_probe(shared: &Arc, python: &PathBuf, ledger_path: &PathBuf) { + let stats = shared.stats(Persona::Python); + let application_name = format!("soak_python_probe_{}", uuid::Uuid::new_v4().simple()); + let phase = Arc::new(AtomicU8::new(Phase::Migrating as u8)); + let mut child = match spawn_process( + shared, + python, + ledger_path, + "migrate_probe", + &application_name, + Some(phase.clone()), + ) { + Ok(child) => child, + Err(error) => { + shared.incomplete( + format!("Python migrate probe could not start: {error}"), + vec![], + ); + return; + } + }; + let pool = match super::super::db::connect_pool( + &shared.config.database_url, + "soak_python_probe_watch", + 1, + ) + .await + { + Ok(pool) => pool, + Err(error) => { + shared.incomplete( + format!("Python probe observer could not connect: {error}"), + vec![], + ); + let _ = child.kill().await; + return; + } + }; + let mut evidence = Evidence::default(); + let deadline = Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(45)); + let exit_success = loop { + match child.try_wait() { + Ok(Some(status)) => { + if !status.success() { + evidence.incomplete(format!("Python migrate probe exited with {status}")); + } + break status.success(); + } + Err(error) => { + evidence.incomplete(format!( + "Python migrate probe exit could not be read: {error}" + )); + let _ = child.kill().await; + break false; + } + Ok(None) if Instant::now() >= deadline => { + evidence + .incomplete("Python migrate-then-query probe did not exit within its deadline"); + let _ = child.kill().await; + break false; + } + Ok(None) => {} + } + // Capture readiness before starting the query: a MIGRATED marker that + // arrives during it must not relabel an earlier migration snapshot. + let before_query = Phase::from_byte(phase.load(Ordering::Acquire)); + evidence.observe( + before_query, + python_probe::sessions(&pool, &application_name).await, + ); + tokio::time::sleep(Duration::from_millis(500)).await; + }; + evidence + .finish_exit(&phase, exit_success, || { + python_probe::sessions(&pool, &application_name) + }) + .await; + stats.add("python_migrate_probes", 1); + stats.set_max("python_probe_peak_sessions", evidence.peak_sessions); + stats.set_max( + "python_probe_migrating_idle_transactions", + evidence.migrating_idle_transactions, + ); + stats.add("python_probe_ready_samples", evidence.ready_samples as i64); + if !evidence.complete() { + shared.incomplete( + format!( + "Python migrate probe {application_name} has incomplete evidence: {evidence:?}" + ), + vec![], + ); + } + if evidence.violation() { + shared.violation( + "python_connections_stranded", + format!("Python migrate probe {application_name}: peak={} against pool4, idle-in-transaction after ready={}, sessions after exit={:?}", evidence.peak_sessions, evidence.ready_idle_transactions, evidence.remaining), + vec![], + ); + } + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Python; + let python = shared + .config + .python_bin + .clone() + .ok_or("the python persona needs AWA_SOAK_PYTHON, the interpreter of a virtualenv built with `maturin develop` (see scripts/persona-soak.sh)")?; + let shared_queue = shared.queues.register(QUEUE_SHARED, persona, false); + shared.queues.register(QUEUE_WORK, persona, false); + shared.queues.register(QUEUE_DEADLINE, persona, false); + + // The Rust runtime that shares the queue with the Python workers. + let rust_shared = shared.clone(); + let factory: super::super::db::ClientFactory = Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE_SHARED, queue(6)) + .register_worker(SharedQueueWorker { + shared: rust_shared.clone(), + }); + attach_events( + tune(builder, &rust_shared), + &rust_shared, + &["py_shared_job"], + ) + .build() + .expect("build python shared-queue client") + }); + let fleet = Fleet::start( + &shared, + "pyrust", + persona, + 1, + 8, + 2, + vec![(shared_queue, 6)], + factory, + ) + .await?; + + std::fs::create_dir_all(&shared.config.report_dir).map_err(|error| error.to_string())?; + let ledger_path = shared.config.report_dir.join("python-ledger.jsonl"); + let _ = std::fs::remove_file(&ledger_path); + let child = spawn(&shared, &python, &ledger_path, "full")?; + let handle = Arc::new(PythonHandle { + process: tokio::sync::Mutex::new(PythonProcess { + child: Some(child), + ledger_path: ledger_path.clone(), + python: python.clone(), + restarts: 0, + }), + }); + shared + .python + .lock() + .expect("python handle") + .replace(handle.clone()); + shared.note("python", format!("{} started; queues {QUEUE_WORK}, {QUEUE_DEADLINE}, {QUEUE_SHARED} (shared with a Rust runtime)", python.display())); + + let tail_pool = + super::super::db::connect_pool(&shared.config.database_url, "soak_python_tail", 1) + .await + .map_err(|error| error.to_string())?; + let mut tail_tick = 0u64; + let mut probes = tokio::task::JoinSet::new(); + let mut importer = Importer { + path: ledger_path.clone(), + offset: 0, + rolled_back: HashSet::new(), + }; + let pace = (shared.config.load.as_secs_f64() / 900.0).clamp(0.1, 1.0); + let mut next_probe = Instant::now() + Duration::from_secs_f64(30.0 * pace.max(0.3)); + let probe_every = Duration::from_secs_f64(150.0 * pace.max(0.3)); + while !shared.load_stop.sleep(Duration::from_secs(1)).await { + importer.import(&shared); + tail_terminal_rows(&shared, &tail_pool, &mut tail_tick).await; + handle.ensure_running(&shared).await; + if Instant::now() >= next_probe { + next_probe = Instant::now() + probe_every; + let (shared, python, ledger_path) = + (shared.clone(), python.clone(), ledger_path.clone()); + probes.spawn(async move { migrate_probe(&shared, &python, &ledger_path).await }); + } + } + + // Stop producing immediately; outstanding migration probes must not delay + // this signal or the liveness/evidence checks during the fixed drain grace. + let errors = begin_drain( + async { + if let Some(pid) = handle + .process + .lock() + .await + .child + .as_ref() + .and_then(Child::id) + { + match std::process::Command::new("kill") + .args(["-USR1", &pid.to_string()]) + .status() + { + Ok(status) if status.success() => {} + result => shared.incomplete( + format!("Python stop-production signal failed: {result:?}"), + vec![], + ), + } + } + }, + &mut probes, + ) + .await; + for error in errors { + shared.incomplete(format!("Python migration probe failed: {error}"), vec![]); + } + shared.note("python", "load stopped; draining"); + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + handle.ensure_running(&shared).await; + importer.import(&shared); + tail_terminal_rows(&shared, &tail_pool, &mut tail_tick).await; + for error in finished_probe_errors(&mut probes) { + shared.incomplete(format!("Python migration probe failed: {error}"), vec![]); + } + } + // Finished tasks retain their result. Only a task actually cancelled at + // shutdown lacks the required evidence; do not add time to the drain grace. + probes.abort_all(); + while let Some(result) = probes.join_next().await { + if let Err(error) = result { + shared.incomplete( + format!("Python migration probe unfinished or failed at drain shutdown: {error}"), + vec![], + ); + } + } + importer.import(&shared); + if let Some(mut child) = handle.process.lock().await.child.take() { + let _ = child.kill().await; + } + shared.python.lock().expect("python handle").take(); + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + importer.import(&shared); + check_rolled_back(&shared, &importer.rolled_back).await?; + Ok(()) +} + +/// Python jobs report no lifecycle events, and ring rotation prunes their terminal rows +/// within seconds, so their terminal outcomes are read from the store as they appear. +async fn tail_terminal_rows(shared: &Arc, pool: &sqlx::PgPool, tick: &mut u64) { + *tick += 1; + if !(*tick).is_multiple_of(3) { + return; + } + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".into()); + let queues = vec![ + QUEUE_WORK.to_string(), + QUEUE_DEADLINE.to_string(), + QUEUE_SHARED.to_string(), + ]; + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, run_lease, state::text FROM {schema}.done_entries WHERE queue = ANY($1) AND finalized_at > now() - interval '25 seconds' + UNION ALL SELECT job_id, run_lease, 'failed' FROM {schema}.dlq_entries WHERE queue = ANY($1) AND dlq_at > now() - interval '25 seconds' + UNION ALL SELECT b.job_id, b.run_lease, 'completed' FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + WHERE c.queue = ANY($1) AND c.finalized_at > now() - interval '25 seconds'" + )) + .bind(&queues) + .fetch_all(pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + shared.incomplete( + format!("Python terminal evidence query failed: {error}"), + vec![], + ); + shared.stats(Persona::Python).add("python_tail_errors", 1); + if shared.stats(Persona::Python).get("python_tail_errors") <= 3 { + shared.note("python", format!("terminal row tail failed: {error}")); + } + return; + } + }; + shared + .stats(Persona::Python) + .add("python_tail_rows", rows.len() as i64); + let now = shared.now_ms() as u32; + for (id, run_lease, state) in rows { + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + u32::MAX, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| { + record.note_observed_terminal(run_lease, &state); + }, + ); + } +} + +/// A transaction rolled back by the Python application must leave no job anywhere. +async fn check_rolled_back(shared: &Arc, ids: &HashSet) -> Result<(), String> { + if ids.is_empty() { + return Ok(()); + } + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".into()); + let pool = super::super::db::connect_pool(&shared.config.database_url, "soak_python_check", 1) + .await + .map_err(|error| error.to_string())?; + let ids: Vec = ids.iter().copied().collect(); + let mut visible: Vec = Vec::new(); + for chunk in ids.chunks(5000) { + let found: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {schema}.ready_entries WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.deferred_jobs WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.done_entries WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.dlq_entries WHERE job_id = ANY($1)" + )) + .bind(chunk) + .fetch_all(&pool) + .await + .map_err(|error| error.to_string())?; + visible.extend(found); + } + let executed: Vec = ids + .iter() + .copied() + .filter(|id| { + shared + .ledger + .get(*id) + .is_some_and(|record| record.executions > 0) + }) + .collect(); + shared + .stats(Persona::Python) + .add("python_rolled_back_checked", ids.len() as i64); + if !visible.is_empty() || !executed.is_empty() { + shared.violation( + "rolled_back_job_visible", + format!("{} job(s) enqueued in rolled-back Python transactions exist in the store, {} of them executed", visible.len(), executed.len()), + visible.iter().take(10).map(|id| format!("{{\"job_id\":{id}}}")).collect(), + ); + } + pool.close().await; + Ok(()) +} + +#[cfg(test)] +mod lifecycle_tests { + use super::*; + + fn finished_child() -> Result { + Command::new(std::env::current_exe().map_err(|error| error.to_string())?) + .arg("--list") + .stdout(Stdio::null()) + .spawn() + .map_err(|error| error.to_string()) + } + + #[tokio::test] + async fn dead_child_is_restarted_for_the_drain_liveness_loop() { + let mut child = Some(finished_child().unwrap()); + child.as_mut().unwrap().wait().await.unwrap(); + let result = restart_if_needed(&mut child, finished_child).unwrap(); + assert!(result.is_some(), "a dead drain worker needs a replacement"); + assert!(child.is_some()); + child.as_mut().unwrap().wait().await.unwrap(); + } + + #[tokio::test] + async fn outstanding_probe_does_not_delay_stop_production_or_worker_recovery() { + let (release, blocked) = tokio::sync::oneshot::channel(); + let (stopped, mut stop_seen) = tokio::sync::oneshot::channel(); + let mut probes = tokio::task::JoinSet::new(); + probes.spawn(async move { + let _ = blocked.await; + }); + let errors = tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(1)), + begin_drain( + async { + stopped.send(()).unwrap(); + }, + &mut probes, + ), + ) + .await + .expect("stop production must precede waiting on a blocked probe"); + assert!(errors.is_empty()); + assert!(stop_seen.try_recv().is_ok()); + assert!( + !probes.is_empty(), + "the probe remains blocked while drain proceeds" + ); + release.send(()).unwrap(); + probes.join_next().await.unwrap().unwrap(); + } + + #[test] + fn failed_restart_is_reported_instead_of_losing_the_process_silently() { + let mut child = None; + let result = restart_if_needed(&mut child, || Err("injected spawn failure".into())); + assert_eq!(result.unwrap_err(), "injected spawn failure"); + assert!(child.is_none()); + } +} diff --git a/awa/tests/persona_soak/personas/sam.rs b/awa/tests/persona_soak/personas/sam.rs new file mode 100644 index 000000000..b70876db0 --- /dev/null +++ b/awa/tests/persona_soak/personas/sam.rs @@ -0,0 +1,359 @@ +//! Sam, ordered streams: a queue with enqueue shards and ordering keys, a striped +//! queue (`queue_stripe_count`), scheduled jobs and a rate-limited queue. Stresses +//! routing, promotion and fairness. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert_many, QueueStorage, QueueStorageConfig}; +use awa::{ + Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, QueueConfig, RateLimit, Worker, +}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +pub const ORDERED_QUEUE: &str = "awa_soak_sam_ord"; +pub const RATE_QUEUE: &str = "awa_soak_sam_rate"; +pub const SCHEDULED_QUEUE: &str = "awa_soak_sam_sched"; +pub const STRIPED_QUEUE: &str = "awa_soak_sam_stripe"; +const KEYS: u64 = 64; +const ENQUEUE_SHARDS: i16 = 4; +const STRIPES: usize = 4; +const RATE_PER_SECOND: f64 = 15.0; +/// Starts this far behind the newest start for the same key count as order inversions. +const INVERSION_WINDOW: u64 = 64; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct SamJob { + pub stream: String, + pub key: u64, + pub seq: u64, +} + +#[derive(Default)] +struct Order { + newest_start: HashMap, + inversions: u64, + worst_distance: u64, + striped_seen: HashSet, + striped_repeats: u64, +} + +struct SamWorker { + shared: Arc, + order: Arc>, +} + +#[async_trait::async_trait] +impl Worker for SamWorker { + fn kind(&self) -> &'static str { + "sam_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Sam, &ctx.job, false); + let args: SamJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + { + let mut order = self.order.lock().expect("order"); + if args.stream == "ordered" { + let newest = order.newest_start.entry(args.key).or_insert(0); + if args.seq + INVERSION_WINDOW < *newest { + let distance = *newest - args.seq; + order.inversions += 1; + order.worst_distance = order.worst_distance.max(distance); + } else { + *newest = (*newest).max(args.seq); + } + } else if args.stream == "striped" && !order.striped_seen.insert(args.seq) { + order.striped_repeats += 1; + } + } + tokio::time::sleep(Duration::from_millis(5 + args.seq % 40)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn plain_factory(shared: Arc, order: Arc>) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(ORDERED_QUEUE, queue(12)) + .queue(SCHEDULED_QUEUE, queue(8)) + .queue( + RATE_QUEUE, + QueueConfig { + max_workers: 8, + rate_limit: Some(RateLimit { + max_rate: RATE_PER_SECOND, + burst: RATE_PER_SECOND as u32, + }), + ..QueueConfig::default() + }, + ) + .register_worker(SamWorker { + shared: shared.clone(), + order: order.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["sam_job"]) + .build() + .expect("build sam client") + }) +} + +fn striped_config() -> QueueStorageConfig { + QueueStorageConfig { + queue_stripe_count: STRIPES, + ..QueueStorageConfig::default() + } +} + +fn striped_factory(shared: Arc, order: Arc>) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(STRIPED_QUEUE, queue(16)) + .queue_storage( + striped_config(), + Duration::from_millis(1_000), + Duration::from_millis(1_000), + ) + .register_worker(SamWorker { + shared: shared.clone(), + order: order.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["sam_job"]) + .build() + .expect("build sam stripe client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Sam; + let ordered = shared.queues.register(ORDERED_QUEUE, persona, false); + let scheduled = shared.queues.register(SCHEDULED_QUEUE, persona, false); + let rated = shared.queues.register(RATE_QUEUE, persona, true); + let striped = shared.queues.register(STRIPED_QUEUE, persona, false); + + let setup_pool = connect_pool(&shared.config.database_url, "soak_sam_setup", 1) + .await + .map_err(|error| error.to_string())?; + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(ORDERED_QUEUE) + .bind(ENQUEUE_SHARDS) + .execute(&setup_pool) + .await + .map_err(|error| format!("set enqueue_shards: {error}"))?; + setup_pool.close().await; + + let order = Arc::new(Mutex::new(Order::default())); + let plain = Fleet::start( + &shared, + "sam", + persona, + 2, + 14, + 2, + vec![(ordered, 12), (scheduled, 8), (rated, 8)], + plain_factory(shared.clone(), order.clone()), + ) + .await?; + let stripe_fleet = Fleet::start( + &shared, + "sam-stripe", + persona, + 2, + 10, + 2, + vec![(striped, 16)], + striped_factory(shared.clone(), order.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_sam_producer", + shared.budget.reserve("sam-producer", 4, 1), + ) + .await + .map_err(|error| error.to_string())?; + let store = QueueStorage::new(striped_config()).map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "sam-producer"); + let scale = shared.config.scale; + let mut key_seq = vec![0u64; KEYS as usize]; + let mut stripe_seq = 0u64; + let mut seq = 0u64; + let mut rate_carry = 0.0_f64; + let tick = Duration::from_millis(100); + shared.note( + "sam", + format!("{ORDERED_QUEUE} ({ENQUEUE_SHARDS} shards, {KEYS} ordering keys), {STRIPED_QUEUE} ({STRIPES} stripes), {SCHEDULED_QUEUE}, {RATE_QUEUE} ({RATE_PER_SECOND}/s)"), + ); + + while !shared.load_stop.sleep(tick).await { + // Ordered stream with ordering keys. + let ordered_count = (4.0 * scale).ceil() as usize; + let mut params = Vec::new(); + for _ in 0..ordered_count { + let key = rng.below(KEYS); + key_seq[key as usize] += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "ordered".into(), + key, + seq: key_seq[key as usize], + }, + InsertOpts { + queue: ORDERED_QUEUE.into(), + ordering_key: Some(key.to_be_bytes().to_vec()), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(persona) + .add("enqueued_ordered", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // Scheduled jobs a few seconds out, and the rate-limited stream. + let mut params = Vec::new(); + for _ in 0..(1.0 * scale).ceil() as usize { + seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "scheduled".into(), + key: 0, + seq, + }, + InsertOpts { + queue: SCHEDULED_QUEUE.into(), + run_at: Some( + Utc::now() + + chrono::Duration::milliseconds(rng.range(1_000, 12_000) as i64), + ), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + // Stay below the queue's rate limit so a backlog does not build by design. + rate_carry += RATE_PER_SECOND * 0.8 * tick.as_secs_f64(); + while rate_carry >= 1.0 { + rate_carry -= 1.0; + seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "rated".into(), + key: 0, + seq, + }, + InsertOpts { + queue: RATE_QUEUE.into(), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(persona) + .add("enqueued_sched_rate", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // Striped queue, enqueued through a store that shares the workers' routing. + let mut params = Vec::new(); + for _ in 0..(3.0 * scale).ceil() as usize { + stripe_seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "striped".into(), + key: 0, + seq: stripe_seq, + }, + InsertOpts { + queue: STRIPED_QUEUE.into(), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match store.enqueue_params_batch(&producer_pool, ¶ms).await { + Ok(count) => shared.stats(persona).add("enqueued_striped", count as i64), + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note( + "sam", + format!("load stopped: {stripe_seq} striped jobs; draining"), + ); + shared.drain_stop.wait().await; + plain.shutdown_all(Duration::from_secs(20), &shared).await; + stripe_fleet + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + + let order = order.lock().expect("order"); + let executed_striped = order.striped_seen.len() as u64; + let missing: Vec = (1..=stripe_seq) + .filter(|seq| !order.striped_seen.contains(seq)) + .take(20) + .collect(); + let stats = shared.stats(persona); + stats.add("striped_executed_distinct", executed_striped as i64); + stats.add("striped_repeat_executions", order.striped_repeats as i64); + stats.add("order_inversions", order.inversions as i64); + stats.set_max("order_worst_distance", order.worst_distance as i64); + if executed_striped < stripe_seq { + let message = format!( + "{} of {stripe_seq} striped jobs never executed", + stripe_seq - executed_striped + ); + let rows = vec![format!( + "{{\"queue\":\"{STRIPED_QUEUE}\",\"missing_seq_sample\":{missing:?}}}" + )]; + if shared.quiesced.load(std::sync::atomic::Ordering::Relaxed) { + shared.violation("lost_job", message, rows); + } else { + shared.finding( + "striped_backlog_at_end", + format!("{message}; the system had not quiesced"), + rows, + ); + } + } + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/shardops.rs b/awa/tests/persona_soak/personas/shardops.rs new file mode 100644 index 000000000..a462c632b --- /dev/null +++ b/awa/tests/persona_soak/personas/shardops.rs @@ -0,0 +1,991 @@ +//! Shard operator: an operator changes `awa.queue_meta.enqueue_shards` on two hot queues +//! while SQL-path producers (`insert_many`, `insert_many_copy`), `QueueStorage` producers, +//! promotion (scheduled jobs and retries) and claimers run, as different non-superuser +//! roles (see `roles.rs`). +//! +//! The operator raises a queue from 1 to 4 to 8 shards, lowers it to 2 and back to 1, +//! sometimes running `prepare_queue` first and sometimes not, restarting some runtimes +//! before others so cached shard counts disagree across the fleet, and, for some raises, +//! holding the producers back until the maintenance leader's runtime has created the new +//! lanes so a different role is the next to use them. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::flag; +use super::super::registry::QueueInfo; +use super::super::rng::Rng; +use super::super::roles; +use super::super::shard_oracle; +use super::super::shared::Shared; +use super::super::stop::Stop; +use super::{attach_events, backoff, queue, register_rows, tune, ANY}; +use awa::model::{insert_many, insert_many_copy_from_pool, QueueStorage, QueueStorageConfig}; +use awa::{Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, Worker}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +pub const HOT_QUEUE: &str = "awa_soak_shardops_hot"; +pub const WARM_QUEUE: &str = "awa_soak_shardops_warm"; +const QUEUES: [&str; 2] = [HOT_QUEUE, WARM_QUEUE]; +const KEYS: u64 = 16; +const WORKERS: u32 = 8; +/// The sequence of shard counts each queue is taken through, repeatedly. +const STEPS: [(i16, i16); 4] = [(1, 4), (4, 8), (8, 2), (2, 1)]; +const MAINTENANCE_LEADER_LOCK_KEY: i64 = 0x4157_415f_4d41_494e; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct ShardOpsJob { + pub seq: u64, + pub mode: String, + pub key: Option, + pub key_seq: u64, +} + +struct ShardOpsWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for ShardOpsWorker { + fn kind(&self) -> &'static str { + "shard_ops_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ShardOps, &ctx.job, false); + let args: ShardOpsJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + self.shared.shardops.note_execution( + args.seq, + ctx.job.id, + args.key.map(|key| (key, args.key_seq)), + ); + tokio::time::sleep(Duration::from_millis(5 + args.seq % 20)).await; + if args.mode == "retry" && ctx.job.attempt <= 1 { + return Ok(JobResult::RetryAfter(Duration::from_millis( + 1_000 + args.seq % 2_000, + ))); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(HOT_QUEUE, queue(WORKERS)) + .queue(WARM_QUEUE, queue(WORKERS)) + .register_worker(ShardOpsWorker { + shared: shared.clone(), + }); + attach_events( + // Retrying election quickly lets these runtimes, which run as application + // roles, take maintenance leadership from the other personas' runtimes. + tune(builder, &shared).leader_election_interval(Duration::from_millis(100)), + &shared, + &["shard_ops_job"], + ) + .build() + .expect("build shard-ops client") + }) +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Mode { + Plain, + Scheduled, + Retry, +} + +impl Mode { + fn roll(rng: &mut Rng) -> Mode { + match rng.below(100) { + 0..=11 => Mode::Retry, + 12..=31 => Mode::Scheduled, + _ => Mode::Plain, + } + } + + fn name(self) -> &'static str { + match self { + Mode::Plain => "plain", + Mode::Scheduled => "scheduled", + Mode::Retry => "retry", + } + } + + fn executions(self) -> u8 { + if self == Mode::Retry { + 2 + } else { + 1 + } + } +} + +struct Built { + params: InsertParams, + seq: u64, + executions: u8, +} + +fn build( + shared: &Shared, + rng: &mut Rng, + queue: &str, + mode: Mode, + key: Option<(u64, u64)>, + run_in: Option, +) -> Result { + let seq = shared.shardops.next_seq(); + let run_at = match (mode, run_in) { + (_, Some(delay)) => { + Some(Utc::now() + chrono::Duration::from_std(delay).unwrap_or_default()) + } + (Mode::Scheduled, None) => { + Some(Utc::now() + chrono::Duration::milliseconds(rng.range(800, 6_000) as i64)) + } + _ => None, + }; + let params = awa::insert::params_with( + &ShardOpsJob { + seq, + mode: mode.name().to_string(), + key: key.map(|(key, _)| key), + key_seq: key.map_or(0, |(_, key_seq)| key_seq), + }, + InsertOpts { + queue: queue.to_string(), + max_attempts: 5, + run_at, + ordering_key: key.map(|(key, _)| key.to_be_bytes().to_vec()), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?; + Ok(Built { + params, + seq, + executions: mode.executions(), + }) +} + +fn pick_queue<'a>(rng: &mut Rng) -> &'a str { + QUEUES[rng.below(QUEUES.len() as u64) as usize] +} + +fn count_enqueued(shared: &Shared, queues: impl Iterator) { + let mut count = 0; + for name in queues { + count += 1; + if let Some(queue) = shared.queues.get(&name) { + queue.enqueued.fetch_add(1, Ordering::Relaxed); + } + } + shared.stats(Persona::ShardOps).add("enqueued", count); +} + +/// Keeps the operator from changing the shard count while an insert is mid-flight +/// during the leader-first hold. +struct InsertGuard<'a>(&'a Shared); + +impl<'a> InsertGuard<'a> { + fn new(shared: &'a Shared) -> Self { + shared + .shardops + .producers_inserting + .fetch_add(1, Ordering::SeqCst); + Self(shared) + } +} + +impl Drop for InsertGuard<'_> { + fn drop(&mut self) { + self.0 + .shardops + .producers_inserting + .fetch_sub(1, Ordering::SeqCst); + } +} + +async fn enqueue_failed( + shared: &Shared, + role: &str, + operation: &str, + error: impl std::fmt::Display, +) { + let message = error.to_string(); + let short: String = message.chars().take(120).collect(); + shared.shardops.violation_once( + shared, + "shardops_enqueue_error", + &format!("{role}:{operation}:{short}"), + format!("role {role} failed {operation}: {message}"), + vec![format!( + "{{\"role\":\"{role}\",\"statement\":\"{operation}\",\"error\":{}}}", + serde_json::Value::String(message.clone()) + )], + ); + backoff(shared, &message, Persona::ShardOps).await; +} + +/// A producer's `QueueStorage` handle. A runtime's cached shard count is only dropped +/// by a restart, so the operator replaces the handle to restart the producer. +struct Stores { + handles: Mutex>>, +} + +impl Stores { + fn new(count: usize) -> Result { + let handles = (0..count) + .map(|_| fresh_store()) + .collect::, _>>()?; + Ok(Self { + handles: Mutex::new(handles), + }) + } + + fn get(&self, index: usize) -> Arc { + self.handles.lock().expect("stores")[index].clone() + } + + fn renew(&self, index: usize) { + if let Ok(store) = fresh_store() { + self.handles.lock().expect("stores")[index] = store; + } + } + + fn len(&self) -> usize { + self.handles.lock().expect("stores").len() + } +} + +fn fresh_store() -> Result, String> { + QueueStorage::new(QueueStorageConfig::default()) + .map(Arc::new) + .map_err(|error| error.to_string()) +} + +/// A SQL-path producer: `insert_many`, or `insert_many_copy` when `copy` is set. +async fn sql_producer(shared: Arc, pool: PgPool, role: &'static str, copy: bool) { + let label = if copy { + "shardops-copy" + } else { + "shardops-sql" + }; + let mut rng = Rng::for_task(shared.config.seed, label); + let scale = shared.config.shardops_scale; + let (tick, per_tick) = if copy { + (Duration::from_secs(1), (40.0 * scale).ceil() as usize) + } else { + (Duration::from_millis(100), (4.0 * scale).ceil() as usize) + }; + while !shared.load_stop.sleep(tick).await { + if shared.shardops.producers_held() { + continue; + } + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let queue = pick_queue(&mut rng); + let mode = Mode::roll(&mut rng); + match build(&shared, &mut rng, queue, mode, None, None) { + Ok(job) => built.push(job), + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let result = if copy { + insert_many_copy_from_pool(&pool, ¶ms).await + } else { + insert_many(&pool, ¶ms).await + }; + match result { + Ok(rows) => { + for job in &built { + shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&shared, rows.iter().map(|row| row.queue.clone())); + register_rows(&shared, Persona::ShardOps, &rows, |_| ANY, 0); + } + Err(error) => { + enqueue_failed( + &shared, + role, + if copy { + "insert_many_copy" + } else { + "insert_many" + }, + error, + ) + .await + } + } + } +} + +/// A producer that enqueues through `QueueStorage`, the runtime's own path. +async fn store_producer( + shared: Arc, + pool: PgPool, + role: &'static str, + stores: Arc, + store_index: usize, +) { + let mut rng = Rng::for_task(shared.config.seed, "shardops-store"); + let per_tick = (3.0 * shared.config.shardops_scale).ceil() as usize; + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + if shared.shardops.producers_held() { + continue; + } + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let queue = pick_queue(&mut rng); + let mode = Mode::roll(&mut rng); + match build(&shared, &mut rng, queue, mode, None, None) { + Ok(job) => built.push(job), + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let store = stores.get(store_index); + match store.enqueue_params_batch(&pool, ¶ms).await { + Ok(_) => { + for job in &built { + shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&shared, params.iter().map(|job| job.opts.queue.clone())); + } + Err(error) => { + enqueue_failed(&shared, role, "QueueStorage::enqueue_params_batch", error).await + } + } + } +} + +/// Jobs with an ordering key, enqueued strictly in sequence per key, alternating between +/// the SQL path as one role and the `QueueStorage` path as the other. +async fn keyed_producer( + shared: Arc, + sql_pool: PgPool, + store_pool: PgPool, + stores: Arc, + store_index: usize, +) { + let mut rng = Rng::for_task(shared.config.seed, "shardops-keyed"); + let per_tick = (3.0 * shared.config.shardops_scale).ceil() as usize; + let mut key_seq = vec![0u64; KEYS as usize]; + let mut tick = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + if shared.shardops.producers_held() { + continue; + } + tick += 1; + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + let mut keys = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let key = rng.below(KEYS); + key_seq[key as usize] += 1; + let queue = QUEUES[(key % 2) as usize]; + match build( + &shared, + &mut rng, + queue, + Mode::Plain, + Some((key, key_seq[key as usize])), + None, + ) { + Ok(job) => { + built.push(job); + keys.push((key, key_seq[key as usize])); + } + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let window_before = shared.shardops.window(); + let outcome = if tick.is_multiple_of(2) { + match insert_many(&sql_pool, ¶ms).await { + Ok(rows) => { + register_rows(&shared, Persona::ShardOps, &rows, |_| ANY, 0); + Ok(rows.len()) + } + Err(error) => Err(( + roles::APP_A, + "insert_many (ordering_key)", + error.to_string(), + )), + } + } else { + let store = stores.get(store_index); + match store.enqueue_params_batch(&store_pool, ¶ms).await { + Ok(count) => Ok(count), + Err(error) => Err(( + roles::APP_B, + "QueueStorage::enqueue_params_batch (ordering_key)", + error.to_string(), + )), + } + }; + match outcome { + Ok(_) => { + let window_after = shared.shardops.window(); + let window = if window_before == window_after { + window_before + } else { + 0 + }; + for (job, (key, seq)) in built.iter().zip(&keys) { + shared.shardops.note_enqueued(job.seq, job.executions); + shared.shardops.note_keyed(*key, *seq, window); + } + shared + .stats(Persona::ShardOps) + .add("enqueued_keyed", built.len() as i64); + count_enqueued(&shared, params.iter().map(|job| job.opts.queue.clone())); + } + Err((role, operation, error)) => enqueue_failed(&shared, role, operation, error).await, + } + } +} + +/// Everything the operator acts through. +struct Operator { + shared: Arc, + fleets: Vec>, + stores: Arc, + owner_pool: PgPool, + migrator_pool: PgPool, + admin_pool: PgPool, + app_pool: PgPool, + rng: Rng, +} + +impl Operator { + fn note(&self, text: impl Into) { + self.shared.note("shards", text); + } + + fn operator_error(&self, what: &str, error: impl std::fmt::Display) { + let message = error.to_string(); + self.shared.shardops.violation_once( + &self.shared, + "shardops_operator_error", + &format!("{what}:{message}"), + format!("{what} failed: {message}"), + vec![format!( + "{{\"statement\":\"{what}\",\"error\":{}}}", + serde_json::Value::String(message.clone()) + )], + ); + } + + async fn prepare(&mut self, queue: &str, shards: i16) { + let as_migrator = self.rng.chance(50); + let (role, pool) = if as_migrator { + (roles::MIGRATOR, &self.migrator_pool) + } else { + (roles::OWNER, &self.owner_pool) + }; + self.note(format!( + "prepare_queue {queue} for {shards} shards as {role}" + )); + let store = match fresh_store() { + Ok(store) => store, + Err(error) => return self.operator_error("prepare_queue", error), + }; + if let Err(error) = store.prepare_queue(pool, queue, shards).await { + self.operator_error(&format!("prepare_queue as {role}"), error); + } + } + + async fn set_shards(&self, queue: &str, shards: i16) { + let result = sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(&self.owner_pool) + .await; + match result { + Ok(_) => { + let now = self.shared.now_ms(); + self.shared.shardops.plans.lock().expect("plans").insert( + queue.to_string(), + shard_oracle::QueuePlan { + configured: shards, + changed_at_ms: now, + settled_at_ms: None, + }, + ); + } + Err(error) => self.operator_error("set enqueue_shards", error), + } + } + + /// Restarts runtime `(fleet, replica)`, which drops its cached shard counts. + async fn restart(&self, fleet: usize, replica: usize) { + let Some(target) = self.fleets.get(fleet) else { + return; + }; + match target.restart(replica, Duration::from_secs(20)).await { + Ok(true) => self + .shared + .stats(Persona::ShardOps) + .add("runtime_restarts", 1), + Ok(false) => {} + Err(error) => self.operator_error("runtime restart", error), + } + } + + async fn runtimes(&self) -> Vec<(usize, usize)> { + let mut all = Vec::new(); + for (index, fleet) in self.fleets.iter().enumerate() { + for replica in fleet.live_replicas().await { + all.push((index, replica)); + } + } + all + } + + fn settle(&self, queue: &str) { + let now = self.shared.now_ms(); + if let Some(plan) = self + .shared + .shardops + .plans + .lock() + .expect("plans") + .get_mut(queue) + { + plan.settled_at_ms = Some(now); + } + self.shared.shardops.open_window(); + } + + /// The runtime holding maintenance leadership, as `(pid, role, application_name)`. + async fn leader(&self) -> Option<(i32, String, String)> { + sqlx::query_as( + "SELECT l.pid, a.usename::text, a.application_name + FROM pg_locks l JOIN pg_stat_activity a ON a.pid = l.pid + WHERE l.locktype = 'advisory' AND l.granted AND a.application_name LIKE 'soak\\_%' + AND ((l.classid::bigint << 32) | l.objid::bigint) = $1 + ORDER BY l.pid LIMIT 1", + ) + .bind(MAINTENANCE_LEADER_LOCK_KEY) + .fetch_optional(&self.admin_pool) + .await + .ok() + .flatten() + } + + /// Makes one of this persona's runtimes the maintenance leader, ending other + /// runtimes' leadership if needed. Returns its role. + async fn lead_with_own_runtime(&self) -> Option { + for _ in 0..10 { + if let Some((pid, role, application)) = self.leader().await { + if application.starts_with("soak_shardops") { + self.note(format!("maintenance leader is {application} ({role})")); + return Some(role); + } + self.note(format!("ending leadership of {application} ({role})")); + let _ = sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(&self.admin_pool) + .await; + } + if self.shared.load_stop.sleep(Duration::from_secs(2)).await { + return None; + } + } + self.note("another persona's runtime kept maintenance leadership"); + None + } + + async fn lane_owners(&self, queue: &str, above: i16) -> String { + let owners: Result, _> = sqlx::query_as( + "SELECT pg_get_userbyid(c.relowner)::text, count(*)::bigint + FROM awa.queue_enqueue_heads h + JOIN pg_class c ON c.relname = h.seq_name + JOIN pg_namespace n ON n.oid = c.relnamespace AND n.nspname = 'awa' + WHERE h.queue = $1 AND h.enqueue_shard >= $2 + GROUP BY 1 ORDER BY 1", + ) + .bind(queue) + .bind(above) + .fetch_all(&self.admin_pool) + .await; + match owners { + Ok(owners) if owners.is_empty() => "none yet".to_string(), + Ok(owners) => owners + .into_iter() + .map(|(role, count)| format!("{role}={count}")) + .collect::>() + .join(","), + Err(error) => format!("lookup failed: {error}"), + } + } + + /// Returns false when the load stopped mid-change. + async fn change(&mut self, queue: &str, from: i16, to: i16) -> bool { + let raising = to > from; + let prepare = + raising && (!self.shared.config.shardops_schema_create || self.rng.chance(50)); + let leader_first = raising && !prepare && self.rng.chance(40); + self.note(format!( + "{queue}: {from} -> {to} shards{}{}", + if prepare { ", prepare_queue first" } else { "" }, + if leader_first { + ", leader creates the new lanes first" + } else { + "" + } + )); + self.shared.stats(Persona::ShardOps).add("shard_changes", 1); + if prepare { + self.prepare(queue, to).await; + } + self.shared.shardops.close_window(); + + if leader_first { + self.shared.shardops.hold_producers(true); + for _ in 0..40 { + if self + .shared + .shardops + .producers_inserting + .load(Ordering::SeqCst) + == 0 + { + break; + } + tokio::time::sleep(Duration::from_millis(250)).await; + } + self.set_shards(queue, to).await; + for (fleet, replica) in self.runtimes().await { + self.restart(fleet, replica).await; + } + for index in 0..self.stores.len() { + self.stores.renew(index); + } + let leader_role = self.lead_with_own_runtime().await; + // Scheduled jobs due shortly: the leader promotes them into new lanes before + // any producer has touched those lanes. + for _ in 0..6 { + let mut built = Vec::new(); + for _ in 0..4 { + let delay = Duration::from_millis(self.rng.range(1_500, 4_000)); + if let Ok(job) = build( + &self.shared, + &mut self.rng, + queue, + Mode::Scheduled, + None, + Some(delay), + ) { + built.push(job); + } + } + let params: Vec = + built.iter().map(|job| job.params.clone()).collect(); + match insert_many(&self.app_pool, ¶ms).await { + Ok(rows) => { + for job in &built { + self.shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&self.shared, rows.iter().map(|row| row.queue.clone())); + register_rows(&self.shared, Persona::ShardOps, &rows, |_| ANY, 0); + } + Err(error) => self.operator_error("seed scheduled jobs", error), + } + if self.shared.load_stop.sleep(Duration::from_secs(1)).await { + self.shared.shardops.hold_producers(false); + return false; + } + } + for _ in 0..15 { + let created: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.queue_enqueue_heads WHERE queue = $1 AND enqueue_shard >= $2", + ) + .bind(queue) + .bind(from) + .fetch_one(&self.admin_pool) + .await + .unwrap_or(0); + if created > 0 { + break; + } + if self.shared.load_stop.sleep(Duration::from_secs(1)).await { + self.shared.shardops.hold_producers(false); + return false; + } + } + let owners = self.lane_owners(queue, from).await; + self.note(format!( + "new lanes of {queue} created before producers resumed; sequence owners: {owners} (leader role {})", + leader_role.as_deref().unwrap_or("not ours") + )); + self.shared.shardops.hold_producers(false); + self.settle(queue); + return true; + } + + self.set_shards(queue, to).await; + if self.shared.load_stop.sleep(Duration::from_secs(4)).await { + return false; + } + // Restart some runtimes and producer stores now and the rest later, so cached + // shard counts disagree across the fleet for a while. + let runtimes = self.runtimes().await; + let mut restart_now: Vec<(usize, usize)> = runtimes + .iter() + .copied() + .filter(|_| self.rng.chance(50)) + .collect(); + if restart_now.is_empty() { + restart_now.extend(runtimes.first().copied()); + } + if restart_now.len() == runtimes.len() && runtimes.len() > 1 { + restart_now.pop(); + } + let renew_now: Vec = (0..self.stores.len()) + .filter(|_| self.rng.chance(50)) + .collect(); + self.note(format!( + "{queue}: restarting {} of {} runtimes and {} of {} producer stores; the others keep the old shard count", + restart_now.len(), + runtimes.len(), + renew_now.len(), + self.stores.len() + )); + for (fleet, replica) in &restart_now { + self.restart(*fleet, *replica).await; + } + for index in &renew_now { + self.stores.renew(*index); + } + if self.shared.load_stop.sleep(Duration::from_secs(18)).await { + return false; + } + for (fleet, replica) in runtimes + .iter() + .filter(|runtime| !restart_now.contains(runtime)) + { + self.restart(*fleet, *replica).await; + } + for index in 0..self.stores.len() { + if !renew_now.contains(&index) { + self.stores.renew(index); + } + } + self.settle(queue); + let owners = self.lane_owners(queue, from.min(to)).await; + self.note(format!( + "{queue}: settled at {to} shards; lane sequence owners above shard {}: {owners}", + from.min(to) + )); + true + } + + async fn run(&mut self, stop: &Stop, period: Duration) { + if stop.sleep(Duration::from_secs(20)).await { + return; + } + let mut step = [0usize; QUEUES.len()]; + let mut turn = 0usize; + loop { + let which = turn % QUEUES.len(); + turn += 1; + let (from, to) = STEPS[step[which] % STEPS.len()]; + step[which] += 1; + let started = Instant::now(); + if !self.change(QUEUES[which], from, to).await { + return; + } + let remaining = period + .saturating_sub(started.elapsed()) + .max(Duration::from_secs(5)); + if stop.sleep(remaining).await { + return; + } + } + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::ShardOps; + let queue_infos: Vec> = QUEUES + .iter() + .map(|name| shared.queues.register(name, persona, false)) + .collect(); + let config = &shared.config; + + let owner_pool = roles::pool_as(config, roles::OWNER, "soak_shardops_owner", 2).await?; + let migrator_pool = + roles::pool_as(config, roles::MIGRATOR, "soak_shardops_migrator", 1).await?; + let admin_pool = connect_pool(&config.database_url, "soak_shardops_admin", 2) + .await + .map_err(|error| error.to_string())?; + + // The schema owner provisions each queue's first shard before any application role runs. + let owner_store = fresh_store()?; + for name in QUEUES { + owner_store + .prepare_queue(&owner_pool, name, 1) + .await + .map_err(|error| format!("prepare_queue {name}: {error}"))?; + shared.shardops.plans.lock().expect("plans").insert( + name.to_string(), + shard_oracle::QueuePlan { + configured: 1, + changed_at_ms: 0, + settled_at_ms: Some(0), + }, + ); + } + shared.shardops.open_window(); + + let served: Vec<(Arc, i64)> = queue_infos + .iter() + .map(|info| (info.clone(), i64::from(WORKERS))) + .collect(); + let mut fleets = Vec::new(); + for (name, role) in [("shardops-a", roles::APP_A), ("shardops-b", roles::APP_B)] { + fleets.push( + Fleet::start_with_url( + &shared, + roles::url_for_role(config, role), + name, + persona, + 2, + 10, + 2, + served.clone(), + client_factory(shared.clone()), + ) + .await?, + ); + } + + let app_a = roles::pool_as( + config, + roles::APP_A, + "soak_shardops_producer_a", + shared.budget.reserve("shardops-producer-a", 6, 2), + ) + .await?; + let app_b = roles::pool_as( + config, + roles::APP_B, + "soak_shardops_producer_b", + shared.budget.reserve("shardops-producer-b", 6, 2), + ) + .await?; + let stores = Arc::new(Stores::new(2)?); + + let period = if config.load >= Duration::from_secs(600) { + Duration::from_secs(60) + } else { + Duration::from_secs(40) + }; + let queue_names: Vec = QUEUES.iter().map(|name| name.to_string()).collect(); + shared.note( + "shards", + format!( + "queues {QUEUES:?} cycle through {STEPS:?} shards every {}s per change; roles: owner {}, migrator {}, applications {} and {}", + period.as_secs(), + roles::OWNER, + roles::MIGRATOR, + roles::APP_A, + roles::APP_B + ), + ); + + let watch_stop = Stop::new(); + let watcher = tokio::spawn(shard_oracle::watch( + shared.clone(), + admin_pool.clone(), + queue_names.clone(), + watch_stop.clone(), + )); + let producers = vec![ + tokio::spawn(sql_producer( + shared.clone(), + app_a.clone(), + roles::APP_A, + false, + )), + tokio::spawn(sql_producer( + shared.clone(), + app_b.clone(), + roles::APP_B, + true, + )), + tokio::spawn(store_producer( + shared.clone(), + app_b.clone(), + roles::APP_B, + stores.clone(), + 0, + )), + tokio::spawn(keyed_producer( + shared.clone(), + app_a.clone(), + app_b.clone(), + stores.clone(), + 1, + )), + ]; + let operator_stop = shared.load_stop.clone(); + let mut operator = Operator { + shared: shared.clone(), + fleets: fleets.clone(), + stores: stores.clone(), + owner_pool: owner_pool.clone(), + migrator_pool, + admin_pool: admin_pool.clone(), + app_pool: app_a.clone(), + rng: Rng::for_task(shared.config.seed, "shardops-operator"), + }; + operator.run(&operator_stop, period).await; + shared.shardops.hold_producers(false); + for producer in producers { + let _ = producer.await; + } + shared.note("shards", "load stopped; draining"); + + shared.drain_stop.wait().await; + shard_oracle::final_check(&shared, &admin_pool, &queue_names).await; + watch_stop.trigger(); + let _ = watcher.await; + for fleet in &fleets { + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + } + for pool in [app_a, app_b, owner_pool, admin_pool] { + pool.close().await; + } + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/tom.rs b/awa/tests/persona_soak/personas/tom.rs new file mode 100644 index 000000000..af963e56c --- /dev/null +++ b/awa/tests/persona_soak/personas/tom.rs @@ -0,0 +1,279 @@ +//! Tom, multi-tenant fan-out: many queues each at a low rate, unique jobs, cron and +//! priority aging. Stresses lane and queue cardinality, per-queue maintenance cost +//! and cron fencing across runtimes. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, register_rows, tune}; +use awa::model::{insert_many, insert_with}; +use awa::{ + AwaError, Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, PeriodicJob, + QueueConfig, UniqueOpts, Worker, +}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +const PREFIX: &str = "awa_soak_tom_"; +const CRON_JOBS: usize = 4; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct TomJob { + pub seq: u64, + pub work_ms: u64, + pub unique_key: Option, +} + +struct TomWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for TomWorker { + fn kind(&self) -> &'static str { + "tom_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Tom, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(50); + tokio::time::sleep(Duration::from_millis(work_ms)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn queue_name(index: usize) -> String { + format!("{PREFIX}{index:03}") +} + +/// Four schedules firing every ten seconds into the first four tenant queues. +pub fn declare_cron(mut builder: awa::ClientBuilder) -> awa::ClientBuilder { + for cron in 0..CRON_JOBS { + let job = PeriodicJob::builder(format!("soak_tom_cron_{cron}"), "*/10 * * * * *") + .queue(queue_name(cron)) + .build(&TomJob { + seq: 0, + work_ms: 20, + unique_key: None, + }) + .expect("cron job"); + builder = builder.periodic(job); + } + builder +} + +fn client_factory(shared: Arc, queue_count: usize) -> ClientFactory { + Arc::new(move |pool| { + let mut builder = Client::builder(pool); + for index in 0..queue_count { + builder = builder.queue( + queue_name(index), + QueueConfig { + max_workers: 2, + priority_aging_interval: Duration::from_secs(5), + ..QueueConfig::default() + }, + ); + } + builder = builder.register_worker(TomWorker { + shared: shared.clone(), + }); + for cron in 0..CRON_JOBS { + let job = PeriodicJob::builder(format!("soak_tom_cron_{cron}"), "*/10 * * * * *") + .queue(queue_name(cron)) + .build(&TomJob { + seq: 0, + work_ms: 20, + unique_key: None, + }) + .expect("cron job"); + builder = builder.periodic(job); + } + attach_events(tune(builder, &shared), &shared, &["tom_job"]) + .build() + .expect("build tom client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Tom; + let queue_count = shared.config.scaled(100) as usize; + let mut served = Vec::new(); + for index in 0..queue_count { + served.push(( + shared.queues.register(&queue_name(index), persona, false), + 2, + )); + } + let fleet = Fleet::start( + &shared, + "tom", + persona, + 2, + 24, + 2, + served, + client_factory(shared.clone(), queue_count), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_tom_producer", + shared.budget.reserve("tom-producer", 4, 1), + ) + .await + .map_err(|error| error.to_string())?; + let tick = Duration::from_millis(100); + let mut rng = Rng::for_task(shared.config.seed, "tom-producer"); + let plain_per_tick = (queue_count as f64 * 0.05).max(1.0) as usize; + shared.note("tom", format!("{queue_count} queues, {plain_per_tick} jobs per 100 ms tick, {CRON_JOBS} cron schedules every 10 s")); + let unique_double_accept = Arc::new(AtomicU64::new(0)); + let mut seq = 0u64; + let mut unique_key = 0u64; + let mut unique_tasks = tokio::task::JoinSet::new(); + + while !shared.load_stop.sleep(tick).await { + let mut params = Vec::with_capacity(plain_per_tick); + for _ in 0..plain_per_tick { + seq += 1; + let scheduled = rng.chance(20); + let run_at = scheduled.then(|| { + Utc::now() + chrono::Duration::milliseconds(rng.range(1_000, 10_000) as i64) + }); + let queue_index = rng.below(queue_count as u64) as usize; + params.push( + awa::insert::params_with( + &TomJob { + seq, + work_ms: rng.range(10, 200), + unique_key: None, + }, + InsertOpts { + queue: queue_name(queue_index), + priority: rng.range(1, 4) as i16, + run_at, + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + shared.stats(persona).add( + if scheduled { + "planned_scheduled" + } else { + "planned_immediate" + }, + 1, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // One logical unique job per second: two concurrent identical inserts. + if seq.is_multiple_of(10) && unique_tasks.len() < 16 { + unique_key += 1; + let queue_index = rng.below(queue_count as u64) as usize; + let pool = producer_pool.clone(); + let shared = shared.clone(); + let double = unique_double_accept.clone(); + let key = unique_key; + let work_ms = rng.range(10, 100); + unique_tasks.spawn(async move { + let args = TomJob { + seq: key, + work_ms, + unique_key: Some(key), + }; + let opts = || InsertOpts { + queue: queue_name(queue_index), + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + by_period: None, + states: 0b0001_1111, + }), + ..Default::default() + }; + let (first, second) = tokio::join!( + insert_with(&pool, &args, opts()), + insert_with(&pool, &args, opts()) + ); + let mut accepted = 0; + for outcome in [first, second] { + match outcome { + Ok(row) => { + accepted += 1; + let queue = shared + .queues + .get(&row.queue) + .map_or(u32::MAX, |queue| queue.index); + shared.ledger.register( + row.id, + JobRecord::new( + Persona::Tom, + queue, + shared.now_ms() as u32, + expect::COMPLETED, + 0, + ), + ); + } + Err(AwaError::UniqueConflict { .. }) => { + shared.stats(Persona::Tom).add("unique_conflicts", 1) + } + Err(error) => shared.stats(Persona::Tom).add( + &format!("unique_error:{error}") + .chars() + .take(60) + .collect::(), + 1, + ), + } + } + shared.stats(Persona::Tom).add("unique_logical", 1); + if accepted > 1 { + double.fetch_add(1, Ordering::Relaxed); + shared.violation( + "unique_double_accept", + format!( + "both concurrent identical unique inserts (key {key}) were accepted" + ), + vec![format!( + "{{\"unique_key\":{key},\"queue\":\"{}\"}}", + queue_name(queue_index) + )], + ); + } + }); + } + while unique_tasks.try_join_next().is_some() {} + } + while unique_tasks.join_next().await.is_some() {} + + shared.note( + "tom", + format!("load stopped after {seq} plain jobs; draining"), + ); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/ui.rs b/awa/tests/persona_soak/personas/ui.rs new file mode 100644 index 000000000..d2148fc43 --- /dev/null +++ b/awa/tests/persona_soak/personas/ui.rs @@ -0,0 +1,881 @@ +//! Admin UI users, driving the real `awa serve` HTTP server with the requests the web +//! UI makes (see `awa-ui/frontend/src/lib/api.ts`) at the UI's own refresh intervals. +//! +//! - Watchers: several dashboard tabs polling overview, queues, counts and timeseries. +//! - Support: searching and paging through jobs, opening details, DLQ, runtime and cron. +//! - UiOps: the UI's mutations (retry, cancel, pause, DLQ retry and purge, bulk actions). +//! - Impatient: requests abandoned mid-flight and bursts of refreshes. + +use super::super::adminserver::APPLICATION_NAME; +use super::super::config::Persona; +use super::super::db::connect_pool; +use super::super::ledger::flag; +use super::super::rng::{percentile, Rng}; +use super::super::shared::Shared; +use super::{alex, mia}; +use reqwest::Method; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +/// Latency samples per endpoint over the run, and the number of requests in flight. +#[derive(Default)] +pub struct UiMetrics { + samples: Mutex>>, + pub inflight: AtomicI64, + base_url: Mutex>, + reported_errors: Mutex>, +} + +impl UiMetrics { + pub fn set_base_url(&self, url: String) { + *self.base_url.lock().expect("ui base url") = Some(url); + } + + fn base_url(&self) -> String { + self.base_url + .lock() + .expect("ui base url") + .clone() + .unwrap_or_default() + } + + fn record(&self, endpoint: &str, at_ms: i64, latency_ms: i64) { + let mut samples = self.samples.lock().expect("ui samples"); + let list = samples.entry(endpoint.to_string()).or_default(); + if list.len() < 200_000 { + list.push((at_ms, latency_ms)); + } + } + + /// True the first few times an error class is seen for an endpoint. + fn first_errors(&self, key: &str) -> bool { + let mut reported = self.reported_errors.lock().expect("ui errors"); + let count = reported.entry(key.to_string()).or_insert(0); + *count += 1; + *count <= 3 + } +} + +pub enum Outcome { + Json(Value), + Status(u16), + Timeout, + Failed, +} + +impl Outcome { + fn json(self) -> Option { + match self { + Outcome::Json(value) => Some(value), + _ => None, + } + } +} + +pub struct Http { + client: reqwest::Client, + shared: Arc, + persona: Persona, +} + +impl Http { + pub fn new(shared: &Arc, persona: Persona) -> Self { + Self { + client: reqwest::Client::builder() + .timeout(shared.config.ui_client_timeout) + .pool_max_idle_per_host(4) + .build() + .expect("http client"), + shared: shared.clone(), + persona, + } + } + + pub async fn get(&self, endpoint: &str, path: &str) -> Outcome { + self.request(Method::GET, endpoint, path, None).await + } + + pub async fn post(&self, endpoint: &str, path: &str, body: Value) -> Outcome { + self.request(Method::POST, endpoint, path, Some(body)).await + } + + pub async fn request( + &self, + method: Method, + endpoint: &str, + path: &str, + body: Option, + ) -> Outcome { + let name = format!("{method} {endpoint}"); + let started = Instant::now(); + let at_ms = self.shared.now_ms(); + self.shared.ui.inflight.fetch_add(1, Ordering::Relaxed); + let mut request = self + .client + .request(method, format!("{}{path}", self.shared.ui.base_url())); + if let Some(body) = body { + request = request.json(&body); + } + let outcome = match request.send().await { + Ok(response) => { + let status = response.status(); + match response.text().await { + Ok(text) if status.is_success() => { + Outcome::Json(serde_json::from_str(&text).unwrap_or(Value::Null)) + } + Ok(text) => { + if status.is_server_error() && self.shared.ui.first_errors(&name) { + self.shared.finding( + "ui_server_error", + format!("{name} answered {status}"), + vec![text.chars().take(300).collect()], + ); + } + Outcome::Status(status.as_u16()) + } + Err(error) if error.is_timeout() => Outcome::Timeout, + Err(_) => Outcome::Failed, + } + } + Err(error) if error.is_timeout() => Outcome::Timeout, + Err(_) => Outcome::Failed, + }; + self.shared.ui.inflight.fetch_sub(1, Ordering::Relaxed); + let latency = started.elapsed().as_millis() as i64; + self.shared.ui.record(&name, at_ms, latency); + let stats = self.shared.stats(self.persona); + stats.observe(&name, latency); + match &outcome { + Outcome::Timeout => { + stats.add("client_timeouts", 1); + if self.shared.ui.first_errors(&format!("timeout {name}")) { + self.shared.finding( + "ui_request_hang", + format!( + "{name} did not answer within the {}s client timeout", + self.shared.config.ui_client_timeout.as_secs() + ), + vec![format!("{{\"path\":{path:?}}}")], + ); + } + } + Outcome::Failed => stats.add("connection_errors", 1), + Outcome::Status(code) => stats.add(&format!("http_{code}"), 1), + Outcome::Json(_) => {} + } + outcome + } + + /// Sends a GET and abandons it after `after`, closing the connection mid-request. + pub async fn get_abandoned(&self, path: &str, after: Duration) -> bool { + let url = format!("{}{path}", self.shared.ui.base_url()); + let send = self.client.get(url).send(); + tokio::select! { + _ = send => false, + _ = tokio::time::sleep(after) => true, + } + } +} + +fn ids_of(value: &Value) -> Vec { + value + .as_array() + .map(|rows| { + rows.iter() + .filter_map(|row| row.get("id").and_then(Value::as_i64)) + .collect() + }) + .unwrap_or_default() +} + +async fn poll_interval(http: &Http) -> Duration { + let ms = http + .get("/api/capabilities", "/api/capabilities") + .await + .json() + .and_then(|value| value.get("poll_interval_ms").and_then(Value::as_u64)) + .unwrap_or(5_000); + Duration::from_millis(ms) +} + +/// Dashboard tabs: each polls what the overview, the shell and (for some tabs) the +/// queues, kinds and cron pages poll, every server-suggested interval. +pub async fn run_watchers(shared: Arc) -> Result<(), String> { + let tabs = shared.config.scaled(4) as usize; + shared.note( + "watchers", + format!("{tabs} dashboard tabs polling the admin UI"), + ); + let mut tasks = tokio::task::JoinSet::new(); + for tab in 0..tabs { + let shared = shared.clone(); + tasks.spawn(async move { + let http = Http::new(&shared, Persona::Watchers); + let mut rng = Rng::for_task(shared.config.seed, &format!("watcher-{tab}")); + let interval = poll_interval(&http).await; + shared + .load_stop + .sleep(Duration::from_millis( + rng.below(interval.as_millis() as u64), + )) + .await; + while !shared.load_stop.is_stopped() { + // The overview page and the shell, in parallel as the browser does. + tokio::join!( + http.get("/api/stats", "/api/stats"), + http.get("/api/stats/timeseries", "/api/stats/timeseries?minutes=60"), + http.get("/api/queues", "/api/queues"), + http.get("/api/jobs", "/api/jobs?state=failed&limit=10"), + http.get("/api/runtime", "/api/runtime"), + http.get("/api/storage", "/api/storage?history=true"), + ); + match tab % 3 { + 1 => { + tokio::join!( + http.get("/api/queues/runtime", "/api/queues/runtime"), + http.get("/api/cron", "/api/cron") + ); + } + 2 => { + http.get("/api/kinds", "/api/kinds").await; + } + _ => {} + } + let jitter = rng.range(90, 110); + shared.load_stop.sleep(interval * jitter as u32 / 100).await; + } + }); + } + while tasks.join_next().await.is_some() {} + Ok(()) +} + +pub async fn run_support(shared: Arc) -> Result<(), String> { + let http = Http::new(&shared, Persona::Support); + let mut rng = Rng::for_task(shared.config.seed, "support"); + let kinds = [ + "mia_job", + "tom_job", + "priya_job", + "sam_job", + "alex_job", + "py_job", + "py_shared_job", + ]; + let states = [ + "available", + "running", + "scheduled", + "retryable", + "completed", + "failed", + "cancelled", + ]; + let mut known: Vec = Vec::new(); + shared.note( + "support", + "searching, paging and opening jobs, DLQ, runtime and cron pages", + ); + while !shared + .load_stop + .sleep(Duration::from_millis(rng.range(600, 1800))) + .await + { + match rng.below(10) { + 0 | 1 => { + let kind = rng.pick(&kinds).expect("kinds"); + known.extend(ids_of( + &http + .get("/api/jobs", &format!("/api/jobs?kind={kind}&limit=50")) + .await + .json() + .unwrap_or(Value::Null), + )); + } + 2 => { + let queues = shared.queues.all(); + if let Some(queue) = rng.pick(&queues) { + known.extend(ids_of( + &http + .get( + "/api/jobs", + &format!("/api/jobs?queue={}&limit=50", queue.name), + ) + .await + .json() + .unwrap_or(Value::Null), + )); + } + } + 3 => { + let state = rng.pick(&states).expect("states"); + known.extend(ids_of( + &http + .get("/api/jobs", &format!("/api/jobs?state={state}&limit=100")) + .await + .json() + .unwrap_or(Value::Null), + )); + } + 4 if rng.chance(30) => { + http.get("/api/jobs", "/api/jobs?tag=soak&limit=20").await; + } + 4 => {} + 5 => { + // Follow one job the way the detail page does: a fast poll for a while. + if let Some(&id) = rng.pick(&known) { + for _ in 0..5 { + http.get("/api/jobs/{id}", &format!("/api/jobs/{id}")).await; + if shared.load_stop.sleep(Duration::from_secs(2)).await { + break; + } + } + } + http.get( + "/api/jobs/{id}", + &format!("/api/jobs/{}", 9_000_000_000_i64 + rng.below(1000) as i64), + ) + .await; + } + 6 => { + // Page through a large result set with the before_id cursor. + let state = *rng.pick(&["completed", "available"]).expect("states"); + let mut before: Option = None; + for page in 0..10 { + let path = match before { + Some(id) => format!("/api/jobs?state={state}&limit=100&before_id={id}"), + None => format!("/api/jobs?state={state}&limit=100"), + }; + let endpoint = if page == 0 { + "/api/jobs (first page)" + } else { + "/api/jobs (later page)" + }; + let ids = ids_of( + &http + .get(endpoint, &path) + .await + .json() + .unwrap_or(Value::Null), + ); + match ids.iter().min() { + Some(&lowest) if ids.len() >= 100 => before = Some(lowest), + _ => break, + } + } + } + 7 => { + let dlq = http + .get("/api/dlq", "/api/dlq?limit=50") + .await + .json() + .unwrap_or(Value::Null); + let ids = ids_of(&dlq); + if let Some(last) = dlq.as_array().and_then(|rows| rows.last()) { + if let (Some(id), Some(at)) = ( + last.get("id").and_then(Value::as_i64), + last.get("dlq_at").and_then(Value::as_str), + ) { + let at = at.replace('+', "%2B"); + http.get( + "/api/dlq (later page)", + &format!("/api/dlq?limit=50&before_id={id}&before_dlq_at={at}"), + ) + .await; + } + } + if let Some(&id) = rng.pick(&ids) { + http.get("/api/dlq/{id}", &format!("/api/dlq/{id}")).await; + } + http.get("/api/dlq/depth", "/api/dlq/depth").await; + } + 8 => { + tokio::join!( + http.get("/api/runtime", "/api/runtime"), + http.get("/api/queues/runtime", "/api/queues/runtime"), + http.get("/api/cron", "/api/cron"), + ); + } + _ => { + tokio::join!( + http.get("/api/kinds", "/api/kinds"), + http.get("/api/stats/kinds", "/api/stats/kinds"), + http.get("/api/stats/queues", "/api/stats/queues"), + ); + } + } + if known.len() > 500 { + known.drain(..known.len() - 500); + } + } + Ok(()) +} + +/// POSTs a resume until it succeeds; a queue or schedule left paused looks like a stall. +async fn resume_until_ok(http: &Http, shared: &Shared, path: &str, endpoint: &str) { + for attempt in 0..8 { + if http.post(endpoint, path, json!({})).await.json().is_some() { + return; + } + tokio::time::sleep(Duration::from_secs(2 + attempt)).await; + } + shared.violation( + "resume_failed", + format!("{path} did not succeed after eight attempts"), + vec![], + ); +} + +pub async fn run_ui_operator(shared: Arc) -> Result<(), String> { + let persona = Persona::UiOps; + let http = Http::new(&shared, persona); + let mut rng = Rng::for_task(shared.config.seed, "uiops"); + shared.note("uiops", "operating through the UI's mutation endpoints"); + while !shared + .load_stop + .sleep(Duration::from_millis(rng.range(2_000, 5_000))) + .await + { + match rng.below(9) { + 0 => { + let mut candidates = shared.queues.of(Persona::Sam); + candidates.extend(shared.queues.of(Persona::Tom)); + if let Some(queue) = rng.pick(&candidates).cloned() { + queue.hold(); + let paused = http + .post( + "/api/queues/{queue}/pause", + &format!("/api/queues/{}/pause", queue.name), + json!({"paused_by": "soak-ui"}), + ) + .await + .json() + .is_some(); + if paused { + shared.stats(persona).add("pauses", 1); + shared + .stop_sleep_load(Duration::from_secs(rng.range(3, 6))) + .await; + } + // A pause whose answer was lost may still have applied, so always resume. + resume_until_ok( + &http, + &shared, + &format!("/api/queues/{}/resume", queue.name), + "/api/queues/{queue}/resume", + ) + .await; + shared.stop_sleep_load(Duration::from_secs(3)).await; + queue.release(); + } + } + 1 | 2 => { + // Cancel live jobs from a list view, singly or in bulk. + let mut candidates = shared.queues.of(Persona::Tom); + candidates.extend(shared.queues.of(Persona::Sam)); + let Some(queue) = rng.pick(&candidates).cloned() else { + continue; + }; + let state = *rng + .pick(&["available", "scheduled", "running"]) + .expect("states"); + let list = http + .get( + "/api/jobs", + &format!("/api/jobs?queue={}&state={state}&limit=20", queue.name), + ) + .await + .json() + .unwrap_or(Value::Null); + let ids: Vec = ids_of(&list).into_iter().take(6).collect(); + if ids.is_empty() { + continue; + } + let now = shared.now_ms() as u32; + if rng.chance(50) { + let completed_before = shared + .ledger + .get(ids[0]) + .is_some_and(|record| record.completed_events > 0); + if http + .post( + "/api/jobs/{id}/cancel", + &format!("/api/jobs/{}/cancel", ids[0]), + json!({}), + ) + .await + .json() + .is_some() + { + shared.record_operator_cancel(persona, ids[0], completed_before, now); + } + } else if let Some(cancelled) = http + .post( + "/api/jobs/bulk-cancel", + "/api/jobs/bulk-cancel", + json!({"ids": ids}), + ) + .await + .json() + { + for id in ids_of(&cancelled) { + shared.record_operator_cancel(persona, id, false, now); + } + } + } + 3 | 4 => { + let queue = if rng.chance(50) { + mia::INFER_QUEUE + } else { + mia::DEADLINE_QUEUE + }; + let dlq = http + .get("/api/dlq", &format!("/api/dlq?queue={queue}&limit=10")) + .await + .json() + .unwrap_or(Value::Null); + let ids = ids_of(&dlq); + let Some(&id) = rng.pick(&ids) else { continue }; + if rng.chance(60) { + shared.ledger.mark(id, flag::DLQ_RETRIED); + if http + .post( + "/api/dlq/{id}/retry", + &format!("/api/dlq/{id}/retry"), + json!({}), + ) + .await + .json() + .is_some() + { + shared.stats(persona).add("dlq_retried", 1); + } + } else { + shared.ledger.mark(id, flag::DLQ_PURGED); + let outcome = http + .request( + Method::DELETE, + "/api/dlq/{id}", + &format!("/api/dlq/{id}"), + None, + ) + .await; + if outcome.json().is_some() { + shared.stats(persona).add("dlq_purged", 1); + } + } + } + 5 => { + // Bulk actions on the DLQ by filter, as the DLQ page's toolbar does. + if rng.chance(50) { + shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_RETRIED, + ); + http.post( + "/api/dlq/bulk-retry", + "/api/dlq/bulk-retry", + json!({"queue": mia::INFER_QUEUE}), + ) + .await; + } else { + shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_PURGED, + ); + http.post( + "/api/dlq/bulk-purge", + "/api/dlq/bulk-purge", + json!({"queue": mia::TRAIN_QUEUE}), + ) + .await; + } + } + 6 => { + let name = format!("soak_tom_cron_{}", rng.below(4)); + let paused = http + .post( + "/api/cron/{name}/pause", + &format!("/api/cron/{name}/pause"), + json!({"paused_by": "soak-ui"}), + ) + .await + .json() + .is_some(); + if paused { + shared.stop_sleep_load(Duration::from_secs(8)).await; + } + resume_until_ok( + &http, + &shared, + &format!("/api/cron/{name}/resume"), + "/api/cron/{name}/resume", + ) + .await; + } + 7 => { + let name = format!("soak_tom_cron_{}", rng.below(4)); + http.post( + "/api/cron/{name}/trigger", + &format!("/api/cron/{name}/trigger"), + json!({}), + ) + .await; + } + _ => { + // Preview then submit a batch operation, as the batch operations page does. + let payload = json!({ + "op_kind": "set_priority", + "spec": { "priority": 2 }, + "filter": { "queue": alex::QUEUE_B, "state": "available" }, + "submitted_by": "soak-ui", + }); + http.post( + "/api/batch-ops/preview", + "/api/batch-ops/preview", + payload.clone(), + ) + .await; + http.post("/api/batch-ops", "/api/batch-ops", payload).await; + http.get("/api/batch-ops", "/api/batch-ops?limit=20").await; + } + } + } + Ok(()) +} + +/// Requests abandoned mid-flight and bursts of refreshes, while sampling the server's +/// database sessions: an abandoned request must not keep holding a connection. +pub async fn run_impatient(shared: Arc) -> Result<(), String> { + let persona = Persona::Impatient; + let http = Arc::new(Http::new(&shared, persona)); + let sampler_pool = connect_pool( + &shared.config.database_url, + "soak_ui_sampler", + shared.budget.reserve("ui-sampler", 1, 1), + ) + .await + .map_err(|error| error.to_string())?; + let statement_timeout = shared.config.ui_statement_timeout; + let sole_ui_persona = [Persona::Watchers, Persona::Support, Persona::UiOps] + .iter() + .all(|other| !shared.config.has(*other)); + let heavy = [ + "/api/kinds", + "/api/queues", + "/api/stats", + "/api/stats/timeseries?minutes=1440", + "/api/jobs?limit=100", + "/api/jobs?state=completed&limit=100", + "/api/dlq?limit=100", + "/api/runtime", + "/api/storage?history=true", + ]; + + let sampler = { + let shared = shared.clone(); + let pool = sampler_pool.clone(); + tokio::spawn(async move { + let mut reported = false; + loop { + if shared.drain_stop.is_stopped() { + return; + } + let row: Result<(i64, i64, f64), _> = sqlx::query_as(&format!( + "SELECT count(*)::bigint, count(*) FILTER (WHERE state = 'active')::bigint, + COALESCE(max(extract(epoch FROM now() - query_start)) FILTER (WHERE state = 'active'), 0)::float8 + FROM pg_stat_activity WHERE application_name = '{APPLICATION_NAME}'" + )) + .fetch_one(&pool) + .await; + if let Ok((sessions, active, oldest)) = row { + shared.run.set_max("ui_server_sessions_peak", sessions); + shared.run.set_max("ui_server_active_sessions_peak", active); + shared + .run + .set_max("ui_server_oldest_active_query_s", oldest as i64); + shared.run.observe("ui_server_sessions", sessions); + if oldest > statement_timeout.as_secs_f64() + 20.0 && !reported { + reported = true; + shared.violation( + "ui_query_outlived_statement_timeout", + format!("an admin server query has been active for {oldest:.0}s with a {}s statement timeout", statement_timeout.as_secs()), + vec![format!("{{\"sessions\":{sessions},\"active\":{active},\"oldest_active_s\":{oldest:.0}}}")], + ); + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + }; + + let mut rng = Rng::for_task(shared.config.seed, "impatient"); + let mut next_burst = Instant::now() + Duration::from_secs(20); + shared.note( + "impatient", + "abandoning requests mid-flight and refreshing in bursts", + ); + while !shared + .load_stop + .sleep(Duration::from_millis(rng.range(200, 900))) + .await + { + let path = *rng.pick(&heavy).expect("paths"); + let after = Duration::from_millis(rng.range(2, 150)); + if http.get_abandoned(path, after).await { + shared.stats(persona).add("abandoned_requests", 1); + } else { + shared.stats(persona).add("completed_before_abandon", 1); + } + if Instant::now() >= next_burst { + next_burst = Instant::now() + Duration::from_secs(rng.range(30, 70)); + let mut burst = tokio::task::JoinSet::new(); + let size = shared.config.scaled(12); + for _ in 0..size { + let http = http.clone(); + let path = *rng.pick(&heavy).expect("paths"); + let abandon = rng.chance(70); + let after = Duration::from_millis(rng.range(5, 120)); + burst.spawn(async move { + if abandon { + http.get_abandoned(path, after).await + } else { + http.get(path, path).await; + false + } + }); + } + let mut abandoned = 0; + while let Some(result) = burst.join_next().await { + if result.unwrap_or(false) { + abandoned += 1; + } + } + shared.stats(persona).add("bursts", 1); + shared.stats(persona).add("burst_abandoned", abandoned); + // With no other UI user running, every query still active after the burst + // belongs to an abandoned request; time how long the server keeps running them. + if sole_ui_persona { + let released_started = Instant::now(); + loop { + let active: i64 = sqlx::query_scalar(&format!( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE application_name = '{APPLICATION_NAME}' AND state = 'active'" + )) + .fetch_one(&sampler_pool) + .await + .unwrap_or(0); + if active == 0 + || released_started.elapsed() > Duration::from_secs(90) + || shared.load_stop.is_stopped() + { + break; + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + let seconds = released_started.elapsed().as_secs() as i64; + shared + .stats(persona) + .observe("release_after_burst_s", seconds); + if seconds >= 5 { + shared.finding( + "ui_abandoned_request_runs_on", + format!("after a burst of abandoned requests the server kept {seconds}s of queries running (statement timeout {}s)", statement_timeout.as_secs()), + vec![], + ); + } + } + } + } + + // With every UI client stopped, the server must release its connections once the + // statements it was running reach their timeout. + let quiet_started = Instant::now(); + while shared.ui.inflight.load(Ordering::Relaxed) > 0 + && quiet_started.elapsed() < Duration::from_secs(60) + { + tokio::time::sleep(Duration::from_millis(500)).await; + } + let deadline = Instant::now() + statement_timeout + Duration::from_secs(15); + let mut active = i64::MAX; + while Instant::now() < deadline { + let row: Result<(i64, i64), _> = sqlx::query_as(&format!( + "SELECT count(*)::bigint, count(*) FILTER (WHERE state = 'active')::bigint + FROM pg_stat_activity WHERE application_name = '{APPLICATION_NAME}'" + )) + .fetch_one(&sampler_pool) + .await; + if let Ok((sessions, now_active)) = row { + active = now_active; + shared.run.add("ui_server_sessions_after_quiet", 0); + shared + .run + .set_max("ui_server_sessions_after_quiet", sessions); + if active == 0 { + break; + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + if active != 0 { + shared.violation( + "ui_sessions_not_released", + format!("{active} admin server database session(s) still running a query {}s after every UI client stopped", (statement_timeout + Duration::from_secs(15)).as_secs()), + vec![], + ); + } + shared.drain_stop.wait().await; + let _ = sampler.await; + sampler_pool.close().await; + Ok(()) +} + +impl Shared { + /// Sleeps unless the load phase ends first. + pub async fn stop_sleep_load(&self, duration: Duration) -> bool { + self.load_stop.sleep(duration).await + } +} + +/// After the run: latency growth and slow-endpoint findings. +pub fn analyse(shared: &Shared) { + let samples = shared.ui.samples.lock().expect("ui samples").clone(); + let run_ms = shared.config.load.as_millis() as i64; + for (endpoint, mut points) in samples { + if points.len() < 20 { + continue; + } + points.sort_by_key(|point| point.0); + let quarter = points.len() / 4; + let median = |slice: &[(i64, i64)]| { + let mut values: Vec = slice.iter().map(|point| point.1).collect(); + values.sort_unstable(); + percentile(&values, 0.5) + }; + let early = median(&points[..quarter]); + let late = median(&points[points.len() - quarter..]); + let mut all: Vec = points.iter().map(|point| point.1).collect(); + all.sort_unstable(); + let p99 = percentile(&all, 0.99); + let max = all.last().copied().unwrap_or(0); + if endpoint.starts_with("GET") && late >= 500 && late >= early.max(1) * 3 { + shared.finding( + "ui_latency_grows", + format!("{endpoint}: median latency {early} ms in the first quarter of its samples, {late} ms in the last (run {}s)", run_ms / 1000), + vec![format!("{{\"endpoint\":\"{endpoint}\",\"samples\":{},\"p99_ms\":{p99},\"max_ms\":{max}}}", points.len())], + ); + } + if p99 >= 10_000 { + shared.finding( + "ui_endpoint_slow", + format!( + "{endpoint}: p99 {p99} ms, max {max} ms over {} requests", + points.len() + ), + vec![], + ); + } + } +} diff --git a/awa/tests/persona_soak/pgserver.rs b/awa/tests/persona_soak/pgserver.rs new file mode 100644 index 000000000..10b018d99 --- /dev/null +++ b/awa/tests/persona_soak/pgserver.rs @@ -0,0 +1,99 @@ +//! Server crash detection: postmaster start time, and a scan of the server log for +//! the lines a crashed backend leaves behind. + +use super::config::SoakConfig; +use sqlx::PgPool; +use std::process::Command; + +const CRASH_PATTERNS: [&str; 6] = [ + "terminated by signal", + "TRAP:", + "server process (PID", + "PANIC:", + "all server processes terminated", + "database system was not properly shut down", +]; + +pub struct PgWatch { + pub started_at: String, + log_start_offset: u64, +} + +impl PgWatch { + pub async fn begin(config: &SoakConfig, pool: &PgPool) -> Result { + let started_at: String = sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await?; + Ok(Self { + started_at, + log_start_offset: current_log_size(config), + }) + } + + pub async fn postmaster_start(pool: &PgPool) -> Result { + sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await + } + + /// Log text written since the run began, or None when no log source is configured. + pub fn log_since_start(&self, config: &SoakConfig) -> Option { + if let Some(path) = &config.pg_log_file { + let bytes = std::fs::read(path).ok()?; + let from = (self.log_start_offset as usize).min(bytes.len()); + return Some(String::from_utf8_lossy(&bytes[from..]).into_owned()); + } + let container = config.pg_container.as_ref()?; + let output = match &config.pg_container_log { + Some(path) => Command::new("docker") + .args([ + "exec", + container, + "sh", + "-c", + &format!("tail -c +{} {path}", self.log_start_offset + 1), + ]) + .output() + .ok()?, + None => Command::new("docker") + .args(["logs", container]) + .output() + .ok()?, + }; + let mut text = String::from_utf8_lossy(&output.stdout).into_owned(); + text.push_str(&String::from_utf8_lossy(&output.stderr)); + Some(text) + } +} + +fn current_log_size(config: &SoakConfig) -> u64 { + if let Some(path) = &config.pg_log_file { + return std::fs::metadata(path).map(|meta| meta.len()).unwrap_or(0); + } + let (Some(container), Some(path)) = (&config.pg_container, &config.pg_container_log) else { + return 0; + }; + Command::new("docker") + .args(["exec", container, "sh", "-c", &format!("wc -c < {path}")]) + .output() + .ok() + .and_then(|output| String::from_utf8_lossy(&output.stdout).trim().parse().ok()) + .unwrap_or(0) +} + +/// Lines of `log` that indicate a crashed server process, each with a few lines of +/// context on either side. +pub fn crash_excerpts(log: &str) -> Vec { + let lines: Vec<&str> = log.lines().collect(); + let mut excerpts = Vec::new(); + let mut next_allowed = 0; + for (index, line) in lines.iter().enumerate() { + if index >= next_allowed && CRASH_PATTERNS.iter().any(|pattern| line.contains(pattern)) { + let from = index.saturating_sub(5); + let to = (index + 8).min(lines.len()); + excerpts.push(lines[from..to].join("\n")); + next_allowed = to; + } + } + excerpts +} diff --git a/awa/tests/persona_soak/python/soak_python.py b/awa/tests/persona_soak/python/soak_python.py new file mode 100644 index 000000000..ff6920c3e --- /dev/null +++ b/awa/tests/persona_soak/python/soak_python.py @@ -0,0 +1,359 @@ +"""Python side of the persona soak: producers, async workers and the migrate path. + +Started and stopped by the Rust harness. Every enqueue, handler start and handler end +is appended to a JSON-lines ledger that the harness folds into its own ledger, so the +oracle checks Python jobs like any others. + +Environment: + SOAK_DATABASE_URL database to use (an awa_soak* database) + SOAK_LEDGER ledger file, appended to + SOAK_SEED, SOAK_SCALE, SOAK_PACE + SOAK_ROLE "full" (default) or "migrate_probe" + SOAK_QUEUE_WORK, SOAK_QUEUE_DEADLINE, SOAK_QUEUE_SHARED queue names + +Signals: SIGUSR1 stops producing (workers keep running); SIGTERM exits at once. +""" + +import asyncio +import json +import os +import random +import signal +import sys +import time +from dataclasses import dataclass + +import awa + +URL = os.environ["SOAK_DATABASE_URL"] +ROLE = os.environ.get("SOAK_ROLE", "full") +SEED = int(os.environ.get("SOAK_SEED", "1")) +SCALE = float(os.environ.get("SOAK_SCALE", "1")) +PACE = float(os.environ.get("SOAK_PACE", "1")) +QUEUE_WORK = os.environ.get("SOAK_QUEUE_WORK", "awa_soak_py_work") +QUEUE_DEADLINE = os.environ.get("SOAK_QUEUE_DEADLINE", "awa_soak_py_deadline") +QUEUE_SHARED = os.environ.get("SOAK_QUEUE_SHARED", "awa_soak_py_shared") +WORK_WORKERS = 8 +DEADLINE_WORKERS = 16 +SHARED_WORKERS = 6 +DEADLINE_MS = 5000 + +# Outcomes a job may end in, matching the harness (1 completed, 2 failed, 4 cancelled). +COMPLETED, FAILED, CANCELLED = 1, 2, 4 + +ledger = open(os.environ["SOAK_LEDGER"], "a", buffering=1) if "SOAK_LEDGER" in os.environ else None + + +def log(**event): + if ledger is not None: + ledger.write(json.dumps(event, separators=(",", ":")) + "\n") + + +@dataclass +class PyJob: + seq: int + behavior: str + work_ms: int + born: float = 0.0 + + +@dataclass +class PySharedJob: + seq: int + behavior: str + work_ms: int + + +active = 0 +max_active = 0 +started = 0 +enqueued = 0 +errors = 0 + + +async def handle(job): + """One handler for every behaviour; see `plan` for what each one is for.""" + global active, max_active, started + args = job.args + attempt = job.attempt + active += 1 + max_active = max(max_active, active) + started += 1 + log(e="start", id=job.id, attempt=attempt, q=job.queue) + result = "completed" + try: + behavior = args.behavior + if behavior == "ok": + await asyncio.sleep(args.work_ms / 1000) + elif behavior == "progress": + for step in range(1, 5): + await asyncio.sleep(args.work_ms / 4000) + job.set_progress(step * 25, f"step {step}") + await job.flush_progress() + elif behavior == "retry": + await asyncio.sleep(0.05) + if attempt < 3: + result = "raise" + raise RuntimeError("transient failure") + elif behavior == "doomed": + result = "raise" + raise RuntimeError("permanent failure") + elif behavior == "retry_after": + if attempt == 1: + result = "retry_after" + return awa.RetryAfter(seconds=1.0) + elif behavior == "snooze": + # A snooze does not count as an attempt, so the wait is keyed on the clock. + if time.time() < args.born + 3.0: + result = "snooze" + return awa.Snooze(seconds=1.0) + elif behavior == "cancel": + result = "cancel" + return awa.Cancel("rejected") + elif behavior == "long": + # Overruns the queue deadline on the first attempt only. + await asyncio.sleep(args.work_ms / 1000 if attempt == 1 else 0.1) + elif behavior == "ignore_deadline": + # Documented behaviour: after the deadline rescue the worker slot is freed + # but this coroutine keeps running, so more handlers than max_workers run. + await asyncio.sleep(args.work_ms / 1000) + else: + await asyncio.sleep(args.work_ms / 1000) + return None + except Exception: + raise + finally: + active -= 1 + log(e="end", id=job.id, attempt=attempt, r=result) + + +def plan(rng, scale_rows=None): + """Returns (queue, behavior, work_ms, max_attempts, allowed_outcomes, kind_class).""" + roll = rng.random() * 100 + if roll < 40: + return QUEUE_WORK, "ok", rng.randint(20, 400), 5, COMPLETED + if roll < 48: + return QUEUE_WORK, "progress", rng.randint(800, 2500), 5, COMPLETED + if roll < 56: + return QUEUE_WORK, "retry", 50, 5, COMPLETED | FAILED + if roll < 60: + return QUEUE_WORK, "doomed", 50, 3, FAILED + if roll < 66: + return QUEUE_WORK, "retry_after", 50, 5, COMPLETED | FAILED + if roll < 72: + return QUEUE_WORK, "snooze", 50, 5, COMPLETED + if roll < 75: + return QUEUE_WORK, "cancel", 50, 3, CANCELLED + if roll < 80: + return QUEUE_DEADLINE, "long", rng.randint(8000, 12000), 3, COMPLETED | FAILED + if roll < 83: + return QUEUE_DEADLINE, "ignore_deadline", rng.randint(9000, 14000), 2, COMPLETED | FAILED + return QUEUE_SHARED, rng.choice(["ok", "ok", "retry", "progress"]), rng.randint(20, 300), 5, COMPLETED | FAILED + + +def kind_for(queue): + return PySharedJob if queue == QUEUE_SHARED else PyJob + + +def make_args(queue, rng, behavior, work_ms): + seq = rng.randint(1, 1 << 40) + if queue == QUEUE_SHARED: + return PySharedJob(seq=seq, behavior=behavior, work_ms=work_ms) + return PyJob(seq=seq, behavior=behavior, work_ms=work_ms, born=time.time()) + + +def kind_name(queue): + return "py_shared_job" if queue == QUEUE_SHARED else "py_job" + + +def record_enqueued(job, queue, behavior, allowed, max_attempts): + global enqueued + enqueued += 1 + log(e="enq", id=job.id, q=queue, b=behavior, x=allowed, ma=max_attempts) + + +async def producer_singles(client, rng, stop): + """Single inserts, one or two at a time.""" + period = 0.25 / max(SCALE, 0.1) + while not stop.is_set(): + for _ in range(rng.randint(1, 2)): + queue, behavior, work_ms, max_attempts, allowed = plan(rng) + args = make_args(queue, rng, behavior, work_ms) + try: + job = await client.insert(args, kind=kind_name(queue), queue=queue, max_attempts=max_attempts) + record_enqueued(job, queue, behavior, allowed, max_attempts) + except Exception as error: # connection killed by chaos, pool timeout + note_error("insert", error) + await asyncio.sleep(period) + + +async def producer_batches(client, rng, stop): + """insert_many above the old 5,461-row bind-parameter limit, and COPY.""" + rows = int(6500 * SCALE) if SCALE >= 1 else 6500 + interval = max(20.0, 75.0 * PACE) + await asyncio.sleep(min(15.0, interval)) + flip = 0 + while not stop.is_set(): + flip += 1 + args = [PyJob(seq=i, behavior="ok", work_ms=0, born=time.time()) for i in range(rows)] + try: + started_at = time.monotonic() + if flip % 3 == 0: + jobs = await client.insert_many_copy(args, kind="py_job", queue=QUEUE_WORK, max_attempts=3) + how = "copy" + else: + # One transaction, one insert_many call far above the old bind-parameter limit. + tx = await client.transaction() + try: + jobs = await tx.insert_many(args, kind="py_job", queue=QUEUE_WORK, max_attempts=3) + await tx.commit() + except Exception: + await tx.rollback() + raise + how = "insert_many" + for job in jobs: + record_enqueued(job, QUEUE_WORK, "ok", COMPLETED, 3) + print(f"BATCH how={how} rows={len(jobs)} seconds={time.monotonic() - started_at:.2f}", flush=True) + except Exception as error: + note_error(f"batch of {rows}", error) + await wait_or_stop(stop, interval) + + +async def producer_transactions(client, rng, stop): + """Enqueue inside an application transaction that sometimes rolls back.""" + while not stop.is_set(): + count = rng.randint(1, 20) + plans = [plan(rng) for _ in range(count)] + rollback = rng.random() < 0.35 + try: + tx = await client.transaction() + try: + created = [] + for queue, behavior, work_ms, max_attempts, allowed in plans: + args = make_args(queue, rng, behavior, work_ms) + job = await tx.insert(args, kind=kind_name(queue), queue=queue, max_attempts=max_attempts) + created.append((job, queue, behavior, allowed, max_attempts)) + ids = [item[0].id for item in created] + if rollback: + await tx.rollback() + log(e="rolled_back", ids=ids) + else: + await tx.commit() + for job, queue, behavior, allowed, max_attempts in created: + record_enqueued(job, queue, behavior, allowed, max_attempts) + except Exception: + try: + await tx.rollback() + except Exception: + pass + raise + except Exception as error: + note_error("transaction", error) + await asyncio.sleep(1.5 / max(SCALE, 0.1)) + + +def note_error(what, error): + global errors + errors += 1 + log(e="error", what=what, error=str(error)[:200]) + if errors <= 5: + print(f"ERROR {what}: {str(error)[:200]}", flush=True) + + +async def wait_or_stop(stop, seconds): + try: + await asyncio.wait_for(stop.wait(), timeout=seconds) + except asyncio.TimeoutError: + pass + + +async def reporter(stop): + while True: + await asyncio.sleep(2) + print(f"STAT active={active} max_active={max_active} started={started} enqueued={enqueued} errors={errors}", flush=True) + + +async def migrate_and_probe(client, label): + """What an application does at startup: migrate, then query.""" + await client.migrate() + stats = await client.queue_stats() + jobs = await client.list_jobs(limit=5) + print(f"MIGRATED {label} queues={len(stats)} sample={len(jobs)}", flush=True) + + +async def migrate_probe_main(): + client = awa.AsyncClient(URL, max_connections=4) + await migrate_and_probe(client, "probe") + # Idle with the pool open: connections stranded by migrate would show up as + # sessions beyond the pool size, or as idle-in-transaction. + await asyncio.sleep(4) + print("IDLE", flush=True) + await asyncio.sleep(4) + await client.close() + print("CLOSED", flush=True) + + +async def exit_with_parent(): + """The harness owns this process; if it dies without cleaning up, so do we.""" + parent = os.getppid() + while True: + await asyncio.sleep(3) + if os.getppid() != parent: + os._exit(1) + + +async def main(): + asyncio.create_task(exit_with_parent()) + if ROLE == "migrate_probe": + await migrate_probe_main() + return + # A chaos restart can still be migrating when the load ends. Remember the + # drain request before startup awaits, then start workers without producers. + stop = asyncio.Event() + if os.environ.get("SOAK_DRAIN") == "1": + stop.set() + loop = asyncio.get_running_loop() + loop.add_signal_handler(signal.SIGUSR1, stop.set) + rng = random.Random(SEED) + client = awa.AsyncClient(URL, max_connections=14) + await migrate_and_probe(client, "startup") + client.task(PyJob, kind="py_job", queue=QUEUE_WORK)(handle) + client.task(PySharedJob, kind="py_shared_job", queue=QUEUE_SHARED)(handle) + await client.start( + [ + {"name": QUEUE_WORK, "max_workers": WORK_WORKERS}, + {"name": QUEUE_DEADLINE, "max_workers": DEADLINE_WORKERS, "deadline_duration_ms": DEADLINE_MS}, + {"name": QUEUE_SHARED, "max_workers": SHARED_WORKERS}, + ], + heartbeat_interval_ms=2000, + heartbeat_staleness_ms=12000, + heartbeat_rescue_interval_ms=3000, + deadline_rescue_interval_ms=2000, + promote_interval_ms=250, + ) + # The deadline queue serves a second kind through the same handler. + print(f"READY pid={os.getpid()}", flush=True) + + tasks = [ + asyncio.create_task(producer_singles(client, rng, stop)), + asyncio.create_task(producer_batches(client, rng, stop)), + asyncio.create_task(producer_transactions(client, rng, stop)), + asyncio.create_task(reporter(stop)), + ] + await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.Event().wait() + + +def exit_now(signum, _frame): + os._exit(128 + signum) + + +if __name__ == "__main__": + # PyO3/tokio teardown can crash while Python unloads modules; exit without finalizers. + signal.signal(signal.SIGTERM, exit_now) + signal.signal(signal.SIGINT, exit_now) + try: + asyncio.run(main()) + finally: + sys.stdout.flush() + os._exit(0) diff --git a/awa/tests/persona_soak/python/test_soak_python.py b/awa/tests/persona_soak/python/test_soak_python.py new file mode 100644 index 000000000..337301719 --- /dev/null +++ b/awa/tests/persona_soak/python/test_soak_python.py @@ -0,0 +1,119 @@ +"""Process-boundary regressions; uses a delayed fake client, never a database.""" + +import os +from pathlib import Path +import select +import signal +import subprocess +import sys +import tempfile +import time +import unittest + + +SCRIPT = Path(os.environ.get("SOAK_PYTHON_SCRIPT", Path(__file__).with_name("soak_python.py"))) +FAKE_AWA = ''' +import asyncio +import os +from pathlib import Path +from types import SimpleNamespace + +class AsyncClient: + def __init__(self, *args, **kwargs): + pass + + async def migrate(self): + print("STARTUP_BLOCKED", flush=True) + while not Path(os.environ["SOAK_TEST_GATE"]).exists(): + await asyncio.sleep(0.01) + + async def queue_stats(self): + return [] + + async def list_jobs(self, **kwargs): + return [] + + def task(self, *args, **kwargs): + return lambda handler: handler + + async def start(self, *args, **kwargs): + Path(os.environ["SOAK_TEST_WORKERS"]).touch() + + async def insert(self, *args, **kwargs): + Path(os.environ["SOAK_TEST_PRODUCED"]).touch() + return SimpleNamespace(id=1) + + async def transaction(self): + Path(os.environ["SOAK_TEST_PRODUCED"]).touch() + raise RuntimeError("unexpected production during drain") +''' + + +class StartupDrainTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory(prefix="awa-python-drain-") + self.root = Path(self.directory.name) + (self.root / "awa.py").write_text(FAKE_AWA) + self.gate = self.root / "gate" + self.workers = self.root / "workers" + self.produced = self.root / "produced" + self.child = None + + def tearDown(self): + if self.child is not None: + if self.child.poll() is None: + self.child.terminate() + self.child.wait(timeout=10) + self.child.stdout.close() + self.child.stderr.close() + self.directory.cleanup() + + def start(self, *, draining=False): + env = os.environ | { + "PYTHONPATH": str(self.root), + "SOAK_DATABASE_URL": "unused-by-the-fake-client", + "SOAK_ROLE": "full", + "SOAK_DRAIN": "1" if draining else "0", + "SOAK_TEST_GATE": str(self.gate), + "SOAK_TEST_WORKERS": str(self.workers), + "SOAK_TEST_PRODUCED": str(self.produced), + } + self.child = subprocess.Popen( + [sys.executable, str(SCRIPT)], env=env, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, + ) + self.wait_for("STARTUP_BLOCKED") + + def wait_for(self, prefix): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + if select.select([self.child.stdout], [], [], 0.1)[0]: + line = self.child.stdout.readline() + if line.startswith(prefix): + return + if not line: + break + self.fail(f"did not see {prefix!r}; child status={self.child.poll()}") + + def assert_draining_workers(self): + self.gate.touch() + self.wait_for("READY pid=") + time.sleep(0.1) + self.assertIsNone(self.child.poll(), "workers must stay alive to drain") + self.assertTrue(self.workers.exists(), "workers must start after migration") + self.assertFalse(self.produced.exists(), "a drain restart must not enqueue new work") + + def test_stop_during_migration_is_remembered_until_workers_start(self): + self.start() + self.child.send_signal(signal.SIGUSR1) + time.sleep(0.05) + self.assertIsNone(self.child.poll(), "SIGUSR1 during startup must not terminate the worker") + self.assert_draining_workers() + + def test_restart_during_drain_starts_workers_without_producers(self): + self.start(draining=True) + self.assert_draining_workers() + + +if __name__ == "__main__": + unittest.main() diff --git a/awa/tests/persona_soak/python_probe.rs b/awa/tests/persona_soak/python_probe.rs new file mode 100644 index 000000000..c91d74247 --- /dev/null +++ b/awa/tests/persona_soak/python_probe.rs @@ -0,0 +1,451 @@ +//! A migrate probe may hold a normal transaction until its ready marker. + +use sqlx::PgPool; +use std::future::Future; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::time::Duration; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[repr(u8)] +pub enum Phase { + #[default] + Migrating, + Ready, + Idle, + Closed, +} + +impl Phase { + pub fn marker(line: &str) -> Option { + if line.starts_with("MIGRATED probe ") { + Some(Self::Ready) + } else { + match line { + "IDLE" => Some(Self::Idle), + "CLOSED" => Some(Self::Closed), + _ => None, + } + } + } + + pub fn from_byte(value: u8) -> Self { + match value { + 1 => Self::Ready, + 2 => Self::Idle, + 3 => Self::Closed, + _ => Self::Migrating, + } + } +} + +/// The unique application name isolates this process, including backends older +/// than a stalled probe's deadline. A failed query is never a zero-session fact. +pub async fn sessions(pool: &PgPool, application_name: &str) -> Result<(i64, i64), String> { + tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(2)), + sqlx::query_as::<_, (i64, i64)>( + "SELECT count(*)::bigint, + count(*) FILTER (WHERE state = 'idle in transaction')::bigint + FROM pg_stat_activity WHERE application_name = $1 AND pid <> pg_backend_pid()", + ) + .bind(application_name) + .fetch_one(pool), + ) + .await + .map_err(|error| format!("session query timed out: {error}"))? + .map_err(|error| format!("session query failed: {error}")) +} + +#[derive(Debug, Default)] +pub struct Evidence { + pub peak_sessions: i64, + pub migrating_idle_transactions: i64, + pub ready_idle_transactions: i64, + pub ready_samples: usize, + pub remaining: Option, + pub last_exit_sample_valid: bool, + pub exit_phase: Option, + pub exit_success: bool, + pub incomplete: Option, +} + +impl Evidence { + pub fn observe(&mut self, phase_before_query: Phase, result: Result<(i64, i64), String>) { + match result { + Ok((sessions, idle)) => { + self.peak_sessions = self.peak_sessions.max(sessions); + if phase_before_query == Phase::Migrating { + self.migrating_idle_transactions = self.migrating_idle_transactions.max(idle); + } else { + if phase_before_query != Phase::Closed { + self.ready_samples += 1; + } + self.ready_idle_transactions = self.ready_idle_transactions.max(idle); + } + } + Err(error) => self.incomplete(error), + } + } + + /// Process exit alone does not prove that an active server query was cancelled. + /// CLOSED is emitted only after the client acknowledged its graceful close. + pub fn record_exit(&mut self, phase: Phase, success: bool) { + self.exit_phase = Some(phase); + self.exit_success = success; + if !self.closed_exit_proven() { + self.incomplete(format!( + "Python migrate probe exited without successful CLOSED evidence (phase={phase:?}, success={success})" + )); + } + } + + fn closed_exit_proven(&self) -> bool { + self.exit_success && self.exit_phase == Some(Phase::Closed) + } + + fn post_exit_observation_complete(&self, phase: Phase, success: bool) -> bool { + self.last_exit_sample_valid + && self.remaining == Some(0) + && (!success || phase == Phase::Closed) + } + + pub fn after_exit(&mut self, result: Result<(i64, i64), String>) { + match result { + Ok((sessions, _)) => { + self.remaining = Some(sessions); + self.last_exit_sample_valid = true; + } + Err(error) => { + self.last_exit_sample_valid = false; + self.incomplete(error); + } + } + } + + /// Reaping may precede stdout consumption. Marker and session observations + /// share the existing post-exit budget, including failed/slow queries. + pub async fn finish_exit(&mut self, phase: &AtomicU8, success: bool, mut sessions: F) + where + F: FnMut() -> Fut, + Fut: Future>, + { + let deadline = + tokio::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(2)); + loop { + if tokio::time::Instant::now() >= deadline { + break; + } + let result = tokio::time::timeout_at(deadline, sessions()) + .await + .unwrap_or_else(|_| { + Err("post-exit session observation exceeded its existing deadline".into()) + }); + let failed = result.is_err(); + self.after_exit(result); + let observed_phase = Phase::from_byte(phase.load(Ordering::Acquire)); + if self.post_exit_observation_complete(observed_phase, success) || failed { + break; + } + tokio::time::sleep_until( + (tokio::time::Instant::now() + Duration::from_millis(100)).min(deadline), + ) + .await; + } + self.record_exit(Phase::from_byte(phase.load(Ordering::Acquire)), success); + } + + pub fn incomplete(&mut self, error: impl Into) { + self.incomplete.get_or_insert_with(|| error.into()); + } + + pub fn violation(&self) -> bool { + self.peak_sessions > 4 + || self.ready_idle_transactions > 0 + || (self.closed_exit_proven() + && self.last_exit_sample_valid + && self.remaining.is_some_and(|sessions| sessions > 0)) + } + + pub fn complete(&self) -> bool { + self.incomplete.is_none() + && self.ready_samples > 0 + && self.closed_exit_proven() + && self.last_exit_sample_valid + && self.remaining.is_some() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::postgres::PgPoolOptions; + + async fn pool(application_name: &str, size: u32) -> PgPool { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".into()); + PgPoolOptions::new() + .max_connections(size) + .connect(&format!( + "{url}{}application_name={application_name}", + if url.contains('?') { '&' } else { '?' } + )) + .await + .expect("probe test pool") + } + + fn name() -> String { + format!("probe_test_{}", uuid::Uuid::new_v4().simple()) + } + + #[tokio::test] + async fn healthy_migration_transaction_is_not_a_stranded_ready_connection() { + let application = name(); + let observer = pool(&name(), 1).await; + let client = pool(&application, 1).await; + let mut tx = client.begin().await.expect("migration transaction"); + sqlx::query("SELECT 1").execute(tx.as_mut()).await.unwrap(); + let mut evidence = Evidence::default(); + let migrating = sessions(&observer, &application).await.unwrap(); + assert_eq!( + migrating, + (1, 1), + "a legitimate transaction between statements" + ); + evidence.observe(Phase::Migrating, Ok(migrating)); + tx.commit().await.unwrap(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + client.close().await; + evidence.record_exit(Phase::Closed, true); + // An acknowledged close can precede the server removing its backend. + let deadline = + std::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(2)); + loop { + let count = sessions(&observer, &application).await.unwrap(); + evidence.after_exit(Ok(count)); + if count.0 == 0 || std::time::Instant::now() >= deadline { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert!(evidence.complete(), "{evidence:?}"); + assert!(!evidence.violation(), "{evidence:?}"); + observer.close().await; + } + + #[tokio::test] + async fn ready_transaction_left_open_is_a_violation() { + let application = name(); + let observer = pool(&name(), 1).await; + let client = pool(&application, 1).await; + let mut tx = client.begin().await.unwrap(); + sqlx::query("SELECT 1").execute(tx.as_mut()).await.unwrap(); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + assert!(evidence.violation(), "{evidence:?}"); + tx.rollback().await.unwrap(); + client.close().await; + observer.close().await; + } + + #[tokio::test] + async fn retained_session_after_exit_is_a_violation_and_other_probes_are_excluded() { + let application = name(); + let observer = pool(&name(), 1).await; + let leaked = pool(&application, 1).await; + let unrelated = pool(&name(), 1).await; + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(sessions(&observer, &application).await); + assert_eq!(evidence.remaining, Some(1)); + assert!(evidence.complete()); + assert!(evidence.violation(), "{evidence:?}"); + leaked.close().await; + unrelated.close().await; + observer.close().await; + } + + #[tokio::test] + async fn session_query_failure_is_inconclusive_and_not_zero() { + let observer = pool(&name(), 1).await; + observer.close().await; + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &name()).await); + evidence.after_exit(sessions(&observer, &name()).await); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, None); + } + + #[test] + fn probe_exit_forced_during_a_read_retains_incomplete_session_evidence() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Ready, false); + evidence.after_exit(Ok((1, 0))); + assert!( + !evidence.violation(), + "a forced exit does not prove graceful-close failure" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_delayed_closed_marker_uses_the_existing_session_window() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.after_exit(Ok((0, 0))); + assert!( + !evidence.post_exit_observation_complete(Phase::Ready, true), + "zero sessions must not stop the reader before its flushed CLOSED marker" + ); + // The detached reader consumes the marker during the same post-exit window. + assert!(evidence.post_exit_observation_complete(Phase::Closed, true)); + evidence.record_exit(Phase::Closed, true); + assert!(evidence.complete()); + assert!(!evidence.violation()); + } + + #[tokio::test] + async fn probe_exit_finish_waits_for_delayed_stdout_closed_marker() { + let phase = std::sync::Arc::new(AtomicU8::new(Phase::Ready as u8)); + let reader_phase = phase.clone(); + let reader = tokio::spawn(async move { + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(10)).await; + reader_phase.store(Phase::Closed as u8, Ordering::Release); + }); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + let mut calls = 0; + evidence + .finish_exit(&phase, true, || { + calls += 1; + async { Ok((0, 0)) } + }) + .await; + reader.await.unwrap(); + assert!(evidence.complete(), "{evidence:?}"); + assert!(!evidence.violation()); + assert_eq!(evidence.exit_phase, Some(Phase::Closed)); + assert!( + calls >= 2, + "the helper must continue past zero sessions before CLOSED" + ); + } + + #[tokio::test] + async fn probe_exit_finish_query_failure_does_not_prove_a_retained_session_leak() { + let phase = AtomicU8::new(Phase::Closed as u8); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + let mut calls = 0; + evidence + .finish_exit(&phase, true, || { + calls += 1; + let result = if calls == 1 { + Ok((1, 0)) + } else { + Err("injected final query failure".into()) + }; + async move { result } + }) + .await; + assert!( + !evidence.violation(), + "an earlier positive count is not a successful final observation" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!( + evidence.remaining, + Some(1), + "retain the last successful count for diagnosis" + ); + assert!(!evidence.last_exit_sample_valid); + } + + #[test] + fn probe_exit_success_without_closed_marker_is_incomplete() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Ready, true); + evidence.after_exit(Ok((1, 0))); + assert!( + !evidence.violation(), + "an exit code is not an acknowledged pool close" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_acknowledged_close_with_retained_session_is_a_violation() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(Ok((1, 0))); + assert!(evidence.violation()); + assert!( + evidence.complete(), + "complete coverage still reports the failure" + ); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_acknowledged_close_and_no_session_is_complete() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(Ok((0, 0))); + assert!(!evidence.violation()); + assert!(evidence.complete()); + } + + #[test] + fn probe_exit_unproven_close_does_not_hide_pool_or_idle_transaction_violations() { + for (phase, sample) in [(Phase::Migrating, (5, 0)), (Phase::Ready, (2, 1))] { + let mut evidence = Evidence::default(); + evidence.observe(phase, Ok(sample)); + evidence.record_exit(phase, false); + evidence.after_exit(Ok((1, 0))); + assert!(evidence.violation(), "{evidence:?}"); + assert!(!evidence.complete()); + assert_eq!(evidence.remaining, Some(1)); + } + } + + #[test] + fn ready_marker_and_missing_ready_coverage_are_explicit() { + assert_eq!( + Phase::marker("MIGRATED probe queues=1 sample=0"), + Some(Phase::Ready) + ); + assert_eq!(Phase::marker("IDLE"), Some(Phase::Idle)); + assert_eq!(Phase::marker("CLOSED"), Some(Phase::Closed)); + assert_eq!(Phase::marker("MIGRATED startup queues=1 sample=0"), None); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Migrating, Ok((4, 1))); + evidence.after_exit(Ok((0, 0))); + assert!( + !evidence.complete(), + "an exit without ready samples proves no idle invariant" + ); + evidence.observe(Phase::Migrating, Ok((5, 0))); + assert!( + evidence.violation(), + "the pool budget applies during migration too" + ); + // Even valid ready/idle/closed markers may all reach the reader before + // the next snapshot. An already closed pool is not live idle coverage. + let mut closed_only = Evidence::default(); + closed_only.observe(Phase::Closed, Ok((0, 0))); + closed_only.record_exit(Phase::Closed, true); + closed_only.after_exit(Ok((0, 0))); + assert!(!closed_only.complete()); + } +} diff --git a/awa/tests/persona_soak/registry.rs b/awa/tests/persona_soak/registry.rs new file mode 100644 index 000000000..a0cc42835 --- /dev/null +++ b/awa/tests/persona_soak/registry.rs @@ -0,0 +1,131 @@ +//! The set of queues the run owns and the live facts the oracle needs about each: +//! which persona serves it, how many workers can run it, how many are busy, and +//! whether it is intentionally not being served right now. + +use super::config::Persona; +use std::collections::HashMap; +use std::sync::atomic::{AtomicI32, AtomicI64, Ordering}; +use std::sync::{Arc, RwLock}; + +impl QueueInfo { + pub fn hold(&self) { + self.pause_holds.fetch_add(1, Ordering::Relaxed); + } + + pub fn release(&self) { + self.pause_holds.fetch_sub(1, Ordering::Relaxed); + } + + pub fn is_held(&self) -> bool { + self.pause_holds.load(Ordering::Relaxed) > 0 + } + + /// Jobs enqueued and not yet started, as far as the producers and workers know. + pub fn backlog_estimate(&self) -> i64 { + self.enqueued.load(Ordering::Relaxed) - self.started.load(Ordering::Relaxed) + } +} + +pub struct QueueInfo { + pub name: String, + pub persona: Persona, + pub index: u32, + /// Workers currently able to run this queue across all live runtimes. + pub capacity: AtomicI64, + pub in_flight: AtomicI64, + /// Jobs producers enqueued here and executions begun, for rough backlog estimates. + pub enqueued: AtomicI64, + pub started: AtomicI64, + /// An operator (or a rate limit) is deliberately holding jobs back. + /// Operators currently holding the queue paused; the stall check skips it while non-zero. + pub pause_holds: AtomicI32, + pub rate_limited: bool, + /// Milliseconds since run start of the most recent operator drain, 0 if never. + pub drained_at_ms: AtomicI64, +} + +#[derive(Default)] +pub struct QueueRegistry { + queues: RwLock>>, + by_name: RwLock>, +} + +impl QueueRegistry { + pub fn register(&self, name: &str, persona: Persona, rate_limited: bool) -> Arc { + let mut queues = self.queues.write().expect("registry lock"); + let mut by_name = self.by_name.write().expect("registry lock"); + if let Some(index) = by_name.get(name) { + return queues[*index as usize].clone(); + } + let index = queues.len() as u32; + let info = Arc::new(QueueInfo { + name: name.to_string(), + persona, + index, + capacity: AtomicI64::new(0), + in_flight: AtomicI64::new(0), + enqueued: AtomicI64::new(0), + started: AtomicI64::new(0), + pause_holds: AtomicI32::new(0), + rate_limited, + drained_at_ms: AtomicI64::new(0), + }); + by_name.insert(name.to_string(), index); + queues.push(info.clone()); + info + } + + /// Looks a queue up by its physical name; striped queues (`name#N`) resolve to + /// their logical queue. + pub fn get(&self, name: &str) -> Option> { + let by_name = self.by_name.read().expect("registry lock"); + let index = by_name.get(name).or_else(|| { + let (logical, suffix) = name.rsplit_once('#')?; + suffix.parse::().ok()?; + by_name.get(logical) + })?; + Some(self.queues.read().expect("registry lock")[*index as usize].clone()) + } + + pub fn by_index(&self, index: u32) -> Option> { + self.queues + .read() + .expect("registry lock") + .get(index as usize) + .cloned() + } + + pub fn all(&self) -> Vec> { + self.queues.read().expect("registry lock").clone() + } + + pub fn of(&self, persona: Persona) -> Vec> { + self.all() + .into_iter() + .filter(|queue| queue.persona == persona) + .collect() + } +} + +/// Counts one running handler against its queue for as long as the guard lives, +/// including when the runtime is killed and the handler future is dropped. +pub struct InFlightGuard { + queue: Option>, +} + +impl InFlightGuard { + pub fn new(queue: Option>) -> Self { + if let Some(queue) = &queue { + queue.in_flight.fetch_add(1, Ordering::Relaxed); + } + Self { queue } + } +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + if let Some(queue) = &self.queue { + queue.in_flight.fetch_sub(1, Ordering::Relaxed); + } + } +} diff --git a/awa/tests/persona_soak/report.rs b/awa/tests/persona_soak/report.rs new file mode 100644 index 000000000..d309bfc94 --- /dev/null +++ b/awa/tests/persona_soak/report.rs @@ -0,0 +1,341 @@ +//! The run report and, for a failed run, the reproducer. + +use super::config::{Persona, SoakConfig}; +use super::pgserver::crash_excerpts; +use super::shared::Shared; +use serde_json::json; +use std::fmt::Write as _; +use std::path::Path; +use std::process::Command; + +/// Printed at the top of every report; also lets a built binary be matched to this harness. +pub const HEADER: &str = "AWA PERSONA SOAK HARNESS"; + +pub fn git_revision() -> String { + Command::new("git") + .args(["rev-parse", "--short=10", "HEAD"]) + .output() + .ok() + .map(|output| String::from_utf8_lossy(&output.stdout).trim().to_string()) + .filter(|revision| !revision.is_empty()) + .unwrap_or_else(|| "unknown".into()) +} + +pub struct Outcome { + pub quiesced: bool, + pub elapsed_seconds: f64, + pub crash_excerpts: Vec, + pub log_tail: Option, +} + +pub fn render(shared: &Shared, outcome: &Outcome) -> String { + let config = &shared.config; + let info = &shared.info; + let violations = shared.violations.lock().expect("violations"); + let mut out = String::new(); + let _ = writeln!(out, "# {HEADER}"); + let _ = writeln!(out); + let _ = writeln!( + out, + "verdict: {}", + if violations + .iter() + .any(|violation| violation.kind != "oracle_incomplete") + { + "FAIL" + } else if shared + .oracle_incomplete + .load(std::sync::atomic::Ordering::Relaxed) + { + "INCONCLUSIVE" + } else { + "PASS" + } + ); + let _ = writeln!( + out, + "engine: {} (schema {})", + info.engine, + info.queue_storage_schema.as_deref().unwrap_or("-") + ); + let _ = writeln!(out, "server: PostgreSQL {}", info.version); + let _ = writeln!( + out, + "max_connections: {} (usable by personas: {})", + info.max_connections, + shared.budget.usable() + ); + let _ = writeln!(out, "connections at start: {}", info.connections_at_start); + let _ = writeln!(out, "awa: {} @ {}", info.awa_version, git_revision()); + let _ = writeln!( + out, + "seed: {} personas: {} load: {}s scale: {} elapsed: {:.0}s quiesced: {}", + config.seed, + config + .personas + .iter() + .map(|persona| persona.name()) + .collect::>() + .join(","), + config.load.as_secs(), + config.scale, + outcome.elapsed_seconds, + outcome.quiesced + ); + let _ = writeln!(out, "ledger jobs: {}", shared.ledger.len()); + let _ = writeln!(out, "NO_EVENTS durable loss detection: job_id % {} = 0; other returned NO_EVENTS jobs are excluded from terminal-commit proof.", super::ledger::TERMINAL_PROOF_STRIDE); + let _ = writeln!(out); + + let _ = writeln!(out, "## Connection budget"); + for (label, count) in shared.budget.table() { + let _ = writeln!(out, "- {label}: {count}"); + } + let _ = writeln!(out); + + for persona in Persona::ALL { + if !config.has(persona) { + continue; + } + let stats = shared.stats(persona); + let counters = stats.counters(); + let _ = writeln!(out, "## {}", persona.name()); + for (name, value) in counters { + let _ = writeln!(out, "- {name}: {value}"); + } + for (name, summary) in stats.summaries() { + let _ = writeln!( + out, + "- {name}: n={} p50={} p99={} max={}", + summary.count, summary.p50, summary.p99, summary.max + ); + } + if persona == Persona::ClaimAbandon { + let _ = writeln!(out, "{}", shared.abandon.render()); + } + let _ = writeln!(out); + } + + let _ = writeln!(out, "## Run"); + if shared.run.get("claim_ring_lock_capture_records") > 0 { + let _ = writeln!( + out, + "Lock evidence: [claim-ring-locks.jsonl](claim-ring-locks.jsonl)" + ); + } + for (name, value) in shared.run.counters() { + let _ = writeln!(out, "- {name}: {value}"); + } + for (name, summary) in shared.run.summaries() { + let _ = writeln!( + out, + "- {name}: n={} p50={} p99={} max={}", + summary.count, summary.p50, summary.p99, summary.max + ); + } + let _ = writeln!(out); + + let findings = shared.findings.lock().expect("findings"); + let _ = writeln!(out, "## Findings, not fatal ({})", findings.len()); + for finding in findings.iter() { + let _ = writeln!( + out, + "- [{:.1}s] {}: {}", + finding.at_ms as f64 / 1000.0, + finding.kind, + finding.summary + ); + for row in finding.rows.iter().take(6) { + let _ = writeln!(out, " {row}"); + } + } + let _ = writeln!(out); + + let _ = writeln!(out, "## Violations ({})", violations.len()); + for violation in violations.iter() { + let _ = writeln!( + out, + "- [{:.1}s{}] {}: {}", + violation.at_ms as f64 / 1000.0, + if violation.during_chaos { + ", during chaos" + } else { + "" + }, + violation.kind, + violation.summary + ); + for row in violation.rows.iter().take(10) { + let _ = writeln!(out, " {row}"); + } + } + if !outcome.crash_excerpts.is_empty() { + let _ = writeln!(out, "\n## Server log: crash lines"); + for excerpt in &outcome.crash_excerpts { + let _ = writeln!(out, "```\n{excerpt}\n```"); + } + } + out +} + +/// Writes `report.md` and `report.json`, and for a failed run `reproducer.md`. +pub fn write(shared: &Shared, outcome: &Outcome) -> std::io::Result<()> { + let dir = &shared.config.report_dir; + std::fs::create_dir_all(dir)?; + let markdown = render(shared, outcome); + println!("\n{markdown}"); + std::fs::write(dir.join("report.md"), &markdown)?; + + let violations = shared.violations.lock().expect("violations").clone(); + let timeline = shared.timeline.lock().expect("timeline").clone(); + let per_persona: serde_json::Map = Persona::ALL + .into_iter() + .filter(|persona| shared.config.has(*persona)) + .map(|persona| { + ( + persona.name().to_string(), + json!({ "counters": shared.stats(persona).counters(), "samples": shared.stats(persona).summaries() }), + ) + }) + .collect(); + let summary = json!({ + "verdict": if violations.iter().any(|violation| violation.kind != "oracle_incomplete") { "fail" } + else if shared.oracle_incomplete.load(std::sync::atomic::Ordering::Relaxed) { "inconclusive" } else { "pass" }, + "terminal_proof_sample_stride": super::ledger::TERMINAL_PROOF_STRIDE, + "oracle_complete": !shared.oracle_incomplete.load(std::sync::atomic::Ordering::Relaxed), + "seed": shared.config.seed, + "personas": shared.config.personas.iter().map(|p| p.name()).collect::>(), + "engine": shared.info.engine, + "server_version": shared.info.version, + "max_connections": shared.info.max_connections, + "awa_version": shared.info.awa_version, + "git": git_revision(), + "elapsed_seconds": outcome.elapsed_seconds, + "quiesced": outcome.quiesced, + "personas_stats": per_persona, + "abandon": shared.abandon.to_json(), + "run": { "counters": shared.run.counters(), "samples": shared.run.summaries() }, + "violations": violations, + "findings": shared.findings.lock().expect("findings").clone(), + "timeline": timeline, + }); + std::fs::write( + dir.join("report.json"), + serde_json::to_string_pretty(&summary).unwrap_or_default(), + )?; + + if !violations.is_empty() { + write_reproducer(dir, &shared.config, shared, outcome)?; + } + Ok(()) +} + +fn write_reproducer( + dir: &Path, + config: &SoakConfig, + shared: &Shared, + outcome: &Outcome, +) -> std::io::Result<()> { + let violations = shared.violations.lock().expect("violations").clone(); + let timeline = shared.timeline.lock().expect("timeline").clone(); + let mut out = String::new(); + let _ = writeln!(out, "# Reproducer\n"); + let _ = writeln!(out, "Re-run with the same seed, personas and duration:\n"); + let _ = writeln!( + out, + "```\nAWA_SOAK_DATABASE_URL= AWA_SOAK_SEED={} AWA_SOAK_PERSONAS={} AWA_SOAK_SECONDS={} AWA_SOAK_SCALE={} \\\n cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture\n```\n", + config.seed, + config.personas.iter().map(|persona| persona.name()).collect::>().join(","), + config.load.as_secs(), + config.scale + ); + let _ = writeln!(out, "Seed fixes the workload mix and the chaos schedule; timing of races still varies between runs.\n"); + let _ = writeln!( + out, + "Engine: {}; PostgreSQL {}; awa {} @ {}\n", + shared.info.engine, + shared.info.version, + shared.info.awa_version, + git_revision() + ); + let _ = writeln!(out, "## Violations\n"); + for violation in &violations { + let persona = violation.kind.as_str(); + let _ = writeln!( + out, + "### {persona} at {:.1}s{}\n{}\n", + violation.at_ms as f64 / 1000.0, + if violation.during_chaos { + " (during chaos)" + } else { + "" + }, + violation.summary + ); + for row in &violation.rows { + let _ = writeln!(out, " {row}"); + } + let nearby: Vec<_> = timeline + .iter() + .filter(|event| (event.at_ms - violation.at_ms).abs() <= 60_000) + .collect(); + if !nearby.is_empty() { + let _ = writeln!(out, "\nTimeline within 60s:"); + for event in nearby { + let _ = writeln!( + out, + "- {:.1}s [{}] {}", + event.at_ms as f64 / 1000.0, + event.source, + event.text + ); + } + } + let _ = writeln!(out); + } + let _ = writeln!(out, "## Full timeline\n"); + for event in &timeline { + let _ = writeln!( + out, + "- {:.1}s [{}] {}", + event.at_ms as f64 / 1000.0, + event.source, + event.text + ); + } + let _ = writeln!(out, "\n## Server log excerpt\n"); + if !outcome.crash_excerpts.is_empty() { + for excerpt in &outcome.crash_excerpts { + let _ = writeln!(out, "```\n{excerpt}\n```"); + } + } else if let Some(tail) = &outcome.log_tail { + let interesting: Vec<&str> = tail + .lines() + .filter(|line| { + line.contains("ERROR") + || line.contains("FATAL") + || line.contains("PANIC") + || line.contains("WARNING") + }) + .collect(); + let start = interesting.len().saturating_sub(80); + let _ = writeln!( + out, + "No crash lines. Last warnings and errors ({} total):\n```", + interesting.len() + ); + for line in &interesting[start..] { + let _ = writeln!(out, "{line}"); + } + let _ = writeln!(out, "```"); + } else { + let _ = writeln!( + out, + "No server log source configured (set AWA_SOAK_PG_LOG or AWA_SOAK_PG_CONTAINER)." + ); + } + std::fs::write(dir.join("reproducer.md"), out) +} + +pub fn excerpts_of(log: &str) -> Vec { + crash_excerpts(log) +} diff --git a/awa/tests/persona_soak/ring_lock_capture.sql b/awa/tests/persona_soak/ring_lock_capture.sql new file mode 100644 index 000000000..165990fe1 --- /dev/null +++ b/awa/tests/persona_soak/ring_lock_capture.sql @@ -0,0 +1,63 @@ +WITH current_db AS ( + SELECT oid FROM pg_catalog.pg_database WHERE datname = current_database() +), +relations AS MATERIALIZED ( + SELECT c.oid, c.relname::text AS relation + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname::text = ANY($2::text[]) +), +activity AS MATERIALIZED ( + SELECT pid, leader_pid, backend_start, application_name, usename, state, + wait_event_type, wait_event, xact_start, query_start, + EXTRACT(EPOCH FROM (clock_timestamp() - xact_start))::float8 AS tx_age_s, + EXTRACT(EPOCH FROM (clock_timestamp() - query_start))::float8 AS query_age_s, + left(query, 512) AS query, + query IS NULL OR query = '' AS visibility_limited + FROM pg_catalog.pg_stat_activity + WHERE datid = (SELECT oid FROM current_db) +), +locks AS MATERIALIZED ( + SELECT l.pid, r.oid AS relation_oid, r.relation, l.mode, l.granted, l.fastpath, + l.virtualtransaction, l.waitstart + FROM pg_catalog.pg_locks l + JOIN relations r ON r.oid = l.relation + LEFT JOIN activity a ON a.pid = l.pid + WHERE l.database = (SELECT oid FROM current_db) + AND l.locktype = 'relation' + AND l.pid IS DISTINCT FROM pg_backend_pid() + ORDER BY l.granted, a.xact_start NULLS LAST, a.query_start NULLS LAST, + l.pid NULLS FIRST, r.relation, l.mode + LIMIT 257 --256 retained rows plus an explicit overflow sentinel +), +interested AS MATERIALIZED ( + SELECT DISTINCT l.pid + FROM locks l + JOIN activity a ON a.pid = l.pid + WHERE a.wait_event_type = 'Lock' + ORDER BY l.pid + LIMIT 32 +), +blockers AS MATERIALIZED ( + SELECT i.pid, pg_catalog.pg_blocking_pids(i.pid) AS pids + FROM interested i +), +lock_rows AS MATERIALIZED ( +SELECT clock_timestamp() AS observed_at, l.*, + a.leader_pid, a.backend_start, a.application_name, a.usename, a.state, + a.wait_event_type, a.wait_event, a.xact_start, a.query_start, + a.tx_age_s, a.query_age_s, a.query, a.visibility_limited, + a.pid IS NULL AS activity_missing, + b.pids[1:16] AS blocking_pids, + b.pid IS NOT NULL AS blocking_pids_observed, + COALESCE(cardinality(b.pids) > 16, false) AS blocking_pids_truncated +FROM locks l +LEFT JOIN activity a ON a.pid = l.pid +LEFT JOIN blockers b ON b.pid = l.pid +) +SELECT jsonb_build_object( + 'server_statement_at', statement_timestamp(), + 'server_observed_at', clock_timestamp(), + 'resolved_relations', COALESCE((SELECT jsonb_agg(to_jsonb(r)) FROM relations r), '[]'::jsonb), + 'locks', COALESCE((SELECT jsonb_agg(to_jsonb(l) ORDER BY l.granted, l.xact_start NULLS LAST, l.pid NULLS FIRST, l.relation, l.mode) FROM lock_rows l), '[]'::jsonb) +); diff --git a/awa/tests/persona_soak/ring_locks.rs b/awa/tests/persona_soak/ring_locks.rs new file mode 100644 index 000000000..2045ed1de --- /dev/null +++ b/awa/tests/persona_soak/ring_locks.rs @@ -0,0 +1,543 @@ +//! Supplemental lock evidence; never changes the claim-ring oracle's verdict. + +use super::abandon_oracle::RingEpisode; +use super::shared::Shared; +use serde_json::{json, Value}; +use sqlx::PgPool; +use std::collections::BTreeSet; +use std::io::Write; +use std::time::Duration; + +const SQL: &str = include_str!("ring_lock_capture.sql"); +const LOCK_ROWS: usize = 256; +const BLOCKER_PIDS: usize = 64; +const CLIENT_DEADLINE: Duration = Duration::from_millis(750); + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Stage { + Suspected, + Violation, +} + +impl Stage { + fn index(self) -> usize { + match self { + Self::Suspected => 0, + Self::Violation => 1, + } + } + + fn name(self) -> &'static str { + match self { + Self::Suspected => "suspected", + Self::Violation => "violation", + } + } +} + +#[derive(Debug, PartialEq)] +pub enum Decision { + Capture(Stage), + Exhausted(Stage), +} + +/// Four slots per stage reserve half the run budget for actual violations. +#[derive(Default)] +pub struct CaptureBudget { + episode: Option<(i64, i64)>, + seen: [bool; 2], + used: [u8; 2], + exhausted_reported: [bool; 2], +} + +impl CaptureBudget { + pub fn plan( + &mut self, + episode: &RingEpisode, + now_ms: i64, + bound_ms: i64, + zero_live: bool, + ) -> Option { + let key = (episode.generation, episode.start_ms); + if self.episode != Some(key) { + self.episode = Some(key); + self.seen = [false; 2]; + } + let stage = if episode.reported { + Stage::Violation + } else if zero_live && now_ms - episode.start_ms >= bound_ms / 2 { + Stage::Suspected + } else { + return None; + }; + let index = stage.index(); + if self.seen[index] { + return None; + } + self.seen[index] = true; + if self.used[index] < 4 { + self.used[index] += 1; + Some(Decision::Capture(stage)) + } else if !self.exhausted_reported[index] { + self.exhausted_reported[index] = true; + Some(Decision::Exhausted(stage)) + } else { + None + } + } +} + +fn names(slot: i32) -> Vec { + [ + format!("lease_claims_{slot}"), + format!("lease_claim_batches_{slot}"), + format!("lease_claim_closures_{slot}"), + format!("lease_claim_closure_batches_{slot}"), + "claim_ring_state".into(), + "claim_ring_slots".into(), + ] + .into() +} + +fn missing(error: impl std::fmt::Display) -> Value { + json!({"status": "missing", "error": error.to_string().chars().take(512).collect::()}) +} + +fn finish(mut value: Value, expected: &[String]) -> Value { + let mut gaps = Vec::new(); + let resolved: BTreeSet<_> = value["resolved_relations"] + .as_array() + .into_iter() + .flatten() + .filter_map(|row| row["relation"].as_str()) + .collect(); + let unresolved: Vec<_> = expected + .iter() + .filter(|name| !resolved.contains(name.as_str())) + .cloned() + .collect(); + if !unresolved.is_empty() { + gaps.push("unresolved_relations"); + } + if let Some(rows) = value["locks"].as_array_mut() { + if rows.len() > LOCK_ROWS { + rows.truncate(LOCK_ROWS); + gaps.push("lock_rows_truncated"); + } + if rows.iter().any(|row| { + row["blocking_pids_truncated"] == true + || (row["wait_event_type"] == "Lock" && row["blocking_pids_observed"] != true) + || row["activity_missing"] == true + || row["visibility_limited"] == true + }) { + gaps.push("incomplete_activity_or_blocker_graph"); + } + } + if value["blocker_pid_lookup_truncated"] == true { + gaps.push("blocker_pid_lookup_truncated"); + } + if value["missing_blocker_activity"] + .as_array() + .is_some_and(|rows| !rows.is_empty()) + { + gaps.push("missing_blocker_activity"); + } + if value["blocker_activity"] + .as_array() + .is_some_and(|rows| rows.iter().any(|row| row["visibility_limited"] == true)) + { + gaps.push("blocker_activity_visibility_limited"); + } + value["status"] = json!(if gaps.is_empty() { "ok" } else { "partial" }); + value["missing"] = json!(gaps); + value["unresolved_relations"] = json!(unresolved); + value +} + +async fn collect(pool: &PgPool, schema: &str, slot: i32, sql: &str) -> Value { + let expected = names(slot); + let result = tokio::time::timeout(CLIENT_DEADLINE, async { + let mut tx = pool.begin().await?; + sqlx::query("SET TRANSACTION READ ONLY") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL statement_timeout = '250ms'") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL lock_timeout = '100ms'") + .execute(tx.as_mut()) + .await?; + let mut value: Value = sqlx::query_scalar(sql) + .bind(schema) + .bind(&expected) + .fetch_one(tx.as_mut()) + .await?; + let pids: BTreeSet = value["locks"] + .as_array() + .into_iter() + .flatten() + .flat_map(|row| row["blocking_pids"].as_array().into_iter().flatten()) + .filter_map(|pid| pid.as_i64().and_then(|pid| i32::try_from(pid).ok())) + .filter(|pid| *pid > 0) + .collect(); + value["blocker_pid_lookup_truncated"] = json!(pids.len() > BLOCKER_PIDS); + let requested: Vec<_> = pids.into_iter().take(BLOCKER_PIDS).collect(); + let activity: Value = sqlx::query_scalar( + "SELECT COALESCE(jsonb_agg(to_jsonb(a)), '[]'::jsonb) FROM ( + SELECT pid, leader_pid, backend_start, application_name, usename, state, + wait_event_type, wait_event, xact_start, query_start, + EXTRACT(EPOCH FROM(clock_timestamp()-xact_start))::float8 AS tx_age_s, + EXTRACT(EPOCH FROM(clock_timestamp()-query_start))::float8 AS query_age_s, + left(query,512) AS query, + query IS NULL OR query = '' AS visibility_limited + FROM pg_catalog.pg_stat_activity + WHERE datid=(SELECT oid FROM pg_catalog.pg_database WHERE datname=current_database()) + AND pid=ANY($1::int4[]) ORDER BY pid LIMIT 65 + ) a", + ) + .bind(&requested) + .fetch_one(tx.as_mut()) + .await?; + let found: BTreeSet<_> = activity + .as_array() + .into_iter() + .flatten() + .filter_map(|row| row["pid"].as_i64()) + .collect(); + value["missing_blocker_activity"] = json!(requested + .iter() + .filter(|pid| !found.contains(&i64::from(**pid))) + .collect::>()); + value["blocker_activity"] = activity; + tx.rollback().await?; + Ok::<_, sqlx::Error>(value) + }) + .await; + match result { + Ok(Ok(value)) => finish(value, &expected), + Ok(Err(error)) => missing(error), + Err(_) => missing("claim-ring metadata capture exceeded its 750ms client deadline"), + } +} + +pub async fn observe( + shared: &Shared, + pool: &PgPool, + schema: &str, + slot_count: i32, + next_slot: Option<(i64, i64)>, + budget: &mut CaptureBudget, +) { + let Some(episode) = shared.abandon.current_ring_episode() else { + return; + }; + let now = shared.now_ms(); + let Some(decision) = budget.plan( + &episode, + now, + shared.config.abandon_bound.as_millis() as i64, + next_slot.is_some_and(|(live, _)| live == 0), + ) else { + return; + }; + let stage = match decision { + Decision::Capture(stage) => stage, + Decision::Exhausted(stage) => { + shared.finding( + "claim_ring_lock_capture_missing", + format!("{} lock-capture budget exhausted", stage.name()), + vec![], + ); + return; + } + }; + let incoming = (episode.slot + 1).rem_euclid(slot_count); + let mut value = collect(pool, schema, incoming, SQL).await; + value["stage"] = json!(stage.name()); + value["started_at_ms"] = json!(now); + value["finished_at_ms"] = json!(shared.now_ms()); + value["current_slot"] = json!(episode.slot); + value["next_slot"] = json!(incoming); + value["generation"] = json!(episode.generation); + value["next_slot_probe"] = json!(next_slot); + shared.run.add("claim_ring_lock_capture_attempts", 1); + if value["status"] != "ok" { + shared.finding( + "claim_ring_lock_capture_missing", + "claim-ring lock evidence is missing or partial", + vec![value.to_string()], + ); + } + let path = shared.config.report_dir.join("claim-ring-locks.jsonl"); + let written = (|| { + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path)?; + writeln!(file, "{value}") + })(); + match written { + Ok(()) => shared.run.add("claim_ring_lock_capture_records", 1), + Err(error) => shared.finding( + "claim_ring_lock_capture_missing", + format!("could not write claim-ring lock evidence: {error}"), + vec![], + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn episode(generation: i64, reported: bool) -> RingEpisode { + RingEpisode { + generation, + reported, + slot: 7, + start_ms: 0, + end_ms: None, + last_live_ms: 0, + expired_in_next_slot: 0, + } + } + + #[test] + fn ring_lock_budget_waits_half_bound_and_reserves_violation_captures() { + let mut budget = CaptureBudget::default(); + assert_eq!(budget.plan(&episode(0, false), 59_999, 120_000, true), None); + assert_eq!( + budget.plan(&episode(0, false), 60_000, 120_000, false), + None + ); + for generation in 1..=4 { + assert_eq!( + budget.plan(&episode(generation, false), 60_000, 120_000, true), + Some(Decision::Capture(Stage::Suspected)) + ); + assert_eq!( + budget.plan(&episode(generation, false), 70_000, 120_000, true), + None + ); + } + assert_eq!( + budget.plan(&episode(5, false), 60_000, 120_000, true), + Some(Decision::Exhausted(Stage::Suspected)) + ); + assert_eq!(budget.plan(&episode(6, false), 60_000, 120_000, true), None); + for generation in 7..=10 { + assert_eq!( + budget.plan(&episode(generation, true), 120_001, 120_000, true), + Some(Decision::Capture(Stage::Violation)) + ); + assert_eq!( + budget.plan(&episode(generation, true), 130_000, 120_000, true), + None + ); + } + assert_eq!( + budget.plan(&episode(11, true), 120_001, 120_000, true), + Some(Decision::Exhausted(Stage::Violation)) + ); + assert_eq!( + budget.plan(&episode(12, true), 120_001, 120_000, true), + None + ); + assert_eq!(budget.used, [4, 4]); + } + + #[test] + fn ring_lock_result_caps_and_missing_relations_stay_explicit() { + let rows = vec![json!({"blocking_pids_truncated": true}); LOCK_ROWS + 1]; + let value = finish(json!({"resolved_relations": [], "locks": rows}), &names(0)); + assert_eq!(value["status"], "partial"); + assert_eq!(value["locks"].as_array().unwrap().len(), LOCK_ROWS); + assert_eq!(value["unresolved_relations"].as_array().unwrap().len(), 6); + let gaps = value["missing"].as_array().unwrap(); + assert!(gaps.contains(&json!("lock_rows_truncated"))); + assert!(gaps.contains(&json!("incomplete_activity_or_blocker_graph"))); + } + + #[test] + fn ring_lock_hidden_second_lookup_is_partial_even_without_selected_locks() { + let expected = names(0); + let inventory: Vec<_> = expected + .iter() + .map(|name| json!({"relation": name})) + .collect(); + let value = finish( + json!({ + "resolved_relations": inventory, + "locks": [], + "blocker_activity": [{"pid": 42, "query": null, "visibility_limited": true}] + }), + &expected, + ); + assert_eq!(value["status"], "partial"); + assert!(value["missing"] + .as_array() + .unwrap() + .contains(&json!("blocker_activity_visibility_limited"))); + } + + struct Database { + config: super::super::config::SoakConfig, + pool: PgPool, + } + + impl Database { + async fn create() -> Self { + use super::super::config::{replace_database_name, SoakConfig}; + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_ring_locks_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + super::super::db::recreate_database(&config) + .await + .expect("isolated database"); + let pool = super::super::db::connect_pool(&config.database_url, "ring_capture_test", 6) + .await + .expect("fixture pool"); + Self { config, pool } + } + + async fn drop(self) { + self.pool.close().await; + let admin = + super::super::db::connect_pool(&self.config.admin_url(), "ring_capture_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(&format!( + "DROP DATABASE {} WITH (FORCE)", + self.config.database_name() + )) + .execute(&admin) + .await + .expect("drop fixture database"); + admin.close().await; + } + } + + #[tokio::test] + async fn ring_lock_capture_identifies_holder_waiter_and_releases_cleanly() { + let db = Database::create().await; + let pool = &db.pool; + let mut holder = pool.begin().await.unwrap(); + sqlx::query("SET LOCAL application_name = 'ring_capture_holder'") + .execute(holder.as_mut()) + .await + .unwrap(); + sqlx::query("LOCK TABLE awa.lease_claims_0 IN ACCESS SHARE MODE") + .execute(holder.as_mut()) + .await + .unwrap(); + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(holder.as_mut()) + .await + .unwrap(); + let mut waiter = pool.begin().await.unwrap(); + sqlx::query("SET LOCAL application_name = 'ring_capture_waiter'") + .execute(waiter.as_mut()) + .await + .unwrap(); + let waiter_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(waiter.as_mut()) + .await + .unwrap(); + let waiting = tokio::spawn(async move { + sqlx::query("LOCK TABLE awa.lease_claims_0 IN ACCESS EXCLUSIVE MODE") + .execute(waiter.as_mut()) + .await + .unwrap(); + waiter.rollback().await.unwrap(); + }); + let deadline = + tokio::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_locks WHERE pid=$1 AND NOT granted)", + ) + .bind(waiter_pid) + .fetch_one(pool) + .await + .unwrap(); + if waiting { + break; + } + assert!( + tokio::time::Instant::now() < deadline, + "exclusive waiter did not queue" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + let value = collect(pool, "awa", 0, SQL).await; + assert_eq!(value["status"], "ok", "{value}"); + assert_eq!(value["resolved_relations"].as_array().unwrap().len(), 6); + let rows = value["locks"].as_array().unwrap(); + let held = rows + .iter() + .find(|row| row["pid"] == holder_pid && row["relation"] == "lease_claims_0") + .expect("holder evidence"); + assert_eq!(held["mode"], "AccessShareLock"); + assert_eq!(held["granted"], true); + assert_eq!(held["application_name"], "ring_capture_holder"); + assert!(held["tx_age_s"].as_f64().unwrap() >= 0.0); + let blocked = rows + .iter() + .find(|row| row["pid"] == waiter_pid && row["granted"] == false) + .expect("waiter evidence"); + assert_eq!(blocked["mode"], "AccessExclusiveLock"); + assert_eq!(blocked["application_name"], "ring_capture_waiter"); + assert_eq!(blocked["wait_event_type"], "Lock"); + assert!(blocked["blocking_pids"] + .as_array() + .unwrap() + .contains(&json!(holder_pid))); + assert!(value["blocker_activity"] + .as_array() + .unwrap() + .iter() + .any(|row| row["pid"] == holder_pid)); + holder.rollback().await.unwrap(); + tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(5)), + waiting, + ) + .await + .unwrap() + .unwrap(); + let clean = collect(pool, "awa", 0, SQL).await; + assert_eq!(clean["status"], "ok", "{clean}"); + assert!(clean["server_statement_at"].as_str().is_some()); + assert!(clean["server_observed_at"].as_str().is_some()); + assert_eq!(clean["resolved_relations"].as_array().unwrap().len(), 6); + assert!(clean["locks"] + .as_array() + .unwrap() + .iter() + .all(|row| row["pid"] != holder_pid && row["pid"] != waiter_pid)); + db.drop().await; + } + + #[tokio::test] + async fn ring_lock_capture_failure_and_unresolved_inventory_are_not_empty_success() { + let db = Database::create().await; + let value = collect(&db.pool, "missing_ring_schema", 0, SQL).await; + assert_eq!(value["status"], "partial"); + assert_eq!(value["resolved_relations"], json!([])); + assert_eq!(value["locks"], json!([])); + assert!(value["server_statement_at"].as_str().is_some()); + assert!(value["server_observed_at"].as_str().is_some()); + assert_eq!(value["unresolved_relations"].as_array().unwrap().len(), 6); + let failed = collect(&db.pool, "awa", 0, + "SELECT jsonb_build_object() FROM pg_sleep(1) WHERE $1::text IS NOT NULL AND $2::text[] IS NOT NULL").await; + assert_eq!(failed["status"], "missing", "{failed}"); + assert!( + failed.get("locks").is_none(), + "failed capture must not look unlocked" + ); + assert!(!failed["error"].as_str().unwrap().is_empty()); + db.drop().await; + } +} diff --git a/awa/tests/persona_soak/rng.rs b/awa/tests/persona_soak/rng.rs new file mode 100644 index 000000000..ca836b27b --- /dev/null +++ b/awa/tests/persona_soak/rng.rs @@ -0,0 +1,64 @@ +//! Deterministic pseudo-random numbers so a failing soak can be replayed from its seed. + +/// SplitMix64: tiny, fast, and good enough for workload shaping. +#[derive(Debug, Clone)] +pub struct Rng { + state: u64, +} + +impl Rng { + pub fn new(seed: u64) -> Self { + Self { state: seed } + } + + /// A generator for one task, derived from the run seed and a stable label so + /// that adding a persona does not shift another persona's sequence. + pub fn for_task(seed: u64, label: &str) -> Self { + let mut hash = seed ^ 0xcbf2_9ce4_8422_2325; + for byte in label.bytes() { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + Self::new(hash) + } + + pub fn next_u64(&mut self) -> u64 { + self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut value = self.state; + value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + value ^ (value >> 31) + } + + /// Uniform in `0..bound` (`bound` must be non-zero). + pub fn below(&mut self, bound: u64) -> u64 { + self.next_u64() % bound.max(1) + } + + /// Uniform in `low..=high`. + pub fn range(&mut self, low: u64, high: u64) -> u64 { + low + self.below(high.saturating_sub(low) + 1) + } + + /// True with probability `percent` / 100. + pub fn chance(&mut self, percent: u64) -> bool { + self.below(100) < percent + } + + pub fn pick<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> { + if items.is_empty() { + None + } else { + Some(&items[self.below(items.len() as u64) as usize]) + } + } +} + +/// Nearest-rank percentile of a sorted slice (`quantile` in `0.0..=1.0`). +pub fn percentile(sorted: &[i64], quantile: f64) -> i64 { + if sorted.is_empty() { + return 0; + } + let index = ((sorted.len() - 1) as f64 * quantile).round() as usize; + sorted[index] +} diff --git a/awa/tests/persona_soak/roles.rs b/awa/tests/persona_soak/roles.rs new file mode 100644 index 000000000..bf819c3a0 --- /dev/null +++ b/awa/tests/persona_soak/roles.rs @@ -0,0 +1,143 @@ +//! The database roles of the shard-ops persona: the shape of a production install that +//! runs its application processes without superuser or schema ownership. +//! +//! * `OWNER` owns the database and the `awa` schema, runs migrations and `prepare_queue`. +//! * `MIGRATOR` is a migration-like role that is a member of `OWNER` and has default +//! privileges for the objects it creates. +//! * `APP_A` and `APP_B` are members of the `RUNTIME_GROUP` role and have no default +//! privileges of their own by default. An explicit comparison profile can share +//! their newly created objects too. They run producers and the maintenance leader. +//! +//! Grants follow the restricted-runtime section of `docs/security.md`. + +use super::config::{replace_credentials, SoakConfig}; +use super::db::connect_pool; +use sqlx::PgPool; + +pub const OWNER: &str = "awa_soak_owner"; +pub const MIGRATOR: &str = "awa_soak_migrator"; +pub const RUNTIME_GROUP: &str = "awa_soak_runtime"; +pub const APP_A: &str = "awa_soak_app_a"; +pub const APP_B: &str = "awa_soak_app_b"; +const PASSWORD: &str = "soak"; + +pub fn url_for_role(config: &SoakConfig, role: &str) -> String { + replace_credentials(&config.database_url, role, PASSWORD) +} + +/// Creates the cluster-level roles. Safe to repeat. +pub async fn create_roles(admin: &PgPool) -> Result<(), String> { + let login = [OWNER, MIGRATOR, APP_A, APP_B]; + let statements: Vec = login + .iter() + .map(|role| { + format!( + "DO $$ BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '{role}') THEN + CREATE ROLE {role} LOGIN PASSWORD '{PASSWORD}' NOSUPERUSER NOCREATEDB NOCREATEROLE; + END IF; + END $$" + ) + }) + .chain([ + format!( + "DO $$ BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '{RUNTIME_GROUP}') THEN + CREATE ROLE {RUNTIME_GROUP} NOLOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE; + END IF; + END $$" + ), + format!("GRANT {OWNER} TO {MIGRATOR}"), + format!("GRANT {RUNTIME_GROUP} TO {APP_A}, {APP_B}"), + ]) + .collect(); + for statement in statements { + sqlx::query(&statement) + .execute(admin) + .await + .map_err(|error| format!("create roles: {error}"))?; + } + Ok(()) +} + +/// Applies the restricted-runtime grants. Run once the schema is migrated, over a +/// superuser connection to the soak database. +pub async fn grant_runtime( + superuser_pool: &PgPool, + database: &str, + schema_create: bool, + mixed_personas: bool, +) -> Result<(), String> { + // Comparison-only setup profile for older binaries whose application roles + // create their own lane sequences. Default off preserves regression coverage. + let shared_creator_grants = match std::env::var("AWA_SOAK_SHARED_CREATOR_GRANTS") { + Ok(value) if value == "1" => true, + Ok(value) if value == "0" => false, + Err(std::env::VarError::NotPresent) => false, + _ => return Err("AWA_SOAK_SHARED_CREATOR_GRANTS must be 0 or 1".into()), + }; + if shared_creator_grants && !schema_create { + return Err("shared creator grants require schema CREATE for the runtime group".into()); + } + let mut statements = vec![ + format!("GRANT USAGE ON SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("REVOKE EXECUTE ON FUNCTION awa.install_queue_storage_substrate(TEXT, INT, INT, INT, BOOLEAN) FROM {RUNTIME_GROUP}"), + format!("GRANT TEMP ON DATABASE {database} TO {RUNTIME_GROUP}"), + ]; + if schema_create { + statements.push(format!("GRANT CREATE ON SCHEMA awa TO {RUNTIME_GROUP}")); + } + let mut creators = vec![OWNER.to_string(), MIGRATOR.to_string()]; + if mixed_personas { + // The other personas keep their superuser runtimes and producers. Their objects + // must stay usable by the application roles, which now lead maintenance for every + // queue. Application-role defaults remain absent unless the explicit + // shared-creator comparison profile below enables them. + let superuser: String = sqlx::query_scalar("SELECT current_user::text") + .fetch_one(superuser_pool) + .await + .map_err(|error| format!("current_user: {error}"))?; + creators.push(superuser); + } + if shared_creator_grants { + creators.extend([APP_A.to_string(), APP_B.to_string()]); + eprintln!("SOAK_SETUP shared_creator_grants=1 schema_create=1 application_creators={APP_A},{APP_B}"); + } + for creator in &creators { + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON TABLES TO {RUNTIME_GROUP}" + )); + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO {RUNTIME_GROUP}" + )); + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT EXECUTE ON FUNCTIONS TO {RUNTIME_GROUP}" + )); + } + for statement in statements { + sqlx::query(&statement) + .execute(superuser_pool) + .await + .map_err(|error| format!("grant runtime privileges ({statement}): {error}"))?; + } + Ok(()) +} + +/// A pool that connects as `role`. +pub async fn pool_as( + config: &SoakConfig, + role: &str, + application_name: &str, + max_connections: u32, +) -> Result { + connect_pool( + &url_for_role(config, role), + application_name, + max_connections, + ) + .await + .map_err(|error| format!("connect as {role}: {error}")) +} diff --git a/awa/tests/persona_soak/run.rs b/awa/tests/persona_soak/run.rs new file mode 100644 index 000000000..a17aaeac3 --- /dev/null +++ b/awa/tests/persona_soak/run.rs @@ -0,0 +1,301 @@ +//! Orchestration: set up the database, start the personas and the oracle, run the +//! load, let the system quiesce, reconcile, and report. + +use super::adminserver::AdminServer; +use super::config::{Persona, SoakConfig}; +use super::db::{self, Budget}; +use super::oracle::Oracle; +use super::personas; +use super::pgserver::PgWatch; +use super::report::{self, Outcome}; +use super::shared::Shared; +use super::stop::Stop; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// Connections kept free of persona pools: oracle, chaos sessions, admin probes. +const HEADROOM: u32 = 14; + +pub async fn run_soak(config: SoakConfig) -> Result, String> { + eprintln!( + "{}: preparing database {}", + report::HEADER, + config.database_name() + ); + db::recreate_database(&config).await?; + let oracle_pool = db::connect_pool(&config.database_url, "soak_oracle", 4) + .await + .map_err(|error| format!("oracle connect: {error}"))?; + let info = db::server_info(&oracle_pool) + .await + .map_err(|error| format!("server info: {error}"))?; + let watch = Arc::new( + PgWatch::begin(&config, &oracle_pool) + .await + .map_err(|error| format!("pg watch: {error}"))?, + ); + let budget = Budget::new(&info, HEADROOM); + eprintln!( + "{}: engine={} server={} max_connections={} connections_now={} usable={} seed={} personas={:?} load={}s scale={}", + report::HEADER, + info.engine, + info.version, + info.max_connections, + info.connections_at_start, + budget.usable(), + config.seed, + config.personas.iter().map(|persona| persona.name()).collect::>(), + config.load.as_secs(), + config.scale + ); + if info.engine != "queue_storage" { + return Err(format!( + "expected the queue_storage engine, found {}", + info.engine + )); + } + let shared = Shared::new(config, info, budget); + + // The UI personas talk to the real admin server, run as a child process. + let ui_personas = [ + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + ]; + let admin_server = if ui_personas + .iter() + .any(|persona| shared.config.has(*persona)) + { + std::fs::create_dir_all(&shared.config.report_dir).map_err(|error| error.to_string())?; + let server = AdminServer::start(&shared.config).await?; + shared.ui.set_base_url(server.base_url.clone()); + shared.note( + "run", + format!( + "admin server at {} (statement timeout {}s, client timeout {}s)", + server.base_url, + server.statement_timeout.as_secs(), + shared.config.ui_client_timeout.as_secs() + ), + ); + Some(server) + } else { + None + }; + + let terminal_observer = tokio::spawn(super::terminal_proof::run( + shared.clone(), + oracle_pool.clone(), + shared.drain_stop.clone(), + )); + + // Personas start their runtimes and then produce until the load stops. Chaos, the + // operator and the UI users reach other personas' fleets and queues, so they start + // once the rest are up. + let mut tasks: Vec<(Persona, tokio::task::JoinHandle>)> = Vec::new(); + for persona in [ + Persona::Mia, + Persona::Tom, + Persona::Priya, + Persona::Sam, + Persona::Python, + Persona::Fanout, + Persona::ShardOps, + Persona::ClaimAbandon, + ] { + if shared.config.has(persona) { + tasks.push((persona, spawn_persona(persona, shared.clone()))); + } + } + tokio::time::sleep(Duration::from_secs(8)).await; + for persona in [ + Persona::Alex, + Persona::Chaos, + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + ] { + if shared.config.has(persona) { + tasks.push((persona, spawn_persona(persona, shared.clone()))); + } + } + + let oracle_stop = Stop::new(); + let oracle_task = { + let oracle = Oracle::new(shared.clone(), oracle_pool.clone(), watch.clone()); + let stop = oracle_stop.clone(); + tokio::spawn(async move { oracle.run(stop).await }) + }; + shared.note( + "run", + format!("load phase started: {}s", shared.config.load.as_secs()), + ); + shared.started_load(Instant::now()); + + let load_deadline = Instant::now() + shared.config.load; + let mut next_progress = Instant::now() + Duration::from_secs(30); + while Instant::now() < load_deadline && !shared.load_stop.is_stopped() { + tokio::time::sleep(Duration::from_secs(1)).await; + if let Some((persona, _)) = tasks.iter().find(|(_, task)| task.is_finished()) { + // A persona only returns early on a setup failure or a panic. + let persona = *persona; + let index = tasks + .iter() + .position(|(p, _)| *p == persona) + .expect("finished task"); + let (_, task) = tasks.remove(index); + match task.await { + Ok(Ok(())) => shared.note("run", format!("{} finished early", persona.name())), + Ok(Err(error)) => shared.violation( + "persona_failed", + format!("{} stopped: {error}", persona.name()), + vec![], + ), + Err(error) => shared.violation( + "persona_panicked", + format!("{} panicked: {error}", persona.name()), + vec![], + ), + } + } + if Instant::now() >= next_progress { + next_progress = Instant::now() + Duration::from_secs(30); + progress(&shared); + } + } + + shared.note("run", "load phase over; letting the system quiesce"); + shared.load_stop.trigger(); + oracle_stop.trigger(); + let mut oracle = oracle_task + .await + .map_err(|error| format!("oracle task: {error}"))?; + let quiesced = oracle.wait_quiescent(shared.config.drain_grace).await; + shared + .quiesced + .store(quiesced, std::sync::atomic::Ordering::Relaxed); + shared.note("run", format!("quiesced: {quiesced}")); + super::terminal_proof::observe(&shared, &oracle_pool).await; + oracle.final_check(quiesced).await; + let excerpts = oracle.crash_scan(); + for excerpt in &excerpts { + shared.violation( + "server_crash_log", + "the Postgres log contains crash lines", + vec![excerpt.clone()], + ); + } + let log_tail = watch.log_since_start(&shared.config); + if shared.config.has(Persona::ShardOps) { + match &log_tail { + Some(log) => super::shard_oracle::scan_server_log(&shared, &oracle_pool, log).await, + None => shared.note( + "run", + "no Postgres log source is configured, so permission errors raised inside the runtimes are not scanned", + ), + } + } + + shared.drain_stop.trigger(); + for (persona, task) in tasks { + match tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(180)), + task, + ) + .await + { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => shared.violation( + "persona_failed", + format!("{} stopped: {error}", persona.name()), + vec![], + ), + Ok(Err(error)) => shared.violation( + "persona_panicked", + format!("{} panicked: {error}", persona.name()), + vec![], + ), + Err(_) => shared.violation( + "shutdown_hang", + format!("{} did not shut down within 180s", persona.name()), + vec![], + ), + } + } + + terminal_observer + .await + .map_err(|error| format!("terminal observer task: {error}"))?; + personas::ui::analyse(&shared); + if let Some(mut server) = admin_server { + if !server.is_running() { + shared.violation( + "admin_server_died", + "the awa serve process exited during the run; see awa-serve.log", + vec![], + ); + } + server.stop(); + } + let outcome = Outcome { + quiesced, + elapsed_seconds: shared.started.elapsed().as_secs_f64(), + crash_excerpts: excerpts, + log_tail, + }; + if let Err(error) = report::write(&shared, &outcome) { + eprintln!("could not write the report: {error}"); + } + oracle_pool.close().await; + Ok(shared) +} + +fn spawn_persona( + persona: Persona, + shared: Arc, +) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + match persona { + Persona::Mia => personas::mia::run(shared).await, + Persona::Tom => personas::tom::run(shared).await, + Persona::Priya => personas::priya::run(shared).await, + Persona::Sam => personas::sam::run(shared).await, + Persona::Alex => personas::alex::run(shared).await, + Persona::Chaos => personas::chaos::run(shared).await, + Persona::Watchers => personas::ui::run_watchers(shared).await, + Persona::Support => personas::ui::run_support(shared).await, + Persona::UiOps => personas::ui::run_ui_operator(shared).await, + Persona::Impatient => personas::ui::run_impatient(shared).await, + Persona::Python => personas::python::run(shared).await, + Persona::Fanout => personas::fanout::run(shared).await, + Persona::ShardOps => personas::shardops::run(shared).await, + Persona::ClaimAbandon => personas::abandon::run(shared).await, + } + }) +} + +fn progress(shared: &Arc) { + let mut line = format!("t={:.0}s", shared.now_ms() as f64 / 1000.0); + for persona in Persona::ALL { + if !shared.config.has(persona) { + continue; + } + let counters = shared.stats(persona).counters(); + let enqueued: i64 = counters + .iter() + .filter(|(name, _)| name.starts_with("enqueued")) + .map(|(_, value)| *value) + .sum(); + line.push_str(&format!( + " {}[enq={} exec={}]", + persona.name(), + enqueued, + counters.get("executions").copied().unwrap_or(0) + )); + } + let violations = shared.violations.lock().expect("violations").len(); + line.push_str(&format!(" violations={violations}")); + shared.note("run", line); +} diff --git a/awa/tests/persona_soak/shard_oracle.rs b/awa/tests/persona_soak/shard_oracle.rs new file mode 100644 index 000000000..ef436c04b --- /dev/null +++ b/awa/tests/persona_soak/shard_oracle.rs @@ -0,0 +1,638 @@ +//! Checks specific to live `enqueue_shards` changes (ADR-025): every shard drains, +//! claim cursors keep moving while their shard has backlog, jobs run exactly once, +//! `ordering_key` jobs stay FIFO within a settled window, and no role ever hits a +//! permission or missing-lane error. + +use super::config::Persona; +use super::ledger::flag; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +const MAX_ROWS: usize = 25; +/// Same-key starts further behind the newest start than this are an order inversion. +const INVERSION_WINDOW: u64 = 8; +/// Seconds a shard above the configured count may keep backlog once every runtime has +/// observed the new count. +const ABOVE_COUNT_DRAIN_SECONDS: i64 = 45; + +#[derive(Debug, Clone, Default)] +pub struct QueuePlan { + pub configured: i16, + pub changed_at_ms: i64, + /// When every runtime and producer store had observed `configured`; None while the + /// fleet disagrees. + pub settled_at_ms: Option, +} + +/// What the shard-ops persona and the oracle share. +#[derive(Default)] +pub struct ShardOpsState { + pub plans: Mutex>, + /// Id of the current settled window; 0 while a change is propagating. + window: AtomicU64, + window_counter: AtomicU64, + producers_held: AtomicBool, + pub producers_inserting: AtomicI64, + seq: AtomicU64, + start_tick: AtomicU64, + /// Executions each enqueued job should have: 1, or 2 for a job that retries once. + expected: Mutex>, + /// Executions per job sequence number, and the job id that ran. + executions: Mutex>, + keyed_window: Mutex>, + keyed_starts: Mutex>, + reported: Mutex>, +} + +impl ShardOpsState { + pub fn next_seq(&self) -> u64 { + self.seq.fetch_add(1, Ordering::Relaxed) + 1 + } + + pub fn window(&self) -> u64 { + self.window.load(Ordering::SeqCst) + } + + pub fn open_window(&self) -> u64 { + let id = self.window_counter.fetch_add(1, Ordering::SeqCst) + 1; + self.window.store(id, Ordering::SeqCst); + id + } + + pub fn close_window(&self) { + self.window.store(0, Ordering::SeqCst); + } + + pub fn hold_producers(&self, held: bool) { + self.producers_held.store(held, Ordering::SeqCst); + } + + pub fn producers_held(&self) -> bool { + self.producers_held.load(Ordering::SeqCst) + } + + pub fn note_enqueued(&self, seq: u64, executions: u8) { + self.expected + .lock() + .expect("expected") + .insert(seq, executions); + } + + pub fn note_keyed(&self, key: u64, key_seq: u64, window: u64) { + self.keyed_window + .lock() + .expect("keyed window") + .insert((key, key_seq), window); + } + + /// Called by the worker at the start of every execution. + pub fn note_execution(&self, seq: u64, job_id: i64, key: Option<(u64, u64)>) { + let mut executions = self.executions.lock().expect("executions"); + let entry = executions.entry(seq).or_insert((0, job_id)); + entry.0 += 1; + drop(executions); + if let Some((key, key_seq)) = key { + let tick = self.start_tick.fetch_add(1, Ordering::SeqCst); + self.keyed_starts + .lock() + .expect("keyed starts") + .push((key, key_seq, tick)); + } + } + + pub fn configured(&self, queue: &str) -> i16 { + self.plans + .lock() + .expect("plans") + .get(queue) + .map_or(1, |plan| plan.configured.max(1)) + } + + /// Reports `summary` as a violation the first time `key` is seen and only counts + /// repeats, so a persistent failure does not bury the report. + pub fn violation_once( + &self, + shared: &Shared, + kind: &str, + key: &str, + summary: String, + rows: Vec, + ) { + let first = { + let mut reported = self.reported.lock().expect("reported"); + let count = reported.entry(format!("{kind}:{key}")).or_insert(0); + *count += 1; + *count == 1 + }; + if first { + shared.violation(kind, summary, rows); + } + shared + .stats(Persona::ShardOps) + .add(&format!("{kind}_total"), 1); + } +} + +/// One lane's cursors and backlog, as read from the store. +#[derive(Debug, Clone)] +struct ShardRow { + queue: String, + priority: i16, + shard: i16, + claim_seq: i64, + next_seq: i64, + available: i64, +} + +impl ShardRow { + fn json(&self, extra: &str) -> String { + format!( + "{{\"queue\":\"{}\",\"priority\":{},\"enqueue_shard\":{},\"claim_seq\":{},\"next_seq\":{},\"lag\":{},\"available\":{}{extra}}}", + self.queue, + self.priority, + self.shard, + self.claim_seq, + self.next_seq, + self.next_seq - self.claim_seq, + self.available + ) + } +} + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Per-shard claim cursor, enqueue cursor and the rows the cursor has not yet claimed +/// (the lag query from ADR-025, with the unclaimed rows counted exactly because aborted +/// transactions leave gaps in the sequences). +async fn shard_rows( + pool: &PgPool, + schema: &str, + queues: &[String], +) -> Result, sqlx::Error> { + let rows: Vec<(String, i16, i16, i64, i64, i64)> = sqlx::query_as(&format!( + "WITH claims AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS claim_seq + FROM {schema}.queue_claim_heads WHERE queue = ANY($1) + ), + enqueues AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS next_seq + FROM {schema}.queue_enqueue_heads WHERE queue = ANY($1) + ), + available AS ( + SELECT r.queue, r.priority, r.enqueue_shard, count(*)::bigint AS n + FROM {schema}.ready_entries r + JOIN claims c ON c.queue = r.queue AND c.priority = r.priority AND c.enqueue_shard = r.enqueue_shard + WHERE r.queue = ANY($1) AND r.lane_seq >= c.claim_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones t + WHERE t.ready_slot = r.ready_slot AND t.ready_generation = r.ready_generation + AND t.queue = r.queue AND t.priority = r.priority + AND t.enqueue_shard = r.enqueue_shard AND t.lane_seq = r.lane_seq + ) + GROUP BY r.queue, r.priority, r.enqueue_shard + ) + SELECT c.queue, c.priority, c.enqueue_shard, c.claim_seq::bigint, e.next_seq::bigint, + COALESCE(a.n, 0)::bigint + FROM claims c + JOIN enqueues e USING (queue, priority, enqueue_shard) + LEFT JOIN available a USING (queue, priority, enqueue_shard) + ORDER BY c.queue, c.priority, c.enqueue_shard" + )) + .bind(queues) + .fetch_all(pool) + .await?; + Ok(rows + .into_iter() + .map( + |(queue, priority, shard, claim_seq, next_seq, available)| ShardRow { + queue, + priority, + shard, + claim_seq, + next_seq, + available, + }, + ) + .collect()) +} + +/// Periodic per-shard checks until `stop`: claim cursors advance while their shard has +/// backlog and idle workers, and shards above a lowered count drain. +pub async fn watch(shared: Arc, pool: PgPool, queues: Vec, stop: Stop) { + let schema = schema(&shared); + // (queue, priority, shard) -> (claim cursor, when it last moved, already reported) + let mut cursors: HashMap<(String, i16, i16), (i64, i64, bool)> = HashMap::new(); + // (queue, priority, shard) -> when backlog above the configured count was first seen + let mut above_since: HashMap<(String, i16, i16), (i64, bool)> = HashMap::new(); + while !stop.sleep(Duration::from_secs(5)).await { + let rows = match shard_rows(&pool, &schema, &queues).await { + Ok(rows) => rows, + Err(error) => { + shared.incomplete(format!("shard cursor query failed: {error}"), vec![]); + continue; + } + }; + let now = shared.now_ms(); + let mut window = shared.config.stall_window; + if shared.chaos_active() { + window += shared.config.chaos_slack; + } + let mut shards_seen: BTreeMap = BTreeMap::new(); + for row in rows { + let key = (row.queue.clone(), row.priority, row.shard); + let seen = shards_seen.entry(row.queue.clone()).or_insert(0); + *seen = (*seen).max(i64::from(row.shard) + 1); + + let entry = cursors + .entry(key.clone()) + .or_insert((row.claim_seq, now, false)); + if entry.0 != row.claim_seq { + *entry = (row.claim_seq, now, false); + } else if let Some(queue) = shared.queues.get(&row.queue) { + let capacity = queue.capacity.load(Ordering::Relaxed); + let in_flight = queue.in_flight.load(Ordering::Relaxed); + let idle_workers = capacity > 0 && in_flight < capacity; + if row.available <= 0 || !idle_workers || queue.is_held() { + entry.1 = now; + } else if !entry.2 && (now - entry.1) as u64 >= window.as_millis() as u64 { + entry.2 = true; + shared.violation( + "shard_cursor_stalled", + format!( + "queue {} priority {} shard {}: {} unclaimed job(s), claim cursor frozen at {} for {}s with {in_flight}/{capacity} workers busy", + row.queue, + row.priority, + row.shard, + row.available, + row.claim_seq, + (now - entry.1) / 1000 + ), + vec![row.json(&format!( + ",\"configured_shards\":{}", + shared.shardops.configured(&row.queue) + ))], + ); + } + } + + let plan = shared + .shardops + .plans + .lock() + .expect("plans") + .get(&row.queue) + .cloned(); + let above = plan.as_ref().is_some_and(|plan| { + i32::from(row.shard) >= i32::from(plan.configured) + && plan + .settled_at_ms + .is_some_and(|at| now - at >= ABOVE_COUNT_DRAIN_SECONDS * 1000) + }); + if above && row.available > 0 { + let since = above_since.entry(key).or_insert((now, false)); + if !since.1 && now - since.0 >= ABOVE_COUNT_DRAIN_SECONDS * 1000 { + since.1 = true; + shared.violation( + "shard_above_count_not_draining", + format!( + "queue {} shard {} is above the configured count {} and still holds {} unclaimed job(s) {}s after every runtime observed the change", + row.queue, + row.shard, + plan.as_ref().map_or(0, |plan| plan.configured), + row.available, + (now - since.0) / 1000 + ), + vec![row.json("")], + ); + } + } else { + above_since.remove(&key); + } + } + for (queue, shards) in shards_seen { + shared + .stats(Persona::ShardOps) + .set_max(&format!("highest_shard_count_{queue}"), shards); + } + } +} + +/// After the system quiesced: every shard of every shard-ops queue is empty, every +/// enqueued job ran exactly the expected number of times, and keyed jobs kept their +/// order within each settled window. +pub async fn final_check(shared: &Arc, pool: &PgPool, queues: &[String]) { + let quiesced = shared.quiesced.load(Ordering::Relaxed); + let schema = schema(shared); + match shard_rows(pool, &schema, queues).await { + Ok(rows) => { + let stuck: Vec = rows + .iter() + .filter(|row| row.available > 0) + .map(|row| { + row.json(&format!( + ",\"configured_shards\":{}", + shared.shardops.configured(&row.queue) + )) + }) + .take(MAX_ROWS) + .collect(); + shared.run.add("shardops_lanes_final", rows.len() as i64); + if !stuck.is_empty() { + let message = format!( + "{} shard(s) still hold unclaimed jobs after the load stopped", + stuck.len() + ); + if quiesced { + shared.violation("shard_not_drained", message, stuck); + } else { + shared.finding("shard_backlog_at_end", message, stuck); + } + } + } + Err(error) => shared.incomplete(format!("final shard query failed: {error}"), vec![]), + } + + let state = &shared.shardops; + let expected = state.expected.lock().expect("expected").clone(); + let executions = state.executions.lock().expect("executions").clone(); + let mut missing = Vec::new(); + let mut repeated = Vec::new(); + let mut rescued = Vec::new(); + let mut short = Vec::new(); + for (seq, want) in &expected { + let (ran, job_id) = executions.get(seq).copied().unwrap_or((0, 0)); + let record = shared.ledger.get(job_id); + let row = format!( + "{{\"seq\":{seq},\"job_id\":{job_id},\"expected_executions\":{want},\"executions\":{ran},\"max_attempt\":{},\"ledger_flags\":{}}}", + record.as_ref().map_or(0, |record| record.max_attempt), + record.as_ref().map_or(0, |record| record.flags) + ); + if ran == 0 { + missing.push(row); + } else if ran > u16::from(*want) { + // A rescue re-runs a job whose attempt lost its lease, which is at-least-once + // delivery working as designed; any other repeat is a duplicate. + let was_rescued = record + .as_ref() + .is_some_and(|record| record.has(flag::RESCUED)); + if was_rescued { + rescued.push(row); + } else { + repeated.push(row); + } + } else if ran < u16::from(*want) { + short.push(row); + } + } + shared + .stats(Persona::ShardOps) + .add("tracked_jobs", expected.len() as i64); + let report = |kind: &str, what: &str, rows: Vec, fatal: bool| { + if rows.is_empty() { + return; + } + let summary = format!("{} job(s) {what}", rows.len()); + let rows = rows.into_iter().take(MAX_ROWS).collect(); + if fatal { + shared.violation(kind, summary, rows); + } else { + shared.finding(kind, summary, rows); + } + }; + report( + "lost_job", + "were enqueued across a shard change and never ran", + missing, + quiesced, + ); + report( + "rescued_reexecution", + "ran again after a stale-heartbeat rescue", + rescued, + false, + ); + report( + "duplicate_execution", + "ran more often than their handler asked for", + repeated, + true, + ); + report( + "retry_not_executed", + "ran fewer times than their handler asked for", + short, + quiesced, + ); + + let windows = state.keyed_window.lock().expect("keyed window").clone(); + let mut starts = state.keyed_starts.lock().expect("keyed starts").clone(); + starts.sort_by_key(|(_, _, tick)| *tick); + let mut newest: HashMap<(u64, u64), u64> = HashMap::new(); + let mut inversions = Vec::new(); + let mut checked = 0u64; + let mut worst = 0u64; + let mut seen_per_key: HashSet<(u64, u64)> = HashSet::new(); + for (key, key_seq, _) in starts { + let window = windows.get(&(key, key_seq)).copied().unwrap_or(0); + if window == 0 { + continue; + } + if !seen_per_key.insert((key, key_seq)) { + continue; + } + checked += 1; + let top = newest.entry((key, window)).or_insert(0); + if key_seq + INVERSION_WINDOW < *top { + worst = worst.max(*top - key_seq); + if inversions.len() < MAX_ROWS { + inversions.push(format!( + "{{\"ordering_key\":{key},\"window\":{window},\"started_seq\":{key_seq},\"newest_started_seq\":{top}}}" + )); + } + } else { + *top = (*top).max(key_seq); + } + } + shared + .stats(Persona::ShardOps) + .add("keyed_starts_checked", checked as i64); + shared + .stats(Persona::ShardOps) + .set_max("keyed_worst_inversion", worst as i64); + if !inversions.is_empty() { + shared.violation( + "ordering_key_inversion", + format!( + "keyed jobs started out of order within a settled shard-count window (worst distance {worst})" + ), + inversions, + ); + } +} + +struct LoggedError { + role: String, + application: String, + message: String, + statement: String, +} + +fn is_lane_error(message: &str) -> bool { + message.contains("permission denied") + || message.contains("must be owner") + || message.contains("has not been provisioned") + || message.contains("provision") + || is_lane_race(message) +} + +/// Two sessions creating the same lane sequence at once: `CREATE SEQUENCE IF NOT EXISTS` +/// loses the race with a unique violation on the catalog. +fn is_lane_race(message: &str) -> bool { + message.contains("pg_class_relname_nsp_index") +} + +/// Permission-denied and missing-lane errors in the server log, one violation per role +/// and error message with the lane, the sequence's owner and the statement. +pub async fn scan_server_log(shared: &Arc, pool: &PgPool, log: &str) { + let mut errors: Vec = Vec::new(); + let lines: Vec<&str> = log.lines().collect(); + let is_entry = |line: &str| { + let bytes = line.as_bytes(); + bytes.len() > 20 && bytes[4] == b'-' && bytes[7] == b'-' && bytes[10] == b' ' + }; + let mut index = 0; + while index < lines.len() { + let line = lines[index]; + index += 1; + let Some((head, message)) = line.split_once(" ERROR: ") else { + continue; + }; + if !is_lane_error(message) { + continue; + } + let identity = head + .rsplit_once("] ") + .map_or("?", |(_, identity)| identity.trim()); + let (role, application) = identity.split_once('@').unwrap_or((identity, "")); + let (role, application) = (role.to_string(), application.to_string()); + let mut statement = String::new(); + while index < lines.len() && !is_entry(lines[index]) { + if let Some(text) = lines[index].trim_start().strip_prefix("STATEMENT:") { + statement = text.trim().chars().take(400).collect(); + } else if !statement.is_empty() { + statement.push(' '); + statement.push_str(lines[index].trim()); + } + index += 1; + } + errors.push(LoggedError { + role, + application, + message: message.trim().to_string(), + statement, + }); + } + if errors.is_empty() { + return; + } + let schema = schema(shared); + let mut groups: BTreeMap<(String, String), (u64, String, BTreeSet)> = BTreeMap::new(); + for error in &errors { + let entry = groups + .entry((error.role.clone(), error.message.clone())) + .or_insert((0, error.statement.clone(), BTreeSet::new())); + entry.0 += 1; + entry.2.insert(error.application.clone()); + } + shared + .run + .add("shard_lane_errors_logged", errors.len() as i64); + for ((role, message), (count, statement, applications)) in groups.into_iter().take(MAX_ROWS) { + let (lane, owner) = describe_sequence(pool, &schema, &message).await; + let applications = applications.into_iter().collect::>().join(","); + let kind = if is_lane_race(&message) { + "lane_creation_race" + } else { + "permission_denied" + }; + shared.violation( + kind, + format!("role {role}: {message} ({count}x; lane {lane}; sequence owner {owner})"), + vec![format!( + "{{\"role\":\"{role}\",\"applications\":{},\"error\":{},\"lane\":{},\"sequence_owner\":{},\"occurrences\":{count},\"statement\":{}}}", + json_string(&applications), + json_string(&message), + json_string(&lane), + json_string(&owner), + json_string(&statement) + )], + ); + } +} + +fn json_string(text: &str) -> String { + serde_json::Value::String(text.to_string()).to_string() +} + +/// The lane and owning role of the sequence an error message names, when it names one. +async fn describe_sequence(pool: &PgPool, schema: &str, message: &str) -> (String, String) { + let Some(name) = message + .split("sequence ") + .nth(1) + .map(|rest| { + rest.split_whitespace() + .next() + .unwrap_or("") + .trim_matches('"') + }) + .map(|name| { + name.rsplit('.') + .next() + .unwrap_or(name) + .trim_matches('"') + .to_string() + }) + .filter(|name| !name.is_empty()) + else { + return ("unknown".into(), "unknown".into()); + }; + let owner: Option = sqlx::query_scalar( + "SELECT pg_get_userbyid(c.relowner) FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname = $2", + ) + .bind(schema) + .bind(&name) + .fetch_optional(pool) + .await + .ok() + .flatten(); + let lane: Option = sqlx::query_scalar(&format!( + "SELECT queue || ' priority ' || priority || ' shard ' || enqueue_shard + FROM {schema}.queue_enqueue_heads WHERE seq_name = $1 + UNION ALL + SELECT queue || ' priority ' || priority || ' shard ' || enqueue_shard + FROM {schema}.queue_claim_heads WHERE seq_name = $1 + LIMIT 1" + )) + .bind(&name) + .fetch_optional(pool) + .await + .ok() + .flatten(); + ( + lane.unwrap_or_else(|| format!("sequence {name}")), + owner.unwrap_or_else(|| "unknown".into()), + ) +} diff --git a/awa/tests/persona_soak/shared.rs b/awa/tests/persona_soak/shared.rs new file mode 100644 index 000000000..2bd038081 --- /dev/null +++ b/awa/tests/persona_soak/shared.rs @@ -0,0 +1,468 @@ +//! State shared by every persona, the oracle and the report. + +use super::config::{Persona, SoakConfig}; +use super::db::{Budget, ServerInfo}; +use super::fleet::Fleet; +use super::ledger::{expect, flag, JobRecord, Ledger}; +use super::registry::{InFlightGuard, QueueRegistry}; +use super::rng::percentile; +use super::stop::Stop; +use awa::{JobRow, UntypedJobEvent}; +use serde::Serialize; +use std::collections::{BTreeMap, HashMap}; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +#[derive(Debug, Clone, Serialize)] +pub struct TimelineEvent { + pub at_ms: i64, + pub source: String, + pub text: String, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Violation { + pub kind: String, + pub at_ms: i64, + pub during_chaos: bool, + pub summary: String, + /// The violating rows, as JSON text. + pub rows: Vec, +} + +#[derive(Default)] +pub struct Stats { + counters: Mutex>, + samples: Mutex>>, +} + +impl Stats { + pub fn add(&self, name: &str, amount: i64) { + *self + .counters + .lock() + .expect("stats") + .entry(name.to_string()) + .or_insert(0) += amount; + } + + pub fn set_max(&self, name: &str, value: i64) { + let mut counters = self.counters.lock().expect("stats"); + let slot = counters.entry(name.to_string()).or_insert(i64::MIN); + *slot = (*slot).max(value); + } + + pub fn get(&self, name: &str) -> i64 { + self.counters + .lock() + .expect("stats") + .get(name) + .copied() + .unwrap_or(0) + } + + /// Records a latency-style sample; kept bounded so long runs stay small. + pub fn observe(&self, name: &str, value: i64) { + let mut samples = self.samples.lock().expect("stats"); + let list = samples.entry(name.to_string()).or_default(); + if list.len() < 400_000 { + list.push(value); + } + } + + pub fn counters(&self) -> BTreeMap { + self.counters.lock().expect("stats").clone() + } + + pub fn summaries(&self) -> BTreeMap { + self.samples + .lock() + .expect("stats") + .iter() + .map(|(name, values)| { + let mut sorted = values.clone(); + sorted.sort_unstable(); + ( + name.clone(), + SampleSummary { + count: sorted.len() as i64, + p50: percentile(&sorted, 0.5), + p99: percentile(&sorted, 0.99), + max: sorted.last().copied().unwrap_or(0), + }, + ) + }) + .collect() + } +} + +#[derive(Debug, Clone, Serialize)] +pub struct SampleSummary { + pub count: i64, + pub p50: i64, + pub p99: i64, + pub max: i64, +} + +pub struct Shared { + pub config: SoakConfig, + pub info: ServerInfo, + pub started: Instant, + pub ledger: Ledger, + pub queues: QueueRegistry, + pub budget: Budget, + pub load_stop: Stop, + pub drain_stop: Stop, + pub timeline: Mutex>, + pub violations: Mutex>, + /// Required oracle evidence was unavailable; the run cannot establish a pass. + pub oracle_incomplete: AtomicBool, + /// Problems that are not invariant violations (slow admin queries); reported, not fatal. + pub findings: Mutex>, + pub stats: HashMap, + pub ui: super::personas::ui::UiMetrics, + /// Handlers currently running in the Python process, from its own reports. + pub python_active: AtomicI64, + /// Set once the post-load drain finished with no live work left. + pub quiesced: std::sync::atomic::AtomicBool, + /// The running Python process, so the chaos persona can kill it. + pub python: Mutex>>, + /// Run-wide measurements (resource use, oracle probes). + pub run: Stats, + pub fleets: Mutex>>, + /// State of the shard-ops persona and the checks that go with it. + pub shardops: super::shard_oracle::ShardOpsState, + /// State of the claim-and-abandon persona and the checks that go with it. + pub abandon: super::abandon_oracle::AbandonState, + chaos_until_ms: AtomicI64, +} + +impl Shared { + pub fn new(config: SoakConfig, info: ServerInfo, budget: Budget) -> Arc { + let stats = Persona::ALL + .into_iter() + .map(|persona| (persona, Stats::default())) + .collect(); + Arc::new(Self { + config, + info, + started: Instant::now(), + ledger: Ledger::new(), + queues: QueueRegistry::default(), + budget, + load_stop: Stop::new(), + drain_stop: Stop::new(), + timeline: Mutex::new(Vec::new()), + violations: Mutex::new(Vec::new()), + oracle_incomplete: AtomicBool::new(false), + findings: Mutex::new(Vec::new()), + stats, + run: Stats::default(), + ui: Default::default(), + python_active: AtomicI64::new(0), + quiesced: std::sync::atomic::AtomicBool::new(false), + python: Mutex::new(None), + fleets: Mutex::new(HashMap::new()), + shardops: Default::default(), + abandon: Default::default(), + chaos_until_ms: AtomicI64::new(0), + }) + } + + pub fn started_load(&self, _at: Instant) {} + + pub fn now_ms(&self) -> i64 { + self.started.elapsed().as_millis() as i64 + } + + pub fn stats(&self, persona: Persona) -> &Stats { + &self.stats[&persona] + } + + pub fn note(&self, source: &str, text: impl Into) { + let text = text.into(); + let at_ms = self.now_ms(); + eprintln!("[{:>7.1}s] {source:<6} {text}", at_ms as f64 / 1000.0); + self.timeline.lock().expect("timeline").push(TimelineEvent { + at_ms, + source: source.to_string(), + text, + }); + } + + /// Marks a disruptive window (kill, failover, connection termination) so that the + /// time-based checks allow the system its recovery time. + pub fn chaos_window(&self, duration: std::time::Duration) { + let until = self.now_ms() + duration.as_millis() as i64; + self.chaos_until_ms.fetch_max(until, Ordering::Relaxed); + } + + pub fn chaos_active(&self) -> bool { + self.now_ms() < self.chaos_until_ms.load(Ordering::Relaxed) + } + + pub fn violation(&self, kind: &str, summary: impl Into, rows: Vec) { + let violation = Violation { + kind: kind.to_string(), + at_ms: self.now_ms(), + during_chaos: self.chaos_active(), + summary: summary.into(), + rows, + }; + eprintln!( + "[{:>7.1}s] VIOLATION {}: {}", + violation.at_ms as f64 / 1000.0, + violation.kind, + violation.summary + ); + self.violations.lock().expect("violations").push(violation); + } + + pub fn incomplete(&self, summary: impl Into, rows: Vec) { + self.run.add("oracle_incomplete_checks", 1); + if !self.oracle_incomplete.swap(true, Ordering::Relaxed) { + self.violation("oracle_incomplete", summary, rows); + } + } + + /// Records a non-fatal finding: something worth a bug report that is not one of the + /// invariants the soak fails on. + pub fn finding(&self, kind: &str, summary: impl Into, rows: Vec) { + let finding = Violation { + kind: kind.to_string(), + at_ms: self.now_ms(), + during_chaos: self.chaos_active(), + summary: summary.into(), + rows, + }; + eprintln!( + "[{:>7.1}s] FINDING {}: {}", + finding.at_ms as f64 / 1000.0, + finding.kind, + finding.summary + ); + self.findings.lock().expect("findings").push(finding); + } + + /// Records a successful operator cancel: the job may legitimately end cancelled, a + /// handler must not start after it, and it must not overwrite a completion. + pub fn record_operator_cancel( + &self, + persona: Persona, + job_id: i64, + completed_before: bool, + call_started_ms: u32, + ) { + let now = self.now_ms() as u32; + self.stats(persona).add("cancelled", 1); + self.ledger.with( + job_id, + || JobRecord::new(persona, u32::MAX, now, expect::ANY_TERMINAL, 0), + |record| { + record.flags |= flag::OPERATOR_REMOVED | flag::OPERATOR_CANCEL_OK; + record.operator_cancel_ms = now; + }, + ); + if completed_before { + self.violation( + "cancel_overwrote_completed", + format!("operator cancel of job {job_id} succeeded after it was reported completed (call began at {call_started_ms} ms)"), + vec![format!("{{\"job_id\":{job_id}}}")], + ); + } + } + + pub fn fleet(&self, name: &str) -> Option> { + self.fleets.lock().expect("fleets").get(name).cloned() + } + + pub fn fleets_of(&self, persona: Persona) -> Vec> { + self.fleets + .lock() + .expect("fleets") + .values() + .filter(|fleet| fleet.persona == persona) + .cloned() + .collect() + } + + /// Records the start of a handler. Hold the returned guard for the duration of + /// `perform`; dropping it, including when the runtime is killed, ends the execution. + pub fn begin_execution( + self: &Arc, + persona: Persona, + job: &JobRow, + no_events: bool, + ) -> ExecGuard { + let queue = self.queues.get(&job.queue); + let queue_index = queue.as_ref().map_or(u32::MAX, |queue| queue.index); + let now = self.now_ms() as u32; + let default_flags = if no_events { flag::NO_EVENTS } else { 0 }; + let (ran_after_cancel, overlapped) = self.ledger.with( + job.id, + || { + JobRecord::new( + persona, + queue_index, + now, + expect::ANY_TERMINAL, + default_flags, + ) + }, + |record| { + let overlapped = record.in_flight > 0; + record.overlapped |= overlapped; + record.in_flight = record.in_flight.saturating_add(1); + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(job.attempt); + record.latest_run_lease = record.latest_run_lease.max(job.run_lease); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + ( + record.operator_cancel_ms != 0 && now > record.operator_cancel_ms, + overlapped, + ) + }, + ); + if ran_after_cancel { + self.violation( + "ran_after_cancel", + format!("job {} started after its operator cancel committed", job.id), + vec![format!( + "{{\"job_id\":{},\"queue\":\"{}\",\"attempt\":{}}}", + job.id, job.queue, job.attempt + )], + ); + } + if overlapped { + self.stats(persona).add("overlapping_executions", 1); + } + if let Some(name) = job + .metadata + .get("cron_name") + .and_then(|value| value.as_str()) + { + let fire = job + .metadata + .get("cron_fire_time") + .and_then(|value| value.as_str()) + .unwrap_or_default(); + // Manually triggered jobs carry the schedule's name but no fire time. + if !fire.is_empty() + && !self + .ledger + .note_cron_fire(name.to_string(), fire.to_string(), job.id) + { + self.violation( + "cron_double_fire", + format!( + "cron {name} fire {fire} produced more than one job (latest {})", + job.id + ), + vec![], + ); + } + } + if let Some(queue) = &queue { + queue.started.fetch_add(1, Ordering::Relaxed); + } + self.stats(persona).add("executions", 1); + ExecGuard { + _flight: InFlightGuard::new(queue), + shared: self.clone(), + job_id: job.id, + } + } + + /// Records a lifecycle event emitted by a runtime after the transition committed. + pub fn on_event(&self, event: &UntypedJobEvent) { + let job = event.job(); + let now = self.now_ms() as u32; + let queue_index = self + .queues + .get(&job.queue) + .map_or(u32::MAX, |queue| queue.index); + let persona = self + .queues + .get(&job.queue) + .map_or(Persona::Chaos, |queue| queue.persona); + let mut double_complete = false; + self.ledger.with( + job.id, + || JobRecord::new(persona, queue_index, now, expect::ANY_TERMINAL, 0), + |record| match event { + UntypedJobEvent::Completed { .. } => { + record.note_terminal(job.run_lease); + record.completed_events = record.completed_events.saturating_add(1); + record.completed_event_ms = now; + record.flags |= flag::COMPLETED_EVENT; + double_complete = record.completed_events > 1 && !record.has(flag::DLQ_RETRIED); + } + UntypedJobEvent::Exhausted { .. } => { + record.note_terminal(job.run_lease); + record.exhausted_events = record.exhausted_events.saturating_add(1); + record.flags |= flag::EXHAUSTED; + } + UntypedJobEvent::Cancelled { .. } => { + record.note_terminal(job.run_lease); + record.cancel_events = record.cancel_events.saturating_add(1); + record.flags |= flag::HANDLER_CANCELLED; + } + UntypedJobEvent::Rescued { .. } => record.flags |= flag::RESCUED, + _ => {} + }, + ); + match event { + UntypedJobEvent::Rescued { reason, .. } => { + if persona == Persona::ShardOps { + self.note( + "events", + format!( + "job {} on {} rescued ({}) at attempt {}", + job.id, + job.queue, + reason.as_str(), + job.attempt + ), + ); + } + self.stats(persona) + .add(&format!("rescued_{}", reason.as_str()), 1); + } + UntypedJobEvent::Exhausted { .. } => self.stats(persona).add("exhausted", 1), + UntypedJobEvent::Completed { .. } => self.stats(persona).add("completed_events", 1), + UntypedJobEvent::Cancelled { .. } => self.stats(persona).add("handler_cancelled", 1), + UntypedJobEvent::Retried { .. } => self.stats(persona).add("retried", 1), + _ => {} + } + if double_complete { + self.violation( + "duplicate_terminal_event", + format!("job {} reported completed more than once", job.id), + vec![format!( + "{{\"job_id\":{},\"queue\":\"{}\"}}", + job.id, job.queue + )], + ); + } + } +} + +pub struct ExecGuard { + _flight: InFlightGuard, + shared: Arc, + job_id: i64, +} + +impl Drop for ExecGuard { + fn drop(&mut self) { + self.shared.ledger.with( + self.job_id, + || JobRecord::new(Persona::Chaos, u32::MAX, 0, expect::ANY_TERMINAL, 0), + |record| record.in_flight = record.in_flight.saturating_sub(1), + ); + } +} diff --git a/awa/tests/persona_soak/stop.rs b/awa/tests/persona_soak/stop.rs new file mode 100644 index 000000000..6264c06e2 --- /dev/null +++ b/awa/tests/persona_soak/stop.rs @@ -0,0 +1,50 @@ +//! A cheap cooperative stop signal shared by every task in the run. + +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::watch; + +#[derive(Clone)] +pub struct Stop { + sender: Arc>, +} + +impl Stop { + pub fn new() -> Self { + Self { + sender: Arc::new(watch::channel(false).0), + } + } + + pub fn trigger(&self) { + self.sender.send_replace(true); + } + + pub fn is_stopped(&self) -> bool { + *self.sender.borrow() + } + + /// Sleeps for `duration` or until stopped. Returns true when stopped. + pub async fn sleep(&self, duration: Duration) -> bool { + if self.is_stopped() { + return true; + } + let mut receiver = self.sender.subscribe(); + tokio::select! { + _ = tokio::time::sleep(duration) => false, + _ = receiver.wait_for(|stopped| *stopped) => true, + } + } + + /// Completes once stopped. + pub async fn wait(&self) { + let mut receiver = self.sender.subscribe(); + let _ = receiver.wait_for(|stopped| *stopped).await; + } +} + +impl Default for Stop { + fn default() -> Self { + Self::new() + } +} diff --git a/awa/tests/persona_soak/terminal_proof.rs b/awa/tests/persona_soak/terminal_proof.rs new file mode 100644 index 000000000..1e9773afc --- /dev/null +++ b/awa/tests/persona_soak/terminal_proof.rs @@ -0,0 +1,122 @@ +//! Durable evidence for returned eventful jobs and the NO_EVENTS sample. +//! Handler returns are not commits; lifecycle hooks are best-effort notifications. + +use super::config::Persona; +use super::ledger::{expect, flag, JobRecord, TERMINAL_PROOF_STRIDE}; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::sync::Arc; +use std::time::Duration; + +const MAX_ROWS: usize = 100_000; + +pub async fn run(shared: Arc, pool: PgPool, stop: Stop) { + let mut tick = tokio::time::interval(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = stop.wait() => return, + _ = tick.tick() => observe(&shared, &pool).await, + } + } +} + +/// One statement gives one snapshot. Explicit eventful attempts have bounded +/// input/result counts and a two-second statement timeout. NO_EVENTS keeps its +/// immutable sample; neither a returned handler nor a missed query proves a commit. +pub async fn observe(shared: &Shared, pool: &PgPool) { + let schema = shared.info.queue_storage_schema.as_deref().unwrap_or("awa"); + let mut pending = Vec::new(); + shared.ledger.for_each(|id, record| { + if record.has(flag::RETURNED_COMPLETED) + && !record.has(flag::NO_EVENTS) + && !record.has_durable_terminal() + && !record.has(flag::OPERATOR_REMOVED) + && !record.has(flag::DLQ_PURGED) + && pending.len() <= MAX_ROWS + { + pending.push((id, record.latest_run_lease)); + } + }); + if pending.len() > MAX_ROWS { + shared.incomplete( + "durable terminal observer exceeded its pending job cap", + vec![], + ); + return; + } + let (pending_ids, pending_leases): (Vec<_>, Vec<_>) = pending.into_iter().unzip(); + let result = async { + let mut tx = pool.begin().await?; + sqlx::query("SET LOCAL statement_timeout = '2s'") + .execute(tx.as_mut()) + .await?; + let rows: Vec<(i64, i64, String)> = sqlx::query_as(&format!( + "WITH pending AS ( + SELECT * FROM unnest($3::bigint[], $4::bigint[]) AS p(job_id, run_lease) + ) + SELECT done.job_id, done.run_lease, done.state::text + FROM {schema}.done_entries AS done + JOIN pending USING (job_id, run_lease) + UNION ALL + SELECT dlq.job_id, dlq.run_lease, 'failed' + FROM {schema}.dlq_entries AS dlq + JOIN pending USING (job_id, run_lease) + UNION ALL + SELECT item.job_id, item.run_lease, 'completed' + FROM {schema}.receipt_completion_batches AS batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + JOIN pending USING (job_id, run_lease) + WHERE batch.job_ids && $3 + UNION ALL + SELECT job_id, run_lease, state::text FROM {schema}.done_entries + WHERE job_id % $1 = 0 AND finalized_at >= now() - interval '30 seconds' + UNION ALL + SELECT job_id, run_lease, 'failed' FROM {schema}.dlq_entries + WHERE job_id % $1 = 0 AND dlq_at >= now() - interval '30 seconds' + UNION ALL + SELECT item.job_id, item.run_lease, 'completed' + FROM {schema}.receipt_completion_batches AS batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + WHERE item.job_id % $1 = 0 AND batch.finalized_at >= now() - interval '30 seconds' + LIMIT $2" + )) + .bind(TERMINAL_PROOF_STRIDE) + .bind(MAX_ROWS as i64 + 1) + .bind(pending_ids) + .bind(pending_leases) + .fetch_all(tx.as_mut()) + .await?; + tx.commit().await?; + Ok::<_, sqlx::Error>(rows) + } + .await; + let rows = match result { + Ok(rows) if rows.len() <= MAX_ROWS => rows, + Ok(_) => { + shared.incomplete("durable terminal observer exceeded its row cap", vec![]); + return; + } + Err(error) => { + shared.incomplete( + format!("durable terminal observer query failed: {error}"), + vec![], + ); + return; + } + }; + shared + .run + .add("terminal_proof_observer_rows", rows.len() as i64); + let now = shared.now_ms() as u32; + for (id, run_lease, state) in rows { + shared.ledger.with( + id, + || JobRecord::new(Persona::Chaos, u32::MAX, now, expect::ANY_TERMINAL, 0), + |record| { + record.note_observed_terminal(run_lease, &state); + }, + ); + } +} diff --git a/awa/tests/persona_soak_test.rs b/awa/tests/persona_soak_test.rs new file mode 100644 index 000000000..1c33ba709 --- /dev/null +++ b/awa/tests/persona_soak_test.rs @@ -0,0 +1,43 @@ +//! Persona soak: several application-shaped workloads (Mia, Tom, Priya, Sam, Alex and +//! a chaos agent) run concurrently against one disposable database on the queue-storage +//! engine while an invariant oracle checks the system. Ignored by default. +//! +//! ```text +//! AWA_SOAK_DATABASE_URL=postgres://postgres:pw@localhost:5432/awa_soak_run \ +//! AWA_SOAK_SECONDS=900 \ +//! cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture +//! ``` +//! +//! See `docs/test-plan.md` for the environment variables and `scripts/persona-soak.sh` +//! for the engine matrix. + +#![allow(dead_code)] +mod ci_timing; +mod persona_soak; + +use persona_soak::config::SoakConfig; + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "drops and recreates an awa_soak* database; see the module docs"] +async fn persona_soak() { + let config = SoakConfig::from_env().unwrap_or_else(|error| panic!("{error}")); + let shared = persona_soak::run::run_soak(config) + .await + .unwrap_or_else(|error| panic!("soak setup failed: {error}")); + let violations = shared.violations.lock().expect("violations").len(); + assert_eq!( + violations, + 0, + "{violations} invariant violation(s); see {}", + shared.config.report_dir.display() + ); +} + +/// The claiming process of the claim-and-abandon persona. The harness runs this test binary +/// again with only this test selected and kills the process with SIGKILL; started any +/// other way it returns at once. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "run by the claim-and-abandon persona as a separate process"] +async fn claim_abandon_child() { + persona_soak::personas::abandon::run_child().await; +} diff --git a/docs/test-plan.md b/docs/test-plan.md index 618ff69a2..54f39beec 100644 --- a/docs/test-plan.md +++ b/docs/test-plan.md @@ -305,6 +305,89 @@ Planned test matrix for the 0.7 cycle, mapped to the roadmap ([`0.7-roadmap.md`] | V28 | Partitioned-queue preset ≥9k jobs/s e2e on the reference 24-CPU harness | #383 Gate 2 | | V29 | Gate A evidence: allocator bake-off (E1) + WAL/job decomposition (E3) run ids recorded on #295 | #295 | +## Persona soak + +Durable loss detection covers every lifecycle-observed job and a deterministic `job_id % 64 = 0` sample of `NO_EVENTS` workloads. The other 63/64 returned `NO_EVENTS` jobs remain in enqueue/execution accounting, but their terminal commits are not proven. A bulk observer polls sampled terminal rows and eventful attempts whose handlers returned without terminal proof once per second (2 s query timeout, 100,000-row cap) and keeps the run lease in the client ledger, so captured proof survives pruning. A handler return is never terminal proof. Missing sampled proof after pruning or a failed invariant query makes the report `INCONCLUSIVE` and fails the test; it does not assert job loss or establish a pass. Reports include sampled, confirmed and excluded counts. + +On CI, assertion waits and heartbeat staleness use `ci_timing.rs` scaling; load durations, kill schedules and measurement windows retain their configured values. + + +`awa/tests/persona_soak_test.rs` (harness in `awa/tests/persona_soak/`, Python side in `awa/tests/persona_soak/python/`) runs several +application-shaped workloads concurrently against one database, the way different +applications share an `awa` schema, and fails the run on any invariant violation. + +| Persona | Workload | Stresses | +| --- | --- | --- | +| Mia, ML pipeline | jobs of 0.1-60 s with per-job deadlines, progress updates, retries into the DLQ, handler cancels, latency spikes | deadline and stale-heartbeat rescues racing heartbeats and completions, DLQ | +| Tom, multi-tenant | 100 queues at a low rate each, unique jobs, scheduled jobs, cron every 10 s on two runtimes, mixed priorities with aging | lane and queue cardinality, per-queue maintenance, cron fencing | +| Priya, firehose | one hot queue, 100k-row `insert_many` bursts, COPY batches, steady small batches, tiny jobs, three runtimes | bind-parameter limit, ring rotation and prune, claim contention | +| Sam, ordered streams | a queue with 4 enqueue shards and 64 ordering keys, a 4-stripe queue (`queue_stripe_count`), scheduled jobs, a rate-limited queue | routing, promotion, fairness | +| Alex, operator | pause/resume, drain, cancel and bulk cancel, move and reprioritise batch operations, DLQ retry and purge, the admin read queries, rolling restarts of other personas' runtimes | admin paths under live load | +| Chaos | kills a runtime or the Python process without a shutdown handshake, terminates the maintenance leader's connection and random runtime connections, exhausts a runtime's pool, holds an idle-in-transaction session | recovery paths | +| Watchers | six dashboard tabs polling overview, queues, counts, timeseries, runtime and storage at the UI's refresh interval, some also the kinds, queue-runtime and cron pages | admin read queries under load, through the real `awa serve` HTTP server | +| Support engineer | searches jobs by kind, queue, state and tag, pages through large result sets with `before_id`, opens job detail and follows a job at the detail page's poll rate, pages the DLQ, opens runtime, cron and kind pages | search and paging cost against retained history and backlog | +| UI operator | the UI's mutations over HTTP: pause and resume queues and cron schedules, cancel and bulk-cancel, DLQ retry and purge (single and by filter), batch operation preview and submit | mutation endpoints racing live work | +| Impatient user | requests abandoned after a few milliseconds, bursts of refreshes, while sampling the server's database sessions | an abandoned request must not keep a pool connection past the statement timeout | +| Fanout, replay fan-out | parent jobs whose handlers batch-insert 100-185 unique child jobs (`insert_many`, unique by queue and args) into their own queue from many concurrent handlers; some parents replay a batch with an already-claimed key, hit the unique violation and fall back to row-by-row inserts | concurrent batch enqueue into one lane, unique claims, inserting from inside a handler | +| Shard operator | an operator changes `enqueue_shards` on two hot queues while producers, promotion and claimers run: 1 → 4 → 8, down to 2 and back to 1, sometimes after `prepare_queue` and sometimes not, restarting some runtimes and producer `QueueStorage` handles before others so cached shard counts disagree. Producers use `insert_many`, `insert_many_copy` and `QueueStorage::enqueue_params_batch`, with scheduled jobs, retries (promotion into shards) and `ordering_key` jobs. The actors run as non-superuser roles: a schema owner that migrates and runs `prepare_queue`, a migrator role with default privileges, and two application roles in a runtime group role with no default privileges, granted as in `docs/security.md`. For some raises the producers are held until the maintenance leader's runtime has created the new lanes, so a different role is the next to use them | live shard-count changes (ADR-025), lane provisioning by different roles, stale per-process shard counts, promotion into new shards | +| Claim and abandon | a separate process (the harness binary started again with one test selected) that claims as much as it can: `max_workers` equal to the target on two queues, one with zero deadlines (compact claim batches) and one with a deadline (one claim row per job), jobs held until the process dies, heartbeats running. It is killed with SIGKILL, repeatedly, and replaced after a few seconds by a process that claims the rescued jobs again; some jobs have `max_attempts` 2 so repeated kills dead-letter them. A surviving in-process runtime leads maintenance and serves a canary queue of short jobs, and the other personas keep working; after the last kill an in-process runtime finishes the rescued jobs | rescue of a large abandoned-claim backlog, claim-ring rotation while slots hold abandoned claims, false rescues of live owners, `max_attempts` and the DLQ under repeated rescue | +| Python users | a separate process using `awa-python`: single inserts, a 6,500-row `insert_many` in one transaction and a COPY batch, enqueues in transactions that sometimes roll back; async workers that complete, raise, `RetryAfter`, `Snooze`, report progress, overrun a queue deadline and ignore their deadline; a queue shared with a Rust runtime; a periodic migrate-then-query probe | Python bindings against live load, mixed Rust and Python runtimes on one lane, the documented slot behaviour of an ignored deadline, connections stranded after `migrate()` | + +The oracle checks, while the load runs and again after it stops: + +- a job attempt with more than one terminal row (`done_entries`, `dlq_entries`, compact receipt completions), a job reported completed twice, and a job that is terminal and still available, deferred or leased +- a cancelled job nobody asked to cancel, an operator cancel that overwrote a completion, and a handler that started after its operator cancel committed +- a queue with claimable jobs and idle workers whose claim cursor stops advancing (stall), and scheduled jobs overdue beyond a threshold (promotion stall) +- a lost job: enqueued, never terminal after the system quiesces, and not removed by an operator action; a cron fire that produced more than one job +- shard operations: every job enqueued before, during and after a shard-count change runs and completes exactly once (a repeat is allowed only after a stale-heartbeat rescue and is reported as a finding), every shard drains including those above a lowered count, claim cursors keep advancing on every shard that has backlog while workers are idle (the per-shard lag query from ADR-025), `ordering_key` jobs start in order per key within a window in which every runtime had observed the shard count, and no role raises a permission-denied or missing-lane error (found in the server log, with the role, lane, sequence owner and statement; needs `AWA_SOAK_PG_CONTAINER` or `AWA_SOAK_PG_LOG`, and `log_line_prefix` containing `%u@%a`, as `scripts/persona-soak.sh` sets) +- claim and abandon: after each SIGKILL the rescue backlog (open claims whose last heartbeat or claim time is older than the staleness window, or whose deadline passed, counted every few seconds from the claim tables and leases) clears within `AWA_SOAK_ABANDON_BOUND_SECONDS`, with the age and count of the oldest unrescued expired claim sampled into `abandon-backlog.csv` and the time to clear reported per kill; the claim ring does not stay still for longer than the bound once the slot it must enter next holds no claim whose owner is alive (a ring pinned by a live process is expected); no job is run again while the process that held it is alive (every start is recorded with the process that ran it, and a repeat is allowed only after that process was killed); no attempt exceeds `max_attempts`, no job is dead-lettered early, and a job abandoned on its last attempt ends in the DLQ +- a transaction the Python application rolled back that left a job in the store, and database sessions left stranded by the Python `migrate()` path +- an admin-server query that outlives its statement timeout, and sessions still running a query after every UI client has stopped +- a server crash: `pg_postmaster_start_time()` changing, or `terminated by signal`, `TRAP:`, `server process (PID` and similar lines in the Postgres log + +Slow or failing admin endpoints, latency that grows with retained history or backlog, and admin queries cancelled by the statement timeout are reported as non-fatal findings in the report, next to the violations. + +Terminal rows are pruned by ring rotation within seconds, so the oracle keeps a +client-side ledger (job ids per persona, executions per job, lifecycle events, +operator actions) and reconciles it with the database. Time-based checks are widened +while a chaos action is recovering. + +Environment: + +| Variable | Meaning | +| --- | --- | +| `AWA_SOAK_DATABASE_URL` | required; database name must match `awa_soak[A-Za-z0-9_]*` and is dropped and recreated | +| `AWA_SOAK_PERSONAS` | comma list of `mia,tom,priya,sam,alex,chaos,watchers,support,uiops,impatient,python,fanout,shardops,abandon`, default all | +| `AWA_SOAK_ABANDON_CLAIMS`, `AWA_SOAK_ABANDON_KILL_SECONDS`, `AWA_SOAK_ABANDON_KILLS` | the claim-and-abandon persona's scale, independent of `AWA_SOAK_SCALE`: claims each killed process holds (default 4000; 25,000 to 70,000 reproduces a large backlog), seconds between kills (default 180, long enough for the process to claim everything and heartbeat for a while) and how many processes are killed (default 4; the last kill is the end of the load phase if that comes first) | +| `AWA_SOAK_ABANDON_BOUND_SECONDS`, `AWA_SOAK_ABANDON_RESTART_SECONDS` | the bound after a kill within which the rescue backlog must clear and the claim ring must be free to rotate (default 120), and the pause before the replacement process starts (default 3) | +| `AWA_SOAK_HEARTBEAT_SECONDS`, `AWA_SOAK_STALENESS_SECONDS` | heartbeat interval (default 2) and heartbeat staleness after which a claim may be rescued (default 12) of every runtime, including the claim-and-abandon process; the library's own defaults are 30 and 90. The claim-and-abandon persona also records the duration of each heartbeat pass of its process against the number of handlers running, read from the timestamps the heartbeats write (`abandon-heartbeat-passes.csv`, and percentiles by handler count in the report) | +| `AWA_SOAK_RESCUE_INTERVAL_SECONDS` | heartbeat and deadline rescue interval of every runtime; unset keeps the harness's 3 s and 2 s, the library default for heartbeat rescue is 30 s | +| `AWA_SOAK_SHARDOPS_SCALE`, `AWA_SOAK_SHARDOPS_SCHEMA_CREATE` | the shard operator's rate multiplier, independent of `AWA_SOAK_SCALE` (default 1), and whether its application roles may create objects in the schema (default 1: lanes are created on the fly; 0 withholds schema `CREATE` and the operator provisions every lane before raising a count) | +| `AWA_SOAK_SHARED_CREATOR_GRANTS` | default `0`; `1` shares application-created objects for comparisons with older binaries and requires schema `CREATE`. Keep it off when checking application-role privileges. | +| `AWA_SOAK_UI_POOL_MAX`, `AWA_SOAK_ADMIN_READ_SECONDS` | the admin server's database pool size (default 10) and the seconds between the operator persona's admin read queries (default 2); lower the load on a slow server by shrinking the pool and stretching the interval | +| `AWA_SOAK_AWA_BIN` | the `awa` binary for the UI personas (`cargo build --release -p awa-cli`); defaults to the one next to the test binary | +| `AWA_SOAK_PYTHON` | interpreter of a virtualenv with `awa-python` installed (`uv venv awa-python/.venv`, `uv pip install maturin`, `maturin develop --release` in `awa-python/`); required by the `python` persona | +| `AWA_SOAK_UI_STATEMENT_TIMEOUT_SECONDS`, `AWA_SOAK_UI_CLIENT_TIMEOUT_SECONDS` | the admin server's statement timeout (default 30) and the UI personas' request timeout (default 45); a request that hits either is reported as a finding | +| `AWA_SOAK_SECONDS` | load phase length, default 900 | +| `AWA_SOAK_SEED` | fixes the workload mix and chaos schedule; printed at start | +| `AWA_SOAK_SCALE` | multiplies rates, burst sizes and queue counts, default 1 | +| `AWA_SOAK_DRAIN_GRACE_SECONDS` | how long to wait for the system to quiesce after load stops, default 120 | +| `AWA_SOAK_PG_LOG` / `AWA_SOAK_PG_CONTAINER` / `AWA_SOAK_PG_CONTAINER_LOG` | where to read the server log for the crash scan: a file on this host, or a Docker container (with an optional log path inside it, otherwise `docker logs`) | +| `AWA_SOAK_REPORT_DIR` | where `report.md`, `report.json` and, for a failed run, `reproducer.md` are written | +| `AWA_SOAK_STALL_SECONDS`, `AWA_SOAK_PROMOTION_OVERDUE_SECONDS`, `AWA_SOAK_CHAOS_SLACK_SECONDS` | oracle thresholds | + +The report records the engine, server version, `max_connections`, connections at +start and the per-persona connection budget. A failed run leaves `reproducer.md` +with the seed, personas and duration, the violating rows, the timeline around each +violation and the relevant server log excerpt. `scripts/persona-soak.sh` runs the +engine matrix (`postgres:17`, `postgres:18`, and an assert-enabled PostgreSQL 18 +image) in private containers on separate ports and removes them afterwards. + +After-exit retained-session findings require a successful `CLOSED` +acknowledgement and a valid final session observation. Forced or unproven exits +remain incomplete evidence. Pool-budget and idle-transaction checks after +readiness remain independent. + ## Running Tests ```bash @@ -329,6 +412,26 @@ DATABASE_URL=postgres://postgres:test@localhost:15432/awa_test \ cargo test --package awa --test chaos_suite_test \ -- --ignored --test-threads=1 --nocapture +# Persona soak: Mia (ML pipeline), Tom (multi-tenant fan-out), Priya (firehose), +# Sam (ordered streams), Alex (operator) and a chaos agent run together against one +# disposable database on the queue-storage engine, checked by an invariant oracle. +# The database name must start with `awa_soak`; it is dropped and recreated. +AWA_SOAK_DATABASE_URL=postgres://postgres:test@localhost:15432/awa_soak_run \ + AWA_SOAK_SECONDS=900 \ + cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture + +# The same soak against throwaway PostgreSQL 17, 18 and assert-enabled 18 containers +scripts/persona-soak.sh # or: scripts/persona-soak.sh pg17 + +# Only the claim-and-abandon persona, with 25,000 claims abandoned by each of four kills; +# the Python bindings and the admin server binary are not needed +SOAK_SKIP_PREPARE=1 AWA_SOAK_PERSONAS=abandon AWA_SOAK_ABANDON_CLAIMS=25000 \ + AWA_SOAK_SECONDS=1500 AWA_SOAK_DRAIN_GRACE_SECONDS=600 scripts/persona-soak.sh pg17 + +# The same harness against another revision of awa (a release tag or a candidate branch): +# a detached worktree of that revision gets this checkout's harness files copied over it +scripts/persona-soak-against.sh v0.6.13 pg17 + # Postgres hot-standby promotion smoke (boots primary + replica + stable proxy endpoint via Docker Compose) cargo test --package awa --test postgres_failover_smoke_test -- --ignored --nocapture diff --git a/scripts/persona-soak-against.sh b/scripts/persona-soak-against.sh new file mode 100755 index 000000000..4e8a86726 --- /dev/null +++ b/scripts/persona-soak-against.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Runs the persona soak against another revision of awa. +# +# scripts/persona-soak-against.sh [engine ...] +# +# The harness lives entirely in awa/tests, so any revision that builds with the same public +# API can be soaked: a detached worktree of is created, this checkout's harness +# files (including uncommitted changes) are copied over it, scripts/persona-soak.sh runs +# there with the arguments after the ref, and the worktree is removed afterwards. Results go +# to this checkout's soak-results/- unless SOAK_RESULTS_DIR is set. +# +# Environment is passed through to scripts/persona-soak.sh. Build outputs go to +# CARGO_TARGET_DIR when it is set, so builds of several revisions share one cache; the +# harness files are copied with a fresh modification time so cargo rebuilds the test binary +# rather than reusing one built from another revision. +set -euo pipefail + +if [ $# -lt 1 ]; then + echo "usage: $0 [engine ...]" >&2 + exit 2 +fi +ref="$1" +shift + +here="$(cd "$(dirname "$0")/.." && pwd)" +label="$(printf '%s' "$ref" | tr -c 'A-Za-z0-9._-' '_')" +export SOAK_RESULTS_DIR="${SOAK_RESULTS_DIR:-$here/soak-results/$(date +%Y%m%d-%H%M%S)-$label}" +worktree="$(mktemp -d "${TMPDIR:-/tmp}/awa-soak-against.XXXXXX")" +cleanup() { git -C "$here" worktree remove --force "$worktree" >/dev/null 2>&1 || rm -rf "$worktree"; } +trap cleanup EXIT + +git -C "$here" worktree add --detach "$worktree" "$ref" >/dev/null +tar -C "$here" -c awa/tests/persona_soak awa/tests/persona_soak_test.rs scripts/persona-soak.sh \ + | tar -C "$worktree" -x -m +echo "== persona soak against $ref ($(git -C "$worktree" rev-parse --short=10 HEAD))" +(cd "$worktree" && scripts/persona-soak.sh "$@") diff --git a/scripts/persona-soak.sh b/scripts/persona-soak.sh new file mode 100755 index 000000000..b1f6a079c --- /dev/null +++ b/scripts/persona-soak.sh @@ -0,0 +1,213 @@ +#!/usr/bin/env bash +# Runs the persona soak (awa/tests/persona_soak_test.rs) against a matrix of +# throwaway PostgreSQL containers, one engine after another. +# +# scripts/persona-soak.sh [engine ...] +# +# Engines: pg17 (postgres:17), pg18 (postgres:18) and pg18-cassert (an +# assert-enabled, -O0 build of PostgreSQL 18; set SOAK_CASSERT_IMAGE). With no +# arguments all three run in that order. +# +# Environment: +# AWA_SOAK_SECONDS load phase length per engine (default 900) +# AWA_SOAK_PERSONAS comma list, default all of mia,tom,priya,sam,alex,chaos, +# watchers,support,uiops,impatient,python,fanout,shardops,abandon +# AWA_SOAK_SEED fixes the workload mix and chaos schedule +# AWA_SOAK_SCALE multiplies rates and queue counts (default 1) +# AWA_SOAK_SHARDOPS_SCALE multiplies the shardops persona's rates on its own (default 1) +# AWA_SOAK_SHARDOPS_SCHEMA_CREATE +# 1 (default) lets the shardops application roles create objects in +# the schema, so lanes are created on the fly; 0 withholds schema +# CREATE and has the operator provision lanes before raising a count +# AWA_SOAK_ABANDON_CLAIMS claims held by each killed process of the abandon persona (default 4000; +# 25000 or more reproduces a large abandoned-claim backlog) +# AWA_SOAK_ABANDON_KILLS, AWA_SOAK_ABANDON_KILL_SECONDS +# how many processes are SIGKILLed (default 4) and the seconds between +# kills (default 180) +# AWA_SOAK_ABANDON_BOUND_SECONDS +# seconds after a kill within which the rescue backlog must clear and the +# claim ring must be free to rotate (default 120) +# AWA_SOAK_RESCUE_INTERVAL_SECONDS +# heartbeat and deadline rescue interval of every runtime (default: the +# harness's 3 s and 2 s; the library default for heartbeat rescue is 30 s) +# AWA_SOAK_HEARTBEAT_SECONDS, AWA_SOAK_STALENESS_SECONDS +# heartbeat interval and heartbeat staleness of every runtime (default 2 s +# and 12 s, so rescues happen often; the library's are 30 s and 90 s) +# SOAK_RESULTS_DIR output directory (default ./soak-results/) +# SOAK_BASE_PORT first host port; engines use consecutive ports (default 15471) +# SOAK_MAX_CONNECTIONS server max_connections (default 400) +# SOAK_PG_CPUS CPU limit for each database container (default: none) +# SOAK_CASSERT_IMAGE image for pg18-cassert (default awa-pg18-cassert:18.6-O0; the +# -epqlog5 build, which logs every partition-pruning setup with its +# query text, is opt-in for diagnosing pruning crashes) +# SOAK_SKIP_PREPARE set to 1 to skip building the awa binary and the Python bindings +# (neither is needed without the UI and python personas) +# SOAK_TEST_BIN a prebuilt harness binary to run instead of `cargo test` +# CARGO_TARGET_DIR cargo target directory, as for any cargo build +# +# The admin UI personas drive the real `awa serve` binary and the python persona runs +# the awa-python bindings in a separate process; both are built once, before the first +# engine, with `cargo build --release -p awa-cli` and `uv` + `maturin develop --release` +# into awa-python/.venv. +# +# Every container is private to the run (name awa-soak--) and is +# removed on exit, including on failure or interruption. Containers that already +# exist are never touched. +set -uo pipefail + +cd "$(dirname "$0")/.." + +SECONDS_PER_ENGINE="${AWA_SOAK_SECONDS:-900}" +BASE_PORT="${SOAK_BASE_PORT:-15471}" +MAX_CONNECTIONS="${SOAK_MAX_CONNECTIONS:-400}" +CASSERT_IMAGE="${SOAK_CASSERT_IMAGE:-awa-pg18-cassert:18.6-O0}" +RESULTS_DIR="${SOAK_RESULTS_DIR:-$PWD/soak-results/$(date +%Y%m%d-%H%M%S)}" +ENGINES=("$@") +if [ ${#ENGINES[@]} -eq 0 ]; then + ENGINES=(pg17 pg18 pg18-cassert) +fi + +mkdir -p "$RESULTS_DIR" + +prepare() { + echo "== building the awa binary and the Python bindings" + cargo build --release -p awa-cli || return 1 + if [ ! -x awa-python/.venv/bin/python ]; then + uv venv --quiet awa-python/.venv || return 1 + uv pip install --quiet --python awa-python/.venv maturin || return 1 + fi + (cd awa-python && VIRTUAL_ENV="$PWD/.venv" PATH="$PWD/.venv/bin:$PATH" maturin develop --release) || return 1 +} +if [ "${SOAK_SKIP_PREPARE:-0}" != 1 ]; then + prepare >"$RESULTS_DIR/prepare.log" 2>&1 || { echo "prepare failed; see $RESULTS_DIR/prepare.log" >&2; exit 2; } +fi +export AWA_SOAK_PYTHON="${AWA_SOAK_PYTHON:-$PWD/awa-python/.venv/bin/python}" +export AWA_SOAK_AWA_BIN="${AWA_SOAK_AWA_BIN:-${CARGO_TARGET_DIR:-$PWD/target}/release/awa}" +CONTAINERS=() +STATS_PID="" + +cleanup() { + [ -n "$STATS_PID" ] && kill "$STATS_PID" 2>/dev/null + for container in "${CONTAINERS[@]:-}"; do + [ -n "$container" ] && docker rm -f -v "$container" >/dev/null 2>&1 + done +} +trap cleanup EXIT +trap 'exit 130' INT TERM + +start_container() { + local engine="$1" name="$2" port="$3" + local cpu_args=() + [ -n "${SOAK_PG_CPUS:-}" ] && cpu_args=(--cpus "$SOAK_PG_CPUS") + local common=(-c "max_connections=$MAX_CONNECTIONS" -c max_locks_per_transaction=256 -c shared_buffers=512MB -c "log_line_prefix=%m [%p] %u@%a ") + case "$engine" in + pg17) + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -e POSTGRES_PASSWORD=test -p "$port:5432" postgres:17 "${common[@]}" >/dev/null ;; + pg18) + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -e POSTGRES_PASSWORD=test -p "$port:5432" postgres:18 "${common[@]}" >/dev/null ;; + pg18-cassert) + # The image ships an initialised, trust-authenticated cluster in /data. Parallel query is + # off: at -O0 a handful of admin reads would otherwise take every core. + local pgargs='-c listen_addresses=0.0.0.0 -c max_connections='"$MAX_CONNECTIONS"' -c max_locks_per_transaction=256 -c shared_buffers=512MB -c track_activity_query_size=65536 -c max_parallel_workers_per_gather=0 -c logging_collector=on -c log_directory=/data/log -c log_filename=pg.log -c log_rotation_size=0' + case "$CASSERT_IMAGE" in + *epqlog*) + # The instrumented build logs an AWAEPQ line, with the full query text, for every + # partition-pruning setup: tens of gigabytes per hour under this workload. The log + # collector writes into a FIFO whose reader drops those messages and keeps the rest + # (assertion failures, crash notices, errors) in /data/log/pg-filtered.log. + SOAK_PG_LOG_FILE=/data/log/pg-filtered.log + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -p "$port:5432" "$CASSERT_IMAGE" sh -c ' + mkdir -p /data/log + rm -f /data/log/pg.log /data/log/pg-filtered.log + mkfifo /data/log/pg.log + awk -W interactive " + /^[0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9] /{ + pid=\"\"; if (match(\$0,/\\[[0-9]+\\]/)) pid=substr(\$0,RSTART,RLENGTH) + if (\$0 ~ /AWAEPQ/) {skip=1; epq[pid]=1} + else if (\$0 ~ /STATEMENT:/ && (pid in epq)) {skip=1; delete epq[pid]} + else {skip=0; delete epq[pid]} + } + !skip{print; fflush()}" \ + >/data/log/pg-filtered.log & + exec postgres -D /data '"$pgargs"'' >/dev/null ;; + *) + SOAK_PG_LOG_FILE=/data/log/pg.log + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -p "$port:5432" "$CASSERT_IMAGE" \ + sh -c 'mkdir -p /data/log && exec postgres -D /data '"$pgargs" >/dev/null ;; + esac ;; + *) + echo "unknown engine: $engine (expected pg17, pg18 or pg18-cassert)" >&2 + return 1 ;; + esac +} + +wait_ready() { + local name="$1" deadline=$((SECONDS + 120)) + until docker exec "$name" psql -U postgres -d postgres -Atc 'select 1' >/dev/null 2>&1; do + if [ "$SECONDS" -ge "$deadline" ]; then + echo "container $name did not become ready" >&2 + docker logs --tail 40 "$name" >&2 + return 1 + fi + sleep 1 + done + sleep 2 +} + +sample_stats() { + local name="$1" out="$2" + echo "time,cpu_percent,mem" >"$out" + while true; do + docker stats --no-stream --format '{{.CPUPerc}},{{.MemUsage}}' "$name" 2>/dev/null | + sed "s/^/$(date +%H:%M:%S),/" >>"$out" + sleep 15 + done +} + +if [ -n "${SOAK_TEST_BIN:-}" ]; then + harness_cmd=("$SOAK_TEST_BIN" persona_soak --ignored --nocapture --exact) +else + harness_cmd=(cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture) +fi + +overall=0 +for index in "${!ENGINES[@]}"; do + engine="${ENGINES[$index]}" + port=$((BASE_PORT + index)) + name="awa-soak-$engine-$$" + out="$RESULTS_DIR/$engine" + mkdir -p "$out" + echo "== $engine on port $port (container $name)" + start_container "$engine" "$name" "$port" || exit 2 + CONTAINERS+=("$name") + wait_ready "$name" || exit 2 + + log_env=(AWA_SOAK_PG_CONTAINER="$name") + [ "$engine" = pg18-cassert ] && log_env+=(AWA_SOAK_PG_CONTAINER_LOG="$SOAK_PG_LOG_FILE") + + sample_stats "$name" "$out/container-stats.csv" & + STATS_PID=$! + + env AWA_SOAK_DATABASE_URL="postgres://postgres:test@localhost:$port/awa_soak_${engine//-/_}" \ + AWA_SOAK_SECONDS="$SECONDS_PER_ENGINE" \ + AWA_SOAK_REPORT_DIR="$out" \ + "${log_env[@]}" \ + "${harness_cmd[@]}" >"$out/run.log" 2>&1 + status=$? + + kill "$STATS_PID" 2>/dev/null + STATS_PID="" + if [ $status -ne 0 ]; then + overall=1 + echo " FAILED (exit $status); see $out/report.md and $out/reproducer.md" + # Keep the server log with the reproducer when the run failed. + docker logs "$name" >"$out/server.log" 2>&1 || true + [ "$engine" = pg18-cassert ] && docker exec "$name" cat "$SOAK_PG_LOG_FILE" >"$out/server.log" 2>/dev/null + else + echo " passed; report in $out/report.md" + fi + docker rm -f -v "$name" >/dev/null 2>&1 +done + +echo "results: $RESULTS_DIR" +exit $overall