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..11a22175a --- /dev/null +++ b/awa/tests/persona_soak/config.rs @@ -0,0 +1,270 @@ +//! Run configuration, read from the environment. + +use std::path::PathBuf; +use std::str::FromStr; +use std::time::Duration; + +#[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, +} + +impl Persona { + pub const ALL: [Persona; 13] = [ + 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, + ]; + + 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", + } + } + + 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") + }) + } +} + +#[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, +} + +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_)", + )?; + 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: 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: Duration::from_secs( + env_parse("AWA_SOAK_PROMOTION_OVERDUE_SECONDS").unwrap_or(10), + ), + stall_window: Duration::from_secs(env_parse("AWA_SOAK_STALL_SECONDS").unwrap_or(30)), + chaos_slack: 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), + }) + } + + 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..e4e331e33 --- /dev/null +++ b/awa/tests/persona_soak/ledger.rs @@ -0,0 +1,225 @@ +//! 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; + +/// 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 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, + 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 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..bb4293fd6 --- /dev/null +++ b/awa/tests/persona_soak/mod.rs @@ -0,0 +1,19 @@ +//! Persona-driven soak harness with an invariant oracle. See the test binary +//! `persona_soak_test.rs` for how to run it. + +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 registry; +pub mod report; +pub mod rng; +pub mod roles; +pub mod run; +pub mod shard_oracle; +pub mod shared; +pub mod stop; diff --git a/awa/tests/persona_soak/oracle.rs b/awa/tests/persona_soak/oracle.rs new file mode 100644 index 000000000..5e16989b8 --- /dev/null +++ b/awa/tests/persona_soak/oracle.rs @@ -0,0 +1,871 @@ +//! 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}; +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 Ok(reset) = reset else { 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.note( + "oracle", + format!("server probe failed repeatedly: {last_error}"), + ); + } + + 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.note( + "oracle", + format!("terminal uniqueness query failed: {error}"), + ), + } + + 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 = still + .map(|rows| rows.into_iter().map(|row| row.0).collect()) + .unwrap_or_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 + .note("oracle", format!("terminal/live query failed: {error}")), + } + 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 Ok(rows) = rows else { 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 Ok(rows) = rows else { 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 Ok(still) = still else { 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 Ok(rows) = rows else { 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 Ok(available) = exact else { 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; + 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); + match live { + Ok(count) if in_flight == 0 && python_active == 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(); + self.shared.ledger.for_each(|id, record| { + if record.has(flag::OPERATOR_REMOVED) || record.has(flag::DLQ_PURGED) { + return; + } + let db_seen = flag::DB_SEEN_COMPLETED | flag::DB_SEEN_FAILED | flag::DB_SEEN_CANCELLED; + let evidence = record.completed_events > 0 + || record.exhausted_events > 0 + || record.cancel_events > 0 + || record.flags & db_seen != 0 + || (record.has(flag::NO_EVENTS) && record.has(flag::RETURNED_COMPLETED)); + if evidence { + let completed = record.completed_events > 0 + || record.has(flag::RETURNED_COMPLETED) + || record.has(flag::DB_SEEN_COMPLETED); + if record.has(flag::DB_SEEN_CANCELLED) + && record.expect & expect::CANCELLED == 0 + && !record.has(flag::HANDLER_CANCELLED) + && unexpected.len() < MAX_ROWS + { + unexpected.push(format!( + "{{\"job_id\":{id},\"outcome\":\"cancelled\",\"allowed\":{}}}", + record.expect + )); + } + if completed + && record.expect & expect::COMPLETED == 0 + && unexpected.len() < MAX_ROWS + { + unexpected.push(format!( + "{{\"job_id\":{id},\"outcome\":\"completed\",\"allowed\":{}}}", + record.expect + )); + } + if (record.exhausted_events > 0 || record.has(flag::DB_SEEN_FAILED)) + && record.expect & expect::FAILED == 0 + && unexpected.len() < MAX_ROWS + { + unexpected.push(format!( + "{{\"job_id\":{id},\"outcome\":\"failed\",\"allowed\":{}}}", + record.expect + )); + } + } 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, ',') FROM ( + SELECT job_id, 'done:' || state::text AS place FROM {schema}.done_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'dlq' FROM {schema}.dlq_entries WHERE job_id = ANY($1) + UNION ALL SELECT b.job_id, 'receipt_completed' + FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids) AS b(job_id) + WHERE c.job_ids && $1 AND b.job_id = ANY($1) + UNION ALL SELECT job_id, 'deferred:' || state::text FROM {schema}.deferred_jobs WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'ready' FROM {schema}.ready_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'leased:' || state::text 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().collect(), + Err(error) => { + self.shared + .note("oracle", format!("final lookup failed: {error}")); + 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()) + ) + }; + match found.get(id) { + Some(place) + if place.contains("done:") + || place.contains("dlq") + || place.contains("receipt_completed") => + { + terminal_in_db += 1; + 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); + } + } + } + } + } + if !forensic_ids.is_empty() { + self.record_forensics(&schema, &forensic_ids).await; + } + 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(), + } + } +} + +fn count_is_idle(count: i64) -> bool { + count <= 0 +} 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..02f434454 --- /dev/null +++ b/awa/tests/persona_soak/personas/chaos.rs @@ -0,0 +1,241 @@ +//! 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 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..4a4c725df --- /dev/null +++ b/awa/tests/persona_soak/personas/mod.rs @@ -0,0 +1,96 @@ +//! The personas. Each one owns its queues, job kinds, runtimes and producers, and +//! reports what it did through [`Shared::stats`]. + +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 + }; + builder + .heartbeat_interval(Duration::from_secs(2)) + .heartbeat_staleness(Duration::from_secs(12)) + .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..f6ba572d4 --- /dev/null +++ b/awa/tests/persona_soak/personas/python.rs @@ -0,0 +1,677 @@ +//! 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::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::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 { + 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 { + '?' + }, + if role == "full" { + "soak_python" + } else { + "soak_python_probe" + } + ), + ) + .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_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(); + 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) +} + +/// 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, +} + +impl PythonHandle { + /// 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 Ok(mut child) = spawn(shared, python, ledger_path, "migrate_probe") else { + return; + }; + let Ok(pool) = + super::super::db::connect_pool(&shared.config.database_url, "soak_python_probe_watch", 1) + .await + else { + return; + }; + let sessions = || async { + sqlx::query_as::<_, (i64, i64)>( + "SELECT count(*)::bigint, count(*) FILTER (WHERE state = 'idle in transaction')::bigint + FROM pg_stat_activity WHERE application_name = 'soak_python_probe' AND pid <> pg_backend_pid() + AND backend_start > now() - interval '30 seconds'", + ) + .fetch_one(&pool) + .await + .unwrap_or((0, 0)) + }; + let mut peak = (0, 0); + let deadline = Instant::now() + Duration::from_secs(45); + loop { + match child.try_wait() { + Ok(Some(_)) | Err(_) => break, + Ok(None) if Instant::now() >= deadline => { + shared.finding( + "python_migrate_probe_hang", + "the Python migrate-then-query probe did not exit within 45s", + vec![], + ); + let _ = child.kill().await; + break; + } + Ok(None) => {} + } + let now = sessions().await; + peak = (peak.0.max(now.0), peak.1.max(now.1)); + tokio::time::sleep(Duration::from_millis(500)).await; + } + stats.add("python_migrate_probes", 1); + stats.set_max("python_probe_peak_sessions", peak.0); + // The probe's pool is capped at 4 connections; more than that, or any session left in a + // transaction, means connections were stranded. + if peak.0 > 4 || peak.1 > 0 { + shared.violation( + "python_connections_stranded", + format!("the Python migrate probe held {} session(s) ({} idle in transaction) against a pool of 4", peak.0, peak.1), + vec![], + ); + } + tokio::time::sleep(Duration::from_secs(2)).await; + let (left, _) = sessions().await; + if left > 0 { + shared.violation( + "python_connections_stranded", + format!("{left} Python session(s) remain after the migrate probe process exited"), + 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; + { + let mut process = handle.process.lock().await; + if let Some(child) = process.child.as_mut() { + if let Ok(Some(status)) = child.try_wait() { + shared.stats(persona).add("python_unexpected_exits", 1); + shared.note( + "python", + format!("python process exited unexpectedly: {status}; restarting"), + ); + process.child = spawn(&shared, &python, &ledger_path, "full").ok(); + } + } + } + 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 }); + } + } + + while probes.join_next().await.is_some() {} + // Stop producing; the workers keep running while the system drains. + if let Some(pid) = handle + .process + .lock() + .await + .child + .as_ref() + .and_then(Child::id) + { + let _ = std::process::Command::new("kill") + .args(["-USR1", &pid.to_string()]) + .status(); + } + shared.note("python", "load stopped; draining"); + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + importer.import(&shared); + tail_terminal_rows(&shared, &tail_pool, &mut tail_tick).await; + } + 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, state::text FROM {schema}.done_entries WHERE queue = ANY($1) AND finalized_at > now() - interval '25 seconds' + UNION ALL SELECT job_id, 'failed' FROM {schema}.dlq_entries WHERE queue = ANY($1) AND dlq_at > now() - interval '25 seconds' + UNION ALL SELECT b.job_id, 'completed' FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids) AS b(job_id) + 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.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, state) in rows { + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + u32::MAX, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| match state.as_str() { + "completed" => record.flags |= flag::DB_SEEN_COMPLETED, + "failed" => record.flags |= flag::DB_SEEN_FAILED, + "cancelled" => record.flags |= flag::DB_SEEN_CANCELLED, + _ => {} + }, + ); + } +} + +/// 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(()) +} 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..73631e039 --- /dev/null +++ b/awa/tests/persona_soak/python/soak_python.py @@ -0,0 +1,355 @@ +"""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 + 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) + + stop = asyncio.Event() + loop = asyncio.get_running_loop() + loop.add_signal_handler(signal.SIGUSR1, stop.set) + 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/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..37e209ac3 --- /dev/null +++ b/awa/tests/persona_soak/report.rs @@ -0,0 +1,319 @@ +//! 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.is_empty() { + "PASS" + } else { + "FAIL" + } + ); + 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); + + 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 + ); + } + let _ = writeln!(out); + } + + let _ = writeln!(out, "## Run"); + 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.is_empty() { "pass" } else { "fail" }, + "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, + "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/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..45c8e7b13 --- /dev/null +++ b/awa/tests/persona_soak/roles.rs @@ -0,0 +1,126 @@ +//! 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. 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> { + 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, without covering anything the application roles create themselves. + 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); + } + 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..0b8d5214e --- /dev/null +++ b/awa/tests/persona_soak/run.rs @@ -0,0 +1,284 @@ +//! 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 + }; + + // 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, + ] { + 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}")); + 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(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![], + ), + } + } + + 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, + } + }) +} + +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..77240ea06 --- /dev/null +++ b/awa/tests/persona_soak/shard_oracle.rs @@ -0,0 +1,634 @@ +//! 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 Ok(rows) = shard_rows(&pool, &schema, &queues).await else { + 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.note("shards", format!("final shard query failed: {error}")), + } + + 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..44a98d49a --- /dev/null +++ b/awa/tests/persona_soak/shared.rs @@ -0,0 +1,451 @@ +//! 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::{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>, + /// 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, + 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()), + 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(), + 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); + } + + /// 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); + 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.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.exhausted_events = record.exhausted_events.saturating_add(1); + record.flags |= flag::EXHAUSTED; + } + UntypedJobEvent::Cancelled { .. } => { + 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_test.rs b/awa/tests/persona_soak_test.rs new file mode 100644 index 000000000..a57b0f5c0 --- /dev/null +++ b/awa/tests/persona_soak_test.rs @@ -0,0 +1,33 @@ +//! 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 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() + ); +} diff --git a/docs/test-plan.md b/docs/test-plan.md index 435b18861..00cfb3b31 100644 --- a/docs/test-plan.md +++ b/docs/test-plan.md @@ -304,6 +304,72 @@ 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 + +`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 | +| 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) +- 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`, default all | +| `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_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. + ## Running Tests ```bash @@ -328,6 +394,17 @@ 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 + # 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.sh b/scripts/persona-soak.sh new file mode 100755 index 000000000..f6bafd0f4 --- /dev/null +++ b/scripts/persona-soak.sh @@ -0,0 +1,193 @@ +#!/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 +# 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 +# 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 +# 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 "$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 +} + +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" + CONTAINERS+=("$name") + + echo "== $engine on port $port (container $name)" + start_container "$engine" "$name" "$port" || exit 2 + 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[@]}" \ + cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture \ + >"$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 "$name" >/dev/null 2>&1 +done + +echo "results: $RESULTS_DIR" +exit $overall