diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8f40a83d4..83079c307 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -285,6 +285,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s @@ -328,6 +329,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s @@ -368,6 +370,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s @@ -402,6 +405,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s @@ -449,6 +453,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s @@ -507,7 +512,7 @@ jobs: POSTGRES_DB: awa_test ports: - 5432:5432 - options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5 + options: --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5 defaults: run: working-directory: awa-python @@ -579,6 +584,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s @@ -685,6 +691,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s @@ -790,6 +797,7 @@ jobs: ports: - 5432:5432 options: >- + --shm-size=256m --health-cmd pg_isready --health-interval 10s --health-timeout 5s diff --git a/CHANGELOG.md b/CHANGELOG.md index db9b15515..4b2a82654 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] +### Forward-ported from 0.6.14 + +- Everything [0.6.14](https://github.com/hardbyte/awa/releases/tag/v0.6.14) shipped that this branch did not already have; see its release notes for the details. Promotion runs each enqueue lane in its own transaction, hands a busy lane to a dedicated waiter, and wakes claimers after the lanes commit; `awa.maintenance.promote.lane_contention` counts the hand-offs. Rescue skips an attempt whose retained ready bodies are ambiguous instead of stopping rescue of every other attempt. Lifecycle events no longer re-read a completed job, and a failed re-read no longer fails a committed callback wait. Admin job lists and kind counts find their candidates in read phases that release partition locks between them, and probe receipt closures more cheaply; a job list still loads each candidate with a point read, which holds `ACCESS SHARE` on every ready, lease and claim partition for the length of that read. `awa migrate` rejects `--sql` combined with `--extract-to`. The persona soak suite (`awa/tests/persona_soak_test.rs`, `scripts/persona-soak.sh`) is added. + ### Added - **The queue-storage prune `lock_timeout` is configurable, and prune diff --git a/Cargo.lock b/Cargo.lock index 2055d2754..a15535ce8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -389,6 +389,7 @@ dependencies = [ "serde_json", "sqlx", "tokio", + "tokio-util", "tonic", "tracing", "tracing-opentelemetry", diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 2a06ecebc..95ba028e8 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -43,7 +43,7 @@ enum Commands { #[arg(long)] extract_to: Option, /// Print migration SQL to stdout instead of applying - #[arg(long)] + #[arg(long, conflicts_with = "extract_to")] sql: bool, /// Only include migrations after this version (exclusive) #[arg(long)] diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 5241d6740..9dc0cfd8e 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -70,6 +70,7 @@ pub mod names { pub const MAINTENANCE_PROMOTE_BATCHES: &str = "awa.maintenance.promote_batches"; pub const MAINTENANCE_PROMOTE_BATCH_SIZE: &str = "awa.maintenance.promote_batch_size"; pub const MAINTENANCE_PROMOTE_DURATION: &str = "awa.maintenance.promote_duration"; + pub const MAINTENANCE_PROMOTE_LANE_CONTENTION: &str = "awa.maintenance.promote.lane_contention"; pub const MAINTENANCE_BRANCH_DURATION: &str = "awa.maintenance.branch.duration"; pub const MAINTENANCE_BRANCH_OVERRUN: &str = "awa.maintenance.branch.overrun"; pub const MAINTENANCE_ROTATE_ATTEMPTS: &str = "awa.maintenance.rotate.attempts"; @@ -192,6 +193,11 @@ pub struct AwaMetrics { pub promotion_batch_size: Histogram, /// Promotion query duration. pub promotion_duration_seconds: Histogram, + /// Promotion lanes whose locks were busy past the leader's wait budget, + /// attributed to (state, queue, outcome): `handed_off` to a dedicated + /// waiter, `waiter_limit` when every waiter was busy, `waiter_timed_out` + /// when a waiter's own budget ran out. + pub promotion_lane_contention: Counter, /// Current in-flight jobs (can go up and down). pub jobs_in_flight: UpDownCounter, /// Total heartbeat batches sent. @@ -482,6 +488,11 @@ impl AwaMetrics { .with_unit("s") .with_boundaries(LATENCY_BUCKETS_SECONDS.to_vec()) .build(), + promotion_lane_contention: meter + .u64_counter(names::MAINTENANCE_PROMOTE_LANE_CONTENTION) + .with_description("Promotion lanes whose locks were busy, by state/queue/outcome") + .with_unit("{lane}") + .build(), jobs_in_flight: meter .i64_up_down_counter(names::JOB_IN_FLIGHT) .with_description("Current number of in-flight jobs") @@ -909,6 +920,23 @@ impl AwaMetrics { .record(duration.as_secs_f64(), &attrs); } + /// Record a promotion lane whose locks were busy past the leader's wait + /// budget, with what became of it (`handed_off`, `waiter_limit`, + /// `waiter_timed_out`). + pub fn record_promotion_lane_contention( + &self, + state: &str, + queue: &str, + outcome: &'static str, + ) { + let attrs = [ + opentelemetry::KeyValue::new("awa.job.state", state.to_string()), + opentelemetry::KeyValue::new("awa.job.queue", queue.to_string()), + opentelemetry::KeyValue::new("awa.promote.outcome", outcome), + ]; + self.promotion_lane_contention.add(1, &attrs); + } + /// Record in-flight change. pub fn record_in_flight_change(&self, queue: &str, delta: i64) { let attrs = [opentelemetry::KeyValue::new( @@ -1462,6 +1490,65 @@ mod tests { assert!(saw_already_pruned); } + #[test] + fn promotion_lane_contention_counts_each_outcome_per_state_and_queue() { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + let metrics = AwaMetrics::new(&provider.meter("awa-metrics-test")); + + metrics.record_promotion_lane_contention("scheduled", "emails", "handed_off"); + metrics.record_promotion_lane_contention("scheduled", "emails", "handed_off"); + metrics.record_promotion_lane_contention("retryable", "emails", "waiter_limit"); + metrics.record_promotion_lane_contention("scheduled", "billing", "waiter_timed_out"); + provider.force_flush().expect("flush metrics"); + let resource_metrics = exporter.get_finished_metrics().expect("read metrics"); + + let mut counts = BTreeMap::new(); + for resource in &resource_metrics { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + if let ( + names::MAINTENANCE_PROMOTE_LANE_CONTENTION, + AggregatedMetrics::U64(MetricData::Sum(sum)), + ) = (metric.name(), metric.data()) + { + for point in sum.data_points() { + let label = |key: &str| { + point + .attributes() + .find(|kv| kv.key.as_str() == key) + .map(|kv| kv.value.as_str().into_owned()) + .unwrap_or_default() + }; + counts.insert( + ( + label("awa.job.state"), + label("awa.job.queue"), + label("awa.promote.outcome"), + ), + point.value(), + ); + } + } + } + } + } + + let key = |state: &str, queue: &str, outcome: &str| { + (state.to_string(), queue.to_string(), outcome.to_string()) + }; + assert_eq!( + counts, + BTreeMap::from([ + (key("retryable", "emails", "waiter_limit"), 1), + (key("scheduled", "billing", "waiter_timed_out"), 1), + (key("scheduled", "emails", "handed_off"), 2), + ]) + ); + } + #[test] fn prune_attempts_count_every_outcome_per_ring() { use awa_model::{PruneDurations, PruneOutcome, SkipReason}; diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 7be7913e0..9dbb5825f 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -10,7 +10,7 @@ use sqlx::Acquire; use sqlx::PgExecutor; use sqlx::PgPool; use std::cmp::max; -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::fmt; use std::str::FromStr; use uuid::Uuid; @@ -684,50 +684,43 @@ fn queue_storage_recent_jobs_cte(schema: &str, since: &str) -> String { ) } -/// The per-kind, per-queue counts behind [`job_kind_overviews`], as a -/// `kind_counts` CTE of `kind`, `job_count`, `queue_count` and -/// `completed_last_hour`. -/// -/// It counts what `current_jobs` holds without expanding it. Each place a job -/// can live is aggregated to `(kind, queue)` on its own, as a `MATERIALIZED` -/// CTE, and only those few rows are combined. Read through `current_jobs`, -/// the receipt-batch join to `ready_entries` and the `UNION ALL` are both -/// estimated at one row whatever the statistics say, so Postgres chooses -/// nested loops that probe once per retained row, and a join on -/// `queue_claim_heads` is only as good as that table's estimate. -/// -/// Available rows are found per lane as a claim finds them (see -/// [`crate::queue_storage::available_by_queue_sql`]). A completion batch -/// member's kind lives on its ready row, so that leg unnests each batch and -/// probes the ready primary key. -fn queue_storage_kind_counts_cte(schema: &str) -> String { +// Per-source kind/queue counts keep the same joins and predicates as the +// combined projection, while allowing each source's relation locks to end +// before another source opens a different set of retained partitions. +#[derive(Clone, Copy, PartialEq, Eq)] +enum KindCountSource { + Available, + Leased, + Deferred, + Done, + Receipts, + Dlq, +} + +// Generate only this source's relations: unused CTEs would still acquire +// relation locks during parse/plan before execution could discard them. +fn queue_storage_kind_count_source_sql(schema: &str, source: KindCountSource) -> String { let available = crate::queue_storage::available_lane_rows_sql(schema, "heads", "ready.kind", "TRUE"); - let open_running = crate::queue_storage::open_receipt_running_claims_sql(schema); - format!( - r#" - heads AS MATERIALIZED ( - SELECT - claims.queue, - claims.priority, - claims.enqueue_shard, - {schema}.sequence_next_value(claims.seq_name) AS claim_seq, - {schema}.sequence_next_value(enqueues.seq_name) AS next_seq - FROM {schema}.queue_claim_heads AS claims - JOIN {schema}.queue_enqueue_heads AS enqueues - ON enqueues.queue = claims.queue - AND enqueues.priority = claims.priority - AND enqueues.enqueue_shard = claims.enqueue_shard - ), - available_by_kind AS MATERIALIZED ( + match source { + KindCountSource::Available => format!( + r#" + WITH heads AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, claim_seq, next_seq + FROM unnest( + $1::text[], $2::smallint[], $3::smallint[], $4::bigint[], $5::bigint[] + ) AS captured(queue, priority, enqueue_shard, claim_seq, next_seq) + ) SELECT lane.kind, heads.queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM heads CROSS JOIN LATERAL ( {available} ) AS lane GROUP BY lane.kind, heads.queue + "# ), - leased_by_kind AS MATERIALIZED ( + KindCountSource::Leased => format!( + r#" SELECT ready.kind, leases.queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM {schema}.leases AS leases JOIN {schema}.ready_entries AS ready @@ -751,27 +744,17 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { AND tomb.lane_seq = ready.lane_seq ) GROUP BY ready.kind, leases.queue + "# ), - deferred_by_kind AS MATERIALIZED ( + KindCountSource::Deferred => format!( + r#" SELECT kind, queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM {schema}.deferred_jobs GROUP BY kind, queue + "# ), - -- Receipt claims that are still open are running; a queue or priority - -- move copies a ready row with its kind, so any one of the job's - -- ready rows gives it. - running_by_kind AS MATERIALIZED ( - SELECT created.kind, open_claims.queue, count(*)::bigint AS n, 0::bigint AS n_recent - FROM ({open_running}) AS open_claims - CROSS JOIN LATERAL ( - SELECT ready.kind - FROM {schema}.ready_entries AS ready - WHERE ready.job_id = open_claims.job_id - LIMIT 1 - ) AS created - GROUP BY created.kind, open_claims.queue - ), - done_by_kind AS MATERIALIZED ( + KindCountSource::Done => format!( + r#" SELECT kind, queue, @@ -782,8 +765,10 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { )::bigint AS n_recent FROM {schema}.done_entries GROUP BY kind, queue + "# ), - receipts_by_kind AS MATERIALIZED ( + KindCountSource::Receipts => format!( + r#" SELECT ready.kind, batch.queue, @@ -797,14 +782,20 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { batch.run_leases, batch.lane_seqs ) AS item(job_id, run_lease, lane_seq) - JOIN {schema}.ready_entries AS ready - ON ready.ready_slot = batch.ready_slot - AND ready.ready_generation = batch.ready_generation - AND ready.queue = batch.queue - AND ready.priority = batch.priority - AND ready.enqueue_shard = batch.enqueue_shard - AND ready.lane_seq = item.lane_seq - AND ready.job_id = item.job_id + CROSS JOIN LATERAL ( + SELECT ready.kind + FROM {schema}.ready_entries AS ready + WHERE ready.ready_slot = batch.ready_slot + AND ready.ready_generation = batch.ready_generation + AND ready.queue = batch.queue + AND ready.priority = batch.priority + AND ready.enqueue_shard = batch.enqueue_shard + AND ready.lane_seq = item.lane_seq + AND ready.job_id = item.job_id + -- Keep the full member key in the parameterized lookup, + -- rather than scanning its retained lane for every member. + OFFSET 0 + ) AS ready WHERE NOT EXISTS ( SELECT 1 FROM {schema}.receipt_completion_tombstones AS tomb WHERE tomb.ready_slot = batch.ready_slot @@ -813,23 +804,25 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { AND tomb.run_lease = item.run_lease ) GROUP BY ready.kind, batch.queue + "# ), - dlq_by_kind AS MATERIALIZED ( + KindCountSource::Dlq => format!( + r#" SELECT kind, queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM {schema}.dlq_entries GROUP BY kind, queue + "# ), + } +} + +// Merge small (kind, queue) aggregates from separately locked read sources. +fn queue_storage_kind_counts_cte() -> &'static str { + r#" kind_queue AS ( SELECT kind, queue, sum(n)::bigint AS n, sum(n_recent)::bigint AS n_recent - FROM ( - SELECT * FROM available_by_kind - UNION ALL SELECT * FROM leased_by_kind - UNION ALL SELECT * FROM deferred_by_kind - UNION ALL SELECT * FROM running_by_kind - UNION ALL SELECT * FROM done_by_kind - UNION ALL SELECT * FROM receipts_by_kind - UNION ALL SELECT * FROM dlq_by_kind - ) AS legs + FROM unnest($1::text[], $2::text[], $3::bigint[], $4::bigint[]) + AS sources(kind, queue, n, n_recent) GROUP BY kind, queue ), kind_counts AS ( @@ -841,8 +834,7 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { FROM kind_queue GROUP BY kind ) - "# - ) + "# } async fn notify_cancellation_tx<'a>( @@ -858,14 +850,21 @@ async fn notify_cancellation_tx<'a>( Ok(()) } -/// Candidate job ids for [`list_jobs`]: the newest `limit` ids below `cursor` -/// matching `filter`'s state, kind and queue, newest first, with the binds the -/// statement takes in order. +#[derive(Clone, Copy, PartialEq, Eq)] +enum JobCandidateSource { + States, + Receipts(usize), +} + +/// Candidate job ids for [`list_jobs`] from one source: the newest `limit` +/// distinct ids below `cursor`, with the statement's binds in order. State +/// sources apply kind/state/queue filters. Receipt sources enumerate open +/// claims in one physical slot and apply kind/queue before their limit. /// -/// It selects what `current_jobs` would with the same filter, but each place a -/// job can live is read and limited on its own, with only the legs that can -/// hold the requested state, and the filter written into the statement rather -/// than as `$n IS NULL OR` terms that a generic plan can't drop. Read through +/// Each place a job can live is read and limited on its own, with only the +/// state legs that can hold the requested state. Known filters are written +/// into the statement rather than as `$n IS NULL OR` terms that a generic plan +/// can't drop. Read through /// `current_jobs`, Postgres joins every retained completion to a scan of the /// ready rows, because that join and the `UNION ALL` are each estimated at one /// row. Available rows are found per lane as a claim finds them, and a @@ -877,13 +876,17 @@ fn queue_storage_job_id_candidates_sql( filter: &ListJobsFilter, cursor: Option, limit: i64, + source: JobCandidateSource, ) -> (String, Vec) { let mut binds = Vec::new(); let mut bind = |value: CandidateBind| { binds.push(value); format!("${}", binds.len()) }; - let state = filter.state.map(|state| bind(CandidateBind::State(state))); + let state = filter + .state + .filter(|_| source == JobCandidateSource::States) + .map(|state| bind(CandidateBind::State(state))); let kind = filter .kind .as_ref() @@ -917,9 +920,10 @@ fn queue_storage_job_id_candidates_sql( sql }; let holds = |states: &[JobState]| { - filter - .state - .is_none_or(|requested| states.contains(&requested)) + source == JobCandidateSource::States + && filter + .state + .is_none_or(|requested| states.contains(&requested)) }; let mut legs: Vec = Vec::new(); @@ -958,7 +962,7 @@ fn queue_storage_job_id_candidates_sql( ); legs.push(format!( r#"( - SELECT lane.job_id + SELECT DISTINCT lane.job_id FROM heads CROSS JOIN LATERAL ({lane_rows}) AS lane ORDER BY lane.job_id DESC @@ -970,7 +974,7 @@ fn queue_storage_job_id_candidates_sql( let predicates = predicates(Some("state"), Some("kind"), "queue", "job_id"); legs.push(format!( r#"( - SELECT job_id + SELECT DISTINCT job_id FROM {schema}.deferred_jobs WHERE TRUE{predicates} ORDER BY job_id DESC @@ -987,7 +991,7 @@ fn queue_storage_job_id_candidates_sql( ); legs.push(format!( r#"( - SELECT leases.job_id + SELECT DISTINCT leases.job_id FROM {schema}.leases AS leases JOIN {schema}.ready_entries AS ready ON ready.ready_slot = leases.ready_slot @@ -1015,8 +1019,9 @@ fn queue_storage_job_id_candidates_sql( )"# )); } - if holds(&[JobState::Running]) { - let open_running = crate::queue_storage::open_receipt_running_claims_sql(schema); + if let JobCandidateSource::Receipts(slot) = source { + let open_running = + crate::queue_storage::open_receipt_running_claims_sql_in_slot(schema, slot); let predicates = predicates( None, Some("created.kind"), @@ -1025,7 +1030,7 @@ fn queue_storage_job_id_candidates_sql( ); legs.push(format!( r#"( - SELECT open_claims.job_id + SELECT DISTINCT open_claims.job_id FROM ({open_running}) AS open_claims CROSS JOIN LATERAL ( SELECT ready.kind @@ -1043,7 +1048,7 @@ fn queue_storage_job_id_candidates_sql( let predicates = predicates(Some("state"), Some("kind"), "queue", "job_id"); legs.push(format!( r#"( - SELECT job_id + SELECT DISTINCT job_id FROM {schema}.done_entries WHERE TRUE{predicates} ORDER BY job_id DESC @@ -1059,7 +1064,7 @@ fn queue_storage_job_id_candidates_sql( .unwrap_or_default(); legs.push(format!( r#"( - SELECT member.job_id + SELECT DISTINCT member.job_id FROM ( SELECT batch.ready_slot, @@ -1108,7 +1113,7 @@ fn queue_storage_job_id_candidates_sql( let predicates = predicates(None, Some("kind"), "queue", "job_id"); legs.push(format!( r#"( - SELECT job_id + SELECT DISTINCT job_id FROM {schema}.dlq_entries WHERE TRUE{predicates} ORDER BY job_id DESC @@ -1120,7 +1125,7 @@ fn queue_storage_job_id_candidates_sql( let sql = format!( r#" {with} - SELECT job_id + SELECT DISTINCT job_id FROM ({legs}) AS candidates ORDER BY job_id DESC LIMIT {limit} @@ -1139,6 +1144,34 @@ enum CandidateBind { Int(i64), } +async fn fetch_queue_storage_job_candidates( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + sql: &str, + binds: Vec, +) -> Result, sqlx::Error> { + sqlx::query("SAVEPOINT awa_list_candidates") + .execute(tx.as_mut()) + .await?; + let mut query = sqlx::query_scalar::<_, i64>(audited_sql(sql)); + for value in binds { + query = match value { + CandidateBind::State(state) => query.bind(state), + CandidateBind::Text(text) => query.bind(text), + CandidateBind::Int(int) => query.bind(int), + }; + } + let ids = query.fetch_all(tx.as_mut()).await?; + // Keep the outer repeatable-read snapshot, but release this read phase's + // relation locks. Destroy the savepoint too rather than stacking names. + sqlx::query("ROLLBACK TO SAVEPOINT awa_list_candidates") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_list_candidates") + .execute(tx.as_mut()) + .await?; + Ok(ids) +} + async fn list_queue_storage_jobs( store: &QueueStorage, pool: &PgPool, @@ -1155,26 +1188,65 @@ async fn list_queue_storage_jobs( let mut cursor = filter.before_id; loop { - let (sql, binds) = - queue_storage_job_id_candidates_sql(store.schema(), filter, cursor, candidate_limit); - let mut query = sqlx::query_scalar::<_, i64>(audited_sql(sql)); - for value in binds { - query = match value { - CandidateBind::State(state) => query.bind(state), - CandidateBind::Text(text) => query.bind(text), - CandidateBind::Int(int) => query.bind(int), - }; - } - // As in `queue_overviews`: compiling costs more than running. + // Every candidate phase sees the same MVCC snapshot. Each phase + // releases its own read locks before the next one opens more tables. let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; - let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + let mut tx = sqlx::Connection::begin_with( + &mut *begin_guard.connection, + "BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY", + ) + .await?; begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; - let ids: Vec = query.fetch_all(tx.as_mut()).await?; + let states = filter.state.map_or_else( + || { + vec![ + JobState::Available, + JobState::Scheduled, + JobState::Retryable, + JobState::Running, + JobState::WaitingExternal, + JobState::Completed, + JobState::Failed, + JobState::Cancelled, + ] + }, + |state| vec![state], + ); + let mut ids = BTreeSet::new(); + for state in states { + let mut phase_filter = filter.clone(); + phase_filter.state = Some(state); + let (sql, binds) = queue_storage_job_id_candidates_sql( + store.schema(), + &phase_filter, + cursor, + candidate_limit, + JobCandidateSource::States, + ); + ids.extend(fetch_queue_storage_job_candidates(&mut tx, &sql, binds).await?); + } + if filter.state.is_none_or(|state| state == JobState::Running) { + for slot in queue_storage_populated_claim_slots(&mut tx, store).await? { + let (sql, binds) = queue_storage_job_id_candidates_sql( + store.schema(), + filter, + cursor, + candidate_limit, + JobCandidateSource::Receipts(slot), + ); + ids.extend(fetch_queue_storage_job_candidates(&mut tx, &sql, binds).await?); + } + } tx.commit().await?; drop(begin_guard); + let ids: Vec = ids + .into_iter() + .rev() + .take(candidate_limit as usize) + .collect(); if ids.is_empty() { break; } @@ -3791,6 +3863,173 @@ pub async fn state_counts(pool: &PgPool) -> Result, AwaEr Ok(rows.into_iter().collect()) } +#[derive(sqlx::FromRow)] +struct KindCountHead { + queue: String, + priority: i16, + enqueue_shard: i16, + claim_seq: i64, + next_seq: i64, +} + +// Sequence values do not obey MVCC. Capture them with the first read's +// snapshot instead of letting later claims advance them during receipt phases. +async fn queue_storage_kind_heads( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, +) -> Result, sqlx::Error> { + let schema = store.schema(); + sqlx::query("SAVEPOINT awa_kind_heads") + .execute(tx.as_mut()) + .await?; + let heads = sqlx::query_as::<_, KindCountHead>(audited_sql(format!( + r#" + SELECT + claims.queue, + claims.priority, + claims.enqueue_shard, + {schema}.sequence_next_value(claims.seq_name) AS claim_seq, + {schema}.sequence_next_value(enqueues.seq_name) AS next_seq + FROM {schema}.queue_claim_heads AS claims + JOIN {schema}.queue_enqueue_heads AS enqueues + ON enqueues.queue = claims.queue + AND enqueues.priority = claims.priority + AND enqueues.enqueue_shard = claims.enqueue_shard + "# + ))) + .fetch_all(tx.as_mut()) + .await?; + sqlx::query("ROLLBACK TO SAVEPOINT awa_kind_heads") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_kind_heads") + .execute(tx.as_mut()) + .await?; + Ok(heads) +} + +async fn queue_storage_base_kind_counts( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, + heads: Vec, +) -> Result, sqlx::Error> { + let mut head_queues = Vec::with_capacity(heads.len()); + let mut head_priorities = Vec::with_capacity(heads.len()); + let mut head_shards = Vec::with_capacity(heads.len()); + let mut claim_seqs = Vec::with_capacity(heads.len()); + let mut next_seqs = Vec::with_capacity(heads.len()); + for head in heads { + head_queues.push(head.queue); + head_priorities.push(head.priority); + head_shards.push(head.enqueue_shard); + claim_seqs.push(head.claim_seq); + next_seqs.push(head.next_seq); + } + let mut counts = Vec::new(); + for source in [ + KindCountSource::Available, + KindCountSource::Leased, + KindCountSource::Deferred, + KindCountSource::Done, + KindCountSource::Receipts, + KindCountSource::Dlq, + ] { + let sql = queue_storage_kind_count_source_sql(store.schema(), source); + sqlx::query("SAVEPOINT awa_kind_source") + .execute(tx.as_mut()) + .await?; + let mut query = sqlx::query_as::<_, (String, String, i64, i64)>(audited_sql(sql)); + if source == KindCountSource::Available { + query = query + .bind(&head_queues) + .bind(&head_priorities) + .bind(&head_shards) + .bind(&claim_seqs) + .bind(&next_seqs); + } + counts.extend(query.fetch_all(tx.as_mut()).await?); + sqlx::query("ROLLBACK TO SAVEPOINT awa_kind_source") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_kind_source") + .execute(tx.as_mut()) + .await?; + } + Ok(counts) +} + +// The result is bounded by the installed claim-ring slot count, not by +// retained job volume. Include all populated slots, including old generations. +async fn queue_storage_populated_claim_slots( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, +) -> Result, AwaError> { + let schema = store.schema(); + sqlx::query("SAVEPOINT awa_claim_slots") + .execute(tx.as_mut()) + .await?; + let slots: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT claim_slot FROM {schema}.lease_claims WHERE closed_at IS NULL \ + UNION SELECT claim_slot FROM {schema}.lease_claim_batches ORDER BY claim_slot" + ))) + .fetch_all(tx.as_mut()) + .await?; + sqlx::query("ROLLBACK TO SAVEPOINT awa_claim_slots") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_claim_slots") + .execute(tx.as_mut()) + .await?; + // A schema-only admin store carries the default slot count, so take the + // installed slot numbers as they are rather than bounding them by it. + slots + .into_iter() + .map(|slot| crate::queue_storage::ring_slot_number(slot, "claim")) + .collect() +} + +// Keep only per-kind, per-queue aggregates in Rust, not every claim identity. +async fn queue_storage_running_kind_counts( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, +) -> Result, AwaError> { + let schema = store.schema(); + let slots = queue_storage_populated_claim_slots(tx, store).await?; + let mut counts = Vec::new(); + for slot in slots { + let open_running = + crate::queue_storage::open_receipt_running_claims_sql_in_slot(schema, slot); + sqlx::query("SAVEPOINT awa_kind_counts") + .execute(tx.as_mut()) + .await?; + let rows: Vec<(String, String, i64)> = sqlx::query_as(audited_sql(format!( + r#" + -- A queue or priority move copies a ready row with its kind, so + -- any one of the job's ready rows gives it. + SELECT created.kind, open_claims.queue, count(*)::bigint + FROM ({open_running}) AS open_claims + CROSS JOIN LATERAL ( + SELECT ready.kind + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = open_claims.job_id + LIMIT 1 + ) AS created + GROUP BY created.kind, open_claims.queue + "# + ))) + .fetch_all(tx.as_mut()) + .await?; + sqlx::query("ROLLBACK TO SAVEPOINT awa_kind_counts") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_kind_counts") + .execute(tx.as_mut()) + .await?; + counts.extend(rows); + } + Ok(counts) +} + /// Get job-kind overviews for all known kinds. pub async fn job_kind_overviews(pool: &PgPool) -> Result, AwaError> { if let Some(store) = active_queue_storage(pool).await? { @@ -3843,17 +4082,43 @@ pub async fn job_kind_overviews(pool: &PgPool) -> Result, A LEFT JOIN live_kind_descriptor_variants kdv ON kdv.kind = k.kind ORDER BY k.kind "#, - kind_counts_cte = queue_storage_kind_counts_cte(store.schema()) + kind_counts_cte = queue_storage_kind_counts_cte() ); // As in `queue_overviews`: compiling costs more than running. let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; - let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + let mut tx = sqlx::Connection::begin_with( + &mut *begin_guard.connection, + "BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY", + ) + .await?; begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; + let heads = queue_storage_kind_heads(&mut tx, &store).await?; + let mut sources = queue_storage_base_kind_counts(&mut tx, &store, heads).await?; + sources.extend( + queue_storage_running_kind_counts(&mut tx, &store) + .await? + .into_iter() + .map(|(kind, queue, n)| (kind, queue, n, 0)), + ); + let mut kinds = Vec::with_capacity(sources.len()); + let mut queues = Vec::with_capacity(sources.len()); + let mut counts = Vec::with_capacity(sources.len()); + let mut recent_counts = Vec::with_capacity(sources.len()); + for (kind, queue, count, recent) in sources { + kinds.push(kind); + queues.push(queue); + counts.push(count); + recent_counts.push(recent); + } let rows = sqlx::query_as::<_, JobKindOverview>(audited_sql(sql)) + .bind(kinds) + .bind(queues) + .bind(counts) + .bind(recent_counts) .fetch_all(tx.as_mut()) .await?; tx.commit().await?; diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 8edc4f099..1d745187d 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -54,9 +54,9 @@ pub use partitioned_queue::{ partition_for_ordering_key, partition_hash64, PartitionedQueue, PartitionedQueueError, }; pub use queue_storage::{ - ClaimedEntry, ClaimedRuntimeJob, PruneDurations, PruneOutcome, QueueCounts, QueueStorage, - QueueStorageConfig, RingLedgerFoldOutcome, RingRotationCadence, RotateOutcome, SkipReason, - TerminalDeltaRollupOutcome, TerminalRollupFoldOutcome, + ClaimedEntry, ClaimedRuntimeJob, PromoteOutcome, PromotionLane, PruneDurations, PruneOutcome, + QueueCounts, QueueStorage, QueueStorageConfig, RingLedgerFoldOutcome, RingRotationCadence, + RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, TerminalRollupFoldOutcome, }; pub use storage::StorageStatus; diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 034e203de..9b0a0813d 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -641,6 +641,29 @@ pub struct RingLedgerFoldOutcome { pub skipped_mvcc_pinned: bool, } +/// A `(queue, priority)` enqueue lane as promotion addresses it. Promotion +/// picks the shard when it re-enqueues a row, so the shard is not part of it. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PromotionLane { + pub queue: String, + pub priority: i16, +} + +/// Outcome of [`QueueStorage::promote_due_lanes`] and +/// [`QueueStorage::promote_due_lane`]. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct PromoteOutcome { + /// Rows moved into the ready ring. + pub promoted: usize, + /// Due rows the batch selected, including those left deferred on a + /// contended lane. A batch that selected its full size may have left + /// further due rows behind. + pub selected: usize, + /// Lanes whose locks were held past `lock_timeout`, so their promotion + /// transaction rolled back. Their rows stay deferred. + pub contended: Vec, +} + /// Discriminator for [`PruneOutcome::SkippedActive`]. /// /// Multiple gates can fire `SkippedActive` for the same ring (e.g. queue @@ -699,7 +722,7 @@ fn is_lock_contention_error(err: &sqlx::Error) -> bool { ) } -async fn set_prune_lock_timeout_tx( +async fn set_lock_timeout_tx( tx: &mut sqlx::Transaction<'_, Postgres>, timeout: Duration, ) -> Result<(), AwaError> { @@ -941,7 +964,7 @@ fn claim_closure_batch_child_name(schema: &str, slot: usize) -> String { format!("{schema}.lease_claim_closure_batches_{slot}") } -fn ring_slot_number(slot: i32, ring: &str) -> Result { +pub(crate) fn ring_slot_number(slot: i32, ring: &str) -> Result { usize::try_from(slot).map_err(|_| { AwaError::Validation(format!( "invalid {ring} ring slot {slot}: slot must be non-negative" @@ -959,42 +982,76 @@ fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result String { - // `receipt_id` is allocated from a global sequence, so it identifies a - // claim without a `claim_slot` predicate. Compact closure batches retain - // the raw receipt array for audit/debugging, but hot membership checks use - // the per-child GiST index on the derived multirange. + receipt_evidence_sql( + schema, + closure_rel, + closure_batch_rel, + claims_alias, + ReceiptEvidence::Closed, + ) +} + +fn receipt_evidence_sql( + schema: &str, + closure_rel: &str, + closure_batch_rel: &str, + claims_alias: &str, + evidence: ReceiptEvidence, +) -> String { + let (null_test, join, exists, range_fence) = match evidence { + ReceiptEvidence::Closed => ("NOT NULL", "OR", "EXISTS", ""), + ReceiptEvidence::Open => ( + "NULL", + "AND", + "NOT EXISTS", + " + -- Keep receipt membership correlated while the other closure + -- exclusions can use full-identity anti joins. + OFFSET 0", + ), + }; + // `receipt_id` is allocated from a global sequence and identifies a claim. + // Closures retain that claim's slot; match it to prune unrelated partitions. + // Compact closure batches retain the raw receipt array for audit/debugging, + // but hot membership checks use the per-child GiST index on the multirange. format!( r#" ( - {claims_alias}.closed_at IS NOT NULL - OR EXISTS ( + {claims_alias}.closed_at IS {null_test} + {join} {exists} ( SELECT 1 FROM {closure_rel} AS closures WHERE closures.claim_slot = {claims_alias}.claim_slot AND closures.job_id = {claims_alias}.job_id AND closures.run_lease = {claims_alias}.run_lease ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {closure_batch_rel} AS closure_batches - WHERE closure_batches.receipt_ranges @> {claims_alias}.receipt_id + WHERE closure_batches.claim_slot = {claims_alias}.claim_slot + AND closure_batches.receipt_ranges @> {claims_alias}.receipt_id{range_fence} ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = {claims_alias}.job_id AND done.run_lease = {claims_alias}.run_lease ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = {claims_alias}.job_id AND deferred.run_lease = {claims_alias}.run_lease ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = {claims_alias}.job_id AND dlq.run_lease = {claims_alias}.run_lease @@ -1420,10 +1477,41 @@ pub fn open_receipt_running_claims_sql(schema: &str) -> String { /// retained batch header is read, but only batches holding one of the jobs /// are unnested and checked. pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option<&str>) -> String { - let closure_rel = format!("{schema}.lease_claim_closures"); - let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); - let row_closed = - receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); + open_receipt_running_claims_sql_from(schema, job_ids, None) +} + +/// The same open-claim projection, scoped to one physical claim slot. +/// Ready and terminal history are deliberately not restricted to that slot. +pub(crate) fn open_receipt_running_claims_sql_in_slot(schema: &str, slot: usize) -> String { + open_receipt_running_claims_sql_from(schema, None, Some(slot)) +} + +fn open_receipt_running_claims_sql_from( + schema: &str, + job_ids: Option<&str>, + slot: Option, +) -> String { + let (claims_rel, batches_rel, closure_rel, closure_batch_rel) = match slot { + Some(slot) => ( + claim_child_name(schema, slot), + claim_batch_child_name(schema, slot), + closure_child_name(schema, slot), + claim_closure_batch_child_name(schema, slot), + ), + None => ( + format!("{schema}.lease_claims"), + format!("{schema}.lease_claim_batches"), + format!("{schema}.lease_claim_closures"), + format!("{schema}.lease_claim_closure_batches"), + ), + }; + let row_open = receipt_evidence_sql( + schema, + &closure_rel, + &closure_batch_rel, + "claims", + ReceiptEvidence::Open, + ); let (row_filter, batch_filter) = match job_ids { Some(ids) => ( format!("claims.job_id = ANY({ids}) AND "), @@ -1434,8 +1522,8 @@ pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option< format!( r#" SELECT claims.job_id, claims.run_lease, claims.queue - FROM {schema}.lease_claims AS claims - WHERE {row_filter}NOT {row_closed} + FROM {claims_rel} AS claims + WHERE {row_filter}{row_open} AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id @@ -1443,7 +1531,7 @@ pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option< ) UNION ALL SELECT items.job_id, items.run_lease, batches.queue - FROM {schema}.lease_claim_batches AS batches + FROM {batches_rel} AS batches CROSS JOIN LATERAL unnest( batches.job_ids, batches.run_leases, @@ -1470,6 +1558,9 @@ pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option< SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches WHERE closure_batches.claim_slot = batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id + -- Keep the range lookup correlated: joining on claim_slot can + -- multiply receipt members by retained closure batches. + OFFSET 0 ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease @@ -2618,6 +2709,7 @@ struct ReadyRowIdentity { run_lease: i64, } +/// Attempts a rescue pass leaves for a later pass, as `(job_id, run_lease)`. struct BusyRetained(Vec<(i64, i64)>); impl BusyRetained { @@ -2633,6 +2725,30 @@ impl BusyRetained { } } +/// A failure hydrating a deleted lease. `Ambiguous` names an attempt whose +/// retained ready bodies rank equally, so none can be chosen for it. +enum HydrateError { + Ambiguous { job_id: i64, run_lease: i64 }, + Awa(AwaError), +} + +impl From for HydrateError { + fn from(err: AwaError) -> Self { + Self::Awa(err) + } +} + +impl From for AwaError { + fn from(err: HydrateError) -> Self { + match err { + HydrateError::Ambiguous { job_id, run_lease } => AwaError::Validation(format!( + "queue storage ready body ambiguous for deleted lease job {job_id} run_lease {run_lease}" + )), + HydrateError::Awa(err) => err, + } + } +} + #[derive(Debug, Clone, sqlx::FromRow)] struct ReadyTransitionRow { ready_slot: i32, @@ -3257,6 +3373,15 @@ struct ExistingReadyRow { payload: serde_json::Value, } +/// Which of a lane's due rows one promotion transaction takes. +#[derive(Debug, Clone, Copy)] +enum DueLaneRows<'a> { + /// These rows, picked by a batch across lanes, where still due. + Selected(&'a [i64]), + /// The oldest due rows, up to this many. + Oldest(i64), +} + #[derive(Debug, Clone, sqlx::FromRow)] struct DeletedLeaseRow { ready_slot: i32, @@ -3300,7 +3425,7 @@ fn select_claimed_ready_body<'a>( candidates: &'a [ReadySnapshotRow], claimed_lanes: &BTreeSet, reference: &DeletedLeaseRow, -) -> Result<(&'a ReadySnapshotRow, bool), AwaError> { +) -> Result<(&'a ReadySnapshotRow, bool), HydrateError> { // Emitted physical-lane evidence wins over a live replacement. // Without it, retain the live-body compatibility fallback; it // cannot identify a legacy claim's original colliding lane. @@ -3332,10 +3457,10 @@ fn select_claimed_ready_body<'a>( let mut best = candidates.iter().filter(|row| rank(row) == strongest); let ready = best.next().expect("ready body rank has a candidate"); if best.next().is_some() { - return Err(AwaError::Validation(format!( - "queue storage ready body ambiguous for deleted lease job {} run_lease {}", - reference.job_id, reference.run_lease - ))); + return Err(HydrateError::Ambiguous { + job_id: reference.job_id, + run_lease: reference.run_lease, + }); } Ok((ready, has_emitted)) } @@ -12730,19 +12855,36 @@ impl QueueStorage { deleted: Vec, write_receipt_closures: bool, ) -> Result, BusyRetained>, AwaError> { - let (hydrated, sources) = self - .hydrate_deleted_leases_with_sources_tx(tx, deleted, write_receipt_closures) - .await?; - if sources.is_empty() { - return Ok(Ok(hydrated)); + let hydrated = async { + let (hydrated, sources) = self + .hydrate_deleted_leases_with_sources_tx(tx, deleted, write_receipt_closures) + .await?; + if sources.is_empty() { + return Ok(Ok(hydrated)); + } + if let Some(busy) = self + .retire_retained_attempt_copies_tx(tx, &sources, true) + .await? + { + return Ok(Err(busy)); + } + Ok(Ok(hydrated)) } - if let Some(busy) = self - .retire_retained_attempt_copies_tx(tx, &sources, true) - .await? - { - return Ok(Err(busy)); + .await; + match hydrated { + Ok(hydrated) => Ok(hydrated), + // One unresolvable attempt must not stop the rescue of the others. + Err(HydrateError::Ambiguous { job_id, run_lease }) => { + tracing::error!( + schema = self.schema(), + job_id, + run_lease, + "Rescue skipped an attempt: queue storage ready body ambiguous for deleted lease; it stays running until its retained bodies are repaired" + ); + Ok(Err(BusyRetained(vec![(job_id, run_lease)]))) + } + Err(HydrateError::Awa(err)) => Err(err), } - Ok(Ok(hydrated)) } async fn hydrate_deleted_leases_with_sources_tx<'a>( @@ -12750,7 +12892,7 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, deleted: Vec, write_receipt_closures: bool, - ) -> Result<(Vec, Vec), AwaError> { + ) -> Result<(Vec, Vec), HydrateError> { if deleted.is_empty() { return Ok((Vec::new(), Vec::new())); } @@ -13008,7 +13150,7 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'_, Postgres>, sources: &[ReadyRowIdentity], skip_busy: bool, - ) -> Result, AwaError> { + ) -> Result, HydrateError> { let schema = self.schema(); let mut expected = BTreeMap::new(); for source in sources { @@ -13016,11 +13158,10 @@ impl QueueStorage { .insert((source.job_id, source.run_lease), source) .is_some() { - return Err(AwaError::Validation(format!( - "queue storage ready body ambiguous for deleted lease job {} run_lease {}", - source.job_id, - source.run_lease + 1 - ))); + return Err(HydrateError::Ambiguous { + job_id: source.job_id, + run_lease: source.run_lease + 1, + }); } } let job_ids: Vec = sources.iter().map(|row| row.job_id).collect(); @@ -13107,7 +13248,8 @@ impl QueueStorage { if copies.is_empty() { return Err(AwaError::Validation( "queue storage ready row missing for deleted lease during retirement".to_string(), - )); + ) + .into()); } let slots: Vec = copies.iter().map(|row| row.ready_slot).collect(); let generations: Vec = copies.iter().map(|row| row.ready_generation).collect(); @@ -13152,11 +13294,10 @@ impl QueueStorage { for row in &emitted { let key = (row.job_id, row.run_lease); if !seen.insert(key) || expected.get(&key).copied() != Some(row) { - return Err(AwaError::Validation(format!( - "queue storage ready body ambiguous for deleted lease job {} run_lease {}", - row.job_id, - row.run_lease + 1 - ))); + return Err(HydrateError::Ambiguous { + job_id: row.job_id, + run_lease: row.run_lease + 1, + }); } } sqlx::query(audited_sql(format!( @@ -18510,32 +18651,224 @@ impl QueueStorage { } } + /// Promote up to `batch_size` due `state` rows into the ready ring. + /// Equivalent to [`Self::promote_due_lanes`] with no lock timeout and no + /// skipped lanes, returning the number of rows promoted. pub async fn promote_due( &self, pool: &PgPool, state: JobState, batch_size: i64, ) -> Result { + Ok(self + .promote_due_lanes(pool, state, batch_size, None, &[]) + .await? + .promoted) + } + + /// Select up to `batch_size` due `state` rows, oldest `run_at` first, + /// outside the lanes in `skip`, and promote each lane's share in its own + /// transaction. + /// + /// The selection skips rows another transaction has locked, such as a + /// lane promotion still waiting for its lane, so they cannot fill every + /// batch of a tick. It runs as its own statement, so its row locks end + /// with it; `deferred_jobs` is not partitioned, so no run-time pruning + /// meets those locks. + /// + /// Re-enqueueing a row takes its lane's enqueue reservation lock, which + /// producers hold until they commit, so a lane under heavy enqueue load + /// can keep its promotion waiting. A transaction per lane confines that + /// wait to the lane's own rows. With `lane_lock_timeout`, each lane's + /// transaction runs under that `lock_timeout`; a lane that hits it rolls + /// back, keeps its rows deferred and is reported in + /// [`PromoteOutcome::contended`], and the remaining lanes go ahead. + pub async fn promote_due_lanes( + &self, + pool: &PgPool, + state: JobState, + batch_size: i64, + lane_lock_timeout: Option, + skip: &[PromotionLane], + ) -> Result { if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { - return Ok(0); + return Ok(PromoteOutcome::default()); } let schema = self.schema(); + let due: Vec<(i64, String, i16)> = sqlx::query_as(audited_sql(format!( + r#" + SELECT job_id, queue, priority + FROM {schema}.deferred_jobs AS deferred + WHERE state = $1 + AND run_at <= clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM awa.queue_meta + WHERE queue = deferred.queue AND paused = TRUE + ) + AND NOT EXISTS ( + SELECT 1 + FROM unnest($3::text[], $4::smallint[]) AS skip (queue, priority) + WHERE skip.queue = deferred.queue AND skip.priority = deferred.priority + ) + ORDER BY run_at ASC, priority ASC, job_id ASC + LIMIT $2 + FOR UPDATE OF deferred SKIP LOCKED + "# + ))) + .bind(state) + .bind(batch_size) + .bind( + skip.iter() + .map(|lane| lane.queue.as_str()) + .collect::>(), + ) + .bind(skip.iter().map(|lane| lane.priority).collect::>()) + .fetch_all(pool) + .await + .map_err(map_sqlx_error)?; + + let mut outcome = PromoteOutcome { + selected: due.len(), + ..PromoteOutcome::default() + }; + let mut lanes: BTreeMap> = BTreeMap::new(); + for (job_id, queue, priority) in due { + lanes + .entry(PromotionLane { queue, priority }) + .or_default() + .push(job_id); + } + let mut promoted_queues = BTreeSet::new(); + for (lane, job_ids) in lanes { + let promoted = self + .promote_lane( + pool, + state, + &lane, + DueLaneRows::Selected(&job_ids), + lane_lock_timeout, + ) + .await; + match promoted { + Ok(Some(promoted)) => { + outcome.promoted += promoted; + if promoted > 0 { + promoted_queues.insert(lane.queue); + } + } + Ok(None) => outcome.contended.push(lane), + Err(err) => { + self.notify_promoted_queues(pool, promoted_queues).await; + return Err(err); + } + } + } + self.notify_promoted_queues(pool, promoted_queues).await; + Ok(outcome) + } + + /// Promote up to `batch_size` due `state` rows of one lane, oldest + /// `run_at` first, in one transaction. With `lane_lock_timeout`, a + /// transaction that hits it rolls back and the lane is returned in + /// [`PromoteOutcome::contended`] with nothing selected. + pub async fn promote_due_lane( + &self, + pool: &PgPool, + state: JobState, + lane: &PromotionLane, + batch_size: i64, + lane_lock_timeout: Option, + ) -> Result { + if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { + return Ok(PromoteOutcome::default()); + } + let promoted = self + .promote_lane( + pool, + state, + lane, + DueLaneRows::Oldest(batch_size), + lane_lock_timeout, + ) + .await?; + if promoted.is_some_and(|promoted| promoted > 0) { + self.notify_promoted_queues(pool, [lane.queue.clone()]) + .await; + } + Ok(match promoted { + Some(promoted) => PromoteOutcome { + promoted, + selected: promoted, + contended: Vec::new(), + }, + None => PromoteOutcome { + contended: vec![lane.clone()], + ..PromoteOutcome::default() + }, + }) + } + + /// One lane's promotion transaction. `None` when its locks hit + /// `lock_timeout` and it rolled back. + async fn promote_lane( + &self, + pool: &PgPool, + state: JobState, + lane: &PromotionLane, + rows: DueLaneRows<'_>, + lock_timeout: Option, + ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let moved: Vec = sqlx::query_as(audited_sql(format!( + match self + .promote_lane_tx(&mut tx, state, lane, rows, lock_timeout) + .await + { + Ok(promoted) => { + tx.commit().await.map_err(map_sqlx_error)?; + Ok(Some(promoted)) + } + Err(AwaError::Database(err)) if is_lock_contention_error(&err) => { + tx.rollback().await.map_err(map_sqlx_error)?; + Ok(None) + } + Err(err) => Err(err), + } + } + + async fn promote_lane_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + state: JobState, + lane: &PromotionLane, + rows: DueLaneRows<'_>, + lock_timeout: Option, + ) -> Result { + if let Some(timeout) = lock_timeout { + set_lock_timeout_tx(tx, timeout).await?; + } + + let schema = self.schema(); + let (row_filter, limit) = match rows { + DueLaneRows::Selected(_) => ("AND job_id = ANY($4)", ""), + DueLaneRows::Oldest(_) => ("", "ORDER BY run_at ASC, job_id ASC LIMIT $4"), + }; + let query = sqlx::query_as::<_, DeferredJobRow>(audited_sql(format!( r#" DELETE FROM {schema}.deferred_jobs WHERE job_id IN ( SELECT job_id - FROM {schema}.deferred_jobs + FROM {schema}.deferred_jobs AS deferred WHERE state = $1 + AND queue = $2 + AND priority = $3 + {row_filter} AND run_at <= clock_timestamp() AND NOT EXISTS ( SELECT 1 FROM awa.queue_meta - WHERE queue = {schema}.deferred_jobs.queue AND paused = TRUE + WHERE queue = deferred.queue AND paused = TRUE ) - ORDER BY run_at ASC, priority ASC, job_id ASC - LIMIT $2 + {limit} FOR UPDATE SKIP LOCKED ) RETURNING @@ -18558,19 +18891,23 @@ impl QueueStorage { "# ))) .bind(state) - .bind(batch_size) - .fetch_all(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - + .bind(&lane.queue) + .bind(lane.priority); + let query = match rows { + DueLaneRows::Selected(job_ids) => query.bind(job_ids), + DueLaneRows::Oldest(batch_size) => query.bind(batch_size), + }; + let mut moved: Vec = + query.fetch_all(tx.as_mut()).await.map_err(map_sqlx_error)?; if moved.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; return Ok(0); } + // DELETE ... RETURNING has no order; lane sequence follows due order. + moved.sort_unstable_by_key(|row| (row.run_at, row.job_id)); + let promoted = moved.len(); let ready_rows: Vec = moved - .iter() - .cloned() + .into_iter() .map(|row| ExistingReadyRow { job_id: row.job_id, kind: row.kind, @@ -18588,15 +18925,45 @@ impl QueueStorage { payload: row.payload, }) .collect(); - let queues = ready_rows - .iter() - .map(|row| row.queue.clone()) - .collect::>(); - self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(state)) + self.insert_existing_ready_rows_tx(tx, ready_rows, Some(state)) .await?; - self.notify_queues_tx(&mut tx, queues).await?; - tx.commit().await.map_err(map_sqlx_error)?; - Ok(moved.len()) + Ok(promoted) + } + + /// Wake claimers of queues whose promotion has committed. + /// + /// Promotion notifies after its lane transactions instead of inside them. + /// A notifying transaction takes PostgreSQL's notify-queue lock at + /// `COMMIT`, which every notifying transaction in the cluster waits its + /// turn for, and a lane transaction would hold its lane's locks, and run + /// under its lock timeout, while it waited. The wakeup is a hint that + /// claimers also poll for, so a failure is logged and not returned. + async fn notify_promoted_queues( + &self, + pool: &PgPool, + queues: impl IntoIterator, + ) { + let channels: Vec = queues + .into_iter() + .map(|queue| format!("awa:{}", self.logical_queue_name(&queue))) + .collect::>() + .into_iter() + .collect(); + if channels.is_empty() { + return; + } + if let Err(err) = + sqlx::query("SELECT pg_notify(channel, '') FROM unnest($1::text[]) AS channel") + .bind(&channels) + .execute(pool) + .await + { + tracing::warn!( + error = %err, + channels = channels.len(), + "Failed to wake claimers after promotion; they pick the rows up on their next poll" + ); + } } async fn relation_has_rows_tx( @@ -19458,7 +19825,7 @@ impl QueueStorage { } let locked: Result = async { - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; sqlx::query(audited_sql(format!( "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" @@ -19916,7 +20283,7 @@ impl QueueStorage { } let locked: Result = async { - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_started = Instant::now(); // Claims read segments before ready rows, while producers write in @@ -20316,7 +20683,7 @@ impl QueueStorage { let lease_child = lease_child_name(schema, slot as usize); let locked: Result = async { - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_started = Instant::now(); sqlx::query(audited_sql(format!( @@ -20681,7 +21048,7 @@ impl QueueStorage { } let locked: Result = async { - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_started = Instant::now(); sqlx::query(audited_sql(format!( diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 37719ff1e..98dcc8bff 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -910,6 +910,27 @@ where } } +/// Re-read a job whose transition has committed, for its lifecycle event. +/// The transition stands whatever the read returns, so a failed read is logged +/// and the caller builds the event from the worker's copy instead. +async fn reload_committed_job( + runtime: &QueueStorageRuntime, + pool: &PgPool, + job_id: i64, +) -> Option { + match runtime.store.load_job(pool, job_id).await { + Ok(job) => job, + Err(err) => { + warn!( + job_id, + error = %err, + "Could not re-read a committed job for its lifecycle event; using the worker's copy" + ); + None + } + } +} + fn is_completion_lock_contention(error: &AwaError) -> bool { matches!(error, AwaError::Database(sqlx::Error::Database(database)) if database.code().as_deref() == Some("55P03")) @@ -2137,20 +2158,13 @@ async fn complete_job_queue_storage( return Ok(CompletionOutcome::IgnoredStale); } if needs_event { - // The output is set from the sealed value rather than re-read, - // which would detoast its compact batch's whole output array. - let mut updated_job = runtime - .store - .load_job_without_output(pool, job.id) - .await? - .unwrap_or_else(|| { - let mut completed_job = job.clone(); - completed_job.state = JobState::Completed; - completed_job.finalized_at = Some(chrono::Utc::now()); - completed_job.progress = None; - completed_job - }); - updated_job.output = job.output.clone(); + // A completed attempt changes only these fields, so the event is + // built from the worker's copy, which carries the sealed output, + // rather than re-reading the job. + let mut updated_job = job.clone(); + updated_job.state = JobState::Completed; + updated_job.finalized_at = Some(chrono::Utc::now()); + updated_job.progress = None; Ok(CompletionOutcome::Applied { event: Some(UntypedJobEvent::Completed { job: updated_job, @@ -2430,10 +2444,8 @@ async fn complete_job_queue_storage( } } let event = if needs_event { - let parked_job = runtime - .store - .load_job(pool, job.id) - .await? + let parked_job = reload_committed_job(runtime, pool, job.id) + .await .unwrap_or_else(|| { let mut parked = job.clone(); parked.state = JobState::WaitingExternal; diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index f1dfe9f79..2c43456a8 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -1,4 +1,5 @@ use crate::executor::DlqPolicy; +use crate::metrics::AwaMetrics; use crate::runtime::InFlightMap; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; use awa_model::audited_sql; @@ -8,7 +9,8 @@ use awa_model::cron::{ #[cfg(test)] use awa_model::SkipReason; use awa_model::{ - JobRow, JobState, PeriodicJob, PruneOutcome, RotateOutcome, TerminalDeltaRollupOutcome, + AwaError, JobRow, JobState, PeriodicJob, PromotionLane, PruneOutcome, QueueStorage, + RotateOutcome, TerminalDeltaRollupOutcome, }; use chrono::Utc; use croner::Cron; @@ -18,7 +20,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use tokio::task::JoinHandle; +use tokio::task::{JoinHandle, JoinSet}; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; use uuid::Uuid; @@ -763,6 +765,232 @@ const CRON_CATCH_UP_LIMIT: usize = 1_000; const TERMINAL_COUNT_ROLLUP_MAX_SLOTS_PER_TICK: usize = 4; type QueueStorageMetricRow = (String, i64, i64, i64, i64, i64, i64, i64, Option); +/// How long a promotion transaction waits on one lane's locks before the +/// leader hands the lane to a dedicated waiter and goes on without it. +const PROMOTE_LANE_LOCK_TIMEOUT: Duration = Duration::from_millis(50); +/// Lock wait of a dedicated lane waiter. PostgreSQL grants a lock to its +/// waiters in arrival order, so a waiter is served after the producers queued +/// ahead of it, however busy the lane stays. +const PROMOTE_LANE_WAITER_LOCK_TIMEOUT: Duration = Duration::from_secs(30); +/// Most lanes with a dedicated waiter at once. A waiter holds a pool +/// connection while it waits, so smaller pools get fewer: one per eight +/// connections, and always one so a busy lane is still promoted. +const PROMOTE_MAX_LANE_WAITERS: usize = 4; + +fn lane_waiter_limit(pool_max_connections: u32) -> usize { + (pool_max_connections as usize / 8).clamp(1, PROMOTE_MAX_LANE_WAITERS) +} + +/// Queue-storage promotion of due scheduled and retryable rows, as the +/// maintenance leader runs it for one store. +/// +/// Each lane's share of a batch is promoted in its own transaction under a +/// short `lock_timeout`, so a lane whose enqueue lock is busy delays only its +/// own rows. A lane that times out keeps its rows deferred and is handed to a +/// dedicated waiter task, which promotes it under a long lock timeout while +/// the leader carries on with other lanes and other maintenance branches. +/// Dropping the promoter aborts its waiters. +#[doc(hidden)] +pub struct QueueStoragePromoter { + pool: PgPool, + store: Arc, + metrics: AwaMetrics, + waiting: Arc>>, + waiters: std::sync::Mutex>, + max_waiters: usize, +} + +impl QueueStoragePromoter { + pub fn new(pool: PgPool, store: Arc, metrics: AwaMetrics) -> Self { + let max_waiters = lane_waiter_limit(pool.options().get_max_connections()); + Self { + max_waiters, + pool, + store, + metrics, + waiting: Arc::default(), + waiters: std::sync::Mutex::new(JoinSet::new()), + } + } + + /// One promotion tick for `state`: batches until one selects fewer rows + /// than the batch size, each skipping lanes that have a waiter or were + /// contended earlier in the tick. Returns the rows promoted by the tick + /// itself, not by waiters. + pub async fn promote_due( + &self, + state: JobState, + cancel: &CancellationToken, + ) -> Result { + { + let mut waiters = self.waiters.lock().expect("lane waiters poisoned"); + while waiters.try_join_next().is_some() {} + } + + let mut contended: Vec = Vec::new(); + let mut promoted = 0usize; + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + if cancel.is_cancelled() { + break; + } + let mut skip = self.waiting_lanes(state); + skip.extend(contended.iter().cloned()); + let batch_started = Instant::now(); + let outcome = self + .store + .promote_due_lanes( + &self.pool, + state, + PROMOTE_BATCH_SIZE, + Some(PROMOTE_LANE_LOCK_TIMEOUT), + &skip, + ) + .await?; + self.metrics.record_promotion_batch( + state.as_str(), + outcome.promoted as u64, + batch_started.elapsed(), + ); + promoted += outcome.promoted; + for lane in outcome.contended { + self.hand_off(state, lane.clone()); + contended.push(lane); + } + if (outcome.selected as i64) < PROMOTE_BATCH_SIZE { + break; + } + } + Ok(promoted) + } + + fn waiting_lanes(&self, state: JobState) -> Vec { + self.waiting + .lock() + .expect("waiting lanes poisoned") + .iter() + .filter(|(waiting_state, _)| *waiting_state == state) + .map(|(_, lane)| lane.clone()) + .collect() + } + + fn hand_off(&self, state: JobState, lane: PromotionLane) { + let key = (state, lane); + { + let mut waiting = self.waiting.lock().expect("waiting lanes poisoned"); + if waiting.len() >= self.max_waiters { + self.metrics.record_promotion_lane_contention( + state.as_str(), + &key.1.queue, + "waiter_limit", + ); + debug!( + queue = %key.1.queue, + priority = key.1.priority, + "Promotion lane contended with every lane waiter busy; retrying next tick" + ); + return; + } + if !waiting.insert(key.clone()) { + return; + } + } + self.metrics + .record_promotion_lane_contention(state.as_str(), &key.1.queue, "handed_off"); + let guard = WaitingLane { + waiting: Arc::clone(&self.waiting), + key, + }; + let pool = self.pool.clone(); + let store = Arc::clone(&self.store); + let metrics = self.metrics.clone(); + self.waiters + .lock() + .expect("lane waiters poisoned") + .spawn(async move { + promote_waiting_lane(&pool, &store, &metrics, &guard.key).await; + drop(guard); + }); + } +} + +/// Holds a lane in [`QueueStoragePromoter`]'s waiting set until its waiter +/// finishes, panics or is aborted. +struct WaitingLane { + waiting: Arc>>, + key: (JobState, PromotionLane), +} + +impl Drop for WaitingLane { + fn drop(&mut self) { + if let Ok(mut waiting) = self.waiting.lock() { + waiting.remove(&self.key); + } + } +} + +async fn promote_waiting_lane( + pool: &PgPool, + store: &QueueStorage, + metrics: &AwaMetrics, + (state, lane): &(JobState, PromotionLane), +) { + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + let batch_started = Instant::now(); + match store + .promote_due_lane( + pool, + *state, + lane, + PROMOTE_BATCH_SIZE, + Some(PROMOTE_LANE_WAITER_LOCK_TIMEOUT), + ) + .await + { + Ok(outcome) if !outcome.contended.is_empty() => { + metrics.record_promotion_lane_contention( + state.as_str(), + &lane.queue, + "waiter_timed_out", + ); + warn!( + queue = %lane.queue, + priority = lane.priority, + lock_timeout_ms = PROMOTE_LANE_WAITER_LOCK_TIMEOUT.as_millis() as u64, + "Promotion of a contended lane timed out waiting for its locks; retrying next tick" + ); + return; + } + Ok(outcome) => { + metrics.record_promotion_batch( + state.as_str(), + outcome.promoted as u64, + batch_started.elapsed(), + ); + if (outcome.selected as i64) < PROMOTE_BATCH_SIZE { + return; + } + } + Err(err) => { + error!( + queue = %lane.queue, + priority = lane.priority, + error = %err, + "Failed to promote a contended lane" + ); + return; + } + } + } +} + +/// The leader's queue-storage promoters: its own store's, and the configured +/// store a canonical-resolved leader promotes once routing has flipped +/// (#456). Scoped to one leadership term. +struct LeaderPromoters { + primary: Option, + standby: Option, +} + impl MaintenanceService { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -1088,6 +1316,22 @@ impl MaintenanceService { // tick. Local to the leader loop so a re-election resets // backoff to zero. let prune_tracker = PruneBackoffTracker::new(); + let promoters = LeaderPromoters { + primary: self.storage.queue_storage().map(|runtime| { + QueueStoragePromoter::new( + self.pool.clone(), + Arc::clone(&runtime.store), + self.metrics.clone(), + ) + }), + standby: self.standby_queue_storage.as_ref().map(|runtime| { + QueueStoragePromoter::new( + self.pool.clone(), + Arc::clone(&runtime.store), + self.metrics.clone(), + ) + }), + }; // Skip the first immediate tick heartbeat_rescue_timer.tick().await; @@ -1153,7 +1397,7 @@ impl MaintenanceService { } _ = promote_timer.tick() => { if let Some(timer) = branch_tracker.try_begin_without_cooldown("promote_scheduled", self.promote_interval, &self.metrics) { - self.promote_scheduled().await; + self.promote_scheduled(&promoters).await; timer.finish(); } } @@ -2541,13 +2785,20 @@ impl MaintenanceService { } /// Promote scheduled jobs that are now due. - #[tracing::instrument(skip(self), name = "maintenance.promote")] - async fn promote_scheduled(&self) { - if let Err(err) = self.promote_due_state("scheduled", "scheduled jobs").await { + #[tracing::instrument(skip(self, promoters), name = "maintenance.promote")] + async fn promote_scheduled(&self, promoters: &LeaderPromoters) { + if let Err(err) = self + .promote_due_state(JobState::Scheduled, "scheduled jobs", promoters) + .await + { error!(error = %err, "Failed to promote scheduled jobs"); } if let Err(err) = self - .promote_due_state("retryable", "retryable jobs (backoff elapsed)") + .promote_due_state( + JobState::Retryable, + "retryable jobs (backoff elapsed)", + promoters, + ) .await { error!(error = %err, "Failed to promote retryable jobs"); @@ -2556,138 +2807,62 @@ impl MaintenanceService { async fn promote_due_state( &self, - state: &'static str, + state: JobState, label: &'static str, - ) -> Result<(), awa_model::AwaError> { + promoters: &LeaderPromoters, + ) -> Result<(), AwaError> { + let canonical_state = match state { + JobState::Scheduled => "scheduled", + JobState::Retryable => "retryable", + other => { + return Err(AwaError::Validation(format!( + "unsupported promote state: {other:?}" + ))); + } + }; let mut promoted_total = 0usize; let mut notified_queues = HashSet::new(); - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - - match &self.storage { - RuntimeStorage::Canonical => { - let (promoted, queues) = self - .promote_due_batch(state) - .await - .map_err(awa_model::AwaError::Database)?; - if promoted == 0 { - break; - } - - promoted_total += promoted; - notified_queues.extend(queues); - - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - RuntimeStorage::QueueStorage(runtime) => { - let job_state = match state { - "scheduled" => awa_model::JobState::Scheduled, - "retryable" => awa_model::JobState::Retryable, - other => { - return Err(awa_model::AwaError::Validation(format!( - "unsupported queue storage promote state: {other}" - ))); - } - }; - let promote_start = std::time::Instant::now(); - let promoted = runtime - .store - .promote_due(&self.pool, job_state, PROMOTE_BATCH_SIZE) - .await?; - self.metrics.record_promotion_batch( - state, - promoted as u64, - promote_start.elapsed(), - ); - if promoted == 0 { - break; - } - - promoted_total += promoted; - - if promoted < PROMOTE_BATCH_SIZE as usize { - break; + match &self.storage { + RuntimeStorage::Canonical => { + promoted_total += self + .promote_canonical_due(canonical_state, &mut notified_queues) + .await?; + + // #456 mirror: a canonical-resolved (pre-flip) leader must + // promote the queue-storage deferred backlog once routing has + // flipped — re-scheduled canonical jobs migrate there, and + // post-flip runtimes that would promote them may not hold + // leadership. The active-schema guard keeps this a no-op until + // the flip. + if let (Some(runtime), Some(promoter)) = + (&self.standby_queue_storage, &promoters.standby) + { + let active_schema: Option = + sqlx::query_scalar("SELECT awa.active_queue_storage_schema()") + .fetch_one(&self.pool) + .await + .map_err(AwaError::Database)?; + if active_schema.as_deref() == Some(runtime.store.schema()) { + promoted_total += promoter.promote_due(state, &self.cancel).await?; } } } - } - - // #456: while a storage transition drains, due jobs can still sit in - // the canonical deferred backlog. Drain-only runtimes execute - // canonical work but stop leading maintenance once a post-flip - // runtime wins the election, so a queue-storage leader must promote - // the canonical plane too. Both queries are cheap no-ops once the - // canonical tables are empty (every finalized or fresh install). - if matches!(&self.storage, RuntimeStorage::QueueStorage(_)) { - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - let (promoted, queues) = self - .promote_due_batch(state) - .await - .map_err(awa_model::AwaError::Database)?; - if promoted == 0 { - break; + RuntimeStorage::QueueStorage(_) => { + if let Some(promoter) = &promoters.primary { + promoted_total += promoter.promote_due(state, &self.cancel).await?; } - promoted_total += promoted; - notified_queues.extend(queues); - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - } - // #456 mirror: a canonical-resolved (pre-flip) leader must promote - // the queue-storage deferred backlog once routing has flipped — - // re-scheduled canonical jobs migrate there, and post-flip runtimes - // that would promote them may not hold leadership. The active-schema - // guard keeps this a no-op until the flip. - if matches!(&self.storage, RuntimeStorage::Canonical) { - if let Some(runtime) = &self.standby_queue_storage { - let active_schema: Option = - sqlx::query_scalar("SELECT awa.active_queue_storage_schema()") - .fetch_one(&self.pool) - .await - .map_err(awa_model::AwaError::Database)?; - if active_schema.as_deref() == Some(runtime.store.schema()) { - let job_state = match state { - "scheduled" => awa_model::JobState::Scheduled, - "retryable" => awa_model::JobState::Retryable, - other => { - return Err(awa_model::AwaError::Validation(format!( - "unsupported queue storage promote state: {other}" - ))); - } - }; - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - let promote_start = std::time::Instant::now(); - let promoted = runtime - .store - .promote_due(&self.pool, job_state, PROMOTE_BATCH_SIZE) - .await?; - self.metrics.record_promotion_batch( - state, - promoted as u64, - promote_start.elapsed(), - ); - if promoted == 0 { - break; - } - promoted_total += promoted; - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - } + // #456: while a storage transition drains, due jobs can still + // sit in the canonical deferred backlog. Drain-only runtimes + // execute canonical work but stop leading maintenance once a + // post-flip runtime wins the election, so a queue-storage + // leader must promote the canonical plane too. This is a + // cheap no-op once the canonical tables are empty (every + // finalized or fresh install). + promoted_total += self + .promote_canonical_due(canonical_state, &mut notified_queues) + .await?; } } @@ -2695,7 +2870,7 @@ impl MaintenanceService { debug!( count = promoted_total, queues = notified_queues.len(), - state, + state = canonical_state, "Promoted {label}" ); } @@ -2703,6 +2878,34 @@ impl MaintenanceService { Ok(()) } + /// Promote due canonical `scheduled_jobs` rows in batches until one comes + /// back short. + async fn promote_canonical_due( + &self, + state: &'static str, + notified_queues: &mut HashSet, + ) -> Result { + let mut promoted_total = 0usize; + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + if self.cancel.is_cancelled() { + break; + } + let (promoted, queues) = self + .promote_due_batch(state) + .await + .map_err(AwaError::Database)?; + if promoted == 0 { + break; + } + promoted_total += promoted; + notified_queues.extend(queues); + if promoted < PROMOTE_BATCH_SIZE as usize { + break; + } + } + Ok(promoted_total) + } + /// SQL template for promotion. The state literal is injected directly /// (not as a parameter) so the planner can match the partial index on /// `(run_at, id) WHERE state = ''`. With a parameter, the planner @@ -3715,6 +3918,15 @@ mod tests { } } + #[test] + fn lane_waiters_scale_with_the_pool_and_never_drop_to_zero() { + assert_eq!(lane_waiter_limit(1), 1); + assert_eq!(lane_waiter_limit(10), 1); + assert_eq!(lane_waiter_limit(16), 2); + assert_eq!(lane_waiter_limit(32), 4); + assert_eq!(lane_waiter_limit(200), 4); + } + #[test] fn compute_fire_times_coalesces_missed_existing_fires_by_default() { let last = Utc.with_ymd_and_hms(2026, 5, 7, 12, 0, 0).unwrap(); diff --git a/awa/Cargo.toml b/awa/Cargo.toml index b928a5a4e..04bb5dee0 100644 --- a/awa/Cargo.toml +++ b/awa/Cargo.toml @@ -25,6 +25,7 @@ sqlx.workspace = true serde.workspace = true serde_json.workspace = true tokio = { workspace = true, features = ["full", "test-util"] } +tokio-util.workspace = true async-trait.workspace = true chrono.workspace = true proptest = "1" diff --git a/awa/tests/enqueue_lane_maintenance_coupling_test.rs b/awa/tests/enqueue_lane_maintenance_coupling_test.rs new file mode 100644 index 000000000..ce7d53040 --- /dev/null +++ b/awa/tests/enqueue_lane_maintenance_coupling_test.rs @@ -0,0 +1,427 @@ +//! How a saturated enqueue lane affects the maintenance leader. +//! +//! Promotion of scheduled and retryable jobs re-enqueues them into their lane, +//! so it queues on the same per-lane advisory lock as producers. The leader +//! runs its branches (promotion, rescue, rotation, ...) one after another on a +//! single task, so a promotion parked behind a hot lane would hold up the due +//! jobs of other queues and every other maintenance branch. The leader's +//! promoter gives each lane its own short transaction and hands a lane whose +//! lock stays busy to a dedicated waiter, so neither should wait on it. +//! +//! `test_maintenance_under_hot_lane` measures, with and without concurrent +//! producers on one hot lane: +//! * how long each promotion tick holds the leader task, and how long a +//! backlog on the hot queue takes to drain, +//! * promotion lag on an unrelated queue (time from a job becoming due to +//! it being visible as ready), with the promotion loop shaped like the +//! leader's (a tick, then batches until one comes back short), +//! * heartbeat rescue of abandoned claims on the hot queue: how long each run +//! takes and how late it starts once its tick is due (1s here, so a +//! blocked promotion shows up as a late start). +//! +//! Tunables (optional): +//! AWA_LANE_MAINT_PRODUCERS producers on the hot lane, default 16 +//! AWA_LANE_MAINT_BACKLOG due scheduled jobs on the hot queue, default 24576 +//! AWA_LANE_MAINT_RESCUE abandoned claims on the hot queue, default 6000 +//! +//! Requires a running Postgres instance. +//! Run with: +//! cargo test --release --package awa --test enqueue_lane_maintenance_coupling_test \ +//! -- --ignored --nocapture --test-threads=1 + +use awa::model::{insert_many, QueueStorage}; +use awa::worker::maintenance::QueueStoragePromoter; +use awa::worker::AwaMetrics; +use awa::{InsertOpts, InsertParams, JobState, QueueStorageConfig}; +use awa_testing::setup; +use sqlx::PgPool; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use uuid::Uuid; + +const PROMOTE_TICK: Duration = Duration::from_millis(250); +const PROBE_PERIOD: Duration = Duration::from_millis(100); +const RESCUE_TICK: Duration = Duration::from_secs(1); + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|raw| raw.parse().ok()) + .unwrap_or(default) +} + +fn percentile(values: &[f64], quantile: f64) -> f64 { + if values.is_empty() { + return 0.0; + } + let mut sorted = values.to_vec(); + sorted.sort_by(|left, right| left.partial_cmp(right).expect("finite")); + sorted[((sorted.len() as f64 - 1.0) * quantile).round() as usize] +} + +fn params( + queue: &str, + tag: &str, + n: u64, + run_at: Option>, +) -> InsertParams { + InsertParams { + kind: "lane_maint_job".to_string(), + args: serde_json::json!({ "tag": tag, "n": n }), + opts: InsertOpts { + queue: queue.to_string(), + run_at, + ..Default::default() + }, + } +} + +async fn seed_due_backlog(pool: &PgPool, queue: &str, count: usize) { + let due = chrono::Utc::now() - chrono::Duration::seconds(60); + let mut inserted = 0usize; + while inserted < count { + let chunk = (count - inserted).min(2_000); + let jobs: Vec = (0..chunk) + .map(|offset| params(queue, "backlog", (inserted + offset) as u64, Some(due))) + .collect(); + insert_many(pool, &jobs).await.expect("seed backlog"); + inserted += chunk; + } +} + +async fn deferred_due(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.deferred_jobs WHERE queue = $1 AND state = 'scheduled'", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("count deferred") +} + +struct Producers { + stop: Arc, + tasks: Vec>, +} + +impl Producers { + fn start(pool: &PgPool, queue: &str, count: usize) -> Self { + let stop = Arc::new(AtomicBool::new(false)); + let tag = Uuid::new_v4().simple().to_string(); + let tasks = (0..count) + .map(|producer| { + let pool = pool.clone(); + let stop = Arc::clone(&stop); + let queue = queue.to_string(); + let tag = tag.clone(); + tokio::spawn(async move { + let mut jobs_inserted = 0u64; + let mut counter = (producer as u64) << 40; + while !stop.load(Ordering::Relaxed) { + let jobs: Vec = (0..185) + .map(|offset| params(&queue, &tag, counter + offset, None)) + .collect(); + counter += 185; + insert_many(&pool, &jobs).await.expect("producer insert"); + jobs_inserted += 185; + } + jobs_inserted + }) + }) + .collect(); + Self { stop, tasks } + } + + async fn stop(self) -> u64 { + self.stop.store(true, Ordering::Relaxed); + let mut total = 0; + for task in self.tasks { + total += task.await.expect("producer task"); + } + total + } +} + +#[derive(Default)] +struct LeaderReport { + /// How long each promotion tick held the leader task. + tick_ms: Vec, + drain_s: f64, + probe_lag_ms: Vec, + /// How long after its tick each heartbeat-rescue run started, and how long + /// the run itself took. + rescue_start_delay_ms: Vec, + rescue_run_ms: Vec, +} + +/// A leader loop shaped like the maintenance leader's: one task selecting over +/// a promotion timer (250ms; each tick runs the leader's promoter) and a +/// heartbeat-rescue timer, running each branch to completion before taking +/// the next tick. Runs until the hot queue's backlog is gone and +/// `min_window` has passed, while a probe feeds one due job per 100ms into an +/// unrelated queue and times how long each takes to become ready. +async fn run_leader( + pool: &PgPool, + store: &Arc, + hot: &str, + other: &str, + min_window: Duration, +) -> LeaderReport { + let mut report = LeaderReport::default(); + let started = Instant::now(); + let stop = Arc::new(AtomicBool::new(false)); + let promoter = + QueueStoragePromoter::new(pool.clone(), Arc::clone(store), AwaMetrics::from_global()); + let cancel = tokio_util::sync::CancellationToken::new(); + + let probe = { + let pool = pool.clone(); + let other = other.to_string(); + let stop = Arc::clone(&stop); + tokio::spawn(async move { + let mut lags = Vec::new(); + let mut n = 0u64; + while !stop.load(Ordering::Relaxed) { + n += 1; + // Older than the hot backlog, so the probe is always in the + // first promotion batch rather than behind the backlog's + // run_at order. + let due_at = chrono::Utc::now() - chrono::Duration::seconds(120); + insert_many(&pool, &[params(&other, "probe", n, Some(due_at))]) + .await + .expect("insert probe"); + let inserted = Instant::now(); + loop { + let visible: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM awa.ready_entries + WHERE queue = $1 AND args->>'tag' = 'probe' AND (args->>'n')::bigint = $2)", + ) + .bind(&other) + .bind(n as i64) + .fetch_one(&pool) + .await + .expect("probe visibility"); + if visible { + lags.push(inserted.elapsed().as_secs_f64() * 1e3); + break; + } + if stop.load(Ordering::Relaxed) { + return lags; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + tokio::time::sleep(PROBE_PERIOD).await; + } + lags + }) + }; + + let mut promote_timer = tokio::time::interval(PROMOTE_TICK); + promote_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut rescue_timer = tokio::time::interval(RESCUE_TICK); + rescue_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + promote_timer.tick().await; + rescue_timer.tick().await; + + let mut backlog_gone_at: Option = None; + loop { + tokio::select! { + tick = promote_timer.tick() => { + let _ = tick; + let tick_started = Instant::now(); + promoter + .promote_due(JobState::Scheduled, &cancel) + .await + .expect("promote_due"); + report.tick_ms.push(tick_started.elapsed().as_secs_f64() * 1e3); + if backlog_gone_at.is_none() && deferred_due(pool, hot).await == 0 { + backlog_gone_at = Some(Instant::now()); + } + if backlog_gone_at.is_some() && started.elapsed() >= min_window { + break; + } + } + tick = rescue_timer.tick() => { + let delay = tokio::time::Instant::now().saturating_duration_since(tick); + let run_started = Instant::now(); + let rescued = store + .rescue_stale_heartbeats(pool, Duration::ZERO) + .await + .expect("rescue_stale_heartbeats"); + if !rescued.is_empty() { + report.rescue_start_delay_ms.push(delay.as_secs_f64() * 1e3); + report.rescue_run_ms.push(run_started.elapsed().as_secs_f64() * 1e3); + } + } + } + } + report.drain_s = backlog_gone_at + .expect("backlog drained") + .duration_since(started) + .as_secs_f64(); + stop.store(true, Ordering::Relaxed); + report.probe_lag_ms = probe.await.expect("probe task"); + report +} + +fn max_of(values: &[f64]) -> f64 { + values.iter().copied().fold(0.0, f64::max) +} + +fn print_leader(label: &str, report: &LeaderReport) { + println!( + "[lane-maint] {label} promote tick on leader ms p50={:.0} p99={:.0} max={:.0} | hot backlog \ + drained in {:.2}s | unrelated-queue promotion lag ms p50={:.0} p99={:.0} max={:.0} ({} probes) \ + | rescue ran {} times, run ms p50={:.0} max={:.0}, started late by ms p50={:.0} max={:.0}", + percentile(&report.tick_ms, 0.5), + percentile(&report.tick_ms, 0.99), + max_of(&report.tick_ms), + report.drain_s, + percentile(&report.probe_lag_ms, 0.5), + percentile(&report.probe_lag_ms, 0.99), + max_of(&report.probe_lag_ms), + report.probe_lag_ms.len(), + report.rescue_run_ms.len(), + percentile(&report.rescue_run_ms, 0.5), + max_of(&report.rescue_run_ms), + percentile(&report.rescue_start_delay_ms, 0.5), + max_of(&report.rescue_start_delay_ms), + ); +} + +/// Enqueue `count` jobs on `queue` and claim every one of them without ever +/// completing or heartbeating, so each becomes an abandoned attempt. +async fn seed_abandoned_claims(pool: &PgPool, store: &QueueStorage, queue: &str, count: usize) { + let tag = Uuid::new_v4().simple().to_string(); + let mut enqueued = 0usize; + while enqueued < count { + let chunk = (count - enqueued).min(2_000); + let jobs: Vec = (0..chunk) + .map(|offset| params(queue, &tag, (enqueued + offset) as u64, None)) + .collect(); + store + .enqueue_params_batch(pool, &jobs) + .await + .expect("enqueue claimable jobs"); + enqueued += chunk; + } + let mut claimed = 0usize; + while claimed < count { + let batch = store + .claim_runtime_batch(pool, queue, 500, Duration::from_secs(3_600)) + .await + .expect("claim batch"); + assert!(!batch.is_empty(), "ran out of jobs to claim"); + claimed += batch.len(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_maintenance_under_hot_lane() { + let producers = env_usize("AWA_LANE_MAINT_PRODUCERS", 16); + let backlog = env_usize("AWA_LANE_MAINT_BACKLOG", 24_576); + let rescue_claims = env_usize("AWA_LANE_MAINT_RESCUE", 6_000); + + let pool = setup::setup(producers as u32 + 12).await; + setup::activate_queue_storage(&pool).await; + let store = Arc::new(QueueStorage::new(QueueStorageConfig::default()).expect("queue storage")); + + for (label, producer_count) in [("idle lane ", 0usize), ("hot lane ", producers)] { + let suffix = Uuid::new_v4().simple().to_string(); + let hot = format!("lane_maint_hot_{}", &suffix[..10]); + let other = format!("lane_maint_other_{}", &suffix[..10]); + + // Warm both lanes so lane provisioning is not measured. + insert_many( + &pool, + &[ + params(&hot, "warm", 0, None), + params(&other, "warm", 0, None), + ], + ) + .await + .expect("warm lanes"); + + seed_due_backlog(&pool, &hot, backlog).await; + seed_abandoned_claims(&pool, &store, &hot, rescue_claims).await; + + let running = (producer_count > 0).then(|| Producers::start(&pool, &hot, producer_count)); + tokio::time::sleep(Duration::from_millis(1_500)).await; + + let report = run_leader(&pool, &store, &hot, &other, Duration::from_secs(12)).await; + print_leader(label, &report); + + if let Some(running) = running { + let produced = running.stop().await; + println!("[lane-maint] {label} producers inserted {produced} jobs meanwhile"); + } + } +} + +/// The mechanism without producer load: one transaction holds the hot lane's +/// reservation lock for a fixed time while the leader loop runs. Promotion of +/// a batch that contains hot-queue rows waits for it, and takes the unrelated +/// queue's due job in the same batch and the leader's other branches with it. +/// Heartbeat rescue called directly does not wait for the lane at all: it +/// writes retryable rows to the deferred table, which promotion later moves. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_leader_behind_a_held_lane_lock() { + let hold = Duration::from_millis(env_usize("AWA_LANE_MAINT_HOLD_MS", 3_000) as u64); + let pool = setup::setup(16).await; + setup::activate_queue_storage(&pool).await; + let store = Arc::new(QueueStorage::new(QueueStorageConfig::default()).expect("queue storage")); + + let suffix = Uuid::new_v4().simple().to_string(); + let hot = format!("lane_maint_hot_{}", &suffix[..10]); + let other = format!("lane_maint_other_{}", &suffix[..10]); + insert_many( + &pool, + &[ + params(&hot, "warm", 0, None), + params(&other, "warm", 0, None), + ], + ) + .await + .expect("warm lanes"); + seed_due_backlog(&pool, &hot, 8_192).await; + seed_abandoned_claims(&pool, &store, &hot, 3_000).await; + + let key: i64 = sqlx::query_scalar( + "SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, 2::smallint, 0::smallint)), + 0)", + ) + .bind(&hot) + .fetch_one(&pool) + .await + .expect("derive lane lock key"); + let mut gate = pool.begin().await.expect("begin gate"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(key) + .execute(&mut *gate) + .await + .expect("take lane lock"); + + let rescue_started = Instant::now(); + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue with the lane lock held"); + println!( + "[lane-maint] rescue of {} claims with the hot lane lock held took {:.0}ms", + rescued.len(), + rescue_started.elapsed().as_secs_f64() * 1e3 + ); + let remaining_hold = hold.saturating_sub(rescue_started.elapsed()); + + let release = tokio::spawn(async move { + tokio::time::sleep(remaining_hold).await; + gate.commit().await.expect("release lane lock"); + }); + let report = run_leader(&pool, &store, &hot, &other, Duration::from_secs(9)).await; + release.await.expect("release task"); + print_leader(&format!("lane held {}ms", hold.as_millis()), &report); +} diff --git a/awa/tests/persona_soak/abandon_oracle.rs b/awa/tests/persona_soak/abandon_oracle.rs new file mode 100644 index 000000000..f28605717 --- /dev/null +++ b/awa/tests/persona_soak/abandon_oracle.rs @@ -0,0 +1,1268 @@ +//! Checks for the claim-and-abandon persona: a process holding thousands of claims is +//! SIGKILLed again and again. After each kill the rescue backlog has to clear within a +//! bound, the claim ring has to keep rotating, no job whose owner is alive may be +//! rescued, and attempts and the dead-letter queue have to follow `max_attempts`. + +use super::config::Persona; +use super::shared::Shared; +use serde::Serialize; +use sqlx::PgPool; +use std::collections::HashMap; +use std::io::Write as _; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; +use std::sync::Mutex; +use std::time::{SystemTime, UNIX_EPOCH}; + +const MAX_ROWS: usize = 25; +/// The ring rotates about once a second; this long without a rotation is a stall. +const RING_STALL_AFTER_MS: i64 = 6_000; +/// Owner of attempts run by the surviving in-process runtime, which is never killed. +pub const FINISHER: u32 = u32::MAX; + +pub fn epoch_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_millis() as i64) + .unwrap_or(0) +} + +#[derive(Debug, Clone, Copy)] +pub struct Attempt { + pub attempt: i16, + pub owner: u32, + pub at_ms: i64, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Kill { + pub number: u32, + pub generation: u32, + /// Milliseconds since the run started, and since the epoch, when SIGKILL was sent. + pub at_run_ms: i64, + pub signalled_epoch_ms: i64, + /// Handlers the process had started, which is what it held when it died. + pub started_by_process: i64, + /// Open claims on the persona's queues right after the kill. + pub open_after_kill: i64, + /// Claims of the process that were rescued before it was killed. + pub rescued_while_alive: i64, + pub expired_peak: i64, + pub oldest_expired_peak_s: f64, + pub cleared_after_s: Option, + pub reported_slow: bool, + pub load_ended: bool, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RingEpisode { + pub start_ms: i64, + pub end_ms: Option, + pub generation: i64, + pub slot: i32, + /// The last time the slot the ring has to enter next held a claim whose owner is + /// alive, which explains a stall. Starts at the beginning of the episode. + pub last_live_ms: i64, + pub expired_in_next_slot: i64, + pub reported: bool, +} + +#[derive(Default)] +struct RingMonitor { + last: Option<(i32, i64)>, + last_change_ms: i64, + episode: Option, + episodes: Vec, +} + +#[derive(Debug, Clone, Copy, Serialize)] +pub struct Sample { + pub at_ms: i64, + pub open: i64, + pub expired: i64, + pub oldest_expired_s: f64, + pub ring_generation: i64, + pub probe_ms: i64, +} + +#[derive(Default)] +pub struct AbandonState { + /// Open claims on the persona's queues at the last probe. + pub open_claims: AtomicI64, + /// A child process is running. + pub child_alive: AtomicBool, + /// Handlers running in the child process at its last report. + pub child_active: AtomicI64, + started_by_process: Mutex>, + /// The running child process and the epoch milliseconds at which it was started. + live: Mutex>, + /// Most claims of a process seen rescued while it was alive, per process. + live_lost: Mutex>, + attempts: Mutex>>, + max_attempts: Mutex>, + /// Epoch milliseconds at which each process generation was killed. + killed_at: Mutex>, + kills: Mutex>, + samples: Mutex>, + ring: Mutex, + csv: Mutex>, + passes_csv: Mutex>, + heartbeat_last: Mutex>, +} + +pub struct Probe { + pub open: i64, + pub expired: i64, + pub oldest_expired_s: f64, + /// Open claims taken at or after the `since` the probe was given. + pub open_since: i64, +} + +/// The process that is alive, and how many handlers it had started before the probe. +#[derive(Debug, Clone, Copy)] +pub struct LiveCheck { + pub generation: u32, + pub started: i64, +} + +impl AbandonState { + pub fn set_csv(&self, path: PathBuf) { + if let Ok(mut file) = std::fs::File::create(&path) { + let _ = writeln!( + file, + "run_s,open_claims,expired_claims,oldest_expired_s,claim_ring_generation,probe_ms" + ); + } + let passes = path.with_file_name("abandon-heartbeat-passes.csv"); + if let Ok(mut file) = std::fs::File::create(&passes) { + let _ = writeln!(file, "run_s,handlers_running,claims_sampled,pass_ms"); + } + *self.passes_csv.lock().expect("passes csv") = Some(passes); + *self.csv.lock().expect("csv") = Some(path); + } + + /// Folds in one reading of the oldest and newest heartbeat (epoch seconds) among the live + /// process's claims. While a pass runs, the oldest heartbeat is still one the previous pass + /// wrote, so it cannot be newer than that pass's newest. The first reading whose oldest + /// heartbeat is newer was taken just after a pass finished, and the spread from oldest to + /// newest is its duration (to within the last chunk). The samples are bucketed by handler + /// count so the report shows how pass time grows with what a process holds. `now` is the + /// database clock at the reading. + pub fn note_heartbeat_sample( + &self, + shared: &Shared, + rows: i64, + oldest: f64, + newest: f64, + now: f64, + active: i64, + ) { + let mut last = self.heartbeat_last.lock().expect("heartbeat last"); + let Some(settled_newest) = *last else { + *last = Some(newest); + return; + }; + if oldest <= settled_newest || active <= 0 { + return; + } + *last = Some(newest); + drop(last); + // A reading taken long after the pass finished may include the next pass's first + // chunks, so only a fresh one measures the pass. + if now - newest > 2.0 { + shared + .stats(Persona::ClaimAbandon) + .add("heartbeat_passes_unmeasured", 1); + return; + } + let ms = ((newest - oldest) * 1000.0).round() as i64; + let stats = shared.stats(Persona::ClaimAbandon); + stats.observe("heartbeat_pass_ms", ms); + let bucket = match active { + ..=4_999 => "under 5k", + 5_000..=9_999 => "5k-10k", + 10_000..=14_999 => "10k-15k", + 15_000..=19_999 => "15k-20k", + _ => "20k and over", + }; + stats.observe(&format!("heartbeat_pass_ms (handlers {bucket})"), ms); + if let Some(path) = self.passes_csv.lock().expect("passes csv").as_ref() { + if let Ok(mut file) = std::fs::OpenOptions::new().append(true).open(path) { + let _ = writeln!( + file, + "{:.1},{active},{rows},{ms}", + shared.now_ms() as f64 / 1000.0 + ); + } + } + } + + pub fn register_job(&self, job_id: i64, max_attempts: i16) { + self.max_attempts + .lock() + .expect("max attempts") + .insert(job_id, max_attempts); + } + + pub fn note_attempt(&self, job_id: i64, attempt: Attempt) { + if attempt.owner != FINISHER { + *self + .started_by_process + .lock() + .expect("started") + .entry(attempt.owner) + .or_insert(0) += 1; + } + self.attempts + .lock() + .expect("attempts") + .entry(job_id) + .or_default() + .push(attempt); + } + + pub fn note_generation_killed(&self, generation: u32, at_epoch_ms: i64) { + self.killed_at + .lock() + .expect("killed at") + .insert(generation, at_epoch_ms); + } + + /// Signal and publish the death before a waiting probe can classify its result. + /// The live-state lock also fences the probe's post-query check; no await occurs + /// between a successful signal and publishing its generation and timestamp. + pub fn signal_generation_killed( + &self, + generation: u32, + signal: impl FnOnce() -> std::io::Result<()>, + ) -> std::io::Result { + let mut live = self.live.lock().expect("live"); + signal()?; + let signalled_at = epoch_ms(); + self.note_generation_killed(generation, signalled_at); + if live.is_some_and(|(current, _)| current == generation) { + *live = None; + } + Ok(signalled_at) + } + + pub fn set_live(&self, live: Option<(u32, i64)>) { + *self.heartbeat_last.lock().expect("heartbeat last") = None; + *self.live.lock().expect("live") = live; + } + + pub fn live(&self) -> Option<(u32, i64)> { + *self.live.lock().expect("live") + } + + pub fn live_lost(&self, generation: u32) -> i64 { + self.live_lost + .lock() + .expect("live lost") + .get(&generation) + .copied() + .unwrap_or(0) + } + + pub fn record_kill(&self, kill: Kill) { + self.kills.lock().expect("kills").push(kill); + } + + pub fn kills(&self) -> Vec { + self.kills.lock().expect("kills").clone() + } + + pub fn started_by(&self, generation: u32) -> i64 { + self.started_by_process + .lock() + .expect("started") + .get(&generation) + .copied() + .unwrap_or(0) + } + + /// Whether the ring has been still for long enough to count as stalled. + pub fn ring_stalled(&self) -> bool { + self.ring.lock().expect("ring").episode.is_some() + } + + pub fn current_ring_episode(&self) -> Option { + self.ring.lock().expect("ring").episode.clone() + } + + pub fn episodes(&self) -> Vec { + let ring = self.ring.lock().expect("ring"); + let mut all = ring.episodes.clone(); + all.extend(ring.episode.clone()); + all + } + + /// Records one probe and advances the per-kill clocks. + pub fn observe( + &self, + shared: &Shared, + probe: &Probe, + live: Option, + ring_generation: i64, + probe_ms: i64, + ) { + let now = shared.now_ms(); + self.open_claims.store(probe.open, Ordering::Relaxed); + let sample = Sample { + at_ms: now, + open: probe.open, + expired: probe.expired, + oldest_expired_s: probe.oldest_expired_s, + ring_generation, + probe_ms, + }; + self.samples.lock().expect("samples").push(sample); + if let Some(path) = self.csv.lock().expect("csv").as_ref() { + if let Ok(mut file) = std::fs::OpenOptions::new().append(true).open(path) { + let _ = writeln!( + file, + "{:.1},{},{},{:.1},{},{}", + now as f64 / 1000.0, + probe.open, + probe.expired, + probe.oldest_expired_s, + ring_generation, + probe_ms + ); + } + } + let stats = shared.stats(Persona::ClaimAbandon); + stats.set_max("peak_open_claims", probe.open); + stats.set_max("peak_expired_claims", probe.expired); + stats.set_max( + "peak_oldest_expired_s", + probe.oldest_expired_s.round() as i64, + ); + stats.observe("probe_ms", probe_ms); + + // A probe can begin while the child is alive and finish after its kill. Only + // a generation still alive after the query proves missing claims were taken + // from a live owner. Keep this check serialized with signaling the kill. + if let Some(live) = live { + let lost = (live.started - probe.open_since).max(0); + let current = self.live.lock().expect("live"); + let still_live = current.is_some_and(|(generation, _)| generation == live.generation) + && !self + .killed_at + .lock() + .expect("killed at") + .contains_key(&live.generation); + if lost > 0 && !still_live { + shared.incomplete( + format!( + "claim probe for process {} crossed its kill or replacement; {lost} missing claim(s) cannot be classified as a live-owner rescue", + live.generation + ), + vec![format!( + "{{\"process\":{},\"started\":{},\"claims_still_open\":{}}}", + live.generation, live.started, probe.open_since + )], + ); + } else { + let mut peaks = self.live_lost.lock().expect("live lost"); + let peak = peaks.entry(live.generation).or_insert(0); + let first = lost > 0 && *peak == 0; + *peak = (*peak).max(lost); + drop(peaks); + stats.set_max("live_claims_rescued_peak", lost); + if first { + shared.violation( + "live_owner_rescued", + format!( + "claims of process {} were rescued while it was alive: it had started {} handler(s) and only {} of its claims were still open ({} claim(s) were eligible for rescue)", + live.generation, live.started, probe.open_since, probe.expired + ), + vec![format!( + "{{\"process\":{},\"started\":{},\"claims_still_open\":{},\"eligible_for_rescue\":{}}}", + live.generation, live.started, probe.open_since, probe.expired + )], + ); + } + } + } + let staleness_ms = shared.config.heartbeat_staleness.as_millis() as i64 + 3_000; + let epoch = epoch_ms(); + let bound = shared.config.abandon_bound; + let mut slow: Vec = Vec::new(); + { + let mut kills = self.kills.lock().expect("kills"); + for kill in kills + .iter_mut() + .filter(|kill| kill.cleared_after_s.is_none()) + { + kill.expired_peak = kill.expired_peak.max(probe.expired); + kill.oldest_expired_peak_s = kill.oldest_expired_peak_s.max(probe.oldest_expired_s); + let since_kill_ms = epoch - kill.signalled_epoch_ms; + // Claims of the dead process only count as stale after the staleness + // window, so a probe sooner than that proves nothing. + if since_kill_ms >= staleness_ms && probe.expired == 0 { + kill.cleared_after_s = Some(since_kill_ms as f64 / 1000.0); + continue; + } + let mut allowed = bound; + if shared.chaos_active() { + allowed += shared.config.chaos_slack; + } + if !kill.reported_slow && since_kill_ms as u64 > allowed.as_millis() as u64 { + kill.reported_slow = true; + slow.push(format!( + "{{\"kill\":{},\"seconds_since_kill\":{:.0},\"open_claims\":{},\"expired_claims\":{},\"oldest_expired_s\":{:.0},\"bound_s\":{}}}", + kill.number, + since_kill_ms as f64 / 1000.0, + probe.open, + probe.expired, + probe.oldest_expired_s, + bound.as_secs() + )); + } + } + } + for row in slow { + shared.violation( + "rescue_backlog_slow", + format!( + "the rescue backlog left by a killed process was not cleared within {}s of the kill ({} expired claim(s) remain, the oldest eligible for {:.0}s)", + bound.as_secs(), + probe.expired, + probe.oldest_expired_s + ), + vec![row], + ); + } + } + + /// Watches the claim ring; called about once a second. + /// + /// `next_slot` is the number of live and expired open claims in the slot the ring + /// enters next, known only while the ring is stalled. + pub fn observe_ring( + &self, + shared: &Shared, + current_slot: i32, + generation: i64, + next_slot: Option<(i64, i64)>, + ) { + let now = shared.now_ms(); + let mut report: Option = None; + { + let mut ring = self.ring.lock().expect("ring"); + let moved = ring.last != Some((current_slot, generation)); + if moved { + ring.last = Some((current_slot, generation)); + ring.last_change_ms = now; + if let Some(mut episode) = ring.episode.take() { + episode.end_ms = Some(now); + ring.episodes.push(episode); + } + return; + } + let mut threshold = crate::ci_timing::scaled_timeout(std::time::Duration::from_millis( + RING_STALL_AFTER_MS as u64, + )) + .as_millis() as i64; + if shared.chaos_active() { + threshold += shared.config.chaos_slack.as_millis() as i64; + } + if now - ring.last_change_ms >= threshold && ring.episode.is_none() { + ring.episode = Some(RingEpisode { + start_ms: ring.last_change_ms, + end_ms: None, + generation, + slot: current_slot, + last_live_ms: now, + expired_in_next_slot: 0, + reported: false, + }); + } + let bound_ms = shared.config.abandon_bound.as_millis() as i64; + if let Some(episode) = ring.episode.as_mut() { + if let Some((live, expired)) = next_slot { + if live > 0 { + episode.last_live_ms = now; + } + episode.expired_in_next_slot = expired; + } + // A stall is expected while a live process holds the next slot; it is a + // violation once nothing alive holds it and the ring still does not move. + if !episode.reported && now - episode.last_live_ms > bound_ms { + episode.reported = true; + report = Some(episode.clone()); + } + } + } + if let Some(episode) = report { + shared.violation( + "claim_ring_stalled", + format!( + "the claim ring has not rotated for more than {}s with no live claim in the slot it has to enter next ({} expired claim(s) there); stuck at slot {}, generation {}", + shared.config.abandon_bound.as_secs(), + episode.expired_in_next_slot, + episode.slot, + episode.generation + ), + vec![format!( + "{{\"stalled_since_s\":{:.1},\"slot\":{},\"generation\":{},\"expired_in_next_slot\":{}}}", + episode.start_ms as f64 / 1000.0, + episode.slot, + episode.generation, + episode.expired_in_next_slot + )], + ); + } + } + + /// Closes the ring episode left open when the run ends. + pub fn finish_ring(&self, shared: &Shared) { + let now = shared.now_ms(); + let mut ring = self.ring.lock().expect("ring"); + if let Some(mut episode) = ring.episode.take() { + episode.end_ms = Some(now); + ring.episodes.push(episode); + } + } + + /// Stall time and resume time of the ring in each kill's window. + fn ring_by_kill(&self) -> Vec<(f64, Option)> { + let kills = self.kills(); + let episodes = self.episodes(); + kills + .iter() + .enumerate() + .map(|(index, kill)| { + let from = kill.at_run_ms; + let until = kills.get(index + 1).map_or(i64::MAX, |next| next.at_run_ms); + let mut stalled = 0i64; + let mut resumed: Option = None; + for episode in &episodes { + let end = episode.end_ms.unwrap_or(i64::MAX); + if end <= from || episode.start_ms >= until { + continue; + } + stalled += end.min(until) - episode.start_ms.max(from); + if let Some(end) = episode.end_ms { + resumed = Some(resumed.map_or(end, |seen| seen.max(end))); + } + } + ( + stalled as f64 / 1000.0, + resumed.map(|end| (end - from) as f64 / 1000.0), + ) + }) + .collect() + } + + /// The per-kill table for the report. + pub fn render(&self) -> String { + let kills = self.kills(); + if kills.is_empty() { + return String::new(); + } + let ring = self.ring_by_kill(); + let samples = self.samples.lock().expect("samples"); + let mut out = String::new(); + out.push_str("### Kills\n\n"); + out.push_str("| kill | held by the process | rescued while alive | open claims after kill | peak expired | oldest expired (s) | time to clear (s) | ring stalled (s) | ring resumed after (s) |\n"); + out.push_str("| -- | -- | -- | -- | -- | -- | -- | -- | -- |\n"); + for (kill, (stalled, resumed)) in kills.iter().zip(ring) { + out.push_str(&format!( + "| {}{} | {} | {} | {} | {} | {:.0} | {} | {:.0} | {} |\n", + kill.number, + if kill.load_ended { " (load end)" } else { "" }, + kill.started_by_process, + kill.rescued_while_alive, + kill.open_after_kill, + kill.expired_peak, + kill.oldest_expired_peak_s, + kill.cleared_after_s + .map_or("not cleared".to_string(), |seconds| format!("{seconds:.0}")), + stalled, + resumed.map_or("-".to_string(), |seconds| format!("{seconds:.0}")), + )); + } + let ring_episodes = self.episodes(); + out.push_str(&format!( + "\nSamples: {}. Ring stall episodes: {}.\n", + samples.len(), + ring_episodes.len() + )); + out + } + + pub fn to_json(&self) -> serde_json::Value { + serde_json::json!({ + "kills": self.kills(), + "ring_episodes": self.episodes(), + "samples": self.samples.lock().expect("samples").len(), + }) + } + + /// Final checks over everything the child processes and the survivors reported. + pub async fn analyse(&self, shared: &Shared, pool: &PgPool, queues: &[String]) { + let killed_at = self.killed_at.lock().expect("killed at").clone(); + let max_attempts = self.max_attempts.lock().expect("max attempts").clone(); + let attempts = self.attempts.lock().expect("attempts").clone(); + let stats = shared.stats(Persona::ClaimAbandon); + + let mut live_rescued: Vec = Vec::new(); + let mut live_rescued_count = 0usize; + let mut duplicates_from_dead = 0i64; + let mut over_max: Vec = Vec::new(); + let mut gaps: Vec = Vec::new(); + let mut exhausted_by_attempts: Vec = Vec::new(); + for (job_id, list) in &attempts { + let mut list = list.clone(); + list.sort_by_key(|attempt| attempt.at_ms); + for pair in list.windows(2) { + let (previous, next) = (pair[0], pair[1]); + let previous_died_at = if previous.owner == FINISHER { + None + } else { + killed_at.get(&previous.owner).copied() + }; + match previous_died_at { + Some(died_at) if next.at_ms >= died_at => duplicates_from_dead += 1, + _ => { + live_rescued_count += 1; + if live_rescued.len() < MAX_ROWS { + live_rescued.push(format!( + "{{\"job_id\":{job_id},\"attempt\":{},\"owner\":\"{}\",\"owner_alive_until_ms\":{},\"next_attempt\":{},\"next_owner\":\"{}\",\"next_started_ms\":{}}}", + previous.attempt, + owner_name(previous.owner), + previous_died_at.map_or("null".to_string(), |at| at.to_string()), + next.attempt, + owner_name(next.owner), + next.at_ms + )); + } + } + } + if next.attempt != previous.attempt + 1 && gaps.len() < MAX_ROWS { + gaps.push(format!( + "{{\"job_id\":{job_id},\"attempt\":{},\"next_attempt\":{}}}", + previous.attempt, next.attempt + )); + } + } + let highest = list + .iter() + .map(|attempt| attempt.attempt) + .max() + .unwrap_or(0); + if let Some(max) = max_attempts.get(job_id) { + if highest > *max && over_max.len() < MAX_ROWS { + over_max.push(format!( + "{{\"job_id\":{job_id},\"attempts_seen\":{highest},\"max_attempts\":{max}}}" + )); + } + let last = list.last().copied(); + if highest == *max + && last.is_some_and(|last| last.owner != FINISHER && last.attempt == *max) + { + exhausted_by_attempts.push(*job_id); + } + } + } + stats.add("duplicate_executions_after_kill", duplicates_from_dead); + stats.add("live_owner_rescues", live_rescued_count as i64); + if live_rescued_count > 0 { + shared.violation( + "live_owner_rescued", + format!("{live_rescued_count} job(s) were run again while the process that held them was still alive"), + live_rescued, + ); + } + if !over_max.is_empty() { + shared.violation( + "attempts_exceeded_max", + format!( + "{} job(s) ran more attempts than max_attempts", + over_max.len() + ), + over_max, + ); + } + if !gaps.is_empty() { + shared.finding( + "abandon_attempt_gap", + format!( + "{} job(s) skipped an attempt number between executions", + gaps.len() + ), + gaps, + ); + } + + // Jobs whose last attempt was abandoned at max_attempts have to be dead-lettered. + let dlq: Result, _> = sqlx::query_as(awa::audited_sql(format!( + "SELECT job_id, attempt, max_attempts FROM {}.dlq_entries WHERE queue = ANY($1)", + schema(shared) + ))) + .bind(queues) + .fetch_all(pool) + .await; + match dlq { + Ok(rows) => { + stats.add("dead_lettered", rows.len() as i64); + let early: Vec = rows + .iter() + .filter(|(_, attempt, max)| attempt < max) + .take(MAX_ROWS) + .map(|(job, attempt, max)| { + format!("{{\"job_id\":{job},\"attempt\":{attempt},\"max_attempts\":{max}}}") + }) + .collect(); + if !early.is_empty() { + shared.violation( + "dead_lettered_early", + format!( + "{} job(s) were dead-lettered before reaching max_attempts", + early.len() + ), + early, + ); + } + let in_dlq: std::collections::HashSet = + rows.iter().map(|(job, _, _)| *job).collect(); + let missing: Vec = exhausted_by_attempts + .iter() + .filter(|job| !in_dlq.contains(job)) + .take(MAX_ROWS) + .map(|job| { + format!( + "{{\"job_id\":{job},\"max_attempts\":{}}}", + max_attempts.get(job).copied().unwrap_or(0) + ) + }) + .collect(); + stats.add("exhausted_after_kill", exhausted_by_attempts.len() as i64); + if !missing.is_empty() && !shared.quiesced.load(Ordering::Relaxed) { + // The rescue was still working when the run ended. + shared.finding( + "abandon_rescue_unfinished", + format!( + "{} job(s) abandoned on their last attempt were not yet dead-lettered when the drain grace ended", + missing.len() + ), + missing, + ); + } else if !missing.is_empty() { + shared.violation( + "exhausted_not_dead_lettered", + format!( + "{} job(s) abandoned on their last attempt are not in the dead-letter queue", + missing.len() + ), + missing, + ); + } + } + Err(error) => shared.incomplete(format!("dead-letter check failed: {error}"), vec![]), + } + } +} + +fn owner_name(owner: u32) -> String { + if owner == FINISHER { + "survivor".to_string() + } else { + format!("process-{owner}") + } +} + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Which claims a backlog probe looks at. +pub enum Scope<'a> { + /// Open receipt claims and running leases of these queues. + Queues(&'a [String]), + /// Open receipt identities in one claim-ring slot, including attempts + /// materialized into running or callback-waiting leases. + Slot(i32), +} + +/// Open claims in `scope`, and the ones the rescue is already entitled to take: claims +/// whose last heartbeat (or claim time, when there never was one) is older than +/// `staleness_secs`, or whose deadline has passed. +/// +/// `since_epoch_ms` splits out the open claims taken at or after that time. +pub async fn probe_backlog( + pool: &PgPool, + schema: &str, + scope: Scope<'_>, + staleness_secs: f64, + since_epoch_ms: i64, +) -> Result { + let (batch_filter, row_filter, receipt_lease_filter, lease_leg) = match &scope { + Scope::Queues(_) => ( + "b.queue = ANY($1)", + "c.queue = ANY($1)", + "l.job_id IS NULL", + format!( + "UNION ALL + SELECT COALESCE(a.heartbeat_at, l.heartbeat_at, l.attempted_at), l.deadline_at, + COALESCE(l.attempted_at, l.heartbeat_at), NULL::timestamptz + FROM {schema}.leases l + LEFT JOIN {schema}.attempt_state a ON a.job_id = l.job_id AND a.run_lease = l.run_lease + WHERE l.queue = ANY($1) AND l.state = 'running'" + ), + ), + Scope::Slot(_) => ( + "b.claim_slot = $4", + "c.claim_slot = $4", + "l.job_id IS NULL OR l.state IN ('running', 'waiting_external')", + String::new(), + ), + }; + let sql = format!( + r#" + WITH bounds AS ( + SELECT clock_timestamp() AS now_ts, + clock_timestamp() - make_interval(secs => $2) AS stale_before + ), + batch_claims AS ( + SELECT b.claim_slot, i.receipt_id, i.job_id, i.run_lease, b.queue, b.claimed_at, b.deadline_at + FROM {schema}.lease_claim_batches b + CROSS JOIN LATERAL unnest(b.job_ids, b.run_leases, b.receipt_ids) AS i(job_id, run_lease, receipt_id) + WHERE {batch_filter} + ), + row_claims AS ( + SELECT c.claim_slot, c.receipt_id, c.job_id, c.run_lease, c.queue, c.claimed_at, c.deadline_at + FROM {schema}.lease_claims c + WHERE {row_filter} AND c.closed_at IS NULL + ), + open_receipts AS ( + SELECT x.* FROM (SELECT * FROM batch_claims UNION ALL SELECT * FROM row_claims) x + WHERE NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closures cl + WHERE cl.claim_slot = x.claim_slot AND cl.job_id = x.job_id AND cl.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closure_batches cb + WHERE cb.claim_slot = x.claim_slot AND cb.receipt_ranges @> x.receipt_id + -- Keep membership correlated: joining by slot can multiply + -- receipt members by retained closure batches. + OFFSET 0) + AND NOT EXISTS (SELECT 1 FROM {schema}.done_entries d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.deferred_jobs d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.dlq_entries d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + ), + holders AS ( + SELECT CASE WHEN l.state = 'waiting_external' THEN NULL + ELSE COALESCE(a.heartbeat_at, l.heartbeat_at, l.attempted_at, x.claimed_at) + END AS beat_at, + CASE WHEN l.job_id IS NULL THEN x.deadline_at + WHEN l.state = 'running' THEN l.deadline_at + END AS deadline_at, + x.claimed_at, + CASE WHEN l.state = 'waiting_external' THEN l.callback_timeout_at END + AS callback_timeout_at + FROM open_receipts x + LEFT JOIN {schema}.leases l ON l.job_id = x.job_id AND l.run_lease = x.run_lease + LEFT JOIN {schema}.attempt_state a ON a.job_id = x.job_id AND a.run_lease = x.run_lease + WHERE {receipt_lease_filter} + {lease_leg} + ), + eligible AS ( + SELECT LEAST( + CASE WHEN beat_at < stale_before THEN beat_at + make_interval(secs => $2) END, + CASE WHEN deadline_at < now_ts THEN deadline_at END, + CASE WHEN callback_timeout_at < now_ts THEN callback_timeout_at END + ) AS eligible_at, + claimed_at >= to_timestamp($3::float8 / 1000.0) AS since_start + FROM holders, bounds + ) + SELECT count(*)::bigint, + count(eligible_at)::bigint, + COALESCE(EXTRACT(EPOCH FROM (clock_timestamp() - min(eligible_at))), 0)::float8, + count(*) FILTER (WHERE since_start)::bigint + FROM eligible + "# + ); + let query = sqlx::query_as::<_, (i64, i64, f64, i64)>(awa::audited_sql(&sql)); + let query = match scope { + Scope::Queues(queues) => query + .bind(queues.to_vec()) + .bind(staleness_secs) + .bind(since_epoch_ms as f64) + .bind(0_i32), + Scope::Slot(slot) => query + .bind(Vec::::new()) + .bind(staleness_secs) + .bind(since_epoch_ms as f64) + .bind(slot), + }; + let (open, expired, oldest, open_since) = query.fetch_one(pool).await?; + Ok(Probe { + open, + expired, + oldest_expired_s: oldest.max(0.0), + open_since, + }) +} + +#[cfg(test)] +mod tests { + use super::{epoch_ms, probe_backlog, LiveCheck, Probe, Scope}; + use crate::persona_soak::config::{replace_database_name, SoakConfig}; + use crate::persona_soak::db::{self, Budget, ServerInfo}; + use crate::persona_soak::shared::Shared; + use awa::model::{ClaimedRuntimeJob, QueueStorage, QueueStorageConfig}; + use sqlx::PgPool; + use std::sync::atomic::Ordering; + use std::sync::Arc; + use std::time::Duration; + + fn live_probe_fixture(generation: u32) -> (Arc, LiveCheck) { + let config = SoakConfig::for_test( + "postgres://postgres:test@localhost/awa_soak_live_probe".to_string(), + ); + let info = ServerInfo::default(); + let budget = Budget::new(&info, 0); + let shared = Shared::new(config, info, budget); + shared.abandon.set_live(Some((generation, epoch_ms()))); + let live = LiveCheck { + generation: shared.abandon.live().expect("live generation").0, + started: 1_000, + }; + (shared, live) + } + + fn observe_live_probe(shared: &Shared, live: LiveCheck, open: i64) { + shared.abandon.observe( + shared, + &Probe { + open, + expired: 0, + oldest_expired_s: 0.0, + open_since: open, + }, + Some(live), + 0, + 1, + ); + } + + fn live_rescue_reports(shared: &Shared) -> usize { + shared + .violations + .lock() + .expect("violations") + .iter() + .filter(|violation| violation.kind == "live_owner_rescued") + .count() + } + + #[test] + fn live_probe_crossing_its_generation_kill_is_incomplete() { + let (shared, captured) = live_probe_fixture(1); + // The query was dispatched with this owner alive, but its snapshot + // boundary relative to the signal is unknown when the result arrives. + shared.abandon.note_generation_killed(1, epoch_ms()); + shared.abandon.set_live(None); + observe_live_probe(&shared, captured, 500); + + assert!(shared.oracle_incomplete.load(Ordering::Relaxed)); + assert_eq!(live_rescue_reports(&shared), 0); + assert_eq!(shared.abandon.live_lost(1), 0); + // The usable backlog evidence is retained even though liveness is unknown. + assert_eq!(shared.abandon.open_claims.load(Ordering::Relaxed), 500); + } + + #[test] + fn live_probe_still_alive_missing_claims_remain_a_violation() { + let (shared, captured) = live_probe_fixture(1); + observe_live_probe(&shared, captured, 500); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(1), 500); + assert!(!shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_other_generation_kill_does_not_hide_live_loss() { + let (shared, captured) = live_probe_fixture(2); + shared.abandon.note_generation_killed(1, epoch_ms()); + observe_live_probe(&shared, captured, 500); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(2), 500); + assert!(!shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_kill_overlap_preserves_prior_proven_loss() { + let (shared, captured) = live_probe_fixture(1); + observe_live_probe(&shared, captured, 500); + assert_eq!(shared.abandon.live_lost(1), 500); + + shared.abandon.note_generation_killed(1, epoch_ms()); + shared.abandon.set_live(None); + observe_live_probe(&shared, captured, 100); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(1), 500); + assert!(shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_successful_signal_publishes_death_before_returning() { + let (shared, _) = live_probe_fixture(1); + let at = shared + .abandon + .signal_generation_killed(1, || Ok(())) + .expect("successful signal"); + + assert_eq!(shared.abandon.live(), None); + assert_eq!( + shared.abandon.killed_at.lock().expect("killed at").get(&1), + Some(&at) + ); + } + + #[test] + fn live_probe_failed_signal_does_not_publish_a_death() { + let (shared, _) = live_probe_fixture(1); + let before = shared.abandon.live(); + let result = shared + .abandon + .signal_generation_killed(1, || Err(std::io::Error::other("signal failed"))); + + assert!(result.is_err()); + assert_eq!(shared.abandon.live(), before); + assert!(shared + .abandon + .killed_at + .lock() + .expect("killed at") + .is_empty()); + } + + /// Compact claims with and without a deadline, and a row-local claim with a deadline. + const CLAIM_SHAPES: [(Duration, bool); 3] = [ + (Duration::ZERO, false), + (Duration::from_secs(3600), false), + (Duration::from_secs(3600), true), + ]; + + struct SlotProbe { + config: SoakConfig, + pool: PgPool, + store: QueueStorage, + } + + impl SlotProbe { + async fn create() -> Self { + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_slot_probe_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + db::recreate_database(&config) + .await + .expect("isolated database"); + let pool = db::connect_pool(&config.database_url, "slot_probe_regression", 4) + .await + .expect("connect"); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("store"); + Self { + config, + pool, + store, + } + } + + /// Claims one job. This runtime writes compact claim batches whatever the + /// deadline; `row_local` rewrites the claim as the `lease_claims` row a 0.6 + /// runtime writes for a deadline-backed claim, which this runtime still reads and + /// closes during a rolling upgrade. + async fn claim(&self, deadline: Duration, row_local: bool) -> ClaimedRuntimeJob { + self.store + .enqueue_batch(&self.pool, "slot_probe", 2, 1) + .await + .expect("enqueue"); + let mut claim = self + .store + .claim_runtime_batch(&self.pool, "slot_probe", 1, deadline) + .await + .expect("claim") + .pop() + .expect("one claim"); + assert!(claim.claim.claim_batch_id.is_some(), "claims are compact"); + if row_local { + sqlx::query("DELETE FROM awa.lease_claim_batches WHERE batch_id = $1") + .bind(claim.claim.claim_batch_id) + .execute(&self.pool) + .await + .expect("remove the compact claim"); + sqlx::query( + "INSERT INTO awa.lease_claims ( + claim_slot, receipt_id, job_id, run_lease, ready_slot, + ready_generation, queue, priority, attempt, max_attempts, + lane_seq, enqueue_shard, deadline_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, + clock_timestamp() + make_interval(secs => $13))", + ) + .bind(claim.claim.claim_slot) + .bind(claim.claim.receipt_id) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .bind(claim.claim.ready_slot) + .bind(claim.claim.ready_generation) + .bind(&claim.claim.queue) + .bind(claim.claim.priority) + .bind(claim.job.attempt) + .bind(claim.job.max_attempts) + .bind(claim.claim.lane_seq) + .bind(claim.claim.enqueue_shard) + .bind(deadline.as_secs_f64()) + .execute(&self.pool) + .await + .expect("write the row-local claim"); + claim.claim.claim_batch_id = None; + claim.claim.claim_batch_index = None; + } + claim + } + + async fn assert_counts(&self, slot: i32, open: i64, expired: i64) { + let probe = probe_backlog(&self.pool, "awa", Scope::Slot(slot), 30.0, 0) + .await + .expect("slot probe"); + assert_eq!((probe.open, probe.expired), (open, expired)); + } + + async fn age_heartbeat(&self, claim: &ClaimedRuntimeJob) { + sqlx::query( + "UPDATE awa.leases + SET heartbeat_at = clock_timestamp() - interval '1 minute' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&self.pool) + .await + .expect("age lease heartbeat"); + sqlx::query( + "INSERT INTO awa.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + VALUES ($1, $2, clock_timestamp() - interval '1 minute', clock_timestamp()) + ON CONFLICT (job_id, run_lease) DO UPDATE + SET heartbeat_at = EXCLUDED.heartbeat_at", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&self.pool) + .await + .expect("age attempt heartbeat"); + } + + async fn drop(self) { + self.pool.close().await; + let admin = db::connect_pool(&self.config.admin_url(), "slot_probe_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(awa::audited_sql(format!( + "DROP DATABASE {} WITH (FORCE)", + self.config.database_name() + ))) + .execute(&admin) + .await + .expect("drop isolated database"); + admin.close().await; + } + } + + #[tokio::test] + async fn slot_probe_keeps_materialized_running_claims_until_durable_completion() { + let run = SlotProbe::create().await; + for (deadline, row_local) in CLAIM_SHAPES { + let claim = run.claim(deadline, row_local).await; + run.store + .register_callback( + &run.pool, + claim.job.id, + claim.job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("materialize the running receipt"); + let slot = claim.claim.claim_slot; + run.assert_counts(slot, 1, 0).await; + run.assert_counts((slot + 1) % 8, 0, 0).await; + run.age_heartbeat(&claim).await; + run.assert_counts(slot, 1, 1).await; + run.store + .heartbeat_batch(&run.pool, &[(claim.job.id, claim.job.run_lease)]) + .await + .expect("refresh authoritative attempt heartbeat"); + // The stale leases.heartbeat_at must not override attempt_state. + run.assert_counts(slot, 1, 0).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&claim)) + .await + .expect("durably complete"); + run.assert_counts(slot, 0, 0).await; + } + run.drop().await; + } + + #[tokio::test] + async fn slot_probe_uses_callback_timeout_for_waiting_materialized_claims() { + let run = SlotProbe::create().await; + for (deadline, row_local) in CLAIM_SHAPES { + let claim = run.claim(deadline, row_local).await; + let callback = run + .store + .register_callback( + &run.pool, + claim.job.id, + claim.job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("materialize the receipt"); + assert!(run + .store + .enter_callback_wait(&run.pool, claim.job.id, claim.job.run_lease, callback) + .await + .expect("wait for callback")); + run.age_heartbeat(&claim).await; + sqlx::query( + "UPDATE awa.leases SET deadline_at = clock_timestamp() - interval '1 second' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&run.pool) + .await + .expect("expire the original job deadline"); + let slot = claim.claim.claim_slot; + // Waiting ignores the old heartbeat and deadline until its callback times out. + run.assert_counts(slot, 1, 0).await; + run.assert_counts((slot + 1) % 8, 0, 0).await; + sqlx::query( + "UPDATE awa.leases SET callback_timeout_at = clock_timestamp() - interval '1 second' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&run.pool) + .await + .expect("expire callback"); + run.assert_counts(slot, 1, 1).await; + run.store + .complete_external(&run.pool, callback, None, Some(claim.job.run_lease), false) + .await + .expect("durably complete callback"); + run.assert_counts(slot, 0, 0).await; + } + run.drop().await; + } +} diff --git a/awa/tests/persona_soak/adminserver.rs b/awa/tests/persona_soak/adminserver.rs new file mode 100644 index 000000000..79e6b0946 --- /dev/null +++ b/awa/tests/persona_soak/adminserver.rs @@ -0,0 +1,132 @@ +//! The real admin HTTP server (`awa serve`), run as a child process against the soak +//! database so the UI personas exercise the same code the web UI talks to. + +use super::config::SoakConfig; +use std::net::TcpListener; +use std::path::PathBuf; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +pub const APPLICATION_NAME: &str = "soak_ui_server"; + +pub struct AdminServer { + child: Child, + pub base_url: String, + pub statement_timeout: Duration, + pub pool_max: u32, +} + +/// Looks for the `awa` binary next to the test executable, where +/// `cargo build --release -p awa-cli` puts it, unless `AWA_SOAK_AWA_BIN` names one. +pub fn locate_binary(config: &SoakConfig) -> Result { + if let Some(path) = &config.awa_bin { + return path + .exists() + .then(|| path.clone()) + .ok_or_else(|| format!("AWA_SOAK_AWA_BIN {} does not exist", path.display())); + } + let exe = std::env::current_exe().map_err(|error| error.to_string())?; + let candidate = exe + .parent() + .and_then(|deps| deps.parent()) + .map(|profile| profile.join("awa")); + match candidate { + Some(path) if path.exists() => Ok(path), + _ => Err( + "the admin UI personas need the awa binary: run `cargo build --release -p awa-cli` \ + or set AWA_SOAK_AWA_BIN" + .to_string(), + ), + } +} + +impl AdminServer { + pub async fn start(config: &SoakConfig) -> Result { + let binary = locate_binary(config)?; + let port = TcpListener::bind("127.0.0.1:0") + .and_then(|listener| listener.local_addr()) + .map_err(|error| format!("pick a port: {error}"))? + .port(); + let separator = if config.database_url.contains('?') { + '&' + } else { + '?' + }; + let url = format!( + "{}{separator}application_name={APPLICATION_NAME}", + config.database_url + ); + let statement_timeout = config.ui_statement_timeout; + let log = std::fs::File::create(config.report_dir.join("awa-serve.log")) + .or_else(|_| { + std::fs::create_dir_all(&config.report_dir)?; + std::fs::File::create(config.report_dir.join("awa-serve.log")) + }) + .map_err(|error| format!("create awa-serve.log: {error}"))?; + let child = Command::new(&binary) + .args([ + "--database-url", + &url, + "serve", + "--host", + "127.0.0.1", + "--port", + &port.to_string(), + ]) + .args([ + "--statement-timeout", + &statement_timeout.as_secs().to_string(), + ]) + .args([ + "--pool-max", + &config.ui_pool_max.to_string(), + "--pool-min", + "2", + ]) + .stdout(Stdio::from( + log.try_clone().map_err(|error| error.to_string())?, + )) + .stderr(Stdio::from(log)) + .spawn() + .map_err(|error| format!("spawn {}: {error}", binary.display()))?; + let mut server = Self { + child, + base_url: format!("http://127.0.0.1:{port}"), + statement_timeout, + pool_max: config.ui_pool_max, + }; + let client = reqwest::Client::new(); + let deadline = Instant::now() + Duration::from_secs(60); + loop { + if let Ok(Some(status)) = server.child.try_wait() { + return Err(format!("awa serve exited during startup: {status}")); + } + if let Ok(response) = client + .get(format!("{}/api/capabilities", server.base_url)) + .send() + .await + { + if response.status().is_success() { + return Ok(server); + } + } + if Instant::now() >= deadline { + return Err("awa serve did not start listening within 60s".to_string()); + } + tokio::time::sleep(Duration::from_millis(300)).await; + } + } + + pub fn is_running(&mut self) -> bool { + matches!(self.child.try_wait(), Ok(None)) + } + + pub fn stop(self) {} +} + +impl Drop for AdminServer { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} diff --git a/awa/tests/persona_soak/config.rs b/awa/tests/persona_soak/config.rs new file mode 100644 index 000000000..c3a62a6f8 --- /dev/null +++ b/awa/tests/persona_soak/config.rs @@ -0,0 +1,327 @@ +//! Run configuration, read from the environment. + +use crate::ci_timing::{scaled_staleness, scaled_timeout}; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::Duration; + +/// The harness's fast heartbeat cadence, so that rescues happen many times in a short run. +pub const DEFAULT_HEARTBEAT_SECONDS: u64 = 2; +pub const DEFAULT_STALENESS_SECONDS: u64 = 12; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum Persona { + Mia, + Tom, + Priya, + Sam, + Alex, + Chaos, + /// Admin UI users, through the real HTTP server. + Watchers, + Support, + UiOps, + Impatient, + /// A separate Python process: producers, async workers and the migrate path. + Python, + /// Handlers that batch-insert unique child jobs into their own queue. + Fanout, + /// An operator changing `enqueue_shards` on hot queues under live load, with the + /// actors running as different non-superuser roles. + ShardOps, + /// A separate process that claims as much as it can and is SIGKILLed again and again. + ClaimAbandon, +} + +impl Persona { + pub const ALL: [Persona; 14] = [ + Persona::Mia, + Persona::Tom, + Persona::Priya, + Persona::Sam, + Persona::Alex, + Persona::Chaos, + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + Persona::Python, + Persona::Fanout, + Persona::ShardOps, + Persona::ClaimAbandon, + ]; + + pub fn name(self) -> &'static str { + match self { + Persona::Mia => "mia", + Persona::Tom => "tom", + Persona::Priya => "priya", + Persona::Sam => "sam", + Persona::Alex => "alex", + Persona::Chaos => "chaos", + Persona::Watchers => "watchers", + Persona::Support => "support", + Persona::UiOps => "uiops", + Persona::Impatient => "impatient", + Persona::Python => "python", + Persona::Fanout => "fanout", + Persona::ShardOps => "shardops", + Persona::ClaimAbandon => "abandon", + } + } + + pub fn index(self) -> u8 { + self as u8 + } + + pub fn from_index(index: u8) -> Persona { + Persona::ALL[(index as usize).min(Persona::ALL.len() - 1)] + } +} + +impl FromStr for Persona { + type Err = String; + + fn from_str(value: &str) -> Result { + Persona::ALL + .into_iter() + .find(|persona| persona.name() == value.trim().to_ascii_lowercase()) + .ok_or_else(|| { + format!("unknown persona {value:?}; expected one of mia,tom,priya,sam,alex,chaos,watchers,support,uiops,impatient,python,fanout,shardops,abandon") + }) + } +} + +#[derive(Debug, Clone)] +pub struct SoakConfig { + /// URL of the disposable soak database. Its database name must start with `awa_soak`. + pub database_url: String, + pub personas: Vec, + /// Length of the load phase. + pub load: Duration, + /// How long the oracle waits for the system to quiesce after load stops. + pub drain_grace: Duration, + pub seed: u64, + /// Multiplies each persona's rates and queue counts. + pub scale: f64, + pub report_dir: PathBuf, + /// Postgres server log file readable from this host. + pub pg_log_file: Option, + /// Docker container running the server, used to read its log. + pub pg_container: Option, + /// Log file path inside `pg_container`; when absent `docker logs` is used. + pub pg_container_log: Option, + /// Seconds a scheduled job may be overdue before the promotion check fails. + pub promotion_overdue: Duration, + /// Seconds a queue may hold available jobs with idle workers before the stall check fails. + pub stall_window: Duration, + /// The `awa` binary that serves the admin UI; located next to the test binary when unset. + pub awa_bin: Option, + /// Python interpreter of a virtualenv with the `awa` extension installed. + pub python_bin: Option, + /// Server-side statement timeout of the admin server, and the matching client timeout. + pub ui_statement_timeout: Duration, + pub ui_client_timeout: Duration, + /// Connections of the admin server's pool, and seconds between the operator's admin reads. + pub ui_pool_max: u32, + pub admin_read_every: Duration, + /// The load phase stops early when the report directory's filesystem has less free space. + pub min_free_disk_gb: i64, + /// Extra tolerance added to the stall and promotion thresholds while chaos is active. + pub chaos_slack: Duration, + /// Multiplies the shard-ops persona's rates independently of `scale`, which is + /// usually lowered to keep the other personas light. + pub shardops_scale: f64, + /// Whether the application roles may create objects in the schema. Without it the + /// operator provisions every lane before the shard count is raised, as the + /// restricted-runtime procedure in `docs/security/database-roles.md` requires. + pub shardops_schema_create: bool, + /// Claim-and-abandon persona: claims held by each killed process, seconds between + /// kills, how many kills, and the bound within which the rescue backlog has to clear + /// and the claim ring has to resume after a kill. Independent of `scale`. + pub abandon_claims: u64, + pub abandon_kill_every: Duration, + pub abandon_kills: u32, + pub abandon_bound: Duration, + /// Pause between a kill and the replacement process. + pub abandon_restart_delay: Duration, + /// Overrides the heartbeat and deadline rescue intervals of every runtime, to run the + /// soak with production's cadence instead of the harness's fast one. + pub rescue_interval: Option, + /// Heartbeat interval and the heartbeat staleness after which a claim may be rescued, + /// for every runtime. The library's own defaults are 30 s and 90 s. + pub heartbeat_interval: Duration, + pub heartbeat_staleness: Duration, +} + +fn env_string(name: &str) -> Option { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) +} + +fn env_parse(name: &str) -> Option { + env_string(name).and_then(|value| value.trim().parse().ok()) +} + +impl SoakConfig { + pub fn from_env() -> Result { + let database_url = env_string("AWA_SOAK_DATABASE_URL").ok_or( + "AWA_SOAK_DATABASE_URL is required (postgres://user:pass@host:port/awa_soak_)", + )?; + Self::from_database_url(database_url) + } + + #[cfg(test)] + pub fn for_test(database_url: String) -> Self { + let mut config = Self::from_database_url(database_url).expect("test database URL"); + config.personas.clear(); + config + } + + fn from_database_url(database_url: String) -> Result { + let database_name = database_name(&database_url); + if !database_name.starts_with("awa_soak") + || !database_name + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || ch == '_') + { + return Err(format!( + "refusing to run: the soak drops and recreates its database, so the name must \ + match awa_soak[A-Za-z0-9_]*, got {database_name:?}" + )); + } + + let personas = match env_string("AWA_SOAK_PERSONAS") { + Some(list) => list + .split(',') + .map(str::parse) + .collect::, _>>()?, + None => Persona::ALL.to_vec(), + }; + let seed = env_parse("AWA_SOAK_SEED").unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|elapsed| elapsed.as_nanos() as u64) + .unwrap_or(1) + }); + let report_dir = env_string("AWA_SOAK_REPORT_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| std::env::temp_dir().join(format!("awa-soak-{seed}"))); + + Ok(Self { + database_url, + personas, + load: Duration::from_secs(env_parse("AWA_SOAK_SECONDS").unwrap_or(900)), + drain_grace: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_DRAIN_GRACE_SECONDS").unwrap_or(120), + )), + seed, + scale: env_parse("AWA_SOAK_SCALE").unwrap_or(1.0), + report_dir, + pg_log_file: env_string("AWA_SOAK_PG_LOG").map(PathBuf::from), + pg_container: env_string("AWA_SOAK_PG_CONTAINER"), + pg_container_log: env_string("AWA_SOAK_PG_CONTAINER_LOG"), + awa_bin: env_string("AWA_SOAK_AWA_BIN").map(PathBuf::from), + python_bin: env_string("AWA_SOAK_PYTHON").map(PathBuf::from), + ui_statement_timeout: Duration::from_secs( + env_parse("AWA_SOAK_UI_STATEMENT_TIMEOUT_SECONDS").unwrap_or(30), + ), + ui_client_timeout: Duration::from_secs( + env_parse("AWA_SOAK_UI_CLIENT_TIMEOUT_SECONDS").unwrap_or(45), + ), + ui_pool_max: env_parse("AWA_SOAK_UI_POOL_MAX").unwrap_or(10), + admin_read_every: Duration::from_secs( + env_parse("AWA_SOAK_ADMIN_READ_SECONDS").unwrap_or(2), + ), + min_free_disk_gb: env_parse("AWA_SOAK_MIN_FREE_DISK_GB").unwrap_or(8), + promotion_overdue: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_PROMOTION_OVERDUE_SECONDS").unwrap_or(10), + )), + stall_window: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_STALL_SECONDS").unwrap_or(30), + )), + chaos_slack: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_CHAOS_SLACK_SECONDS").unwrap_or(45), + )), + shardops_scale: env_parse("AWA_SOAK_SHARDOPS_SCALE").unwrap_or(1.0), + shardops_schema_create: env_parse::("AWA_SOAK_SHARDOPS_SCHEMA_CREATE") + .map(|value| value != 0) + .unwrap_or(true), + abandon_claims: env_parse("AWA_SOAK_ABANDON_CLAIMS").unwrap_or(4_000), + abandon_kill_every: Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_KILL_SECONDS").unwrap_or(180), + ), + abandon_kills: env_parse("AWA_SOAK_ABANDON_KILLS").unwrap_or(4), + abandon_bound: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_BOUND_SECONDS").unwrap_or(120), + )), + abandon_restart_delay: Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_RESTART_SECONDS").unwrap_or(3), + ), + rescue_interval: env_parse::("AWA_SOAK_RESCUE_INTERVAL_SECONDS") + .map(Duration::from_secs), + heartbeat_interval: Duration::from_secs( + env_parse("AWA_SOAK_HEARTBEAT_SECONDS").unwrap_or(DEFAULT_HEARTBEAT_SECONDS), + ), + heartbeat_staleness: scaled_staleness(Duration::from_secs( + env_parse("AWA_SOAK_STALENESS_SECONDS").unwrap_or(DEFAULT_STALENESS_SECONDS), + )), + }) + } + + pub fn has(&self, persona: Persona) -> bool { + self.personas.contains(&persona) + } + + /// Scales a count, never below one. + pub fn scaled(&self, base: u64) -> u64 { + ((base as f64 * self.scale).round() as u64).max(1) + } + + pub fn database_name(&self) -> String { + database_name(&self.database_url) + } + + /// The same server, connected to the maintenance database. + pub fn admin_url(&self) -> String { + replace_database_name(&self.database_url, "postgres") + } +} + +pub fn database_name(url: &str) -> String { + let without_query = url.split_once('?').map_or(url, |(prefix, _)| prefix); + without_query + .rsplit_once('/') + .map(|(_, name)| name.to_string()) + .unwrap_or_default() +} + +pub fn replace_database_name(url: &str, database_name: &str) -> String { + let (without_query, query) = match url.split_once('?') { + Some((prefix, query)) => (prefix, Some(query)), + None => (url, None), + }; + let base = without_query + .rsplit_once('/') + .map_or(without_query, |(base, _)| base); + match query { + Some(query) => format!("{base}/{database_name}?{query}"), + None => format!("{base}/{database_name}"), + } +} + +/// `url` with its user and password replaced. +pub fn replace_credentials(url: &str, user: &str, password: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return url.to_string(); + }; + let authority_end = rest.find('/').unwrap_or(rest.len()); + let host = rest[..authority_end] + .rsplit_once('@') + .map_or(&rest[..authority_end], |(_, host)| host); + format!( + "{scheme}://{user}:{password}@{host}{}", + &rest[authority_end..] + ) +} diff --git a/awa/tests/persona_soak/db.rs b/awa/tests/persona_soak/db.rs new file mode 100644 index 000000000..0bd0b8e33 --- /dev/null +++ b/awa/tests/persona_soak/db.rs @@ -0,0 +1,298 @@ +//! 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(awa::audited_sql(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(awa::audited_sql(format!( + "CREATE DATABASE {name} OWNER {}", + roles::OWNER + ))) + .execute(&admin) + .await + .map_err(|error| format!("create database: {error}"))?; + } else { + sqlx::query(awa::audited_sql(format!("CREATE DATABASE {name}"))) + .execute(&admin) + .await + .map_err(|error| format!("create database: {error}"))?; + } + admin.close().await; + + // With the shard-ops persona the schema belongs to a non-superuser owner role, as in + // a production install, and the application roles get the documented grants. + let setup_url = if restricted_roles { + roles::url_for_role(config, roles::OWNER) + } else { + config.database_url.clone() + }; + let pool = connect_pool(&setup_url, "awa_soak_setup", 2) + .await + .map_err(|error| format!("connect to soak database: {error}"))?; + awa::model::migrations::run(&pool) + .await + .map_err(|error| format!("migrate: {error}"))?; + awa_testing::setup::activate_queue_storage(&pool).await; + pool.close().await; + if restricted_roles { + let superuser = connect_pool(&config.database_url, "awa_soak_setup_grants", 1) + .await + .map_err(|error| format!("connect to soak database: {error}"))?; + let mixed_personas = config + .personas + .iter() + .any(|persona| *persona != Persona::ShardOps); + roles::grant_runtime( + &superuser, + &name, + config.shardops_schema_create, + mixed_personas, + ) + .await?; + superuser.close().await; + } + Ok(()) +} + +#[derive(Debug, Clone, Default)] +pub struct ServerInfo { + pub version: String, + pub max_connections: u32, + pub superuser_reserved: u32, + pub connections_at_start: i64, + pub postmaster_start: String, + pub engine: String, + pub queue_storage_schema: Option, + pub awa_version: String, +} + +pub async fn server_info(pool: &PgPool) -> Result { + let version: String = sqlx::query_scalar("SHOW server_version") + .fetch_one(pool) + .await?; + let max_connections: String = sqlx::query_scalar("SHOW max_connections") + .fetch_one(pool) + .await?; + let reserved: String = sqlx::query_scalar("SHOW superuser_reserved_connections") + .fetch_one(pool) + .await?; + let connections: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE datname IS NOT NULL", + ) + .fetch_one(pool) + .await?; + let postmaster_start: String = sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await?; + let schema: Option = sqlx::query_scalar("SELECT awa.active_queue_storage_schema()") + .fetch_one(pool) + .await?; + Ok(ServerInfo { + version, + max_connections: max_connections.parse().unwrap_or(100), + superuser_reserved: reserved.parse().unwrap_or(3), + connections_at_start: connections, + postmaster_start, + engine: if schema.is_some() { + "queue_storage" + } else { + "canonical" + } + .to_string(), + queue_storage_schema: schema, + awa_version: env!("CARGO_PKG_VERSION").to_string(), + }) +} + +/// Hands out connections to the personas so that all of them together stay inside +/// the server's `max_connections`. +pub struct Budget { + usable: u32, + allocations: Mutex>, +} + +impl Budget { + /// `headroom` connections are kept free for the oracle, operator probes and chaos + /// sessions on top of the server's superuser reservation. + pub fn new(info: &ServerInfo, headroom: u32) -> Self { + let usable = info + .max_connections + .saturating_sub(info.superuser_reserved) + .saturating_sub(headroom) + .saturating_sub(info.connections_at_start.max(0) as u32); + Self { + usable, + allocations: Mutex::new(BTreeMap::new()), + } + } + + /// Reserves up to `want` connections under `label`, never fewer than `min`. + /// Returns the amount granted; panics when even `min` does not fit, because the + /// run would exhaust the server before it tested anything. + pub fn reserve(&self, label: &str, want: u32, min: u32) -> u32 { + let mut allocations = self.allocations.lock().expect("budget lock"); + let used: u32 = allocations.values().sum(); + let available = self.usable.saturating_sub(used); + assert!( + available >= min, + "connection budget exhausted: {label} needs at least {min}, {available} of {} left; \ + raise max_connections or lower AWA_SOAK_SCALE / persona count. Allocations: {allocations:?}", + self.usable + ); + let granted = want.min(available).max(min); + allocations.insert(label.to_string(), granted); + granted + } + + pub fn release(&self, label: &str) { + self.allocations.lock().expect("budget lock").remove(label); + } + + pub fn usable(&self) -> u32 { + self.usable + } + + pub fn table(&self) -> Vec<(String, u32)> { + self.allocations + .lock() + .expect("budget lock") + .iter() + .map(|(label, count)| (label.clone(), *count)) + .collect() + } +} + +/// A function that builds one runtime's client over its own pool. +pub type ClientFactory = Arc Client + Send + Sync>; + +/// One awa runtime living on its own tokio runtime, so it can be stopped gracefully +/// or killed the way a crashed process dies: tasks vanish mid-job without any +/// shutdown handshake and the connections drop. +pub struct ClientHost { + pub label: String, + runtime: tokio::runtime::Runtime, + client: Arc, + pool: PgPool, +} + +impl ClientHost { + pub async fn spawn( + url: String, + label: String, + pool_size: u32, + worker_threads: usize, + factory: ClientFactory, + ) -> Result { + let thread_label = label.clone(); + tokio::task::spawn_blocking(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .thread_name(format!("soak-{thread_label}")) + .enable_all() + .build() + .map_err(|error| format!("build runtime: {error}"))?; + let (pool, client) = runtime.block_on(async { + let pool = connect_pool(&url, &format!("soak_{thread_label}"), pool_size) + .await + .map_err(|error| format!("connect pool: {error}"))?; + let client = factory(pool.clone()); + client + .start() + .await + .map_err(|error| format!("start client: {error}"))?; + Ok::<_, String>((pool, client)) + })?; + Ok(Self { + label: thread_label, + runtime, + client: Arc::new(client), + pool, + }) + }) + .await + .map_err(|error| format!("spawn host task: {error}"))? + } + + pub fn client(&self) -> &Client { + &self.client + } + + pub fn pool(&self) -> &PgPool { + &self.pool + } + + /// Graceful stop: stop claiming, drain in-flight jobs, then close the pool. + pub async fn shutdown(self, timeout: Duration) { + let client = self.client.clone(); + let pool = self.pool.clone(); + let _ = self + .runtime + .handle() + .spawn(async move { + client.shutdown(timeout).await; + pool.close().await; + }) + .await; + self.runtime.shutdown_background(); + } + + /// Abrupt stop: no shutdown handshake, as when a process is killed. + pub fn kill(self) { + self.runtime.shutdown_background(); + } +} + +pub fn url_for_database(config: &SoakConfig, name: &str) -> String { + replace_database_name(&config.database_url, name) +} diff --git a/awa/tests/persona_soak/fleet.rs b/awa/tests/persona_soak/fleet.rs new file mode 100644 index 000000000..b7e709ca1 --- /dev/null +++ b/awa/tests/persona_soak/fleet.rs @@ -0,0 +1,192 @@ +//! A persona's group of runtimes, with the capacity bookkeeping the stall check +//! needs and the lifecycle operations the operator and chaos personas use. + +use super::config::Persona; +use super::db::{ClientFactory, ClientHost}; +use super::registry::QueueInfo; +use super::shared::Shared; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Mutex; + +pub struct Fleet { + pub name: String, + pub persona: Persona, + /// Every runtime in the fleet serves these queues with this many workers. + served: Vec<(Arc, i64)>, + factory: ClientFactory, + pool_size: u32, + worker_threads: usize, + url: String, + hosts: Mutex>>, + budget_label: String, +} + +impl Fleet { + #[allow(clippy::too_many_arguments)] + pub async fn start( + shared: &Arc, + name: &str, + persona: Persona, + replicas: usize, + pool_size: u32, + worker_threads: usize, + served: Vec<(Arc, i64)>, + factory: ClientFactory, + ) -> Result, String> { + let url = shared.config.database_url.clone(); + Self::start_with_url( + shared, + url, + name, + persona, + replicas, + pool_size, + worker_threads, + served, + factory, + ) + .await + } + + /// As [`Fleet::start`], with every runtime connecting through `url`, for a fleet that + /// runs as a role other than the one in the run's database URL. + #[allow(clippy::too_many_arguments)] + pub async fn start_with_url( + shared: &Arc, + url: String, + name: &str, + persona: Persona, + replicas: usize, + pool_size: u32, + worker_threads: usize, + served: Vec<(Arc, i64)>, + factory: ClientFactory, + ) -> Result, String> { + let budget_label = name.to_string(); + let granted = shared.budget.reserve( + &budget_label, + pool_size * replicas as u32, + 2 * replicas as u32, + ); + let pool_size = (granted / replicas as u32).max(2); + let fleet = Arc::new(Self { + name: name.to_string(), + persona, + served, + factory, + pool_size, + worker_threads, + url, + hosts: Mutex::new(Vec::new()), + budget_label, + }); + for index in 0..replicas { + let host = fleet.spawn_host(index).await?; + fleet.hosts.lock().await.push(Some(host)); + } + shared + .fleets + .lock() + .expect("fleets") + .insert(name.to_string(), fleet.clone()); + Ok(fleet) + } + + async fn spawn_host(&self, index: usize) -> Result { + let host = ClientHost::spawn( + self.url.clone(), + format!("{}_{index}", self.name), + self.pool_size, + self.worker_threads, + self.factory.clone(), + ) + .await?; + for (queue, workers) in &self.served { + queue.capacity.fetch_add(*workers, Ordering::Relaxed); + } + Ok(host) + } + + fn release_capacity(&self) { + for (queue, workers) in &self.served { + queue.capacity.fetch_sub(*workers, Ordering::Relaxed); + } + } + + pub async fn replicas(&self) -> usize { + self.hosts.lock().await.len() + } + + pub async fn live_replicas(&self) -> Vec { + self.hosts + .lock() + .await + .iter() + .enumerate() + .filter_map(|(index, host)| host.is_some().then_some(index)) + .collect() + } + + pub async fn pool_of(&self, index: usize) -> Option { + self.hosts + .lock() + .await + .get(index)? + .as_ref() + .map(|host| host.pool().clone()) + } + + /// Stops runtime `index` gracefully and starts a fresh one in its place. + pub async fn restart(&self, index: usize, timeout: Duration) -> Result { + let old = self + .hosts + .lock() + .await + .get_mut(index) + .and_then(Option::take); + let Some(old) = old else { return Ok(false) }; + self.release_capacity(); + old.shutdown(timeout).await; + let fresh = self.spawn_host(index).await?; + self.hosts.lock().await[index] = Some(fresh); + Ok(true) + } + + /// Kills runtime `index` with no shutdown handshake, then (optionally) starts a + /// replacement after `downtime`. + pub async fn kill(&self, index: usize, downtime: Option) -> Result { + let old = self + .hosts + .lock() + .await + .get_mut(index) + .and_then(Option::take); + let Some(old) = old else { return Ok(false) }; + self.release_capacity(); + old.kill(); + if let Some(downtime) = downtime { + tokio::time::sleep(downtime).await; + let fresh = self.spawn_host(index).await?; + self.hosts.lock().await[index] = Some(fresh); + } + Ok(true) + } + + pub async fn shutdown_all(&self, timeout: Duration, shared: &Shared) { + let hosts: Vec = self + .hosts + .lock() + .await + .iter_mut() + .filter_map(Option::take) + .collect(); + for host in hosts { + self.release_capacity(); + host.shutdown(timeout).await; + } + shared.budget.release(&self.budget_label); + } +} diff --git a/awa/tests/persona_soak/ledger.rs b/awa/tests/persona_soak/ledger.rs new file mode 100644 index 000000000..a1c8ca236 --- /dev/null +++ b/awa/tests/persona_soak/ledger.rs @@ -0,0 +1,266 @@ +//! The client-side ledger: every job a persona enqueued, what the workers saw, and +//! what the system reported back. The database may prune terminal rows long before +//! the run ends, so the ledger is the oracle's memory of outcomes. + +use super::config::Persona; +use std::collections::HashMap; +use std::sync::Mutex; + +const SHARDS: usize = 64; + +/// NO_EVENTS workloads keep their fast path; durable loss detection covers +/// this deterministic sample, chosen from the immutable assigned job ID. +pub const TERMINAL_PROOF_STRIDE: i64 = 64; + +pub fn samples_terminal_proof(job_id: i64) -> bool { + job_id.rem_euclid(TERMINAL_PROOF_STRIDE) == 0 +} + +/// What a job is allowed to end as. +pub mod expect { + pub const COMPLETED: u8 = 1; + /// Failed, which with the DLQ enabled means a DLQ row. + pub const FAILED: u8 = 2; + pub const CANCELLED: u8 = 4; + pub const ANY_TERMINAL: u8 = COMPLETED | FAILED | CANCELLED; +} + +/// Facts about a job, accumulated as the run proceeds. +pub mod flag { + /// An operator cancel, drain or purge may legitimately have removed the job. + pub const OPERATOR_REMOVED: u32 = 1 << 0; + pub const OPERATOR_CANCEL_OK: u32 = 1 << 1; + pub const OPERATOR_MOVED: u32 = 1 << 2; + pub const DLQ_RETRIED: u32 = 1 << 3; + pub const HANDLER_CANCELLED: u32 = 1 << 4; + pub const RESCUED: u32 = 1 << 5; + pub const EXHAUSTED: u32 = 1 << 6; + pub const COMPLETED_EVENT: u32 = 1 << 7; + pub const RETURNED_COMPLETED: u32 = 1 << 8; + /// The job is served by a firehose queue where lifecycle events are not collected. + pub const NO_EVENTS: u32 = 1 << 9; + /// First seen at execution time (cron-created); not enqueued by a producer. + pub const UNREGISTERED: u32 = 1 << 10; + /// A DLQ row for this job was purged by an operator. + pub const DLQ_PURGED: u32 = 1 << 11; + pub const CHAOS_VICTIM: u32 = 1 << 12; + /// Terminal rows read from the store for jobs whose runtime reports no lifecycle events. + pub const DB_SEEN_COMPLETED: u32 = 1 << 13; + pub const DB_SEEN_FAILED: u32 = 1 << 14; + pub const DB_SEEN_CANCELLED: u32 = 1 << 15; + /// How a firehose job was enqueued. + pub const SRC_BURST: u32 = 1 << 16; + pub const SRC_COPY: u32 = 1 << 17; + pub const SRC_STEADY: u32 = 1 << 18; +} + +#[derive(Debug, Clone)] +pub struct JobRecord { + pub persona: Persona, + pub queue: u32, + pub enqueued_ms: u32, + pub expect: u8, + pub flags: u32, + pub executions: u16, + pub in_flight: u8, + pub overlapped: bool, + pub max_attempt: i16, + pub latest_run_lease: i64, + pub terminal_run_lease: Option, + pub first_start_ms: u32, + pub last_start_ms: u32, + pub completed_events: u8, + pub exhausted_events: u8, + pub cancel_events: u8, + pub completed_event_ms: u32, + /// When an operator cancel on this job committed, 0 if none. + pub operator_cancel_ms: u32, +} + +impl JobRecord { + pub fn new(persona: Persona, queue: u32, enqueued_ms: u32, expect: u8, flags: u32) -> Self { + Self { + persona, + queue, + enqueued_ms, + expect, + flags, + executions: 0, + in_flight: 0, + overlapped: false, + max_attempt: 0, + latest_run_lease: 0, + terminal_run_lease: None, + first_start_ms: 0, + last_start_ms: 0, + completed_events: 0, + exhausted_events: 0, + cancel_events: 0, + completed_event_ms: 0, + operator_cancel_ms: 0, + } + } + + pub fn note_terminal(&mut self, run_lease: i64) { + if self.terminal_run_lease.is_none_or(|old| run_lease > old) { + // DB outcome bits belong to the newest proven attempt. A later + // lifecycle proof also supersedes an older observed DB state. + self.flags &= + !(flag::DB_SEEN_COMPLETED | flag::DB_SEEN_FAILED | flag::DB_SEEN_CANCELLED); + self.terminal_run_lease = Some(run_lease); + } + } + + pub fn note_observed_terminal(&mut self, run_lease: i64, state: &str) { + self.note_terminal(run_lease); + if self.terminal_run_lease == Some(run_lease) { + // The query is unordered: an older row can follow a newer proof. + // Different states at the same lease remain visible as a conflict. + self.flags |= match state { + "completed" => flag::DB_SEEN_COMPLETED, + "failed" => flag::DB_SEEN_FAILED, + "cancelled" => flag::DB_SEEN_CANCELLED, + _ => 0, + }; + } + } + + pub fn has_durable_terminal(&self) -> bool { + self.terminal_run_lease + .is_some_and(|lease| lease >= self.latest_run_lease) + } + + pub fn has(&self, flag: u32) -> bool { + self.flags & flag != 0 + } +} + +pub struct Ledger { + shards: Vec>>, + /// `(cron name, fire time)` -> job ids executed for it; more than one is a fencing failure. + cron_fires: Mutex>>, +} + +impl Default for Ledger { + fn default() -> Self { + Self::new() + } +} + +impl Ledger { + pub fn new() -> Self { + Self { + shards: (0..SHARDS).map(|_| Mutex::new(HashMap::new())).collect(), + cron_fires: Mutex::new(HashMap::new()), + } + } + + fn shard(&self, id: i64) -> &Mutex> { + &self.shards[(id as u64 % SHARDS as u64) as usize] + } + + /// Records an enqueue. A fast worker can run the job before its producer gets here, so + /// an existing record keeps what the workers already reported. + pub fn register(&self, id: i64, record: JobRecord) { + let mut shard = self.shard(id).lock().expect("ledger"); + match shard.get_mut(&id) { + Some(existing) => { + existing.persona = record.persona; + existing.queue = record.queue; + existing.enqueued_ms = record.enqueued_ms; + existing.expect = record.expect; + existing.flags = (existing.flags | record.flags) & !flag::UNREGISTERED; + } + None => { + shard.insert(id, record); + } + } + } + + /// Applies `update` to the job's record, creating an unregistered record when the + /// job was not enqueued by a producer (cron fires, operator moves). + pub fn with( + &self, + id: i64, + default: impl FnOnce() -> JobRecord, + update: impl FnOnce(&mut JobRecord) -> R, + ) -> R { + let mut shard = self.shard(id).lock().expect("ledger"); + let record = shard.entry(id).or_insert_with(|| { + let mut record = default(); + record.flags |= flag::UNREGISTERED; + record + }); + update(record) + } + + pub fn mark(&self, id: i64, flags: u32) { + if let Some(record) = self.shard(id).lock().expect("ledger").get_mut(&id) { + record.flags |= flags; + } + } + + pub fn mark_many(&self, ids: &[i64], flags: u32) { + for id in ids { + self.mark(*id, flags); + } + } + + /// Adds `flags` to every record for which `predicate` holds. + pub fn mark_where(&self, predicate: impl Fn(&JobRecord) -> bool, flags: u32) -> usize { + let mut marked = 0; + for shard in &self.shards { + for record in shard.lock().expect("ledger").values_mut() { + if predicate(record) { + record.flags |= flags; + marked += 1; + } + } + } + marked + } + + pub fn get(&self, id: i64) -> Option { + self.shard(id).lock().expect("ledger").get(&id).cloned() + } + + pub fn len(&self) -> usize { + self.shards + .iter() + .map(|shard| shard.lock().expect("ledger").len()) + .sum() + } + + /// Calls `visit` for every record. Shards are locked one at a time. + pub fn for_each(&self, mut visit: impl FnMut(i64, &JobRecord)) { + for shard in &self.shards { + let shard = shard.lock().expect("ledger"); + for (id, record) in shard.iter() { + visit(*id, record); + } + } + } + + pub fn note_cron_fire(&self, name: String, fire_time: String, job_id: i64) -> bool { + let mut fires = self.cron_fires.lock().expect("cron fires"); + let ids = fires.entry((name, fire_time)).or_default(); + if !ids.contains(&job_id) { + ids.push(job_id); + } + ids.len() == 1 + } + + pub fn cron_duplicates(&self) -> Vec<((String, String), Vec)> { + self.cron_fires + .lock() + .expect("cron fires") + .iter() + .filter(|(_, ids)| ids.len() > 1) + .map(|(key, ids)| (key.clone(), ids.clone())) + .collect() + } + + pub fn cron_fire_count(&self) -> usize { + self.cron_fires.lock().expect("cron fires").len() + } +} diff --git a/awa/tests/persona_soak/mod.rs b/awa/tests/persona_soak/mod.rs new file mode 100644 index 000000000..2bccae9a0 --- /dev/null +++ b/awa/tests/persona_soak/mod.rs @@ -0,0 +1,24 @@ +//! Persona-driven soak harness with an invariant oracle. See the test binary +//! `persona_soak_test.rs` for how to run it. + +pub mod abandon_oracle; +pub mod adminserver; +pub mod config; +pub mod db; +pub mod fleet; +pub mod ledger; +pub mod oracle; +pub mod personas; +pub mod pgserver; +pub mod python_probe; +pub mod registry; +pub mod report; +pub mod ring_locks; +pub mod rng; +pub mod roles; +pub mod run; +pub mod shard_oracle; +pub mod shared; +pub mod stop; +pub mod terminal_proof; +pub mod worker_errors; diff --git a/awa/tests/persona_soak/oracle.rs b/awa/tests/persona_soak/oracle.rs new file mode 100644 index 000000000..5c6840c3a --- /dev/null +++ b/awa/tests/persona_soak/oracle.rs @@ -0,0 +1,1393 @@ +//! The invariant oracle: periodic checks while the load runs and a final +//! reconciliation of the client ledger against the database. + +use super::ledger::{expect, flag, samples_terminal_proof, JobRecord}; +use super::pgserver::PgWatch; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +const MAX_ROWS: usize = 25; + +/// count, min and max job id, min and max lane sequence. +type ForensicRow = (i64, Option, Option, Option, Option); + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Every terminal outcome the store records, one row per `(job, attempt)`. +fn terminal_cte(schema: &str) -> String { + format!( + "terminal AS ( + SELECT job_id, run_lease, 'done:' || state::text AS src FROM {schema}.done_entries + UNION ALL SELECT job_id, run_lease, 'dlq' FROM {schema}.dlq_entries + UNION ALL SELECT b.job_id, b.run_lease, 'receipt_completed' + FROM {schema}.receipt_completion_batches c, + LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + )" + ) +} + +/// Jobs the lane cursors still consider claimable: at or after the claim cursor and +/// not tombstoned. +fn available_cte(schema: &str) -> String { + format!( + "heads AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS claim_seq + FROM {schema}.queue_claim_heads + ), + available AS MATERIALIZED ( + SELECT r.job_id, r.run_lease, r.queue + FROM {schema}.ready_entries r + JOIN heads h ON h.queue = r.queue AND h.priority = r.priority AND h.enqueue_shard = r.enqueue_shard + WHERE r.lane_seq >= h.claim_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones t + WHERE t.ready_slot = r.ready_slot AND t.ready_generation = r.ready_generation + AND t.queue = r.queue AND t.priority = r.priority + AND t.enqueue_shard = r.enqueue_shard AND t.lane_seq = r.lane_seq + ) + )" + ) +} + +#[derive(Default)] +struct StallTracker { + /// queue -> (claim cursor, millis when the cursor last moved, already reported) + cursors: HashMap, +} + +pub struct Oracle { + shared: Arc, + pool: PgPool, + watch: Arc, + stall: StallTracker, + last_cancel_scan_ms: i64, + reported_duplicates: std::collections::HashSet<(i64, i64)>, + reported_cancel: std::collections::HashSet, + initial_stats_reset: std::sync::Mutex>>, +} + +impl Oracle { + pub fn new(shared: Arc, pool: PgPool, watch: Arc) -> Self { + Self { + shared, + pool, + watch, + stall: StallTracker::default(), + last_cancel_scan_ms: 0, + reported_duplicates: Default::default(), + reported_cancel: Default::default(), + initial_stats_reset: Default::default(), + } + } + + /// Runs the periodic checks until `stop`. + pub async fn run(mut self, stop: Stop) -> Self { + let mut tick = 0u64; + loop { + if stop.sleep(Duration::from_secs(5)).await { + return self; + } + tick += 1; + self.probe_server().await; + self.check_stall().await; + self.check_promotion().await; + if tick.is_multiple_of(2) { + self.check_terminal_uniqueness().await; + self.check_unexplained_cancels().await; + } + if tick.is_multiple_of(3) { + self.sample_resources().await; + } + } + } + + /// Crash recovery discards the cumulative statistics, so a changed reset time means the + /// server restarted its processes even though the postmaster, and its start time, survived. + async fn probe_stats_reset(&self) { + let reset: Result, _> = + sqlx::query_scalar("SELECT stats_reset::text FROM pg_stat_bgwriter") + .fetch_one(&self.pool) + .await; + let reset = match reset { + Ok(reset) => reset, + Err(error) => { + self.shared + .incomplete(format!("crash-recovery probe failed: {error}"), vec![]); + return; + } + }; + let mut initial = self.initial_stats_reset.lock().expect("stats reset"); + match initial.as_ref() { + None => *initial = Some(reset), + Some(first) if *first != reset => { + let message = format!("pg_stat_bgwriter.stats_reset changed from {first:?} to {reset:?}: the server went through crash recovery"); + *initial = Some(reset); + drop(initial); + self.shared + .violation("server_crash_recovery", message, vec![]); + } + Some(_) => {} + } + } + + async fn probe_server(&self) { + self.probe_stats_reset().await; + let mut last_error = String::new(); + for attempt in 0..4 { + match PgWatch::postmaster_start(&self.pool).await { + Ok(start) if start == self.watch.started_at => return, + Ok(start) => { + self.shared.violation( + "server_restart", + format!( + "pg_postmaster_start_time changed from {} to {start}", + self.watch.started_at + ), + vec![], + ); + return; + } + Err(error) => { + last_error = error.to_string(); + tokio::time::sleep(Duration::from_secs(1 + attempt)).await; + } + } + } + self.shared.run.add("server_probe_failures", 1); + self.shared.incomplete( + format!("server probe failed repeatedly: {last_error}"), + vec![], + ); + } + + async fn check_terminal_uniqueness(&mut self) { + let schema = schema(&self.shared); + let terminal = terminal_cte(&schema); + let started = std::time::Instant::now(); + + let duplicates: Result)>, _> = + sqlx::query_as(awa::audited_sql(format!( + "WITH {terminal} + SELECT job_id, run_lease, array_agg(src ORDER BY src) + FROM terminal GROUP BY job_id, run_lease HAVING count(*) > 1 + ORDER BY job_id LIMIT {MAX_ROWS}" + ))) + .fetch_all(&self.pool) + .await; + match duplicates { + Ok(rows) => { + let fresh: Vec<_> = rows + .iter() + .filter(|(job_id, run_lease, _)| { + self.reported_duplicates.insert((*job_id, *run_lease)) + }) + .collect(); + if !fresh.is_empty() { + self.shared.violation( + "duplicate_terminal_rows", + format!("{} job attempt(s) have more than one terminal row", fresh.len()), + fresh + .iter() + .map(|(job_id, run_lease, sources)| { + format!("{{\"job_id\":{job_id},\"run_lease\":{run_lease},\"terminal_rows\":{sources:?}}}") + }) + .collect(), + ); + } + } + Err(error) => self + .shared + .incomplete(format!("terminal uniqueness query failed: {error}"), vec![]), + } + + let available = available_cte(&schema); + let overlap: Result, _> = sqlx::query_as(awa::audited_sql(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(awa::audited_sql(format!( + "WITH {terminal}, {available} + SELECT t.job_id, 'available'::text FROM terminal t JOIN available a ON a.job_id = t.job_id AND a.run_lease <= t.run_lease WHERE t.job_id = ANY($1) + UNION ALL SELECT t.job_id, 'deferred' FROM terminal t JOIN {schema}.deferred_jobs d ON d.job_id = t.job_id AND d.run_lease <= t.run_lease WHERE t.job_id = ANY($1) + UNION ALL SELECT t.job_id, 'leased' FROM terminal t JOIN {schema}.leases l ON l.job_id = t.job_id AND l.run_lease <= t.run_lease WHERE t.job_id = ANY($1)" + ))) + .bind(&ids) + .fetch_all(&self.pool) + .await; + let confirmed: std::collections::HashSet = match still { + Ok(rows) => rows.into_iter().map(|row| row.0).collect(), + Err(error) => { + self.shared.incomplete( + format!("terminal/live confirmation failed: {error}"), + vec![], + ); + Default::default() + } + }; + let confirmed_rows: Vec = rows + .iter() + .filter(|row| confirmed.contains(&row.0)) + .filter(|row| self.reported_duplicates.insert((row.0, -row.2 - 1))) + .map(|(job_id, place, terminal_lease, live_lease)| { + format!("{{\"job_id\":{job_id},\"also_in\":\"{place}\",\"terminal_run_lease\":{terminal_lease},\"live_run_lease\":{live_lease}}}") + }) + .collect(); + if !confirmed_rows.is_empty() { + self.shared.violation( + "terminal_and_live", + format!( + "{} job(s) are terminal and still available, deferred or leased", + confirmed_rows.len() + ), + confirmed_rows, + ); + } + } + Ok(_) => {} + Err(error) => self + .shared + .incomplete(format!("terminal/live query failed: {error}"), vec![]), + } + self.shared.run.observe( + "oracle_terminal_check_ms", + started.elapsed().as_millis() as i64, + ); + } + + /// A `cancelled` terminal row nobody asked for: no operator cancel, drain or purge + /// touched the job and its handler did not cancel it. + async fn check_unexplained_cancels(&mut self) { + let schema = schema(&self.shared); + let since_ms = self.last_cancel_scan_ms; + let scan_started_ms = self.shared.now_ms(); + let rows: Result, _> = sqlx::query_as(awa::audited_sql(format!( + "SELECT job_id, queue, run_lease FROM {schema}.done_entries + WHERE state = 'cancelled' AND finalized_at > now() - make_interval(secs => $1) + LIMIT 2000" + ))) + .bind(((scan_started_ms - since_ms) as f64 / 1000.0 + 15.0).max(20.0)) + .fetch_all(&self.pool) + .await; + self.last_cancel_scan_ms = scan_started_ms; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + let mut unexplained = Vec::new(); + for (job_id, queue_name, run_lease) in rows { + let Some(record) = self.shared.ledger.get(job_id) else { + continue; + }; + let queue = self.shared.queues.get(&queue_name); + let drained = queue.as_ref().is_some_and(|queue| { + let drained_at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + drained_at > 0 && i64::from(record.enqueued_ms) <= drained_at + 1000 + }); + let explained = record.has(flag::OPERATOR_REMOVED) + || record.has(flag::HANDLER_CANCELLED) + || record.has(flag::OPERATOR_MOVED) + || drained; + if !explained { + // An operator flag is set just after its cancel commits; look again before accusing. + let lag = if self + .shared + .ledger + .get(job_id) + .is_some_and(|record| record.persona == super::config::Persona::Python) + { + 4500 + } else { + 1500 + }; + tokio::time::sleep(Duration::from_millis(lag)).await; + } + let explained = explained + || self.shared.ledger.get(job_id).is_some_and(|record| { + record.has(flag::OPERATOR_REMOVED) || record.has(flag::HANDLER_CANCELLED) + }); + if !explained && self.reported_cancel.insert(job_id) { + unexplained.push(format!( + "{{\"job_id\":{job_id},\"queue\":\"{queue_name}\",\"run_lease\":{run_lease},\"executions\":{},\"running_now\":{},\"completed_events\":{}}}", + record.executions, record.in_flight, record.completed_events + )); + } + } + if !unexplained.is_empty() { + self.shared.violation( + "unexplained_cancel", + format!( + "{} job(s) recorded as cancelled with no operator or handler cancel", + unexplained.len() + ), + unexplained.into_iter().take(MAX_ROWS).collect(), + ); + } + } + + async fn check_promotion(&self) { + let schema = schema(&self.shared); + let mut threshold = self.shared.config.promotion_overdue; + if self.shared.chaos_active() { + threshold += self.shared.config.chaos_slack; + } + let rows: Result, _> = sqlx::query_as(awa::audited_sql(format!( + "SELECT job_id, queue, extract(epoch FROM now() - run_at)::float8 + FROM {schema}.deferred_jobs + WHERE state IN ('scheduled', 'retryable') AND run_at < now() - make_interval(secs => $1) + ORDER BY run_at LIMIT {MAX_ROWS}" + ))) + .bind(threshold.as_secs_f64()) + .fetch_all(&self.pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + if rows.is_empty() { + return; + } + // Re-read so that a job promoted a moment ago is not reported. + tokio::time::sleep(Duration::from_secs(2)).await; + let ids: Vec = rows.iter().map(|row| row.0).collect(); + let still: Result, _> = sqlx::query_scalar(awa::audited_sql(format!( + "SELECT job_id FROM {schema}.deferred_jobs WHERE job_id = ANY($1) AND state IN ('scheduled','retryable')" + ))) + .bind(&ids) + .fetch_all(&self.pool) + .await; + let still = match still { + Ok(still) => still, + Err(error) => { + self.shared + .incomplete(format!("promotion confirmation failed: {error}"), vec![]); + return; + } + }; + let stuck: Vec<_> = rows.iter().filter(|row| still.contains(&row.0)).collect(); + if stuck.is_empty() { + return; + } + let worst = stuck.iter().map(|row| row.2).fold(0.0, f64::max); + self.shared + .run + .set_max("promotion_overdue_max_s", worst as i64); + self.shared.violation( + "promotion_stall", + format!( + "{} scheduled job(s) overdue beyond {}s, worst {worst:.1}s", + stuck.len(), + threshold.as_secs() + ), + stuck + .iter() + .map(|(job_id, queue, overdue)| { + format!( + "{{\"job_id\":{job_id},\"queue\":\"{queue}\",\"overdue_s\":{overdue:.1}}}" + ) + }) + .collect(), + ); + } + + async fn check_stall(&mut self) { + let schema = schema(&self.shared); + let rows: Result, _> = sqlx::query_as(awa::audited_sql(format!( + "SELECT c.queue, + COALESCE(sum({schema}.sequence_next_value(e.seq_name) - {schema}.sequence_next_value(c.seq_name)), 0)::bigint, + COALESCE(sum({schema}.sequence_next_value(c.seq_name)), 0)::bigint + FROM {schema}.queue_claim_heads c + JOIN {schema}.queue_enqueue_heads e + ON e.queue = c.queue AND e.priority = c.priority AND e.enqueue_shard = c.enqueue_shard + GROUP BY c.queue" + ))) + .fetch_all(&self.pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + let now = self.shared.now_ms(); + let mut window = self.shared.config.stall_window; + if self.shared.chaos_active() { + window += self.shared.config.chaos_slack; + } + let mut candidates = Vec::new(); + for (queue_name, backlog, cursor) in rows { + let entry = self + .stall + .cursors + .entry(queue_name.clone()) + .or_insert((cursor, now, false)); + if entry.0 != cursor { + *entry = (cursor, now, false); + continue; + } + let Some(queue) = self.shared.queues.get(&queue_name) else { + continue; + }; + let capacity = queue.capacity.load(std::sync::atomic::Ordering::Relaxed); + let in_flight = queue.in_flight.load(std::sync::atomic::Ordering::Relaxed); + let held = queue.is_held() || queue.rate_limited; + let idle_workers = capacity > 0 && in_flight < capacity; + if !entry.2 + && backlog > 0 + && idle_workers + && !held + && (now - entry.1) as u64 >= window.as_millis() as u64 + { + candidates.push((queue_name, backlog, cursor, entry.1, capacity, in_flight)); + } else if !idle_workers || held || backlog <= 0 { + // Busy or deliberately held: the stall clock restarts. + entry.1 = now; + } + } + for (queue_name, backlog, cursor, since, capacity, in_flight) in candidates { + let exact: Result = sqlx::query_scalar(awa::audited_sql(format!( + "WITH {} SELECT count(*)::bigint FROM available WHERE queue = $1", + available_cte(&schema) + ))) + .bind(&queue_name) + .fetch_one(&self.pool) + .await; + let available = match exact { + Ok(available) => available, + Err(error) => { + self.shared + .incomplete(format!("exact claimable-row check failed: {error}"), vec![]); + continue; + } + }; + if available > 0 { + if let Some(entry) = self.stall.cursors.get_mut(&queue_name) { + entry.2 = true; + } + self.shared.run.add("stalls", 1); + self.shared.violation( + "stall", + format!( + "queue {queue_name}: {available} available job(s), claim cursor frozen at {cursor} for {}s with {}/{} workers busy", + (now - since) / 1000, + in_flight, + capacity + ), + vec![format!( + "{{\"queue\":\"{queue_name}\",\"available\":{available},\"lane_backlog\":{backlog},\"claim_cursor\":{cursor},\"workers_busy\":{in_flight},\"worker_capacity\":{capacity}}}" + )], + ); + } else if let Some(entry) = self.stall.cursors.get_mut(&queue_name) { + entry.1 = now; + } + } + } + + /// Stops the run before it can fill the disk it shares with other work. + fn check_disk(&self) { + let output = std::process::Command::new("df") + .args(["--output=avail", "-BM"]) + .arg(&self.shared.config.report_dir) + .output(); + let Ok(output) = output else { return }; + let text = String::from_utf8_lossy(&output.stdout); + let Some(free_mb) = text + .lines() + .nth(1) + .and_then(|line| line.trim().trim_end_matches('M').parse::().ok()) + else { + return; + }; + self.shared.run.observe("free_disk_gb", free_mb / 1024); + if free_mb < self.shared.config.min_free_disk_gb * 1024 + && !self.shared.load_stop.is_stopped() + { + self.shared.violation( + "disk_low", + format!( + "only {} GB free on the report directory's filesystem; stopping the load phase", + free_mb / 1024 + ), + vec![], + ); + self.shared.load_stop.trigger(); + } + } + + async fn sample_resources(&self) { + self.check_disk(); + let connections: Result = sqlx::query_scalar( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE backend_type = 'client backend'", + ) + .fetch_one(&self.pool) + .await; + if let Ok(connections) = connections { + self.shared.run.set_max("peak_connections", connections); + self.shared.run.observe("connections", connections); + } + let size: Result = + sqlx::query_scalar("SELECT pg_database_size(current_database())") + .fetch_one(&self.pool) + .await; + if let Ok(size) = size { + self.shared.run.set_max("peak_db_mb", size / (1024 * 1024)); + self.shared.run.observe("db_mb", size / (1024 * 1024)); + } + if let Ok(status) = std::fs::read_to_string("/proc/self/status") { + if let Some(kb) = status + .lines() + .find_map(|line| line.strip_prefix("VmRSS:")) + .and_then(|rest| { + rest.trim() + .trim_end_matches("kB") + .trim() + .parse::() + .ok() + }) + { + self.shared.run.set_max("peak_harness_rss_mb", kb / 1024); + } + } + if let Ok(load) = std::fs::read_to_string("/proc/loadavg") { + if let Some(one_minute) = load + .split_whitespace() + .next() + .and_then(|v| v.parse::().ok()) + { + self.shared + .run + .observe("loadavg_x100", (one_minute * 100.0) as i64); + } + } + } + + /// Waits for the system to quiesce after load stops. Returns true when no live + /// work remains. + pub async fn wait_quiescent(&mut self, timeout: Duration) -> bool { + let schema = schema(&self.shared); + let deadline = std::time::Instant::now() + timeout; + let mut quiet_polls = 0; + loop { + let live: Result = sqlx::query_scalar(awa::audited_sql(format!( + "WITH {} SELECT (SELECT count(*) FROM available) + (SELECT count(*) FROM {schema}.leases) + + (SELECT count(*) FROM {schema}.deferred_jobs WHERE run_at < now() + interval '1 second')", + available_cte(&schema) + ))) + .fetch_one(&self.pool) + .await; + if let Err(error) = &live { + self.shared + .incomplete(format!("quiescence query failed: {error}"), vec![]); + } + let mut pending_terminal_proof = 0usize; + self.shared.ledger.for_each(|id, record| { + if record.has(flag::RETURNED_COMPLETED) + && !record.has_durable_terminal() + && !record.has(flag::OPERATOR_REMOVED) + && !record.has(flag::DLQ_PURGED) + && (!record.has(flag::NO_EVENTS) || samples_terminal_proof(id)) + { + pending_terminal_proof += 1; + } + }); + let in_flight: i64 = self + .shared + .queues + .all() + .iter() + .map(|queue| queue.in_flight.load(std::sync::atomic::Ordering::Relaxed)) + .sum(); + let python_active = self + .shared + .python_active + .load(std::sync::atomic::Ordering::Relaxed); + // Claims abandoned by a killed process sit in no table the live count reads. + let abandoned_open = self + .shared + .abandon + .open_claims + .load(std::sync::atomic::Ordering::Relaxed); + match live { + Ok(count) + if in_flight == 0 + && python_active == 0 + && abandoned_open == 0 + && pending_terminal_proof == 0 + && count_is_idle(count) => + { + quiet_polls += 1 + } + _ => quiet_polls = 0, + } + if quiet_polls >= 3 { + return true; + } + if std::time::Instant::now() >= deadline { + self.shared.note( + "oracle", + format!( + "not quiescent after {}s: handlers running={in_flight}, live rows={live:?}", + timeout.as_secs() + ), + ); + return false; + } + self.probe_server().await; + self.check_stall().await; + tokio::time::sleep(Duration::from_secs(2)).await; + } + } + + /// Final reconciliation of the ledger against the database. + pub async fn final_check(&mut self, quiesced: bool) { + let schema = schema(&self.shared); + self.check_terminal_uniqueness().await; + self.check_unexplained_cancels().await; + self.check_promotion().await; + self.probe_server().await; + + self.shared.run.add( + "cron_fires_executed", + self.shared.ledger.cron_fire_count() as i64, + ); + for ((name, fire), ids) in self.shared.ledger.cron_duplicates() { + self.shared.violation( + "cron_double_fire", + format!("cron {name} fire {fire} produced jobs {ids:?}"), + vec![], + ); + } + + // Jobs with no terminal evidence on the client side have to be found in the database. + let mut unresolved: Vec = Vec::new(); + let mut unexpected: Vec = Vec::new(); + let mut sampled = 0i64; + let mut confirmed = 0i64; + let mut excluded = 0i64; + self.shared.ledger.for_each(|id, record| { + if record.has(flag::OPERATOR_REMOVED) || record.has(flag::DLQ_PURGED) { + return; + } + let sampled_no_events = record.has(flag::NO_EVENTS) && samples_terminal_proof(id); + sampled += i64::from(sampled_no_events); + if record.has_durable_terminal() { + confirmed += i64::from(sampled_no_events); + for outcome in unexpected_terminal_outcomes(record) { + if unexpected.len() >= MAX_ROWS { + break; + } + let queue = self.shared.queues.by_index(record.queue); + unexpected.push(serde_json::json!({ + "job_id": id, + "outcome": outcome, + "allowed": record.expect, + "persona": record.persona.name(), + "queue": queue.as_ref().map(|queue| queue.name.chars().take(256).collect::()), + "flags": record.flags, + "latest_run_lease": record.latest_run_lease, + "terminal_run_lease": record.terminal_run_lease, + "executions": record.executions, + "in_flight": record.in_flight, + "max_attempt": record.max_attempt, + "completed_events": record.completed_events, + "exhausted_events": record.exhausted_events, + "cancel_events": record.cancel_events, + "enqueued_ms": record.enqueued_ms, + "first_start_ms": record.first_start_ms, + "last_start_ms": record.last_start_ms, + "completed_event_ms": record.completed_event_ms, + "operator_cancel_ms": record.operator_cancel_ms, + "queue_drained_at_ms": queue.map(|queue| queue.drained_at_ms.load(std::sync::atomic::Ordering::Relaxed)), + }).to_string()); + } + } else if record.has(flag::NO_EVENTS) + && record.has(flag::RETURNED_COMPLETED) + && !sampled_no_events + { + // This optimized workload is sampled, not proven per job. + // A return is execution evidence only, never commit evidence. + excluded += 1; + } else { + unresolved.push(id); + } + }); + if !unexpected.is_empty() { + self.shared.violation( + "unexpected_outcome", + format!( + "{} job(s) ended in an outcome their persona did not allow", + unexpected.len() + ), + unexpected, + ); + } + + let mut live_stuck = Vec::new(); + let mut forensic_ids: Vec = Vec::new(); + let mut lost = Vec::new(); + let mut lost_after_run = Vec::new(); + let mut terminal_in_db = 0usize; + for chunk in unresolved.chunks(5000) { + let found: Result)>, _> = sqlx::query_as(awa::audited_sql(format!( + "SELECT job_id, string_agg(place, ','), max(terminal_run_lease) FROM ( + SELECT job_id, 'done:' || state::text AS place, run_lease AS terminal_run_lease FROM {schema}.done_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'dlq', run_lease FROM {schema}.dlq_entries WHERE job_id = ANY($1) + UNION ALL SELECT b.job_id, 'receipt_completed', b.run_lease + FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + WHERE c.job_ids && $1 AND b.job_id = ANY($1) + UNION ALL SELECT job_id, 'deferred:' || state::text, NULL::bigint FROM {schema}.deferred_jobs WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'ready', NULL::bigint FROM {schema}.ready_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'leased:' || state::text, NULL::bigint FROM {schema}.leases WHERE job_id = ANY($1) + ) places GROUP BY job_id" + ))) + .bind(chunk) + .fetch_all(&self.pool) + .await; + let found: HashMap)> = match found { + Ok(rows) => rows + .into_iter() + .map(|(id, place, lease)| (id, (place, lease))) + .collect(), + Err(error) => { + self.shared.run.add("final_lookup_errors", 1); + self.shared.incomplete( + format!( + "final lookup failed for {} unresolved jobs: {error}", + chunk.len() + ), + vec![], + ); + continue; + } + }; + for id in chunk { + let record = self.shared.ledger.get(*id); + let describe = |place: &str| { + let record = record.clone(); + format!( + "{{\"job_id\":{id},\"found\":\"{place}\",\"executions\":{},\"persona\":\"{}\",\"flags\":{},\"enqueued_at_s\":{:.1},\"queue\":\"{}\"}}", + record.as_ref().map_or(0, |r| r.executions), + record.as_ref().map_or("?", |r| r.persona.name()), + record.as_ref().map_or(0, |r| r.flags), + record.as_ref().map_or(0.0, |r| f64::from(r.enqueued_ms) / 1000.0), + record + .as_ref() + .and_then(|r| self.shared.queues.by_index(r.queue)) + .map_or("?".to_string(), |queue| queue.name.clone()) + ) + }; + if record + .as_ref() + .is_some_and(|record| record.has_durable_terminal()) + { + confirmed += i64::from( + record + .as_ref() + .is_some_and(|record| record.has(flag::NO_EVENTS)) + && samples_terminal_proof(*id), + ); + continue; + } + match found.get(id) { + Some((place, terminal_lease)) + if terminal_lease.is_some_and(|lease| { + lease >= record.as_ref().map_or(0, |record| record.latest_run_lease) + }) => + { + terminal_in_db += 1; + confirmed += i64::from( + record + .as_ref() + .is_some_and(|record| record.has(flag::NO_EVENTS)) + && samples_terminal_proof(*id), + ); + if place.contains("done:cancelled") { + if let Some(record) = &record { + let queue = self.shared.queues.by_index(record.queue); + let drained = queue.is_some_and(|queue| { + let at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + at > 0 && i64::from(record.enqueued_ms) <= at + 1000 + }); + if !(record.has(flag::OPERATOR_REMOVED) || drained) + && record.expect & expect::CANCELLED == 0 + { + lost_after_run.push(describe(place)); + forensic_ids.push(*id); + } + } + } + } + Some((place, _)) => live_stuck.push(describe(place)), + None => { + let executed = record.as_ref().is_some_and(|r| r.executions > 0); + let returned = record + .as_ref() + .is_some_and(|r| r.has(flag::RETURNED_COMPLETED)); + let drained = record.as_ref().is_some_and(|r| { + self.shared.queues.by_index(r.queue).is_some_and(|queue| { + let at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + at > 0 && i64::from(r.enqueued_ms) <= at + 1000 + }) + }); + if drained { + continue; + } + if !executed { + lost.push(describe("absent")); + forensic_ids.push(*id); + } else if !returned { + lost_after_run.push(describe("absent")); + forensic_ids.push(*id); + } else { + self.shared.incomplete( + "a returned handler has no durable terminal proof; a missed pruned outcome cannot be distinguished from a lost finalize", + vec![describe("absent without durable terminal evidence")], + ); + } + } + } + } + } + if !forensic_ids.is_empty() { + self.record_forensics(&schema, &forensic_ids).await; + } + self.shared + .run + .add("terminal_proof_no_events_sampled", sampled); + self.shared + .run + .add("terminal_proof_no_events_confirmed", confirmed); + self.shared + .run + .add("terminal_proof_no_events_excluded", excluded); + self.shared + .run + .add("final_unresolved_ids", unresolved.len() as i64); + self.shared + .run + .add("final_terminal_in_db", terminal_in_db as i64); + let report = |kind: &str, what: &str, rows: Vec| { + if !rows.is_empty() { + let count = rows.len(); + self.shared.violation( + kind, + format!("{count} job(s) {what}"), + rows.into_iter().take(MAX_ROWS).collect(), + ); + } + }; + if quiesced { + report( + "stuck_job", + "never reached a terminal state although the system went quiet (available/deferred/leased)", + live_stuck, + ); + } else if !live_stuck.is_empty() { + // Work still queued when the grace ended is backlog, not loss: the run was + // sized beyond what the server could clear. + self.shared.finding( + "backlog_at_end", + format!( + "{} job(s) were still queued when the drain grace ended", + live_stuck.len() + ), + live_stuck.into_iter().take(MAX_ROWS).collect(), + ); + } + report( + "lost_job", + "were enqueued, never ran and are in no table", + lost, + ); + report("lost_job_after_run", "ran without a recorded terminal outcome and are in no table, or ended cancelled without cause", lost_after_run); + } + + /// For jobs that look lost: everywhere the store could still mention them, written to + /// the timeline so the reproducer carries it. + async fn record_forensics(&self, schema: &str, ids: &[i64]) { + let sample: Vec = ids.iter().copied().take(200).collect(); + let tables: [(&str, String); 8] = [ + ("ready_entries", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.ready_entries WHERE job_id = ANY($1)")), + ("ready_tombstones", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.ready_tombstones WHERE job_id = ANY($1)")), + ("lease_claims", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.lease_claims WHERE job_id = ANY($1)")), + ("lease_claim_batches", format!("SELECT count(*), min(j), max(j), 0::bigint, 0::bigint FROM {schema}.lease_claim_batches b, LATERAL unnest(b.job_ids) AS j WHERE b.job_ids && $1 AND j = ANY($1)")), + ("lease_claim_closures", format!("SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM {schema}.lease_claim_closures WHERE job_id = ANY($1)")), + ("attempt_state", format!("SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM {schema}.attempt_state WHERE job_id = ANY($1)")), + ("receipt_completion_tombstones", format!("SELECT count(*), 0::bigint, 0::bigint, 0::bigint, 0::bigint FROM {schema}.receipt_completion_tombstones")), + ("job_unique_claims", "SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM awa.job_unique_claims WHERE job_id = ANY($1)".to_string()), + ]; + let mut lines = Vec::new(); + for (name, sql) in tables { + let row: Result = sqlx::query_as(awa::audited_sql(&sql)) + .bind(&sample) + .fetch_one(&self.pool) + .await; + match row { + Ok((count, min_id, max_id, min_seq, max_seq)) => lines.push(format!("{name}: {count} row(s) for the sampled ids (job ids {min_id:?}..{max_id:?}, lane_seq {min_seq:?}..{max_seq:?})")), + Err(error) => lines.push(format!("{name}: query failed: {error}")), + } + } + self.shared.note( + "oracle", + format!( + "forensics for {} lost-looking job(s), first {:?}: {}", + ids.len(), + &ids[..ids.len().min(5)], + lines.join("; ") + ), + ); + } + + pub fn crash_scan(&self) -> Vec { + match self.watch.log_since_start(&self.shared.config) { + Some(log) => super::pgserver::crash_excerpts(&log), + None => Vec::new(), + } + } +} + +/// The final verdict's outcome predicates, shared by deterministic lease-order controls. +fn unexpected_terminal_outcomes(record: &JobRecord) -> Vec<&'static str> { + if !record.has_durable_terminal() { + return Vec::new(); + } + let completed = record.completed_events > 0 + || record.has(flag::RETURNED_COMPLETED) + || record.has(flag::DB_SEEN_COMPLETED); + [ + ( + "cancelled", + record.has(flag::DB_SEEN_CANCELLED) + && record.expect & expect::CANCELLED == 0 + && !record.has(flag::HANDLER_CANCELLED), + ), + ( + "completed", + completed && record.expect & expect::COMPLETED == 0, + ), + ( + "failed", + (record.exhausted_events > 0 || record.has(flag::DB_SEEN_FAILED)) + && record.expect & expect::FAILED == 0, + ), + ] + .into_iter() + .filter_map(|(outcome, unexpected)| unexpected.then_some(outcome)) + .collect() +} + +fn count_is_idle(count: i64) -> bool { + count <= 0 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::persona_soak::config::{replace_database_name, Persona, SoakConfig}; + use crate::persona_soak::db::{self, Budget}; + use awa::model::{ClaimedRuntimeJob, QueueStorage, QueueStorageConfig}; + use std::sync::atomic::Ordering; + + fn current_terminal_record() -> JobRecord { + let mut record = JobRecord::new(Persona::Tom, 0, 0, expect::COMPLETED, 0); + record.latest_run_lease = 1; + record + } + + #[test] + fn terminal_state_old_cancel_then_new_completion_is_completed() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "cancelled"); + record.note_observed_terminal(1, "completed"); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert_eq!(record.terminal_run_lease, Some(1)); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + } + + #[test] + fn terminal_state_new_completion_then_old_cancel_is_completed() { + let mut record = current_terminal_record(); + record.note_observed_terminal(1, "completed"); + record.note_observed_terminal(0, "cancelled"); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + } + + #[test] + fn terminal_state_lifecycle_completion_supersedes_old_db_cancel() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "cancelled"); + record.note_terminal(1); + record.completed_events += 1; + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + record.note_observed_terminal(0, "cancelled"); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + } + + #[test] + fn terminal_state_new_cancel_remains_an_unexpected_outcome() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "completed"); + record.note_observed_terminal(1, "cancelled"); + assert!(record.has(flag::DB_SEEN_CANCELLED)); + assert_eq!(unexpected_terminal_outcomes(&record), vec!["cancelled"]); + } + + #[test] + fn terminal_state_same_lease_conflicting_states_remain_visible() { + for states in [["cancelled", "completed"], ["completed", "cancelled"]] { + let mut record = current_terminal_record(); + for state in states { + record.note_observed_terminal(1, state); + } + assert!(record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + assert_eq!(unexpected_terminal_outcomes(&record), vec!["cancelled"]); + } + } + + struct TestRun { + shared: Arc, + pool: PgPool, + store: QueueStorage, + claim: ClaimedRuntimeJob, + oracle: Oracle, + } + + impl TestRun { + async fn create() -> Self { + Self::create_with(true, true).await + } + + async fn create_with(no_events: bool, sampled: bool) -> Self { + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_oracle_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + db::recreate_database(&config) + .await + .expect("create isolated oracle database"); + let pool = db::connect_pool(&config.database_url, "soak_oracle_regression", 4) + .await + .expect("connect oracle"); + let info = db::server_info(&pool).await.expect("server info"); + let watch = Arc::new(PgWatch::begin(&config, &pool).await.expect("watch server")); + let budget = Budget::new(&info, 0); + let shared = Shared::new(config, info, budget); + let queue = shared + .queues + .register("oracle_proof", Persona::Priya, false); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("store"); + store + .enqueue_batch(&pool, "oracle_proof", 2, 64) + .await + .expect("enqueue sample candidates"); + let claimed = store + .claim_runtime_batch(&pool, "oracle_proof", 64, Duration::ZERO) + .await + .expect("claim sample"); + let claim = claimed + .into_iter() + .find(|claim| samples_terminal_proof(claim.job.id) == sampled) + .expect("immutable job ID selects the requested sample status"); + shared.ledger.register( + claim.job.id, + JobRecord::new( + Persona::Priya, + queue.index, + 0, + expect::COMPLETED, + if no_events { flag::NO_EVENTS } else { 0 }, + ), + ); + // The handler returned, but its caller has not finalized the claim. + let execution = shared.begin_execution(Persona::Priya, &claim.job, no_events); + shared.ledger.mark(claim.job.id, flag::RETURNED_COMPLETED); + drop(execution); + let oracle = Oracle::new(shared.clone(), pool.clone(), watch); + Self { + shared, + pool, + store, + claim, + oracle, + } + } + + async fn drop(self) { + self.pool.close().await; + let admin = db::connect_pool(&self.shared.config.admin_url(), "oracle_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(awa::audited_sql(format!( + "DROP DATABASE {} WITH (FORCE)", + self.shared.config.database_name() + ))) + .execute(&admin) + .await + .expect("drop oracle database"); + admin.close().await; + } + } + + #[tokio::test] + async fn returned_handler_without_finalization_cannot_pass_the_oracle() { + let mut run = TestRun::create().await; + run.oracle.final_check(true).await; + assert!( + run.shared + .violations + .lock() + .expect("violations") + .iter() + .any(|violation| violation.kind == "stuck_job"), + "a handler return must not stand in for a terminal commit" + ); + run.drop().await; + } + + #[tokio::test] + async fn sampled_terminal_pruned_before_observation_is_inconclusive() { + let mut run = TestRun::create().await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + run.store + .reset(&run.pool) + .await + .expect("remove retained evidence"); + run.oracle.final_check(true).await; + assert!( + run.shared.oracle_incomplete.load(Ordering::Relaxed), + "an unobserved pruned result cannot establish a pass or a lost-job claim" + ); + run.drop().await; + } + + #[tokio::test] + async fn observed_terminal_survives_pruning_in_the_client_ledger() { + let mut run = TestRun::create().await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(run + .shared + .ledger + .get(run.claim.job.id) + .expect("sample record") + .has_durable_terminal()); + run.store + .reset(&run.pool) + .await + .expect("prune retained evidence"); + run.oracle.final_check(true).await; + assert!( + run.shared.violations.lock().expect("violations").is_empty(), + "captured durable proof must survive pruning" + ); + run.drop().await; + } + + #[tokio::test] + async fn eventful_observed_terminal_without_callback_survives_pruning() { + let mut run = TestRun::create_with(false, false).await; + assert!(!samples_terminal_proof(run.claim.job.id)); + let completed = run + .store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion without a lifecycle callback"); + assert_eq!(completed, vec![(run.claim.job.id, run.claim.job.run_lease)]); + let durable: bool = sqlx::query_scalar( + "SELECT EXISTS ( + SELECT 1 FROM awa.receipt_completion_batches AS batch, + LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + WHERE item.job_id = $1 AND item.run_lease = $2 + )", + ) + .bind(run.claim.job.id) + .bind(run.claim.job.run_lease) + .fetch_one(&run.pool) + .await + .expect("read committed completion independently of the observer"); + assert!(durable); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + let record = run + .shared + .ledger + .get(run.claim.job.id) + .expect("eventful record"); + assert!(record.has_durable_terminal()); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + assert_eq!(record.completed_events, 0); + run.store + .reset(&run.pool) + .await + .expect("prune retained evidence"); + run.oracle.final_check(true).await; + assert!(run.shared.violations.lock().expect("violations").is_empty()); + run.drop().await; + } + + #[tokio::test] + async fn eventful_return_without_finalization_cannot_gain_terminal_proof() { + let mut run = TestRun::create_with(false, false).await; + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.oracle.final_check(true).await; + assert!(run + .shared + .violations + .lock() + .expect("violations") + .iter() + .any(|violation| violation.kind == "stuck_job")); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_pruned_before_observation_stays_inconclusive() { + let mut run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + run.store + .reset(&run.pool) + .await + .expect("remove evidence before observation"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + run.oracle.final_check(true).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_observer_query_failure_stays_inconclusive() { + let run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("commit completion"); + sqlx::query("ALTER TABLE awa.receipt_completion_batches RENAME TO receipt_completion_batches_unavailable") + .execute(&run.pool).await.expect("inject observer query failure"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_observer_requires_the_exact_attempt() { + let run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("complete old attempt"); + run.shared.ledger.with( + run.claim.job.id, + || unreachable!("registered"), + |record| { + record.latest_run_lease += 1; + }, + ); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + let record = run.shared.ledger.get(run.claim.job.id).expect("record"); + assert!(!record.has_durable_terminal()); + assert!(!record.has(flag::DB_SEEN_COMPLETED)); + run.drop().await; + } + + #[tokio::test] + async fn unsampled_no_events_job_does_not_gain_periodic_terminal_proof() { + let run = TestRun::create_with(true, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("commit completion"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.drop().await; + } + + #[tokio::test] + async fn failed_final_lookup_cannot_omit_unchecked_jobs_and_pass() { + let mut run = TestRun::create().await; + sqlx::query("ALTER TABLE awa.ready_entries RENAME TO ready_entries_unavailable") + .execute(&run.pool) + .await + .expect("inject a lookup failure"); + run.oracle.final_check(true).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + assert_eq!( + run.shared.run.get("final_lookup_errors"), + 1, + "the failed lookup chunk must be accounted for" + ); + run.drop().await; + } +} diff --git a/awa/tests/persona_soak/personas/abandon.rs b/awa/tests/persona_soak/personas/abandon.rs new file mode 100644 index 000000000..2a5a00f6f --- /dev/null +++ b/awa/tests/persona_soak/personas/abandon.rs @@ -0,0 +1,1096 @@ +//! Claim and abandon: a runtime that claims as much as it can, holds thousands of jobs +//! with deadlines and heartbeats, and is SIGKILLed again and again. +//! +//! The runtime is a separate process: the test binary started again with only the +//! `claim_abandon_child` test selected, so that a kill leaves its claims genuinely +//! abandoned (no shutdown handshake, connections dropped by the kernel). Between kills a +//! surviving in-process runtime and the other personas keep working on the same +//! database. After the last kill an in-process runtime takes over the abandoned queues +//! and finishes whatever the rescue returns to them. +//! +//! The checks are in `abandon_oracle.rs`: how long the rescue backlog takes to clear +//! after each kill, whether the claim ring keeps rotating, whether any job was rescued +//! while its owner was alive, and whether attempts and the dead-letter queue follow +//! `max_attempts`. + +use super::super::abandon_oracle::{ + epoch_ms, probe_backlog, Attempt, Kill, LiveCheck, Scope, FINISHER, +}; +use super::super::config::{Persona, DEFAULT_HEARTBEAT_SECONDS, DEFAULT_STALENESS_SECONDS}; +use super::super::db::{connect_pool, connect_pool_with_timeout, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{backoff, queue, register_rows, timings, tune}; +use awa::model::{insert_many, insert_many_copy_from_pool}; +use awa::{ + Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, QueueConfig, Worker, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::{Child, Command}; + +/// Zero-deadline claims. +pub const QUEUE_PLAIN: &str = "awa_soak_abandon_plain"; +/// Claims with a deadline. Both queues' claims are written as compact claim batches. +pub const QUEUE_DEADLINE: &str = "awa_soak_abandon_deadline"; +/// Short jobs served by the survivors, to see claim latency while the backlog drains. +pub const QUEUE_CANARY: &str = "awa_soak_abandon_canary"; + +/// The name of the test the child process runs. +pub const CHILD_TEST: &str = "claim_abandon_child"; +const ENV_URL: &str = "AWA_SOAK_ABANDON_URL"; +const ENV_GENERATION: &str = "AWA_SOAK_ABANDON_GENERATION"; +const ENV_LEDGER: &str = "AWA_SOAK_ABANDON_LEDGER"; +const ENV_WORKERS: &str = "AWA_SOAK_ABANDON_WORKERS"; +const ENV_POOL: &str = "AWA_SOAK_ABANDON_POOL"; +const ENV_HEARTBEAT: &str = "AWA_SOAK_HEARTBEAT_SECONDS"; +const ENV_STALENESS: &str = "AWA_SOAK_STALENESS_SECONDS"; + +/// Longer than any process lives, so a job is only ever finished by the survivors. +const HOLD: Duration = Duration::from_secs(20 * 60); +const DEADLINE: Duration = Duration::from_secs(15 * 60); +const CHILD_POOL: u32 = 16; +const FINISHER_REPLICAS: usize = 2; +const FINISHER_POOL: u32 = 16; +/// Percentage of jobs that may only be attempted twice, so that repeated kills +/// dead-letter them. +const SHORT_RETRY_PERCENT: u64 = 25; +/// Consecutive heartbeat samples that may find no open claim while the process reports +/// running handlers. Covers a handler count read before the process's claims were visible. +const EMPTY_SAMPLES_BEFORE_INCOMPLETE: u32 = 5; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AbandonJob { + pub seq: u64, + pub cycle: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AbandonCanaryJob { + pub seq: u64, +} + +/// The worker in the child process: records every start in the child ledger and then +/// holds the job until the process is killed. +struct ChildWorker { + generation: u32, + ledger: Mutex, + active: Arc, + peak: Arc, +} + +#[async_trait::async_trait] +impl Worker for ChildWorker { + fn kind(&self) -> &'static str { + "abandon_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let line = format!( + "{{\"e\":\"start\",\"id\":{},\"attempt\":{},\"gen\":{},\"q\":\"{}\",\"t\":{}}}\n", + ctx.job.id, + ctx.job.attempt, + self.generation, + ctx.job.queue, + epoch_ms() + ); + // One write per start, straight to the file: nothing is buffered that a kill could + // lose, and the parent never sees a half line it cannot skip. + if let Ok(mut file) = self.ledger.lock() { + let _ = file.write_all(line.as_bytes()); + } + let now = self.active.fetch_add(1, Ordering::Relaxed) + 1; + self.peak.fetch_max(now, Ordering::Relaxed); + tokio::time::sleep(HOLD).await; + self.active.fetch_sub(1, Ordering::Relaxed); + Ok(JobResult::Completed) + } +} + +/// Entry point of the child process; see [`CHILD_TEST`]. Returns at once when the +/// environment is not set up by the parent, so a plain `--ignored` run is unaffected. +pub async fn run_child() { + let (Ok(url), Ok(generation), Ok(ledger), Ok(workers)) = ( + std::env::var(ENV_URL), + std::env::var(ENV_GENERATION), + std::env::var(ENV_LEDGER), + std::env::var(ENV_WORKERS), + ) else { + return; + }; + super::super::worker_errors::install(true).expect("child worker error tracing"); + let generation: u32 = generation.parse().expect("generation"); + let workers: u32 = workers.parse().expect("workers"); + let pool_size: u32 = std::env::var(ENV_POOL) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(CHILD_POOL); + + // Exit with the harness, so a harness that died leaves no process behind. + let parent = std::os::unix::process::parent_id(); + std::thread::spawn(move || loop { + std::thread::sleep(Duration::from_secs(1)); + if std::os::unix::process::parent_id() != parent { + std::process::exit(0); + } + }); + + let pool = connect_pool(&url, &format!("soak_abandon_{generation}"), pool_size) + .await + .expect("child connect"); + let file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&ledger) + .expect("child ledger"); + let active = Arc::new(AtomicI64::new(0)); + let peak = Arc::new(AtomicI64::new(0)); + let claiming = |deadline_duration: Duration| QueueConfig { + max_workers: workers, + poll_interval: Duration::from_millis(50), + deadline_duration, + claimers: 2, + claim_batch_size: 1_000, + ..QueueConfig::default() + }; + // The child keeps away from maintenance leadership: the survivors do the rescuing, as + // the other replicas of an application would. + let seconds = |name: &str, default: u64| { + Duration::from_secs( + std::env::var(name) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(default), + ) + }; + let client = timings( + Client::builder(pool) + .queue(QUEUE_PLAIN, claiming(Duration::ZERO)) + .queue(QUEUE_DEADLINE, claiming(DEADLINE)) + .register_worker(ChildWorker { + generation, + ledger: Mutex::new(file), + active: active.clone(), + peak: peak.clone(), + }), + seconds(ENV_HEARTBEAT, DEFAULT_HEARTBEAT_SECONDS), + seconds(ENV_STALENESS, DEFAULT_STALENESS_SECONDS), + ) + .leader_election_interval(Duration::from_secs(3600)) + .build() + .expect("build child client"); + client.start().await.expect("start child client"); + println!("READY gen={generation}"); + loop { + tokio::time::sleep(Duration::from_secs(2)).await; + println!( + "STAT gen={generation} active={} peak={}", + active.load(Ordering::Relaxed), + peak.load(Ordering::Relaxed) + ); + } +} + +/// Finishes the abandoned jobs once the child processes are gone. +struct FinisherWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for FinisherWorker { + fn kind(&self) -> &'static str { + "abandon_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ClaimAbandon, &ctx.job, true); + self.shared.abandon.note_attempt( + ctx.job.id, + Attempt { + attempt: ctx.job.attempt, + owner: FINISHER, + at_ms: epoch_ms(), + }, + ); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +struct CanaryWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for CanaryWorker { + fn kind(&self) -> &'static str { + "abandon_canary_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ClaimAbandon, &ctx.job, true); + let waited = (chrono::Utc::now() - ctx.job.created_at).num_milliseconds(); + self.shared + .stats(Persona::ClaimAbandon) + .observe("canary_wait_ms", waited); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +/// Copies the running test binary into the report directory. Every process is started +/// from the copy, so a rebuild of the binary during a long run cannot change what is +/// started, or leave nothing to start. +fn copy_executable(directory: &std::path::Path) -> Result { + let copy = directory.join("claim-abandon-child"); + let _ = std::fs::remove_file(©); + let source = std::path::Path::new("/proc/self/exe"); + let source = if source.exists() { + source.to_path_buf() + } else { + std::env::current_exe().map_err(|error| format!("current_exe: {error}"))? + }; + std::fs::copy(&source, ©) + .map_err(|error| format!("copy {} to {}: {error}", source.display(), copy.display()))?; + Ok(copy) +} + +fn spawn_child( + shared: &Arc, + executable: &PathBuf, + generation: u32, + ledger: &PathBuf, + workers_per_queue: u32, + pool_size: u32, +) -> Result { + let url = &shared.config.database_url; + let mut child = Command::new(executable) + .args([ + "--ignored", + "--exact", + CHILD_TEST, + "--nocapture", + "--test-threads=1", + ]) + .env(ENV_URL, url) + .env(ENV_GENERATION, generation.to_string()) + .env(ENV_LEDGER, ledger) + .env(ENV_WORKERS, workers_per_queue.to_string()) + .env(ENV_POOL, pool_size.to_string()) + .env( + ENV_HEARTBEAT, + shared.config.heartbeat_interval.as_secs().to_string(), + ) + .env( + ENV_STALENESS, + shared.config.heartbeat_staleness.as_secs().to_string(), + ) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|error| format!("spawn the claiming process: {error}"))?; + for stream in [ + child + .stdout + .take() + .map(|s| Box::new(s) as Box), + child + .stderr + .take() + .map(|s| Box::new(s) as Box), + ] + .into_iter() + .flatten() + { + let shared = shared.clone(); + let output_guard = super::super::worker_errors::child_output_guard(); + tokio::spawn(async move { + let _output_guard = output_guard; + let mut lines = BufReader::new(stream).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => absorb_output(&shared, &line), + Ok(None) => break, + Err(error) => { + shared.incomplete( + format!("claiming-child output read failed: {error}"), + vec![], + ); + break; + } + } + } + }); + } + Ok(child) +} + +fn absorb_output(shared: &Shared, line: &str) { + if super::super::worker_errors::absorb_child_marker(line) { + return; + } + let stats = shared.stats(Persona::ClaimAbandon); + let field = |name: &str| { + line.split_whitespace() + .find_map(|part| part.strip_prefix(&format!("{name}="))) + .and_then(|value| value.parse::().ok()) + }; + if line.starts_with("STAT ") { + if let Some(active) = field("active") { + shared.abandon.child_active.store(active, Ordering::Relaxed); + } + if let Some(peak) = field("peak") { + stats.set_max("process_peak_held", peak); + } + } else if line.contains("panicked") || line.contains("Error") { + if stats.get("process_error_lines") < 5 { + shared.note("abandon", line.chars().take(200).collect::()); + } + stats.add("process_error_lines", 1); + } +} + +/// Folds the new lines of the child ledger into the harness ledger. +struct Importer { + path: PathBuf, + offset: u64, +} + +impl Importer { + fn import(&mut self, shared: &Shared) { + let Ok(mut file) = std::fs::File::open(&self.path) else { + return; + }; + if file.seek(SeekFrom::Start(self.offset)).is_err() { + return; + } + let mut text = String::new(); + if file.read_to_string(&mut text).is_err() { + return; + } + for line in text.split_inclusive('\n') { + if !line.ends_with('\n') { + break; + } + self.offset += line.len() as u64; + if let Ok(event) = serde_json::from_str::(line) { + self.apply(shared, &event); + } + } + } + + fn apply(&self, shared: &Shared, event: &Value) { + if event.get("e").and_then(Value::as_str) != Some("start") { + return; + } + let (Some(id), Some(attempt), Some(generation), Some(at_ms)) = ( + event.get("id").and_then(Value::as_i64), + event.get("attempt").and_then(Value::as_i64), + event.get("gen").and_then(Value::as_u64), + event.get("t").and_then(Value::as_i64), + ) else { + return; + }; + let queue = event + .get("q") + .and_then(Value::as_str) + .and_then(|name| shared.queues.get(name)); + let queue_index = queue.as_ref().map_or(u32::MAX, |queue| queue.index); + if let Some(queue) = &queue { + queue.started.fetch_add(1, Ordering::Relaxed); + } + let now = shared.now_ms() as u32; + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::ClaimAbandon, + queue_index, + now, + expect::COMPLETED | expect::FAILED, + flag::NO_EVENTS, + ) + }, + |record| { + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(attempt as i16); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + }, + ); + shared.abandon.note_attempt( + id, + Attempt { + attempt: attempt as i16, + owner: generation as u32, + at_ms, + }, + ); + shared + .stats(Persona::ClaimAbandon) + .add("process_executions", 1); + } +} + +fn params(rng: &mut Rng, seq: &mut u64, cycle: u32, count: usize) -> Vec { + (0..count) + .map(|index| { + *seq += 1; + awa::insert::params_with( + &AbandonJob { seq: *seq, cycle }, + InsertOpts { + queue: if index % 2 == 0 { + QUEUE_PLAIN + } else { + QUEUE_DEADLINE + } + .to_string(), + max_attempts: if rng.chance(SHORT_RETRY_PERCENT) { + 2 + } else { + 25 + }, + ..Default::default() + }, + ) + .expect("abandon params") + }) + .collect() +} + +/// Enqueues `count` jobs for the next process to claim, in chunks that COPY handles +/// comfortably. +async fn enqueue_cycle( + shared: &Arc, + pool: &sqlx::PgPool, + rng: &mut Rng, + seq: &mut u64, + cycle: u32, + count: usize, +) -> Result { + let mut enqueued = 0; + while enqueued < count { + let chunk = (count - enqueued).min(10_000); + let batch = params(rng, seq, cycle, chunk); + let rows = insert_many_copy_from_pool(pool, &batch) + .await + .map_err(|error| format!("enqueue: {error}"))?; + for row in &rows { + shared.abandon.register_job(row.id, row.max_attempts); + if let Some(queue) = shared.queues.get(&row.queue) { + queue.enqueued.fetch_add(1, Ordering::Relaxed); + } + } + register_rows( + shared, + Persona::ClaimAbandon, + &rows, + |_| expect::COMPLETED | expect::FAILED, + flag::NO_EVENTS, + ); + shared + .stats(Persona::ClaimAbandon) + .add("enqueued", rows.len() as i64); + enqueued += rows.len(); + } + Ok(enqueued) +} + +async fn canary_producer(shared: Arc, pool: sqlx::PgPool) { + let mut seq = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(200)).await { + let batch: Vec = (0..2) + .map(|_| { + seq += 1; + awa::insert::params_with( + &AbandonCanaryJob { seq }, + InsertOpts { + queue: QUEUE_CANARY.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .expect("canary params") + }) + .collect(); + match insert_many(&pool, &batch).await { + Ok(rows) => { + shared + .stats(Persona::ClaimAbandon) + .add("canary_enqueued", rows.len() as i64); + register_rows( + &shared, + Persona::ClaimAbandon, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS, + ); + } + Err(error) => backoff(&shared, error, Persona::ClaimAbandon).await, + } + } +} + +/// Measures the live process's heartbeat passes from the timestamps its heartbeats write. +/// A pass updates every in-flight claim with the time of its own chunk, so once a pass has +/// finished the oldest heartbeat is newer than anything the previous pass wrote, and the +/// spread from oldest to newest is the pass's duration (to within the last chunk). Only +/// open claims on the persona's queues taken since the process started are read, so the +/// dead processes' open claims, whose heartbeats stopped, do not count. +/// +/// Returns false when the sample saw no open claim at all, which is only possible while the +/// process holds jobs if the query does not read where claims are written. +async fn sample_heartbeats( + shared: &Arc, + pool: &sqlx::PgPool, + schema: &str, + since_epoch_ms: i64, +) -> bool { + let row: Result<(i64, Option, Option, f64), _> = + sqlx::query_as(awa::audited_sql(format!( + "WITH open_claims AS ( + SELECT b.claim_slot, i.job_id, i.run_lease, i.receipt_id + FROM {schema}.lease_claim_batches b + CROSS JOIN LATERAL unnest(b.job_ids, b.run_leases, b.receipt_ids) + AS i(job_id, run_lease, receipt_id) + WHERE b.queue = ANY($1) + AND b.claimed_at >= to_timestamp($2::float8 / 1000.0) + AND NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closures cl + WHERE cl.claim_slot = b.claim_slot + AND cl.job_id = i.job_id AND cl.run_lease = i.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closure_batches cb + WHERE cb.claim_slot = b.claim_slot + AND cb.receipt_ranges @> i.receipt_id + OFFSET 0) + ) + SELECT count(*)::bigint, + EXTRACT(EPOCH FROM min(a.heartbeat_at))::float8, + EXTRACT(EPOCH FROM max(a.heartbeat_at))::float8, + EXTRACT(EPOCH FROM clock_timestamp())::float8 + FROM open_claims c + LEFT JOIN {schema}.attempt_state a ON a.job_id = c.job_id AND a.run_lease = c.run_lease" + ))) + .bind(vec![QUEUE_PLAIN, QUEUE_DEADLINE]) + .bind(since_epoch_ms as f64) + .fetch_one(pool) + .await; + let (rows, oldest, newest, now) = match row { + Ok(row) => row, + Err(error) => { + shared.incomplete(format!("heartbeat sample query failed: {error}"), vec![]); + return true; + } + }; + if rows == 0 { + return false; + } + // Claimed, but not yet reached by a heartbeat pass. + let (Some(oldest), Some(newest)) = (oldest, newest) else { + return true; + }; + let active = shared.abandon.child_active.load(Ordering::Relaxed); + shared + .abandon + .note_heartbeat_sample(shared, rows, oldest, newest, now, active); + true +} + +/// Probes the backlog every few seconds and the claim ring every second, until the +/// run is over. +async fn monitor(shared: Arc, queues: Vec) { + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()); + let Ok(pool) = connect_pool_with_timeout( + &shared.config.database_url, + "soak_abandon_monitor", + 2, + Some(Duration::from_secs(60)), + ) + .await + else { + shared.violation( + "persona_failed", + "abandon monitor could not connect", + vec![], + ); + return; + }; + let staleness_secs = shared.config.heartbeat_staleness.as_secs_f64(); + let mut tick = 0u64; + let mut empty_samples = 0u32; + let mut lock_captures = super::super::ring_locks::CaptureBudget::default(); + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + tick += 1; + // The cursor lives where `ring_cursor_authority` says: the latest ledger row, or + // the state singleton's columns before the authority flip. + let ring: Result<(i32, i64, i32), _> = sqlx::query_as(awa::audited_sql(format!( + "SELECT CASE WHEN authority.authority = 'ledger' THEN ledger.slot + ELSE state.current_slot END, + CASE WHEN authority.authority = 'ledger' THEN ledger.generation + ELSE state.generation END, + state.slot_count + FROM {schema}.claim_ring_state AS state + CROSS JOIN {schema}.ring_cursor_authority AS authority + LEFT JOIN LATERAL ( + SELECT slot, generation FROM {schema}.claim_ring_rotations + ORDER BY generation DESC LIMIT 1 + ) AS ledger ON TRUE + WHERE state.singleton AND authority.singleton" + ))) + .fetch_one(&pool) + .await; + let (slot, ring_generation, slot_count) = match ring { + Ok(ring) => ring, + Err(error) => { + shared.incomplete(format!("claim-ring monitor query failed: {error}"), vec![]); + continue; + } + }; + if let Some((generation, spawned)) = shared.abandon.live() { + let holding = shared.abandon.child_active.load(Ordering::Relaxed); + if sample_heartbeats(&shared, &pool, &schema, spawned).await || holding <= 0 { + empty_samples = 0; + } else { + empty_samples += 1; + if empty_samples == EMPTY_SAMPLES_BEFORE_INCOMPLETE { + shared.incomplete( + format!( + "the heartbeat sample found no open claim of process {generation} in {empty_samples} consecutive samples while it reported {holding} running handler(s)" + ), + vec![], + ); + } + } + } else { + empty_samples = 0; + } + // While the ring is still, look at the slot it has to enter next: live claims + // there explain the stall. + let next_slot = if shared.abandon.ring_stalled() && tick.is_multiple_of(2) { + probe_backlog( + &pool, + &schema, + Scope::Slot((slot + 1).rem_euclid(slot_count)), + staleness_secs, + 0, + ) + .await + .map(|probe| (probe.open - probe.expired, probe.expired)) + .map_err(|error| { + shared.incomplete(format!("claim-slot monitor query failed: {error}"), vec![]) + }) + .ok() + } else { + None + }; + shared + .abandon + .observe_ring(&shared, slot, ring_generation, next_slot); + super::super::ring_locks::observe( + &shared, + &pool, + &schema, + slot_count, + next_slot, + &mut lock_captures, + ) + .await; + if tick.is_multiple_of(3) { + // Read what the live process has started before the probe, so a claim started + // meanwhile cannot make a claim look missing. + let live = shared.abandon.live(); + let live_check = live.map(|(generation, _)| LiveCheck { + generation, + started: shared.abandon.started_by(generation), + }); + let started = Instant::now(); + match probe_backlog( + &pool, + &schema, + Scope::Queues(&queues), + staleness_secs, + live.map_or(0, |(_, spawned)| spawned), + ) + .await + { + Ok(probe) => shared.abandon.observe( + &shared, + &probe, + live_check, + ring_generation, + started.elapsed().as_millis() as i64, + ), + Err(error) => { + shared.incomplete( + format!("abandoned-claim backlog query failed: {error}"), + vec![], + ); + shared.stats(Persona::ClaimAbandon).add("probe_errors", 1); + if shared.stats(Persona::ClaimAbandon).get("probe_errors") <= 3 { + shared.note("abandon", format!("backlog probe failed: {error}")); + } + } + } + } + } + shared.abandon.finish_ring(&shared); + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::ClaimAbandon; + let config = shared.config.clone(); + let staleness_secs = config.heartbeat_staleness.as_secs_f64(); + let plain = shared.queues.register(QUEUE_PLAIN, persona, false); + let deadline = shared.queues.register(QUEUE_DEADLINE, persona, false); + let canary = shared.queues.register(QUEUE_CANARY, persona, false); + let queues = vec![QUEUE_PLAIN.to_string(), QUEUE_DEADLINE.to_string()]; + if config.abandon_kill_every >= DEADLINE { + return Err(format!( + "AWA_SOAK_ABANDON_KILL_SECONDS must stay below the {}s job deadline, or the owner's own deadline rescue would count as a false rescue", + DEADLINE.as_secs() + )); + } + let target = config.abandon_claims.max(2); + let workers_per_queue = (target / 2).max(1) as u32; + + std::fs::create_dir_all(&config.report_dir).map_err(|error| error.to_string())?; + let ledger_path = config.report_dir.join("abandon-child-ledger.jsonl"); + let _ = std::fs::remove_file(&ledger_path); + let executable = copy_executable(&config.report_dir)?; + shared + .abandon + .set_csv(config.report_dir.join("abandon-backlog.csv")); + + let child_pool = shared.budget.reserve("abandon-child", CHILD_POOL, 4); + // Held for the runtime that takes over at the end; Fleet::start replaces it. + shared.budget.reserve( + "abandon-finisher", + FINISHER_POOL * FINISHER_REPLICAS as u32, + FINISHER_REPLICAS as u32 * 2, + ); + + // The survivors: other replicas of the application. They lead maintenance, so they do + // the rescuing, and they serve a canary queue. + let canary_shared = shared.clone(); + let survivor_factory: ClientFactory = Arc::new(move |pool| { + tune( + Client::builder(pool) + .queue(QUEUE_CANARY, queue(4)) + .register_worker(CanaryWorker { + shared: canary_shared.clone(), + }), + &canary_shared, + ) + .build() + .expect("build abandon survivor client") + }); + let survivors = Fleet::start( + &shared, + "abandon-survivor", + persona, + 2, + 8, + 2, + vec![(canary, 4)], + survivor_factory, + ) + .await?; + + let producer_pool = connect_pool( + &config.database_url, + "soak_abandon_producer", + shared.budget.reserve("abandon-producer", 4, 2), + ) + .await + .map_err(|error| error.to_string())?; + let canary_task = tokio::spawn(canary_producer(shared.clone(), producer_pool.clone())); + let monitor_task = tokio::spawn(monitor(shared.clone(), queues.clone())); + let backlog_pool = connect_pool_with_timeout( + &config.database_url, + "soak_abandon_control", + 2, + Some(Duration::from_secs(60)), + ) + .await + .map_err(|error| error.to_string())?; + + let mut rng = Rng::for_task(config.seed, "abandon"); + let mut importer = Importer { + path: ledger_path.clone(), + offset: 0, + }; + let mut seq = 0u64; + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()); + shared.note( + "abandon", + format!( + "armed: {} kill(s) every {}s, each leaving about {target} open claims ({} plain, {} with a deadline); rescue must clear within {}s", + config.abandon_kills, + config.abandon_kill_every.as_secs(), + target / 2, + target - target / 2, + config.abandon_bound.as_secs() + ), + ); + + let mut generation = 0u32; + let mut load_ended = false; + for number in 1..=config.abandon_kills { + if shared.load_stop.is_stopped() { + break; + } + // The first process finds a full queue. Later ones find the rescued jobs of the + // earlier kills and only need fresh jobs to top their capacity up, so that jobs + // are run, and abandoned, more than once. + let fresh = if number == 1 { target } else { target / 2 }; + match enqueue_cycle( + &shared, + &producer_pool, + &mut rng, + &mut seq, + number, + fresh as usize, + ) + .await + { + Ok(count) => shared.note( + "abandon", + format!("enqueued {count} fresh job(s) for process {number}"), + ), + Err(error) => { + backoff(&shared, error, persona).await; + continue; + } + } + generation += 1; + let mut child = spawn_child( + &shared, + &executable, + generation, + &ledger_path, + workers_per_queue, + child_pool, + )?; + shared.abandon.child_alive.store(true, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + shared + .abandon + .set_live(Some((generation, epoch_ms() - 1_000))); + let spawned = Instant::now(); + let mut ramped = false; + shared.note( + "abandon", + format!("process {generation} started with {workers_per_queue} workers on each queue"), + ); + let mut unexpected_exit = false; + loop { + importer.import(&shared); + if shared.load_stop.sleep(Duration::from_secs(1)).await { + load_ended = true; + break; + } + if let Ok(Some(status)) = child.try_wait() { + shared.violation( + "abandon_process_exited", + format!("the claiming process {generation} exited on its own: {status}"), + vec![], + ); + unexpected_exit = true; + break; + } + if !ramped + && shared.abandon.child_active.load(Ordering::Relaxed) >= (target as i64) * 95 / 100 + { + ramped = true; + shared.note( + "abandon", + format!( + "process {generation} holds {} jobs after {:.0}s", + shared.abandon.child_active.load(Ordering::Relaxed), + spawned.elapsed().as_secs_f64() + ), + ); + shared + .stats(persona) + .observe("ramp_seconds", spawned.elapsed().as_secs() as i64); + } + if spawned.elapsed() >= config.abandon_kill_every { + break; + } + } + + let held = shared.abandon.child_active.load(Ordering::Relaxed); + let signalled_epoch_ms = if unexpected_exit { + let at = epoch_ms(); + shared.abandon.note_generation_killed(generation, at); + shared.abandon.set_live(None); + at + } else { + // SIGKILL: no shutdown handshake, no chance to release a claim. + match shared + .abandon + .signal_generation_killed(generation, || child.start_kill()) + { + Ok(at) => at, + Err(error) => { + shared.incomplete( + format!("could not signal claiming process {generation}: {error}"), + vec![], + ); + // No confirmed signal means no valid death boundary. Stop this + // persona and clean up its owned tasks instead of inventing one. + shared.abandon.set_live(None); + shared.abandon.child_alive.store(false, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + drop(child); // kill_on_drop remains a best-effort cleanup. + canary_task.abort(); + monitor_task.abort(); + let _ = canary_task.await; + let _ = monitor_task.await; + survivors + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + backlog_pool.close().await; + return Ok(()); + } + } + }; + let signalled_run_ms = shared.now_ms(); + shared.abandon.child_alive.store(false, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + if !unexpected_exit { + if let Err(error) = child.wait().await { + shared.incomplete( + format!( + "could not reap claiming process {generation} after signaling it: {error}" + ), + vec![], + ); + } + } + importer.import(&shared); + let open_after_kill = probe_backlog( + &backlog_pool, + &schema, + Scope::Queues(&queues), + staleness_secs, + 0, + ) + .await + .map(|probe| probe.open) + .unwrap_or_else(|error| { + shared.incomplete(format!("post-kill backlog query failed: {error}"), vec![]); + -1 + }); + let started = shared.abandon.started_by(generation); + shared.abandon.record_kill(Kill { + number, + generation, + at_run_ms: signalled_run_ms, + signalled_epoch_ms, + started_by_process: started, + rescued_while_alive: shared.abandon.live_lost(generation), + open_after_kill, + expired_peak: 0, + oldest_expired_peak_s: 0.0, + cleared_after_s: None, + reported_slow: false, + load_ended, + }); + shared.stats(persona).add("kills", 1); + shared.stats(persona).add("claims_abandoned", started); + shared.note( + "abandon", + format!( + "SIGKILL process {generation} (kill {number}): {started} handlers started, {held} running at the last report, {open_after_kill} open claims on its queues" + ), + ); + if load_ended { + break; + } + if shared.load_stop.sleep(config.abandon_restart_delay).await { + break; + } + } + + // Everything the kills left behind is finished by an in-process runtime that is + // never killed. + let finisher_shared = shared.clone(); + let finisher_factory: ClientFactory = Arc::new(move |pool| { + tune( + Client::builder(pool) + .queue(QUEUE_PLAIN, queue(256)) + .queue(QUEUE_DEADLINE, queue(256)) + .register_worker(FinisherWorker { + shared: finisher_shared.clone(), + }), + &finisher_shared, + ) + .build() + .expect("build abandon finisher client") + }); + shared.note( + "abandon", + "kills over; starting the runtime that finishes the rescued jobs", + ); + let finisher = Fleet::start( + &shared, + "abandon-finisher", + persona, + FINISHER_REPLICAS, + FINISHER_POOL, + 4, + vec![(plain, 256), (deadline, 256)], + finisher_factory, + ) + .await?; + + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + importer.import(&shared); + } + importer.import(&shared); + canary_task.abort(); + let _ = monitor_task.await; + shared + .abandon + .analyse(&shared, &backlog_pool, &queues) + .await; + finisher + .shutdown_all(Duration::from_secs(20), &shared) + .await; + survivors + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + backlog_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/alex.rs b/awa/tests/persona_soak/personas/alex.rs new file mode 100644 index 000000000..88498cf04 --- /dev/null +++ b/awa/tests/persona_soak/personas/alex.rs @@ -0,0 +1,756 @@ +//! Alex, the operator: under live load, pauses and resumes queues, drains, cancels, +//! moves and reprioritises jobs, retries and purges DLQ rows, runs the admin read +//! queries the UI issues, and rolls a client restart through another persona's fleet. + +use super::super::config::Persona; +use super::super::db::{connect_pool, connect_pool_with_timeout, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::registry::QueueInfo; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::admin::{self, ListJobsFilter}; +use awa::model::batch_operations::{ + get_batch_operation, submit_batch_operation, BatchOperationFilter, BatchOperationSpec, + BatchOperationState, SubmitBatchOperation, +}; +use awa::model::dlq::{self, ListDlqFilter}; +use awa::model::{insert_many, JobState, RetryFromDlqOpts}; +use awa::{Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const QUEUE_A: &str = "awa_soak_alex_a"; +pub const QUEUE_B: &str = "awa_soak_alex_b"; +pub const QUEUE_DRAIN: &str = "awa_soak_alex_drain"; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AlexJob { + pub seq: u64, + pub work_ms: u64, +} + +struct AlexWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for AlexWorker { + fn kind(&self) -> &'static str { + "alex_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Alex, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(100); + tokio::time::sleep(Duration::from_millis(work_ms)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE_A, queue(4)) + .queue(QUEUE_B, queue(4)) + .queue(QUEUE_DRAIN, queue(4)) + .register_worker(AlexWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["alex_job"]) + .build() + .expect("build alex client") + }) +} + +struct Context { + shared: Arc, + pool: PgPool, + diagnostics: PgPool, + stalled_rows: std::sync::Mutex>, + rng: Rng, +} + +/// pid, application, state, wait event, seconds running, blocking pids, query. +type BackendRow = (i32, String, String, String, i32, String, String); + +/// An admin call this slow under load is a finding in its own right. +const SLOW_CALL: Duration = Duration::from_secs(30); +/// Server-side limit on the operator's statements; hitting it is reported as a finding. +const ADMIN_STATEMENT_TIMEOUT: Duration = Duration::from_secs(150); +const STALLED_CALL: Duration = Duration::from_secs(120); + +impl Context { + fn stats(&self) -> &super::super::shared::Stats { + self.shared.stats(Persona::Alex) + } + + async fn timed( + &self, + name: &str, + call: impl std::future::Future>, + ) -> Option { + let started = Instant::now(); + let mut call = std::pin::pin!(call); + let result = tokio::select! { + result = &mut call => result, + _ = tokio::time::sleep(SLOW_CALL) => { + self.diagnose_slow_call(name).await; + call.await + } + }; + let ms = started.elapsed().as_millis() as i64; + if started.elapsed() > STALLED_CALL { + self.shared.finding( + "admin_call_stalled", + format!("admin call {name} took {ms} ms under load"), + self.stalled_rows.lock().expect("stalled rows").clone(), + ); + } + self.stats().observe(name, ms); + if ms > 5_000 { + self.stats().add("slow_admin_calls_over_5s", 1); + self.shared.note("alex", format!("{name} took {ms} ms")); + } + match result { + Ok(value) => Some(value), + Err(error) => { + if error.to_string().contains("statement timeout") { + self.shared.finding( + "admin_call_timeout", + format!( + "admin call {name} was cancelled by the {}s statement timeout", + ADMIN_STATEMENT_TIMEOUT.as_secs() + ), + self.stalled_rows.lock().expect("stalled rows").clone(), + ); + } + self.stats().add(&format!("{name}_errors"), 1); + if self.stats().get(&format!("{name}_errors")) <= 3 { + self.shared.note("alex", format!("{name} failed: {error}")); + } + None + } + } + } + + /// Records what the operator's backend is waiting on and who blocks it. + async fn diagnose_slow_call(&self, name: &str) { + let rows: Result, _> = sqlx::query_as( + "SELECT a.pid, a.application_name, COALESCE(a.state, ''), + COALESCE(a.wait_event_type, '') || ':' || COALESCE(a.wait_event, ''), + COALESCE(extract(epoch FROM now() - a.query_start)::int, 0), + pg_blocking_pids(a.pid)::text, left(a.query, 200) + FROM pg_stat_activity a + WHERE a.pid <> pg_backend_pid() + AND (a.application_name = 'soak_alex_operator' + OR a.pid IN (SELECT unnest(pg_blocking_pids(op.pid)) FROM pg_stat_activity op + WHERE op.application_name = 'soak_alex_operator'))", + ) + .fetch_all(&self.diagnostics) + .await; + let lines: Vec = match rows { + Ok(rows) => rows + .into_iter() + .map(|(pid, app, state, wait, secs, blockers, query)| { + format!("{{\"pid\":{pid},\"application\":\"{app}\",\"state\":\"{state}\",\"wait\":\"{wait}\",\"running_s\":{secs},\"blocked_by\":\"{blockers}\",\"query\":{query:?}}}") + }) + .collect(), + Err(error) => vec![format!("diagnostic query failed: {error}")], + }; + self.shared.note( + "alex", + format!( + "{name} still running after {}s; backend state: {}", + SLOW_CALL.as_secs(), + lines.join(" | ") + ), + ); + *self.stalled_rows.lock().expect("stalled rows") = lines; + } + + /// A random queue of a persona, other than Alex's own. + fn queue_of(&mut self, persona: Persona) -> Option> { + let queues = self.shared.queues.of(persona); + self.rng.pick(&queues).cloned() + } + + async fn admin_reads(&mut self) { + let pool = self.pool.clone(); + self.timed("admin_queue_overviews", admin::queue_overviews(&pool)) + .await; + self.timed("admin_state_counts", admin::state_counts(&pool)) + .await; + self.timed("admin_job_kind_overviews", admin::job_kind_overviews(&pool)) + .await; + self.timed("admin_runtime_overview", admin::runtime_overview(&pool)) + .await; + self.timed( + "admin_queue_runtime_summary", + admin::queue_runtime_summary(&pool), + ) + .await; + self.timed("admin_dlq_depth_by_queue", dlq::dlq_depth_by_queue(&pool)) + .await; + if let Some(queue) = self.queue_of(Persona::Tom) { + let filter = ListJobsFilter { + queue: Some(queue.name.clone()), + limit: Some(50), + ..Default::default() + }; + self.timed("admin_list_jobs", admin::list_jobs(&pool, &filter)) + .await; + } + } + + /// Resumes a queue, retrying: a queue left paused would look like a stall. + async fn resume_with_retry(&self, queue: &str) { + for attempt in 0..8 { + if self + .timed("resume_queue", admin::resume_queue(&self.pool, queue)) + .await + .is_some() + { + return; + } + tokio::time::sleep(Duration::from_secs(2 + attempt)).await; + } + self.shared.violation( + "resume_failed", + format!("queue {queue} could not be resumed after eight attempts"), + vec![], + ); + } + + async fn pause_resume(&mut self) { + let target = match self.rng.below(3) { + 0 => self.shared.queues.get(QUEUE_B), + 1 => self.queue_of(Persona::Tom), + _ => self + .shared + .queues + .get(super::priya::HOT_QUEUE) + .or_else(|| self.queue_of(Persona::Mia)), + }; + let Some(queue) = target else { return }; + queue.hold(); + let pool = self.pool.clone(); + let paused = self + .timed( + "pause_queue", + admin::pause_queue(&pool, &queue.name, Some("soak")), + ) + .await + .is_some(); + if paused { + self.stats().add("pauses", 1); + self.shared + .stop_sleep(Duration::from_secs(self.rng.range(3, 8))) + .await; + } + // A pause whose answer was lost may still have applied, so always resume. + self.resume_with_retry(&queue.name).await; + // Give the claimers a moment to pick the queue back up before stall checks resume. + self.shared.stop_sleep(Duration::from_secs(3)).await; + queue.release(); + } + + async fn drain(&mut self) { + let target = if self.rng.chance(50) { + self.shared.queues.get(QUEUE_DRAIN) + } else { + self.queue_of(Persona::Tom) + }; + let Some(queue) = target else { return }; + // Until the drain returns, any job enqueued so far may legitimately be cancelled. + queue.drained_at_ms.store(i64::MAX / 4, Ordering::Relaxed); + let pool = self.pool.clone(); + if let Some(count) = self + .timed("drain_queue", admin::drain_queue(&pool, &queue.name)) + .await + { + self.stats().add("drains", 1); + self.shared + .note("alex", format!("drained {} ({count} job(s))", queue.name)); + self.stats().add("drained_jobs", count as i64); + } + queue + .drained_at_ms + .store(self.shared.now_ms(), Ordering::Relaxed); + } + + /// Cancels live jobs of other personas, some of them already running. + async fn cancel(&mut self) { + let persona = *self + .rng + .pick(&[Persona::Mia, Persona::Tom, Persona::Sam, Persona::Priya]) + .expect("personas"); + let Some(queue) = self.queue_of(persona) else { + return; + }; + let states = [ + JobState::Available, + JobState::Running, + JobState::Scheduled, + JobState::Retryable, + ]; + let state = *self.rng.pick(&states).expect("states"); + let pool = self.pool.clone(); + let filter = ListJobsFilter { + queue: Some(queue.name.clone()), + state: Some(state), + limit: Some(30), + ..Default::default() + }; + let Some(rows) = self + .timed("admin_list_jobs", admin::list_jobs(&pool, &filter)) + .await + else { + return; + }; + let ids: Vec = rows.iter().map(|row| row.id).take(8).collect(); + if ids.is_empty() { + return; + } + self.shared.note( + "alex", + format!("cancelling jobs {ids:?} ({state:?}) in {}", queue.name), + ); + if self.rng.chance(50) { + let started_ms = self.shared.now_ms() as u32; + let first = ids[0]; + let completed_before = self + .shared + .ledger + .get(first) + .is_some_and(|record| record.completed_events > 0); + if let Some(Some(job)) = self.timed("cancel", admin::cancel(&pool, first)).await { + self.after_cancel(job.id, completed_before, started_ms); + } + } else if let Some(cancelled) = self + .timed("bulk_cancel", admin::bulk_cancel(&pool, &ids)) + .await + { + let now = self.shared.now_ms() as u32; + for job in &cancelled { + self.after_cancel(job.id, false, now); + } + } + } + + fn after_cancel(&self, job_id: i64, completed_before: bool, call_started_ms: u32) { + self.shared.record_operator_cancel( + Persona::Alex, + job_id, + completed_before, + call_started_ms, + ); + } + + async fn dlq_operations(&mut self) { + let pool = self.pool.clone(); + let filter = ListDlqFilter { + queue: Some(super::mia::INFER_QUEUE.to_string()), + limit: Some(20), + ..Default::default() + }; + let Some(rows) = self.timed("dlq_list", dlq::list_dlq(&pool, &filter)).await else { + return; + }; + let ids: Vec = rows.iter().map(|row| row.job.id).collect(); + self.stats().set_max("dlq_listed_max", ids.len() as i64); + match self.rng.below(4) { + 0 => { + if let Some(&id) = ids.first() { + self.shared.ledger.mark(id, flag::DLQ_RETRIED); + let opts = RetryFromDlqOpts::default(); + if self + .timed("dlq_retry", dlq::retry_from_dlq(&pool, id, &opts)) + .await + .is_some() + { + self.stats().add("dlq_retried", 1); + } + } + } + 1 => { + for id in ids.iter().take(3) { + self.shared.ledger.mark(*id, flag::DLQ_PURGED); + if self + .timed("dlq_purge_job", dlq::purge_dlq_job(&pool, *id)) + .await + == Some(true) + { + self.stats().add("dlq_purged", 1); + } + } + } + 2 => { + // Bulk retry: every exhausted Mia job may come back. + self.shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_RETRIED, + ); + let bulk = ListDlqFilter { + queue: Some(super::mia::INFER_QUEUE.to_string()), + ..Default::default() + }; + if let Some(count) = self + .timed( + "dlq_bulk_retry", + dlq::bulk_retry_from_dlq(&pool, &bulk, false), + ) + .await + { + self.stats().add("dlq_bulk_retried", count as i64); + } + } + _ => { + self.shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_PURGED, + ); + let bulk = ListDlqFilter { + queue: Some(super::mia::TRAIN_QUEUE.to_string()), + ..Default::default() + }; + if let Some(count) = self + .timed("dlq_bulk_purge", dlq::purge_dlq(&pool, &bulk, false)) + .await + { + self.stats().add("dlq_bulk_purged", count as i64); + } + } + } + } + + async fn batch_operations(&mut self) { + let (spec, filter, label, source) = if self.rng.chance(50) { + ( + BatchOperationSpec::MoveQueue { + queue: QUEUE_B.to_string(), + priority: None, + }, + BatchOperationFilter { + queue: Some(QUEUE_A.to_string()), + state: Some(JobState::Available), + ..Default::default() + }, + "move_queue", + QUEUE_A, + ) + } else { + ( + BatchOperationSpec::SetPriority { priority: 1 }, + BatchOperationFilter { + queue: Some(QUEUE_B.to_string()), + state: Some(JobState::Available), + ..Default::default() + }, + "set_priority", + QUEUE_B, + ) + }; + // Hold the source queue so a backlog builds for the operation to act on. + let pool = self.pool.clone(); + let held = self.shared.queues.get(source); + if let Some(queue) = &held { + queue.hold(); + let _ = self + .timed( + "pause_queue", + admin::pause_queue(&pool, &queue.name, Some("soak")), + ) + .await; + self.shared.stop_sleep(Duration::from_secs(6)).await; + } + let request = SubmitBatchOperation { + spec, + filter, + submitted_by: Some("soak".into()), + allow_all: false, + }; + self.run_batch(request, label).await; + if let Some(queue) = &held { + self.resume_with_retry(&queue.name).await; + self.shared.stop_sleep(Duration::from_secs(3)).await; + queue.release(); + } + } + + async fn run_batch(&mut self, request: SubmitBatchOperation, label: &str) { + let pool = self.pool.clone(); + let Some(operation) = self + .timed("batch_submit", submit_batch_operation(&pool, request)) + .await + else { + return; + }; + let started = Instant::now(); + loop { + let Some(current) = self + .timed("batch_get", get_batch_operation(&pool, operation.id)) + .await + else { + return; + }; + match current.state { + BatchOperationState::Completed | BatchOperationState::Cancelled => { + self.stats().add(&format!("batch_{label}"), 1); + self.stats() + .add(&format!("batch_{label}_processed"), current.processed); + self.stats() + .observe("batch_operation_ms", started.elapsed().as_millis() as i64); + return; + } + BatchOperationState::Failed => { + self.shared.violation( + "batch_operation_failed", + format!( + "{label} batch operation {} failed: {:?}", + operation.id, current.last_error + ), + vec![], + ); + return; + } + _ => {} + } + if started.elapsed() > Duration::from_secs(90) { + self.stats().add("batch_operations_incomplete", 1); + self.shared.note( + "alex", + format!( + "{label} batch operation {} still {:?} after 90s", + operation.id, current.state + ), + ); + return; + } + if self.shared.stop_sleep(Duration::from_secs(1)).await { + return; + } + } + } + + async fn rolling_restart(&mut self) { + let names = ["mia", "tom", "priya", "sam", "sam-stripe"]; + let name = *self.rng.pick(&names).expect("fleet names"); + let Some(fleet) = self.shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + let Some(&index) = self.rng.pick(&live) else { + return; + }; + self.shared + .note("alex", format!("rolling restart: {name} replica {index}")); + self.shared.chaos_window(Duration::from_secs(20)); + let started = Instant::now(); + let restarted = fleet.restart(index, Duration::from_secs(30)).await; + // If the restarted runtime led, the next election can take a full interval. + self.shared.chaos_window(Duration::from_secs(45)); + match restarted { + Ok(_) => { + self.stats().add("rolling_restarts", 1); + self.stats() + .observe("rolling_restart_ms", started.elapsed().as_millis() as i64); + } + Err(error) => { + self.shared.violation( + "restart_failed", + format!("restarting {name} replica {index}: {error}"), + vec![], + ); + } + } + } +} + +impl Shared { + /// Sleeps unless the load phase ends first; true when it ended. + pub async fn stop_sleep(&self, duration: Duration) -> bool { + self.load_stop.sleep(duration).await + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Alex; + let queue_a = shared.queues.register(QUEUE_A, persona, false); + let queue_b = shared.queues.register(QUEUE_B, persona, false); + let queue_drain = shared.queues.register(QUEUE_DRAIN, persona, false); + let fleet = Fleet::start( + &shared, + "alex", + persona, + 2, + 10, + 2, + vec![(queue_a, 4), (queue_b, 4), (queue_drain, 4)], + client_factory(shared.clone()), + ) + .await?; + + let pool = connect_pool_with_timeout( + &shared.config.database_url, + "soak_alex_operator", + shared.budget.reserve("alex-operator", 4, 2), + Some(ADMIN_STATEMENT_TIMEOUT), + ) + .await + .map_err(|error| error.to_string())?; + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_alex_producer", + shared.budget.reserve("alex-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + + // Alex's own traffic: a stream the move and reprioritise operations act on, and a + // stream that gets drained while it is still being produced. + let producer = { + let shared = shared.clone(); + tokio::spawn(async move { + let mut rng = Rng::for_task(shared.config.seed, "alex-producer"); + let mut seq = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(200)).await { + let mut params = Vec::new(); + for queue in [ + QUEUE_A, + QUEUE_A, + QUEUE_A, + QUEUE_A, + QUEUE_B, + QUEUE_DRAIN, + QUEUE_DRAIN, + QUEUE_DRAIN, + ] { + seq += 1; + params.push( + awa::insert::params_with( + &AlexJob { + seq, + work_ms: rng.range(50, 300), + }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .expect("alex params"), + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(Persona::Alex) + .add("enqueued", rows.len() as i64); + register_rows( + &shared, + Persona::Alex, + &rows, + |_| expect::ANY_TERMINAL, + flag::OPERATOR_MOVED, + ); + } + Err(error) => backoff(&shared, error, Persona::Alex).await, + } + } + producer_pool.close().await; + }) + }; + + let diagnostics = connect_pool( + &shared.config.database_url, + "soak_alex_diagnostics", + shared.budget.reserve("alex-diagnostics", 1, 1), + ) + .await + .map_err(|error| error.to_string())?; + let mut context = Context { + shared: shared.clone(), + pool, + diagnostics, + stalled_rows: Default::default(), + rng: Rng::for_task(shared.config.seed, "alex-operator"), + }; + let load = shared.config.load.as_secs_f64(); + // Scale the cadence down for short runs so each action still happens a few times. + let pace = (load / 900.0).clamp(0.1, 1.0); + let mut schedule: Vec<(&str, f64, Instant)> = vec![ + ( + "reads", + shared.config.admin_read_every.as_secs_f64(), + Instant::now(), + ), + ( + "pause", + 60.0 * pace, + Instant::now() + Duration::from_secs_f64(20.0 * pace), + ), + ( + "drain", + 75.0 * pace, + Instant::now() + Duration::from_secs_f64(35.0 * pace), + ), + ( + "cancel", + 10.0 * pace, + Instant::now() + Duration::from_secs_f64(15.0 * pace), + ), + ( + "dlq", + 45.0 * pace, + Instant::now() + Duration::from_secs_f64(60.0 * pace), + ), + ( + "batch", + 40.0 * pace, + Instant::now() + Duration::from_secs_f64(25.0 * pace), + ), + ( + "restart", + 150.0 * pace, + Instant::now() + Duration::from_secs_f64(90.0 * pace), + ), + ]; + shared.note("alex", "operating: reads, pause/resume, drain, cancel, move/reprioritise, DLQ retry/purge, rolling restarts"); + while !shared.load_stop.is_stopped() { + let now = Instant::now(); + for (name, interval, due) in schedule.iter_mut() { + if *due > now || shared.load_stop.is_stopped() { + continue; + } + *due = Instant::now() + Duration::from_secs_f64(*interval); + match *name { + "reads" => context.admin_reads().await, + "pause" => context.pause_resume().await, + "drain" => context.drain().await, + "cancel" => context.cancel().await, + "dlq" => context.dlq_operations().await, + "batch" => context.batch_operations().await, + "restart" => context.rolling_restart().await, + _ => {} + } + } + shared.load_stop.sleep(Duration::from_millis(250)).await; + } + + let _ = producer.await; + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + context.pool.close().await; + context.diagnostics.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/chaos.rs b/awa/tests/persona_soak/personas/chaos.rs new file mode 100644 index 000000000..fee1c63eb --- /dev/null +++ b/awa/tests/persona_soak/personas/chaos.rs @@ -0,0 +1,242 @@ +//! Chaos: kills a runtime without a shutdown handshake, terminates the maintenance +//! leader's connection, terminates random runtime connections, exhausts a runtime's +//! pool and holds a long idle-in-transaction session. + +use super::super::config::Persona; +use super::super::db::connect_pool; +use super::super::rng::Rng; +use super::super::shared::Shared; +use sqlx::PgPool; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// The advisory lock the maintenance leader holds for as long as it leads. +const MAINTENANCE_LEADER_LOCK_KEY: i64 = 0x4157_415f_4d41_494e; +const RUNTIME_FLEETS: [&str; 6] = ["mia", "tom", "priya", "sam", "sam-stripe", "fanout"]; + +async fn kill_runtime(shared: &Arc, rng: &mut Rng) { + let Some(&name) = rng.pick(&RUNTIME_FLEETS) else { + return; + }; + let Some(fleet) = shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + // Never take the last runtime of a fleet down; the point is a survivor that rescues. + if live.len() < 2 { + return; + } + let Some(&index) = rng.pick(&live) else { + return; + }; + shared.note( + "chaos", + format!("killing {name} replica {index} without shutdown"), + ); + shared.chaos_window(Duration::from_secs(45)); + match fleet.kill(index, Some(Duration::from_secs(4))).await { + Ok(_) => shared.stats(Persona::Chaos).add("runtime_kills", 1), + Err(error) => shared.violation( + "restart_failed", + format!("restarting killed {name} replica {index}: {error}"), + vec![], + ), + } +} + +async fn kill_python(shared: &Arc) { + let handle = shared.python.lock().expect("python handle").clone(); + let Some(handle) = handle else { return }; + shared.note("chaos", "killing the Python process without shutdown"); + shared.chaos_window(Duration::from_secs(45)); + if let Err(error) = handle + .kill_and_restart(shared, Duration::from_secs(4)) + .await + { + shared.violation( + "restart_failed", + format!("restarting the Python process: {error}"), + vec![], + ); + } +} + +async fn terminate_leader(shared: &Arc, pool: &PgPool) { + // Maintenance leadership is an advisory lock held by one runtime connection. + let holders: Result, _> = sqlx::query_as( + "SELECT l.pid, COALESCE(a.application_name, '') + FROM pg_locks l JOIN pg_stat_activity a ON a.pid = l.pid + WHERE l.locktype = 'advisory' AND l.granted AND a.application_name LIKE 'soak\\_%' + AND ((l.classid::bigint << 32) | l.objid::bigint) = $1 + ORDER BY l.pid", + ) + .bind(MAINTENANCE_LEADER_LOCK_KEY) + .fetch_all(pool) + .await; + let Ok(holders) = holders else { return }; + let Some((pid, application)) = holders.first().cloned() else { + shared.note("chaos", "no advisory-lock holder found to terminate"); + return; + }; + shared.note( + "chaos", + format!("terminating maintenance leader backend {pid} ({application})"), + ); + shared.chaos_window(Duration::from_secs(40)); + let _ = sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(pool) + .await; + shared.stats(Persona::Chaos).add("leader_terminations", 1); +} + +async fn terminate_runtime_connections(shared: &Arc, pool: &PgPool, rng: &mut Rng) { + let count = rng.range(3, 12) as i64; + let terminated: Result, _> = sqlx::query_as( + "SELECT pid FROM ( + SELECT pid FROM pg_stat_activity + WHERE application_name LIKE 'soak\\_%' AND application_name NOT LIKE '%producer%' + AND application_name NOT LIKE '%operator%' AND application_name NOT LIKE '%oracle%' + AND application_name NOT LIKE '%chaos%' AND application_name NOT LIKE 'soak\\_abandon\\_%' + AND backend_type = 'client backend' + AND pid <> pg_backend_pid() + ORDER BY random() LIMIT $1 + ) victims WHERE pg_terminate_backend(pid)", + ) + .bind(count) + .fetch_all(pool) + .await; + let killed = terminated.map(|rows| rows.len()).unwrap_or(0); + shared.note( + "chaos", + format!("pg_terminate_backend on {killed} runtime connection(s)"), + ); + shared.chaos_window(Duration::from_secs(30)); + shared + .stats(Persona::Chaos) + .add("connections_terminated", killed as i64); +} + +/// Takes every connection of one runtime's pool for `hold`, starving its heartbeats +/// and claims the way a connection-pool exhaustion bug would. +async fn exhaust_pool(shared: &Arc, rng: &mut Rng, hold: Duration) { + let Some(&name) = rng.pick(&RUNTIME_FLEETS) else { + return; + }; + let Some(fleet) = shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + let Some(&index) = rng.pick(&live) else { + return; + }; + let Some(pool) = fleet.pool_of(index).await else { + return; + }; + let size = pool.options().get_max_connections() as usize; + shared.note( + "chaos", + format!( + "exhausting the {size}-connection pool of {name} replica {index} for {}s", + hold.as_secs() + ), + ); + shared.chaos_window(hold + Duration::from_secs(40)); + let mut held = Vec::new(); + let deadline = Instant::now() + Duration::from_secs(20); + while held.len() < size && Instant::now() < deadline { + if let Ok(Ok(connection)) = + tokio::time::timeout(Duration::from_secs(2), pool.acquire()).await + { + held.push(connection); + } + } + shared.stats(Persona::Chaos).add("pool_exhaustions", 1); + shared + .stats(Persona::Chaos) + .add("pool_connections_held", held.len() as i64); + shared.load_stop.sleep(hold).await; + drop(held); +} + +/// One session that begins a transaction, takes a snapshot and sits idle, pinning the +/// xmin horizon so vacuum and prune cannot reclaim anything newer. +async fn idle_in_transaction(shared: Arc, hold: Duration) { + let Ok(pool) = connect_pool(&shared.config.database_url, "soak_chaos_idle", 1).await else { + return; + }; + let Ok(mut transaction) = pool.begin().await else { + return; + }; + if sqlx::query("SELECT txid_current()") + .execute(&mut *transaction) + .await + .is_err() + { + return; + } + shared.note( + "chaos", + format!( + "holding an idle-in-transaction session for {}s", + hold.as_secs() + ), + ); + shared + .stats(Persona::Chaos) + .add("idle_in_transaction_holds", 1); + shared.load_stop.sleep(hold).await; + let _ = transaction.rollback().await; + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let pool = connect_pool( + &shared.config.database_url, + "soak_chaos_control", + shared.budget.reserve("chaos", 3, 2), + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "chaos"); + let load = shared.config.load.as_secs_f64(); + let pace = (load / 900.0).clamp(0.15, 1.0); + let second = |seconds: f64| Duration::from_secs_f64(seconds * pace); + let mut schedule: Vec<(&str, Duration, Instant)> = vec![ + ("kill", second(120.0), Instant::now() + second(60.0)), + ("python", second(200.0), Instant::now() + second(110.0)), + ("leader", second(150.0), Instant::now() + second(90.0)), + ("terminate", second(75.0), Instant::now() + second(45.0)), + ("pool", second(240.0), Instant::now() + second(150.0)), + ("idle", second(300.0), Instant::now() + second(30.0)), + ]; + shared.note("chaos", "armed: runtime kill, leader termination, connection termination, pool exhaustion, idle transaction"); + while !shared.load_stop.is_stopped() { + let now = Instant::now(); + for (name, interval, due) in schedule.iter_mut() { + if *due > now || shared.load_stop.is_stopped() { + continue; + } + *due = Instant::now() + *interval; + match *name { + "kill" => kill_runtime(&shared, &mut rng).await, + "python" => kill_python(&shared).await, + "leader" => terminate_leader(&shared, &pool).await, + "terminate" => terminate_runtime_connections(&shared, &pool, &mut rng).await, + "pool" => { + exhaust_pool(&shared, &mut rng, Duration::from_secs(20).max(second(20.0))).await + } + "idle" => { + tokio::spawn(idle_in_transaction( + shared.clone(), + second(90.0).max(Duration::from_secs(30)), + )); + } + _ => {} + } + } + shared.load_stop.sleep(Duration::from_millis(500)).await; + } + pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/fanout.rs b/awa/tests/persona_soak/personas/fanout.rs new file mode 100644 index 000000000..8b3146300 --- /dev/null +++ b/awa/tests/persona_soak/personas/fanout.rs @@ -0,0 +1,278 @@ +//! Fanout, replay-style fan-out: each parent job batch-inserts a hundred or so unique +//! child jobs into the queue it runs on, from many concurrent handlers. Some parents +//! then replay a batch that reuses an already-claimed unique key, fail with a unique +//! violation and fall back to inserting that batch row by row. Stresses concurrent +//! batch enqueue into one lane, unique claims and the insert-from-handler path. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert_many, insert_with}; +use awa::{ + AwaError, Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, JobRow, + UniqueOpts, Worker, +}; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const QUEUE: &str = "awa_soak_fanout"; +const MIN_CHILDREN: u64 = 100; +const MAX_CHILDREN: u64 = 185; +const REPLAY_PERCENT: u64 = 10; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct FanoutParent { + pub seq: u64, + pub children: u64, + pub replay: bool, +} + +/// Distinct `(parent_seq, index)` per child, so the unique key (queue + args) is +/// shared only by a deliberate replay. +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct FanoutChild { + pub parent_seq: u64, + pub index: u64, +} + +fn child_opts() -> InsertOpts { + InsertOpts { + queue: QUEUE.to_string(), + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + ..UniqueOpts::default() + }), + ..Default::default() + } +} + +fn is_unique_violation(error: &AwaError) -> bool { + match error { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(database)) => { + database.code().as_deref() == Some("23505") + } + _ => false, + } +} + +struct ParentWorker { + shared: Arc, +} + +impl ParentWorker { + /// Records rows the handler enqueued. A child can run before this is reached; the + /// ledger merges the two. + fn record(&self, rows: &[JobRow]) { + let persona = Persona::Fanout; + if let Some(queue) = self.shared.queues.get(QUEUE) { + queue + .enqueued + .fetch_add(rows.len() as i64, Ordering::Relaxed); + } + self.shared + .stats(persona) + .add("enqueued_children", rows.len() as i64); + register_rows( + &self.shared, + persona, + rows, + |_| expect::COMPLETED, + flag::NO_EVENTS, + ); + } + + async fn insert_batch( + &self, + pool: &PgPool, + parent_seq: u64, + indexes: impl Iterator + Clone, + ) -> Result<(), JobError> { + let stats = self.shared.stats(Persona::Fanout); + let params: Vec = indexes + .clone() + .map(|index| { + awa::insert::params_with(&FanoutChild { parent_seq, index }, child_opts()) + .map_err(JobError::retryable) + }) + .collect::>()?; + let started = Instant::now(); + match insert_many(pool, ¶ms).await { + Ok(rows) => { + stats.observe("batch_insert_ms", started.elapsed().as_millis() as i64); + self.record(&rows); + Ok(()) + } + Err(error) if is_unique_violation(&error) => { + stats.add("batch_unique_conflicts", 1); + for index in indexes { + match insert_with(pool, &FanoutChild { parent_seq, index }, child_opts()).await + { + Ok(row) => { + stats.add("fallback_rows_inserted", 1); + self.record(&[row]); + } + Err(error) if is_unique_violation(&error) => { + stats.add("fallback_rows_skipped", 1) + } + Err(error) => return Err(JobError::retryable(error)), + } + } + Ok(()) + } + Err(error) => { + stats.add("batch_insert_errors", 1); + Err(JobError::retryable(error)) + } + } + } +} + +#[async_trait::async_trait] +impl Worker for ParentWorker { + fn kind(&self) -> &'static str { + "fanout_parent" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::Fanout, &ctx.job, false); + let args: FanoutParent = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + self.insert_batch(ctx.pool(), args.seq, 0..args.children) + .await?; + if args.replay { + // Index 0 is already claimed, so the batch fails as a whole. + let fresh = args.children..args.children * 2 - 1; + self.insert_batch(ctx.pool(), args.seq, std::iter::once(0).chain(fresh)) + .await?; + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +struct ChildWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for ChildWorker { + fn kind(&self) -> &'static str { + "fanout_child" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + // No lifecycle hooks: a hook per child would add a read per job. + let _guard = self.shared.begin_execution(Persona::Fanout, &ctx.job, true); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE, queue(16)) + .register_worker(ParentWorker { + shared: shared.clone(), + }) + .register_worker(ChildWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["fanout_parent"]) + .build() + .expect("build fanout client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Fanout; + let queue_info = shared.queues.register(QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "fanout", + persona, + 2, + 12, + 2, + vec![(queue_info.clone(), 16)], + client_factory(shared.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_fanout_producer", + shared.budget.reserve("fanout-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "fanout-producer"); + let tick = Duration::from_millis(500); + let parents_per_tick = shared.config.scaled(2) as usize; + let mut seq = 0u64; + shared.note( + "fanout", + format!( + "{parents_per_tick} parents per {} ms tick, {MIN_CHILDREN}-{MAX_CHILDREN} unique children each, {REPLAY_PERCENT}% replay a used key", + tick.as_millis() + ), + ); + + while !shared.load_stop.sleep(tick).await { + let mut params = Vec::with_capacity(parents_per_tick); + for _ in 0..parents_per_tick { + seq += 1; + params.push( + awa::insert::params_with( + &FanoutParent { + seq, + children: rng.range(MIN_CHILDREN, MAX_CHILDREN), + replay: rng.chance(REPLAY_PERCENT), + }, + InsertOpts { + queue: QUEUE.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + queue_info + .enqueued + .fetch_add(rows.len() as i64, Ordering::Relaxed); + shared + .stats(persona) + .add("enqueued_parents", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note( + "fanout", + format!("load stopped after {seq} parents; draining"), + ); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/mia.rs b/awa/tests/persona_soak/personas/mia.rs new file mode 100644 index 000000000..07d77f276 --- /dev/null +++ b/awa/tests/persona_soak/personas/mia.rs @@ -0,0 +1,260 @@ +//! Mia, an ML pipeline: jobs from seconds to about a minute (compressed), deadlines +//! that sometimes expire, heartbeats and progress, retries into the DLQ and latency +//! spikes. Stresses deadline and stale rescues racing heartbeats, completions and +//! the DLQ. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert, insert_many}; +use awa::{Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::Duration; + +pub const TRAIN_QUEUE: &str = "awa_soak_mia_train"; +pub const INFER_QUEUE: &str = "awa_soak_mia_infer"; +pub const DEADLINE_QUEUE: &str = "awa_soak_mia_deadline"; +/// Deadline given to jobs in the deadline queue; they run longer than this. +const DEADLINE: chrono::Duration = chrono::Duration::seconds(6); + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct MiaJob { + pub seq: u64, + pub behavior: String, + pub work_ms: u64, +} + +struct MiaWorker { + shared: Arc, +} + +/// Sleeps in slices, reporting progress and stopping early when the job is cancelled +/// or the runtime is shutting down. +async fn work(ctx: &JobContext, total_ms: u64) -> bool { + let slice = 500; + let mut elapsed = 0; + while elapsed < total_ms { + if ctx.is_cancelled() { + return false; + } + let step = slice.min(total_ms - elapsed); + tokio::time::sleep(Duration::from_millis(step)).await; + elapsed += step; + if total_ms >= 4_000 && elapsed % 2_000 < slice { + ctx.set_progress(((elapsed * 100) / total_ms).min(99) as u8, "training"); + let _ = ctx.flush_progress().await; + } + } + true +} + +#[async_trait::async_trait] +impl Worker for MiaWorker { + fn kind(&self) -> &'static str { + "mia_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Mia, &ctx.job, false); + let args: MiaJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + let attempt = ctx.job.attempt; + let finished = match args.behavior.as_str() { + "flaky" if attempt < 3 => { + work(ctx, args.work_ms.min(1_500)).await; + return Err(JobError::retryable_msg("transient model server error")); + } + "doomed" => { + work(ctx, args.work_ms.min(1_500)).await; + return Err(JobError::retryable_msg("model artefact missing")); + } + "cancel" => { + work(ctx, args.work_ms.min(1_000)).await; + return Ok(JobResult::Cancel("rejected by validation".into())); + } + // The first attempt overruns its deadline; later attempts often recover. + "overrun" if attempt >= 2 && !args.seq.is_multiple_of(3) => work(ctx, 1_000).await, + _ => work(ctx, args.work_ms).await, + }; + if !finished { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(TRAIN_QUEUE, queue(24)) + .queue(INFER_QUEUE, queue(16)) + .queue(DEADLINE_QUEUE, queue(12)) + .register_worker(MiaWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["mia_job"]) + .build() + .expect("build mia client") + }) +} + +/// One enqueued job: queue, behaviour, how long it works, deadline, max attempts and +/// the outcomes it may end in. +struct Plan { + queue: &'static str, + behavior: &'static str, + work_ms: u64, + max_attempts: i16, + expect: u8, +} + +fn plan(rng: &mut Rng) -> Plan { + let roll = rng.below(100); + if roll < 45 { + Plan { + queue: INFER_QUEUE, + behavior: "fast", + work_ms: rng.range(100, 2_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 70 { + Plan { + queue: TRAIN_QUEUE, + behavior: "train", + work_ms: rng.range(2_000, 20_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 78 { + // Latency spike: a long job that heartbeats and reports progress throughout. + Plan { + queue: TRAIN_QUEUE, + behavior: "spike", + work_ms: rng.range(30_000, 60_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 88 { + Plan { + queue: DEADLINE_QUEUE, + behavior: "overrun", + work_ms: rng.range(10_000, 25_000), + max_attempts: 3, + expect: expect::COMPLETED | expect::FAILED, + } + } else if roll < 94 { + Plan { + queue: INFER_QUEUE, + behavior: "flaky", + work_ms: 500, + max_attempts: 5, + expect: expect::COMPLETED | expect::FAILED, + } + } else if roll < 98 { + Plan { + queue: INFER_QUEUE, + behavior: "doomed", + work_ms: 500, + max_attempts: 3, + expect: expect::FAILED, + } + } else { + Plan { + queue: INFER_QUEUE, + behavior: "cancel", + work_ms: 500, + max_attempts: 3, + expect: expect::CANCELLED, + } + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Mia; + let train = shared.queues.register(TRAIN_QUEUE, persona, false); + let infer = shared.queues.register(INFER_QUEUE, persona, false); + let deadline = shared.queues.register(DEADLINE_QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "mia", + persona, + 2, + 14, + 2, + vec![(train, 24), (infer, 16), (deadline, 12)], + client_factory(shared.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_mia_producer", + shared.budget.reserve("mia-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + let rate_per_second = (6.0 * shared.config.scale).max(0.5); + let tick = Duration::from_millis(250); + let mut rng = Rng::for_task(shared.config.seed, "mia-producer"); + let mut seq = 0u64; + let mut carry = 0.0; + shared.note("mia", format!("producing about {rate_per_second:.1} jobs/s across {TRAIN_QUEUE}, {INFER_QUEUE} and {DEADLINE_QUEUE}")); + + while !shared.load_stop.sleep(tick).await { + carry += rate_per_second * tick.as_secs_f64(); + let count = carry as usize; + carry -= count as f64; + if count == 0 { + continue; + } + let mut params = Vec::with_capacity(count); + let mut expects = Vec::with_capacity(count); + for _ in 0..count { + seq += 1; + let plan = plan(&mut rng); + let args = MiaJob { + seq, + behavior: plan.behavior.to_string(), + work_ms: plan.work_ms, + }; + params.push( + insert::params_with( + &args, + InsertOpts { + queue: plan.queue.to_string(), + max_attempts: plan.max_attempts, + deadline_duration: (plan.queue == DEADLINE_QUEUE).then_some(DEADLINE), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + expects.push(plan.expect); + shared + .stats(persona) + .add(&format!("planned_{}", plan.behavior), 1); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows(&shared, persona, &rows, |index| expects[index], 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note("mia", format!("load stopped after {seq} jobs; draining")); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(30), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/mod.rs b/awa/tests/persona_soak/personas/mod.rs new file mode 100644 index 000000000..bfebd3462 --- /dev/null +++ b/awa/tests/persona_soak/personas/mod.rs @@ -0,0 +1,117 @@ +//! The personas. Each one owns its queues, job kinds, runtimes and producers, and +//! reports what it did through [`Shared::stats`]. + +pub mod abandon; +pub mod alex; +pub mod chaos; +pub mod fanout; +pub mod mia; +pub mod priya; +pub mod python; +pub mod sam; +pub mod shardops; +pub mod tom; +pub mod ui; + +use super::config::Persona; +use super::ledger::{expect, JobRecord}; +use super::shared::Shared; +use awa::{ClientBuilder, JobRow, QueueConfig}; +use std::sync::Arc; +use std::time::Duration; + +/// Registers the lifecycle event hook for each kind so the ledger learns about +/// terminal outcomes as they commit. +pub fn attach_events( + mut builder: ClientBuilder, + shared: &Arc, + kinds: &[&str], +) -> ClientBuilder { + for kind in kinds { + let shared = shared.clone(); + builder = builder.on_event_kind(*kind, move |event| { + let shared = shared.clone(); + async move { shared.on_event(&event) } + }); + } + builder +} + +/// Runtime timings shared by every fleet: fast enough that rescues, promotion and +/// rotation happen many times in a short run, but with realistic ratios between them. +pub fn tune(builder: ClientBuilder, shared: &Shared) -> ClientBuilder { + // The maintenance leader evaluates only the schedules its own runtime declares, and + // any runtime can lead, so every fleet declares the same set. + let builder = if shared.config.has(Persona::Tom) { + tom::declare_cron(builder) + } else { + builder + }; + let builder = timings( + builder, + shared.config.heartbeat_interval, + shared.config.heartbeat_staleness, + ); + match shared.config.rescue_interval { + Some(interval) => builder + .heartbeat_rescue_interval(interval) + .deadline_rescue_interval(interval), + None => builder, + } +} + +/// The timings of [`tune`] that do not depend on the run, for runtimes that live in a +/// process of their own. +pub fn timings( + builder: ClientBuilder, + heartbeat_interval: Duration, + heartbeat_staleness: Duration, +) -> ClientBuilder { + builder + .heartbeat_interval(heartbeat_interval) + .heartbeat_staleness(heartbeat_staleness) + .heartbeat_rescue_interval(Duration::from_secs(3)) + .deadline_rescue_interval(Duration::from_secs(2)) + .promote_interval(Duration::from_millis(250)) + .runtime_snapshot_interval(Duration::from_secs(5)) + .dlq_enabled_by_default(true) +} + +pub fn queue(max_workers: u32) -> QueueConfig { + QueueConfig { + max_workers, + ..QueueConfig::default() + } +} + +/// Records freshly enqueued rows in the ledger. +pub fn register_rows( + shared: &Shared, + persona: Persona, + rows: &[JobRow], + expect_for: impl Fn(usize) -> u8, + flags: u32, +) { + let now = shared.now_ms() as u32; + for (index, row) in rows.iter().enumerate() { + let queue = shared + .queues + .get(&row.queue) + .map_or(u32::MAX, |queue| queue.index); + shared.ledger.register( + row.id, + JobRecord::new(persona, queue, now, expect_for(index), flags), + ); + } +} + +pub const ANY: u8 = expect::ANY_TERMINAL; + +/// Short pause used by producers after an error, so a dead connection does not spin. +pub async fn backoff(shared: &Shared, error: impl std::fmt::Display, persona: Persona) { + shared.stats(persona).add("enqueue_errors", 1); + if shared.stats(persona).get("enqueue_errors") <= 5 { + shared.note(persona.name(), format!("enqueue error: {error}")); + } + shared.load_stop.sleep(Duration::from_millis(500)).await; +} diff --git a/awa/tests/persona_soak/personas/priya.rs b/awa/tests/persona_soak/personas/priya.rs new file mode 100644 index 000000000..e674d9a31 --- /dev/null +++ b/awa/tests/persona_soak/personas/priya.rs @@ -0,0 +1,203 @@ +//! Priya, firehose: one hot queue, large batch inserts (above the bind-parameter +//! limit), COPY inserts, tiny jobs and several clients. Stresses bind limits, ring +//! rotation and prune under load and claim contention. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{backoff, queue, register_rows, tune}; +use awa::model::{insert_many, insert_many_copy_from_pool}; +use awa::{Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const HOT_QUEUE: &str = "awa_soak_priya_hot"; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct PriyaJob { + pub seq: u64, + pub source: String, +} + +struct PriyaWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for PriyaWorker { + fn kind(&self) -> &'static str { + "priya_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Priya, &ctx.job, true); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + // No lifecycle hooks: a hook per completion would add a read per job and + // change the load this persona exists to apply. + tune(Client::builder(pool).queue(HOT_QUEUE, queue(48)), &shared) + .register_worker(PriyaWorker { + shared: shared.clone(), + }) + .build() + .expect("build priya client") + }) +} + +fn batch(seq: &mut u64, count: usize, source: &str) -> Vec { + (0..count) + .map(|_| { + *seq += 1; + awa::insert::params_with( + &PriyaJob { + seq: *seq, + source: source.to_string(), + }, + InsertOpts { + queue: HOT_QUEUE.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .expect("priya params") + }) + .collect() +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Priya; + let hot = shared.queues.register(HOT_QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "priya", + persona, + 3, + 12, + 3, + vec![(hot, 48)], + client_factory(shared.clone()), + ) + .await?; + + let pool_size = shared.budget.reserve("priya-producer", 6, 2); + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_priya_producer", + pool_size, + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "priya-producer"); + let burst_size = shared.config.scaled(100_000) as usize; + let copy_size = shared.config.scaled(10_000) as usize; + let steady_per_tick = shared.config.scaled(30) as usize; + let burst_every = Duration::from_secs(180); + let copy_every = Duration::from_secs(45); + let mut next_burst = Instant::now() + Duration::from_secs(10); + let mut next_copy = Instant::now() + Duration::from_secs(25); + let mut seq = 0u64; + shared.note( + "priya", + format!("steady {} jobs/s, {burst_size}-row insert_many burst every {}s, {copy_size}-row COPY every {}s", steady_per_tick * 10, burst_every.as_secs(), copy_every.as_secs()), + ); + + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + let now = Instant::now(); + if now >= next_burst { + next_burst = now + burst_every; + // Vary the size around the target so chunk boundaries move between bursts. + let size = burst_size + rng.below(5_000) as usize; + let params = batch(&mut seq, size, "burst"); + let started = Instant::now(); + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + shared.stats(persona).add("bursts", 1); + shared + .stats(persona) + .observe("burst_insert_ms", started.elapsed().as_millis() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_BURST, + ); + shared.note( + "priya", + format!( + "burst of {} rows inserted in {:?}", + rows.len(), + started.elapsed() + ), + ); + } + Err(error) if shared.chaos_active() => { + backoff(&shared, format!("burst of {size}: {error}"), persona).await + } + Err(error) => { + shared.violation( + "insert_many_failed", + format!("insert_many of {size} rows failed: {error}"), + vec![], + ); + } + } + } + if now >= next_copy { + next_copy = now + copy_every; + let params = batch(&mut seq, copy_size, "copy"); + let started = Instant::now(); + match insert_many_copy_from_pool(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + shared.stats(persona).add("copy_batches", 1); + shared + .stats(persona) + .observe("copy_insert_ms", started.elapsed().as_millis() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_COPY, + ); + } + Err(error) => { + backoff(&shared, format!("copy of {copy_size}: {error}"), persona).await + } + } + } + let params = batch(&mut seq, steady_per_tick, "steady"); + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_STEADY, + ); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note("priya", format!("load stopped after {seq} jobs; draining")); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/python.rs b/awa/tests/persona_soak/personas/python.rs new file mode 100644 index 000000000..4c0808cfe --- /dev/null +++ b/awa/tests/persona_soak/personas/python.rs @@ -0,0 +1,887 @@ +//! Python users: a separate process built on the `awa-python` bindings, started and +//! stopped by the harness. It produces (single inserts, `insert_many` above the old +//! bind-parameter limit, COPY, and enqueues inside transactions that sometimes roll +//! back), runs async workers (completing, retryable, `RetryAfter`, `Snooze`, progress, +//! deadline overruns and a handler that ignores its deadline) and exercises the +//! migrate-then-query startup path. One queue is shared with a Rust runtime so Rust and +//! Python workers claim from the same lanes. The process writes its own per-job ledger, +//! which this module folds into the harness ledger for the oracle. + +use super::super::config::Persona; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::python_probe::{self, Evidence, Phase}; +use super::super::shared::Shared; +use super::{attach_events, queue, tune}; +use awa::{Client, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashSet; +use std::io::{BufRead, Seek, SeekFrom}; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::{Child, Command}; + +pub const QUEUE_WORK: &str = "awa_soak_py_work"; +pub const QUEUE_DEADLINE: &str = "awa_soak_py_deadline"; +pub const QUEUE_SHARED: &str = "awa_soak_py_shared"; +const SCRIPT: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/persona_soak/python/soak_python.py" +); + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct PySharedJob { + pub seq: u64, + pub behavior: String, + pub work_ms: u64, +} + +struct SharedQueueWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for SharedQueueWorker { + fn kind(&self) -> &'static str { + "py_shared_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::Python, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(50); + match ctx.job.args["behavior"].as_str().unwrap_or("ok") { + "retry" if ctx.job.attempt < 3 => { + tokio::time::sleep(Duration::from_millis(50)).await; + return Err(JobError::retryable_msg("transient failure")); + } + "progress" => { + for step in 1..=4u8 { + tokio::time::sleep(Duration::from_millis(work_ms / 4)).await; + ctx.set_progress(step * 25, "step"); + let _ = ctx.flush_progress().await; + } + } + _ => tokio::time::sleep(Duration::from_millis(work_ms)).await, + } + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +/// Handle to the running Python process, shared with the chaos persona. +pub struct PythonProcess { + child: Option, + ledger_path: PathBuf, + python: PathBuf, + restarts: u32, +} + +fn spawn( + shared: &Arc, + python: &PathBuf, + ledger_path: &PathBuf, + role: &str, +) -> Result { + spawn_process(shared, python, ledger_path, role, "soak_python", None) +} + +fn spawn_process( + shared: &Arc, + python: &PathBuf, + ledger_path: &PathBuf, + role: &str, + application_name: &str, + probe_phase: Option>, +) -> Result { + let mut child = Command::new(python) + .arg(SCRIPT) + .env( + "SOAK_DATABASE_URL", + format!( + "{}{}application_name={}", + shared.config.database_url, + if shared.config.database_url.contains('?') { + '&' + } else { + '?' + }, + application_name + ), + ) + .env("SOAK_LEDGER", ledger_path) + .env("SOAK_SEED", shared.config.seed.to_string()) + .env("SOAK_SCALE", shared.config.scale.to_string()) + .env( + "SOAK_PACE", + (shared.config.load.as_secs_f64() / 900.0) + .clamp(0.1, 1.0) + .to_string(), + ) + .env("SOAK_ROLE", role) + .env( + "SOAK_DRAIN", + if shared.load_stop.is_stopped() { + "1" + } else { + "0" + }, + ) + .env("SOAK_QUEUE_WORK", QUEUE_WORK) + .env("SOAK_QUEUE_DEADLINE", QUEUE_DEADLINE) + .env("SOAK_QUEUE_SHARED", QUEUE_SHARED) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|error| format!("spawn {}: {error}", python.display()))?; + for stream in [ + child + .stdout + .take() + .map(|s| Box::new(s) as Box), + child + .stderr + .take() + .map(|s| Box::new(s) as Box), + ] + .into_iter() + .flatten() + { + let shared = shared.clone(); + let probe_phase = probe_phase.clone(); + let application_name = application_name.to_string(); + tokio::spawn(async move { + let mut lines = BufReader::new(stream).lines(); + while let Ok(Some(line)) = lines.next_line().await { + if let (Some(phase), Some(marker)) = (&probe_phase, Phase::marker(&line)) { + // A close marker alone does not prove migrate-and-query + // reached readiness. IDLE also follows that completed call. + if marker != Phase::Closed || phase.load(Ordering::Acquire) > 0 { + phase.fetch_max(marker as u8, Ordering::Release); + } + shared.note("python", format!("{application_name}: {line}")); + } + absorb_output(&shared, &line); + } + }); + } + Ok(child) +} + +/// Turns the interesting stdout lines of the Python process into report counters. +fn absorb_output(shared: &Shared, line: &str) { + let stats = shared.stats(Persona::Python); + let field = |name: &str| { + line.split_whitespace() + .find_map(|part| part.strip_prefix(&format!("{name}="))) + .and_then(|value| value.parse::().ok()) + }; + if line.starts_with("STAT ") { + if let Some(active) = field("active") { + shared + .python_active + .store(active as i64, std::sync::atomic::Ordering::Relaxed); + } + if let Some(max_active) = field("max_active") { + stats.set_max("python_max_concurrent_handlers", max_active as i64); + } + if let Some(errors) = field("errors") { + stats.set_max("python_errors", errors as i64); + } + } else if line.starts_with("BATCH ") { + if let (Some(rows), Some(seconds)) = (field("rows"), field("seconds")) { + stats.add("python_batches", 1); + stats.add("python_batch_rows", rows as i64); + stats.observe("python_batch_insert_ms", (seconds * 1000.0) as i64); + } + } else if line.starts_with("MIGRATED") { + stats.add("python_migrations", 1); + } else if line.starts_with("ERROR") || line.contains("Traceback") { + if stats.get("python_error_lines") < 5 { + shared.note("python", line.chars().take(200).collect::()); + } + stats.add("python_error_lines", 1); + } +} + +pub struct PythonHandle { + process: tokio::sync::Mutex, +} + +// Used in both phases: after load stops a dead process still needs a worker +// replacement, but spawn_process disables its producers through SOAK_DRAIN. +fn restart_if_needed( + child: &mut Option, + start: impl FnOnce() -> Result, +) -> Result, String> { + let reason = match child.as_mut() { + Some(child) => match child.try_wait().map_err(|error| error.to_string())? { + Some(status) => format!("process exited with {status}"), + None => return Ok(None), + }, + None => "process was absent".to_string(), + }; + child.take(); + *child = Some(start()?); + Ok(Some(reason)) +} + +fn finished_probe_errors(probes: &mut tokio::task::JoinSet<()>) -> Vec { + let mut errors = Vec::new(); + while let Some(result) = probes.try_join_next() { + if let Err(error) = result { + errors.push(error.to_string()); + } + } + errors +} + +async fn begin_drain( + stop_producers: impl std::future::Future, + probes: &mut tokio::task::JoinSet<()>, +) -> Vec { + // Startup probes may still be blocked. Stop production before waiting for + // them, and leave their observation tasks running beside worker recovery. + stop_producers.await; + finished_probe_errors(probes) +} + +impl PythonHandle { + async fn ensure_running(&self, shared: &Arc) { + let mut process = self.process.lock().await; + let python = process.python.clone(); + let ledger_path = process.ledger_path.clone(); + match restart_if_needed(&mut process.child, || { + spawn(shared, &python, &ledger_path, "full") + }) { + Ok(Some(reason)) => { + process.restarts += 1; + shared + .stats(Persona::Python) + .add("python_unexpected_exits", 1); + shared.note( + "python", + format!( + "{reason}; restarted (draining={})", + shared.load_stop.is_stopped() + ), + ); + } + Ok(None) => {} + Err(error) => shared.incomplete( + format!("Python worker liveness/restart check failed: {error}"), + vec![], + ), + } + } + /// Kills the process without a shutdown handshake and starts a new one after `downtime`. + pub async fn kill_and_restart( + &self, + shared: &Arc, + downtime: Duration, + ) -> Result<(), String> { + let mut process = self.process.lock().await; + if let Some(mut child) = process.child.take() { + let _ = child.kill().await; + } + tokio::time::sleep(downtime).await; + shared + .python_active + .store(0, std::sync::atomic::Ordering::Relaxed); + process.restarts += 1; + let child = spawn( + shared, + &process.python.clone(), + &process.ledger_path.clone(), + "full", + )?; + process.child = Some(child); + shared.stats(Persona::Python).add("python_kills", 1); + Ok(()) + } +} + +/// Folds new lines of the Python ledger into the harness ledger. +struct Importer { + path: PathBuf, + offset: u64, + rolled_back: HashSet, +} + +impl Importer { + fn import(&mut self, shared: &Shared) { + let Ok(mut file) = std::fs::File::open(&self.path) else { + return; + }; + if file.seek(SeekFrom::Start(self.offset)).is_err() { + return; + } + let mut reader = std::io::BufReader::new(file); + let mut line = String::new(); + loop { + line.clear(); + match reader.read_line(&mut line) { + Ok(0) | Err(_) => break, + Ok(read) => { + if !line.ends_with('\n') { + break; + } + self.offset += read as u64; + if let Ok(event) = serde_json::from_str::(&line) { + self.apply(shared, &event); + } + } + } + } + } + + fn apply(&mut self, shared: &Shared, event: &Value) { + let now = shared.now_ms() as u32; + let id = event.get("id").and_then(Value::as_i64); + let queue_of = |name: Option<&str>| { + name.and_then(|name| shared.queues.get(name)) + .map_or(u32::MAX, |queue| queue.index) + }; + match event.get("e").and_then(Value::as_str) { + Some("enq") => { + let Some(id) = id else { return }; + let allowed = event.get("x").and_then(Value::as_u64).unwrap_or(7) as u8; + let queue = shared + .queues + .get(event.get("q").and_then(Value::as_str).unwrap_or("")); + if let Some(queue) = &queue { + queue + .enqueued + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + let index = queue.map_or(u32::MAX, |queue| queue.index); + let mut record = + JobRecord::new(Persona::Python, index, now, allowed, flag::NO_EVENTS); + record.max_attempt = 0; + // Keep what the Rust side may have seen already (shared queue). + shared.ledger.with( + id, + || record.clone(), + |existing| { + existing.expect = allowed; + existing.flags &= !flag::UNREGISTERED; + existing.persona = Persona::Python; + }, + ); + shared.stats(Persona::Python).add("python_enqueued", 1); + } + Some("start") => { + let Some(id) = id else { return }; + let queue = queue_of(event.get("q").and_then(Value::as_str)); + let attempt = event.get("attempt").and_then(Value::as_i64).unwrap_or(0) as i16; + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + queue, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| { + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(attempt); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + }, + ); + if let Some(queue) = event + .get("q") + .and_then(Value::as_str) + .and_then(|name| shared.queues.get(name)) + { + queue + .started + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + shared.stats(Persona::Python).add("python_executions", 1); + } + Some("end") => { + let Some(id) = id else { return }; + let result = event.get("r").and_then(Value::as_str).unwrap_or(""); + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + u32::MAX, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| match result { + // A completed return is evidence only for jobs that must complete; for + // the rest the outcome is read from the database, since a late return + // after a deadline rescue is discarded by the store. + "completed" if record.expect == expect::COMPLETED => { + record.flags |= flag::RETURNED_COMPLETED + } + "cancel" => { + record.flags |= flag::HANDLER_CANCELLED; + record.cancel_events = record.cancel_events.saturating_add(1); + } + _ => {} + }, + ); + } + Some("rolled_back") => { + if let Some(ids) = event.get("ids").and_then(Value::as_array) { + for id in ids.iter().filter_map(Value::as_i64) { + self.rolled_back.insert(id); + } + shared + .stats(Persona::Python) + .add("python_rolled_back_jobs", ids.len() as i64); + } + } + Some("error") => shared.stats(Persona::Python).add("python_logged_errors", 1), + _ => {} + } + } +} + +/// A migrate-then-query run in a throwaway process; its database sessions must not be +/// left stranded while it idles, and must be gone once it exits. +async fn migrate_probe(shared: &Arc, python: &PathBuf, ledger_path: &PathBuf) { + let stats = shared.stats(Persona::Python); + let application_name = format!("soak_python_probe_{}", uuid::Uuid::new_v4().simple()); + let phase = Arc::new(AtomicU8::new(Phase::Migrating as u8)); + let mut child = match spawn_process( + shared, + python, + ledger_path, + "migrate_probe", + &application_name, + Some(phase.clone()), + ) { + Ok(child) => child, + Err(error) => { + shared.incomplete( + format!("Python migrate probe could not start: {error}"), + vec![], + ); + return; + } + }; + let pool = match super::super::db::connect_pool( + &shared.config.database_url, + "soak_python_probe_watch", + 1, + ) + .await + { + Ok(pool) => pool, + Err(error) => { + shared.incomplete( + format!("Python probe observer could not connect: {error}"), + vec![], + ); + let _ = child.kill().await; + return; + } + }; + let mut evidence = Evidence::default(); + let deadline = Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(45)); + let exit_success = loop { + match child.try_wait() { + Ok(Some(status)) => { + if !status.success() { + evidence.incomplete(format!("Python migrate probe exited with {status}")); + } + break status.success(); + } + Err(error) => { + evidence.incomplete(format!( + "Python migrate probe exit could not be read: {error}" + )); + let _ = child.kill().await; + break false; + } + Ok(None) if Instant::now() >= deadline => { + evidence + .incomplete("Python migrate-then-query probe did not exit within its deadline"); + let _ = child.kill().await; + break false; + } + Ok(None) => {} + } + // Capture readiness before starting the query: a MIGRATED marker that + // arrives during it must not relabel an earlier migration snapshot. + let before_query = Phase::from_byte(phase.load(Ordering::Acquire)); + evidence.observe( + before_query, + python_probe::sessions(&pool, &application_name).await, + ); + tokio::time::sleep(Duration::from_millis(500)).await; + }; + evidence + .finish_exit(&phase, exit_success, || { + python_probe::sessions(&pool, &application_name) + }) + .await; + stats.add("python_migrate_probes", 1); + stats.set_max("python_probe_peak_sessions", evidence.peak_sessions); + stats.set_max( + "python_probe_migrating_idle_transactions", + evidence.migrating_idle_transactions, + ); + stats.add("python_probe_ready_samples", evidence.ready_samples as i64); + if !evidence.complete() { + shared.incomplete( + format!( + "Python migrate probe {application_name} has incomplete evidence: {evidence:?}" + ), + vec![], + ); + } + if evidence.violation() { + shared.violation( + "python_connections_stranded", + format!("Python migrate probe {application_name}: peak={} against pool4, idle-in-transaction after ready={}, sessions after exit={:?}", evidence.peak_sessions, evidence.ready_idle_transactions, evidence.remaining), + vec![], + ); + } + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Python; + let python = shared + .config + .python_bin + .clone() + .ok_or("the python persona needs AWA_SOAK_PYTHON, the interpreter of a virtualenv built with `maturin develop` (see scripts/persona-soak.sh)")?; + let shared_queue = shared.queues.register(QUEUE_SHARED, persona, false); + shared.queues.register(QUEUE_WORK, persona, false); + shared.queues.register(QUEUE_DEADLINE, persona, false); + + // The Rust runtime that shares the queue with the Python workers. + let rust_shared = shared.clone(); + let factory: super::super::db::ClientFactory = Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE_SHARED, queue(6)) + .register_worker(SharedQueueWorker { + shared: rust_shared.clone(), + }); + attach_events( + tune(builder, &rust_shared), + &rust_shared, + &["py_shared_job"], + ) + .build() + .expect("build python shared-queue client") + }); + let fleet = Fleet::start( + &shared, + "pyrust", + persona, + 1, + 8, + 2, + vec![(shared_queue, 6)], + factory, + ) + .await?; + + std::fs::create_dir_all(&shared.config.report_dir).map_err(|error| error.to_string())?; + let ledger_path = shared.config.report_dir.join("python-ledger.jsonl"); + let _ = std::fs::remove_file(&ledger_path); + let child = spawn(&shared, &python, &ledger_path, "full")?; + let handle = Arc::new(PythonHandle { + process: tokio::sync::Mutex::new(PythonProcess { + child: Some(child), + ledger_path: ledger_path.clone(), + python: python.clone(), + restarts: 0, + }), + }); + shared + .python + .lock() + .expect("python handle") + .replace(handle.clone()); + shared.note("python", format!("{} started; queues {QUEUE_WORK}, {QUEUE_DEADLINE}, {QUEUE_SHARED} (shared with a Rust runtime)", python.display())); + + let tail_pool = + super::super::db::connect_pool(&shared.config.database_url, "soak_python_tail", 1) + .await + .map_err(|error| error.to_string())?; + let mut tail_tick = 0u64; + let mut probes = tokio::task::JoinSet::new(); + let mut importer = Importer { + path: ledger_path.clone(), + offset: 0, + rolled_back: HashSet::new(), + }; + let pace = (shared.config.load.as_secs_f64() / 900.0).clamp(0.1, 1.0); + let mut next_probe = Instant::now() + Duration::from_secs_f64(30.0 * pace.max(0.3)); + let probe_every = Duration::from_secs_f64(150.0 * pace.max(0.3)); + while !shared.load_stop.sleep(Duration::from_secs(1)).await { + importer.import(&shared); + tail_terminal_rows(&shared, &tail_pool, &mut tail_tick).await; + handle.ensure_running(&shared).await; + if Instant::now() >= next_probe { + next_probe = Instant::now() + probe_every; + let (shared, python, ledger_path) = + (shared.clone(), python.clone(), ledger_path.clone()); + probes.spawn(async move { migrate_probe(&shared, &python, &ledger_path).await }); + } + } + + // Stop producing immediately; outstanding migration probes must not delay + // this signal or the liveness/evidence checks during the fixed drain grace. + let errors = begin_drain( + async { + if let Some(pid) = handle + .process + .lock() + .await + .child + .as_ref() + .and_then(Child::id) + { + match std::process::Command::new("kill") + .args(["-USR1", &pid.to_string()]) + .status() + { + Ok(status) if status.success() => {} + result => shared.incomplete( + format!("Python stop-production signal failed: {result:?}"), + vec![], + ), + } + } + }, + &mut probes, + ) + .await; + for error in errors { + shared.incomplete(format!("Python migration probe failed: {error}"), vec![]); + } + shared.note("python", "load stopped; draining"); + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + handle.ensure_running(&shared).await; + importer.import(&shared); + tail_terminal_rows(&shared, &tail_pool, &mut tail_tick).await; + for error in finished_probe_errors(&mut probes) { + shared.incomplete(format!("Python migration probe failed: {error}"), vec![]); + } + } + // Finished tasks retain their result. Only a task actually cancelled at + // shutdown lacks the required evidence; do not add time to the drain grace. + probes.abort_all(); + while let Some(result) = probes.join_next().await { + if let Err(error) = result { + shared.incomplete( + format!("Python migration probe unfinished or failed at drain shutdown: {error}"), + vec![], + ); + } + } + importer.import(&shared); + if let Some(mut child) = handle.process.lock().await.child.take() { + let _ = child.kill().await; + } + shared.python.lock().expect("python handle").take(); + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + importer.import(&shared); + check_rolled_back(&shared, &importer.rolled_back).await?; + Ok(()) +} + +/// Python jobs report no lifecycle events, and ring rotation prunes their terminal rows +/// within seconds, so their terminal outcomes are read from the store as they appear. +async fn tail_terminal_rows(shared: &Arc, pool: &sqlx::PgPool, tick: &mut u64) { + *tick += 1; + if !(*tick).is_multiple_of(3) { + return; + } + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".into()); + let queues = vec![ + QUEUE_WORK.to_string(), + QUEUE_DEADLINE.to_string(), + QUEUE_SHARED.to_string(), + ]; + let rows: Result, _> = sqlx::query_as(awa::audited_sql(format!( + "SELECT job_id, run_lease, state::text FROM {schema}.done_entries WHERE queue = ANY($1) AND finalized_at > now() - interval '25 seconds' + UNION ALL SELECT job_id, run_lease, 'failed' FROM {schema}.dlq_entries WHERE queue = ANY($1) AND dlq_at > now() - interval '25 seconds' + UNION ALL SELECT b.job_id, b.run_lease, 'completed' FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + WHERE c.queue = ANY($1) AND c.finalized_at > now() - interval '25 seconds'" + ))) + .bind(&queues) + .fetch_all(pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + shared.incomplete( + format!("Python terminal evidence query failed: {error}"), + vec![], + ); + shared.stats(Persona::Python).add("python_tail_errors", 1); + if shared.stats(Persona::Python).get("python_tail_errors") <= 3 { + shared.note("python", format!("terminal row tail failed: {error}")); + } + return; + } + }; + shared + .stats(Persona::Python) + .add("python_tail_rows", rows.len() as i64); + let now = shared.now_ms() as u32; + for (id, run_lease, state) in rows { + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + u32::MAX, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| { + record.note_observed_terminal(run_lease, &state); + }, + ); + } +} + +/// A transaction rolled back by the Python application must leave no job anywhere. +async fn check_rolled_back(shared: &Arc, ids: &HashSet) -> Result<(), String> { + if ids.is_empty() { + return Ok(()); + } + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".into()); + let pool = super::super::db::connect_pool(&shared.config.database_url, "soak_python_check", 1) + .await + .map_err(|error| error.to_string())?; + let ids: Vec = ids.iter().copied().collect(); + let mut visible: Vec = Vec::new(); + for chunk in ids.chunks(5000) { + let found: Vec = sqlx::query_scalar(awa::audited_sql(format!( + "SELECT job_id FROM {schema}.ready_entries WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.deferred_jobs WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.done_entries WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.dlq_entries WHERE job_id = ANY($1)" + ))) + .bind(chunk) + .fetch_all(&pool) + .await + .map_err(|error| error.to_string())?; + visible.extend(found); + } + let executed: Vec = ids + .iter() + .copied() + .filter(|id| { + shared + .ledger + .get(*id) + .is_some_and(|record| record.executions > 0) + }) + .collect(); + shared + .stats(Persona::Python) + .add("python_rolled_back_checked", ids.len() as i64); + if !visible.is_empty() || !executed.is_empty() { + shared.violation( + "rolled_back_job_visible", + format!("{} job(s) enqueued in rolled-back Python transactions exist in the store, {} of them executed", visible.len(), executed.len()), + visible.iter().take(10).map(|id| format!("{{\"job_id\":{id}}}")).collect(), + ); + } + pool.close().await; + Ok(()) +} + +#[cfg(test)] +mod lifecycle_tests { + use super::*; + + fn finished_child() -> Result { + Command::new(std::env::current_exe().map_err(|error| error.to_string())?) + .arg("--list") + .stdout(Stdio::null()) + .spawn() + .map_err(|error| error.to_string()) + } + + #[tokio::test] + async fn dead_child_is_restarted_for_the_drain_liveness_loop() { + let mut child = Some(finished_child().unwrap()); + child.as_mut().unwrap().wait().await.unwrap(); + let result = restart_if_needed(&mut child, finished_child).unwrap(); + assert!(result.is_some(), "a dead drain worker needs a replacement"); + assert!(child.is_some()); + child.as_mut().unwrap().wait().await.unwrap(); + } + + #[tokio::test] + async fn outstanding_probe_does_not_delay_stop_production_or_worker_recovery() { + let (release, blocked) = tokio::sync::oneshot::channel(); + let (stopped, mut stop_seen) = tokio::sync::oneshot::channel(); + let mut probes = tokio::task::JoinSet::new(); + probes.spawn(async move { + let _ = blocked.await; + }); + let errors = tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(1)), + begin_drain( + async { + stopped.send(()).unwrap(); + }, + &mut probes, + ), + ) + .await + .expect("stop production must precede waiting on a blocked probe"); + assert!(errors.is_empty()); + assert!(stop_seen.try_recv().is_ok()); + assert!( + !probes.is_empty(), + "the probe remains blocked while drain proceeds" + ); + release.send(()).unwrap(); + probes.join_next().await.unwrap().unwrap(); + } + + #[test] + fn failed_restart_is_reported_instead_of_losing_the_process_silently() { + let mut child = None; + let result = restart_if_needed(&mut child, || Err("injected spawn failure".into())); + assert_eq!(result.unwrap_err(), "injected spawn failure"); + assert!(child.is_none()); + } +} diff --git a/awa/tests/persona_soak/personas/sam.rs b/awa/tests/persona_soak/personas/sam.rs new file mode 100644 index 000000000..b70876db0 --- /dev/null +++ b/awa/tests/persona_soak/personas/sam.rs @@ -0,0 +1,359 @@ +//! Sam, ordered streams: a queue with enqueue shards and ordering keys, a striped +//! queue (`queue_stripe_count`), scheduled jobs and a rate-limited queue. Stresses +//! routing, promotion and fairness. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert_many, QueueStorage, QueueStorageConfig}; +use awa::{ + Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, QueueConfig, RateLimit, Worker, +}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +pub const ORDERED_QUEUE: &str = "awa_soak_sam_ord"; +pub const RATE_QUEUE: &str = "awa_soak_sam_rate"; +pub const SCHEDULED_QUEUE: &str = "awa_soak_sam_sched"; +pub const STRIPED_QUEUE: &str = "awa_soak_sam_stripe"; +const KEYS: u64 = 64; +const ENQUEUE_SHARDS: i16 = 4; +const STRIPES: usize = 4; +const RATE_PER_SECOND: f64 = 15.0; +/// Starts this far behind the newest start for the same key count as order inversions. +const INVERSION_WINDOW: u64 = 64; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct SamJob { + pub stream: String, + pub key: u64, + pub seq: u64, +} + +#[derive(Default)] +struct Order { + newest_start: HashMap, + inversions: u64, + worst_distance: u64, + striped_seen: HashSet, + striped_repeats: u64, +} + +struct SamWorker { + shared: Arc, + order: Arc>, +} + +#[async_trait::async_trait] +impl Worker for SamWorker { + fn kind(&self) -> &'static str { + "sam_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Sam, &ctx.job, false); + let args: SamJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + { + let mut order = self.order.lock().expect("order"); + if args.stream == "ordered" { + let newest = order.newest_start.entry(args.key).or_insert(0); + if args.seq + INVERSION_WINDOW < *newest { + let distance = *newest - args.seq; + order.inversions += 1; + order.worst_distance = order.worst_distance.max(distance); + } else { + *newest = (*newest).max(args.seq); + } + } else if args.stream == "striped" && !order.striped_seen.insert(args.seq) { + order.striped_repeats += 1; + } + } + tokio::time::sleep(Duration::from_millis(5 + args.seq % 40)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn plain_factory(shared: Arc, order: Arc>) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(ORDERED_QUEUE, queue(12)) + .queue(SCHEDULED_QUEUE, queue(8)) + .queue( + RATE_QUEUE, + QueueConfig { + max_workers: 8, + rate_limit: Some(RateLimit { + max_rate: RATE_PER_SECOND, + burst: RATE_PER_SECOND as u32, + }), + ..QueueConfig::default() + }, + ) + .register_worker(SamWorker { + shared: shared.clone(), + order: order.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["sam_job"]) + .build() + .expect("build sam client") + }) +} + +fn striped_config() -> QueueStorageConfig { + QueueStorageConfig { + queue_stripe_count: STRIPES, + ..QueueStorageConfig::default() + } +} + +fn striped_factory(shared: Arc, order: Arc>) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(STRIPED_QUEUE, queue(16)) + .queue_storage( + striped_config(), + Duration::from_millis(1_000), + Duration::from_millis(1_000), + ) + .register_worker(SamWorker { + shared: shared.clone(), + order: order.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["sam_job"]) + .build() + .expect("build sam stripe client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Sam; + let ordered = shared.queues.register(ORDERED_QUEUE, persona, false); + let scheduled = shared.queues.register(SCHEDULED_QUEUE, persona, false); + let rated = shared.queues.register(RATE_QUEUE, persona, true); + let striped = shared.queues.register(STRIPED_QUEUE, persona, false); + + let setup_pool = connect_pool(&shared.config.database_url, "soak_sam_setup", 1) + .await + .map_err(|error| error.to_string())?; + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(ORDERED_QUEUE) + .bind(ENQUEUE_SHARDS) + .execute(&setup_pool) + .await + .map_err(|error| format!("set enqueue_shards: {error}"))?; + setup_pool.close().await; + + let order = Arc::new(Mutex::new(Order::default())); + let plain = Fleet::start( + &shared, + "sam", + persona, + 2, + 14, + 2, + vec![(ordered, 12), (scheduled, 8), (rated, 8)], + plain_factory(shared.clone(), order.clone()), + ) + .await?; + let stripe_fleet = Fleet::start( + &shared, + "sam-stripe", + persona, + 2, + 10, + 2, + vec![(striped, 16)], + striped_factory(shared.clone(), order.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_sam_producer", + shared.budget.reserve("sam-producer", 4, 1), + ) + .await + .map_err(|error| error.to_string())?; + let store = QueueStorage::new(striped_config()).map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "sam-producer"); + let scale = shared.config.scale; + let mut key_seq = vec![0u64; KEYS as usize]; + let mut stripe_seq = 0u64; + let mut seq = 0u64; + let mut rate_carry = 0.0_f64; + let tick = Duration::from_millis(100); + shared.note( + "sam", + format!("{ORDERED_QUEUE} ({ENQUEUE_SHARDS} shards, {KEYS} ordering keys), {STRIPED_QUEUE} ({STRIPES} stripes), {SCHEDULED_QUEUE}, {RATE_QUEUE} ({RATE_PER_SECOND}/s)"), + ); + + while !shared.load_stop.sleep(tick).await { + // Ordered stream with ordering keys. + let ordered_count = (4.0 * scale).ceil() as usize; + let mut params = Vec::new(); + for _ in 0..ordered_count { + let key = rng.below(KEYS); + key_seq[key as usize] += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "ordered".into(), + key, + seq: key_seq[key as usize], + }, + InsertOpts { + queue: ORDERED_QUEUE.into(), + ordering_key: Some(key.to_be_bytes().to_vec()), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(persona) + .add("enqueued_ordered", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // Scheduled jobs a few seconds out, and the rate-limited stream. + let mut params = Vec::new(); + for _ in 0..(1.0 * scale).ceil() as usize { + seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "scheduled".into(), + key: 0, + seq, + }, + InsertOpts { + queue: SCHEDULED_QUEUE.into(), + run_at: Some( + Utc::now() + + chrono::Duration::milliseconds(rng.range(1_000, 12_000) as i64), + ), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + // Stay below the queue's rate limit so a backlog does not build by design. + rate_carry += RATE_PER_SECOND * 0.8 * tick.as_secs_f64(); + while rate_carry >= 1.0 { + rate_carry -= 1.0; + seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "rated".into(), + key: 0, + seq, + }, + InsertOpts { + queue: RATE_QUEUE.into(), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(persona) + .add("enqueued_sched_rate", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // Striped queue, enqueued through a store that shares the workers' routing. + let mut params = Vec::new(); + for _ in 0..(3.0 * scale).ceil() as usize { + stripe_seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "striped".into(), + key: 0, + seq: stripe_seq, + }, + InsertOpts { + queue: STRIPED_QUEUE.into(), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match store.enqueue_params_batch(&producer_pool, ¶ms).await { + Ok(count) => shared.stats(persona).add("enqueued_striped", count as i64), + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note( + "sam", + format!("load stopped: {stripe_seq} striped jobs; draining"), + ); + shared.drain_stop.wait().await; + plain.shutdown_all(Duration::from_secs(20), &shared).await; + stripe_fleet + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + + let order = order.lock().expect("order"); + let executed_striped = order.striped_seen.len() as u64; + let missing: Vec = (1..=stripe_seq) + .filter(|seq| !order.striped_seen.contains(seq)) + .take(20) + .collect(); + let stats = shared.stats(persona); + stats.add("striped_executed_distinct", executed_striped as i64); + stats.add("striped_repeat_executions", order.striped_repeats as i64); + stats.add("order_inversions", order.inversions as i64); + stats.set_max("order_worst_distance", order.worst_distance as i64); + if executed_striped < stripe_seq { + let message = format!( + "{} of {stripe_seq} striped jobs never executed", + stripe_seq - executed_striped + ); + let rows = vec![format!( + "{{\"queue\":\"{STRIPED_QUEUE}\",\"missing_seq_sample\":{missing:?}}}" + )]; + if shared.quiesced.load(std::sync::atomic::Ordering::Relaxed) { + shared.violation("lost_job", message, rows); + } else { + shared.finding( + "striped_backlog_at_end", + format!("{message}; the system had not quiesced"), + rows, + ); + } + } + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/shardops.rs b/awa/tests/persona_soak/personas/shardops.rs new file mode 100644 index 000000000..a462c632b --- /dev/null +++ b/awa/tests/persona_soak/personas/shardops.rs @@ -0,0 +1,991 @@ +//! Shard operator: an operator changes `awa.queue_meta.enqueue_shards` on two hot queues +//! while SQL-path producers (`insert_many`, `insert_many_copy`), `QueueStorage` producers, +//! promotion (scheduled jobs and retries) and claimers run, as different non-superuser +//! roles (see `roles.rs`). +//! +//! The operator raises a queue from 1 to 4 to 8 shards, lowers it to 2 and back to 1, +//! sometimes running `prepare_queue` first and sometimes not, restarting some runtimes +//! before others so cached shard counts disagree across the fleet, and, for some raises, +//! holding the producers back until the maintenance leader's runtime has created the new +//! lanes so a different role is the next to use them. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::flag; +use super::super::registry::QueueInfo; +use super::super::rng::Rng; +use super::super::roles; +use super::super::shard_oracle; +use super::super::shared::Shared; +use super::super::stop::Stop; +use super::{attach_events, backoff, queue, register_rows, tune, ANY}; +use awa::model::{insert_many, insert_many_copy_from_pool, QueueStorage, QueueStorageConfig}; +use awa::{Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, Worker}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +pub const HOT_QUEUE: &str = "awa_soak_shardops_hot"; +pub const WARM_QUEUE: &str = "awa_soak_shardops_warm"; +const QUEUES: [&str; 2] = [HOT_QUEUE, WARM_QUEUE]; +const KEYS: u64 = 16; +const WORKERS: u32 = 8; +/// The sequence of shard counts each queue is taken through, repeatedly. +const STEPS: [(i16, i16); 4] = [(1, 4), (4, 8), (8, 2), (2, 1)]; +const MAINTENANCE_LEADER_LOCK_KEY: i64 = 0x4157_415f_4d41_494e; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct ShardOpsJob { + pub seq: u64, + pub mode: String, + pub key: Option, + pub key_seq: u64, +} + +struct ShardOpsWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for ShardOpsWorker { + fn kind(&self) -> &'static str { + "shard_ops_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ShardOps, &ctx.job, false); + let args: ShardOpsJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + self.shared.shardops.note_execution( + args.seq, + ctx.job.id, + args.key.map(|key| (key, args.key_seq)), + ); + tokio::time::sleep(Duration::from_millis(5 + args.seq % 20)).await; + if args.mode == "retry" && ctx.job.attempt <= 1 { + return Ok(JobResult::RetryAfter(Duration::from_millis( + 1_000 + args.seq % 2_000, + ))); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(HOT_QUEUE, queue(WORKERS)) + .queue(WARM_QUEUE, queue(WORKERS)) + .register_worker(ShardOpsWorker { + shared: shared.clone(), + }); + attach_events( + // Retrying election quickly lets these runtimes, which run as application + // roles, take maintenance leadership from the other personas' runtimes. + tune(builder, &shared).leader_election_interval(Duration::from_millis(100)), + &shared, + &["shard_ops_job"], + ) + .build() + .expect("build shard-ops client") + }) +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Mode { + Plain, + Scheduled, + Retry, +} + +impl Mode { + fn roll(rng: &mut Rng) -> Mode { + match rng.below(100) { + 0..=11 => Mode::Retry, + 12..=31 => Mode::Scheduled, + _ => Mode::Plain, + } + } + + fn name(self) -> &'static str { + match self { + Mode::Plain => "plain", + Mode::Scheduled => "scheduled", + Mode::Retry => "retry", + } + } + + fn executions(self) -> u8 { + if self == Mode::Retry { + 2 + } else { + 1 + } + } +} + +struct Built { + params: InsertParams, + seq: u64, + executions: u8, +} + +fn build( + shared: &Shared, + rng: &mut Rng, + queue: &str, + mode: Mode, + key: Option<(u64, u64)>, + run_in: Option, +) -> Result { + let seq = shared.shardops.next_seq(); + let run_at = match (mode, run_in) { + (_, Some(delay)) => { + Some(Utc::now() + chrono::Duration::from_std(delay).unwrap_or_default()) + } + (Mode::Scheduled, None) => { + Some(Utc::now() + chrono::Duration::milliseconds(rng.range(800, 6_000) as i64)) + } + _ => None, + }; + let params = awa::insert::params_with( + &ShardOpsJob { + seq, + mode: mode.name().to_string(), + key: key.map(|(key, _)| key), + key_seq: key.map_or(0, |(_, key_seq)| key_seq), + }, + InsertOpts { + queue: queue.to_string(), + max_attempts: 5, + run_at, + ordering_key: key.map(|(key, _)| key.to_be_bytes().to_vec()), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?; + Ok(Built { + params, + seq, + executions: mode.executions(), + }) +} + +fn pick_queue<'a>(rng: &mut Rng) -> &'a str { + QUEUES[rng.below(QUEUES.len() as u64) as usize] +} + +fn count_enqueued(shared: &Shared, queues: impl Iterator) { + let mut count = 0; + for name in queues { + count += 1; + if let Some(queue) = shared.queues.get(&name) { + queue.enqueued.fetch_add(1, Ordering::Relaxed); + } + } + shared.stats(Persona::ShardOps).add("enqueued", count); +} + +/// Keeps the operator from changing the shard count while an insert is mid-flight +/// during the leader-first hold. +struct InsertGuard<'a>(&'a Shared); + +impl<'a> InsertGuard<'a> { + fn new(shared: &'a Shared) -> Self { + shared + .shardops + .producers_inserting + .fetch_add(1, Ordering::SeqCst); + Self(shared) + } +} + +impl Drop for InsertGuard<'_> { + fn drop(&mut self) { + self.0 + .shardops + .producers_inserting + .fetch_sub(1, Ordering::SeqCst); + } +} + +async fn enqueue_failed( + shared: &Shared, + role: &str, + operation: &str, + error: impl std::fmt::Display, +) { + let message = error.to_string(); + let short: String = message.chars().take(120).collect(); + shared.shardops.violation_once( + shared, + "shardops_enqueue_error", + &format!("{role}:{operation}:{short}"), + format!("role {role} failed {operation}: {message}"), + vec![format!( + "{{\"role\":\"{role}\",\"statement\":\"{operation}\",\"error\":{}}}", + serde_json::Value::String(message.clone()) + )], + ); + backoff(shared, &message, Persona::ShardOps).await; +} + +/// A producer's `QueueStorage` handle. A runtime's cached shard count is only dropped +/// by a restart, so the operator replaces the handle to restart the producer. +struct Stores { + handles: Mutex>>, +} + +impl Stores { + fn new(count: usize) -> Result { + let handles = (0..count) + .map(|_| fresh_store()) + .collect::, _>>()?; + Ok(Self { + handles: Mutex::new(handles), + }) + } + + fn get(&self, index: usize) -> Arc { + self.handles.lock().expect("stores")[index].clone() + } + + fn renew(&self, index: usize) { + if let Ok(store) = fresh_store() { + self.handles.lock().expect("stores")[index] = store; + } + } + + fn len(&self) -> usize { + self.handles.lock().expect("stores").len() + } +} + +fn fresh_store() -> Result, String> { + QueueStorage::new(QueueStorageConfig::default()) + .map(Arc::new) + .map_err(|error| error.to_string()) +} + +/// A SQL-path producer: `insert_many`, or `insert_many_copy` when `copy` is set. +async fn sql_producer(shared: Arc, pool: PgPool, role: &'static str, copy: bool) { + let label = if copy { + "shardops-copy" + } else { + "shardops-sql" + }; + let mut rng = Rng::for_task(shared.config.seed, label); + let scale = shared.config.shardops_scale; + let (tick, per_tick) = if copy { + (Duration::from_secs(1), (40.0 * scale).ceil() as usize) + } else { + (Duration::from_millis(100), (4.0 * scale).ceil() as usize) + }; + while !shared.load_stop.sleep(tick).await { + if shared.shardops.producers_held() { + continue; + } + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let queue = pick_queue(&mut rng); + let mode = Mode::roll(&mut rng); + match build(&shared, &mut rng, queue, mode, None, None) { + Ok(job) => built.push(job), + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let result = if copy { + insert_many_copy_from_pool(&pool, ¶ms).await + } else { + insert_many(&pool, ¶ms).await + }; + match result { + Ok(rows) => { + for job in &built { + shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&shared, rows.iter().map(|row| row.queue.clone())); + register_rows(&shared, Persona::ShardOps, &rows, |_| ANY, 0); + } + Err(error) => { + enqueue_failed( + &shared, + role, + if copy { + "insert_many_copy" + } else { + "insert_many" + }, + error, + ) + .await + } + } + } +} + +/// A producer that enqueues through `QueueStorage`, the runtime's own path. +async fn store_producer( + shared: Arc, + pool: PgPool, + role: &'static str, + stores: Arc, + store_index: usize, +) { + let mut rng = Rng::for_task(shared.config.seed, "shardops-store"); + let per_tick = (3.0 * shared.config.shardops_scale).ceil() as usize; + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + if shared.shardops.producers_held() { + continue; + } + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let queue = pick_queue(&mut rng); + let mode = Mode::roll(&mut rng); + match build(&shared, &mut rng, queue, mode, None, None) { + Ok(job) => built.push(job), + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let store = stores.get(store_index); + match store.enqueue_params_batch(&pool, ¶ms).await { + Ok(_) => { + for job in &built { + shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&shared, params.iter().map(|job| job.opts.queue.clone())); + } + Err(error) => { + enqueue_failed(&shared, role, "QueueStorage::enqueue_params_batch", error).await + } + } + } +} + +/// Jobs with an ordering key, enqueued strictly in sequence per key, alternating between +/// the SQL path as one role and the `QueueStorage` path as the other. +async fn keyed_producer( + shared: Arc, + sql_pool: PgPool, + store_pool: PgPool, + stores: Arc, + store_index: usize, +) { + let mut rng = Rng::for_task(shared.config.seed, "shardops-keyed"); + let per_tick = (3.0 * shared.config.shardops_scale).ceil() as usize; + let mut key_seq = vec![0u64; KEYS as usize]; + let mut tick = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + if shared.shardops.producers_held() { + continue; + } + tick += 1; + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + let mut keys = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let key = rng.below(KEYS); + key_seq[key as usize] += 1; + let queue = QUEUES[(key % 2) as usize]; + match build( + &shared, + &mut rng, + queue, + Mode::Plain, + Some((key, key_seq[key as usize])), + None, + ) { + Ok(job) => { + built.push(job); + keys.push((key, key_seq[key as usize])); + } + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let window_before = shared.shardops.window(); + let outcome = if tick.is_multiple_of(2) { + match insert_many(&sql_pool, ¶ms).await { + Ok(rows) => { + register_rows(&shared, Persona::ShardOps, &rows, |_| ANY, 0); + Ok(rows.len()) + } + Err(error) => Err(( + roles::APP_A, + "insert_many (ordering_key)", + error.to_string(), + )), + } + } else { + let store = stores.get(store_index); + match store.enqueue_params_batch(&store_pool, ¶ms).await { + Ok(count) => Ok(count), + Err(error) => Err(( + roles::APP_B, + "QueueStorage::enqueue_params_batch (ordering_key)", + error.to_string(), + )), + } + }; + match outcome { + Ok(_) => { + let window_after = shared.shardops.window(); + let window = if window_before == window_after { + window_before + } else { + 0 + }; + for (job, (key, seq)) in built.iter().zip(&keys) { + shared.shardops.note_enqueued(job.seq, job.executions); + shared.shardops.note_keyed(*key, *seq, window); + } + shared + .stats(Persona::ShardOps) + .add("enqueued_keyed", built.len() as i64); + count_enqueued(&shared, params.iter().map(|job| job.opts.queue.clone())); + } + Err((role, operation, error)) => enqueue_failed(&shared, role, operation, error).await, + } + } +} + +/// Everything the operator acts through. +struct Operator { + shared: Arc, + fleets: Vec>, + stores: Arc, + owner_pool: PgPool, + migrator_pool: PgPool, + admin_pool: PgPool, + app_pool: PgPool, + rng: Rng, +} + +impl Operator { + fn note(&self, text: impl Into) { + self.shared.note("shards", text); + } + + fn operator_error(&self, what: &str, error: impl std::fmt::Display) { + let message = error.to_string(); + self.shared.shardops.violation_once( + &self.shared, + "shardops_operator_error", + &format!("{what}:{message}"), + format!("{what} failed: {message}"), + vec![format!( + "{{\"statement\":\"{what}\",\"error\":{}}}", + serde_json::Value::String(message.clone()) + )], + ); + } + + async fn prepare(&mut self, queue: &str, shards: i16) { + let as_migrator = self.rng.chance(50); + let (role, pool) = if as_migrator { + (roles::MIGRATOR, &self.migrator_pool) + } else { + (roles::OWNER, &self.owner_pool) + }; + self.note(format!( + "prepare_queue {queue} for {shards} shards as {role}" + )); + let store = match fresh_store() { + Ok(store) => store, + Err(error) => return self.operator_error("prepare_queue", error), + }; + if let Err(error) = store.prepare_queue(pool, queue, shards).await { + self.operator_error(&format!("prepare_queue as {role}"), error); + } + } + + async fn set_shards(&self, queue: &str, shards: i16) { + let result = sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(&self.owner_pool) + .await; + match result { + Ok(_) => { + let now = self.shared.now_ms(); + self.shared.shardops.plans.lock().expect("plans").insert( + queue.to_string(), + shard_oracle::QueuePlan { + configured: shards, + changed_at_ms: now, + settled_at_ms: None, + }, + ); + } + Err(error) => self.operator_error("set enqueue_shards", error), + } + } + + /// Restarts runtime `(fleet, replica)`, which drops its cached shard counts. + async fn restart(&self, fleet: usize, replica: usize) { + let Some(target) = self.fleets.get(fleet) else { + return; + }; + match target.restart(replica, Duration::from_secs(20)).await { + Ok(true) => self + .shared + .stats(Persona::ShardOps) + .add("runtime_restarts", 1), + Ok(false) => {} + Err(error) => self.operator_error("runtime restart", error), + } + } + + async fn runtimes(&self) -> Vec<(usize, usize)> { + let mut all = Vec::new(); + for (index, fleet) in self.fleets.iter().enumerate() { + for replica in fleet.live_replicas().await { + all.push((index, replica)); + } + } + all + } + + fn settle(&self, queue: &str) { + let now = self.shared.now_ms(); + if let Some(plan) = self + .shared + .shardops + .plans + .lock() + .expect("plans") + .get_mut(queue) + { + plan.settled_at_ms = Some(now); + } + self.shared.shardops.open_window(); + } + + /// The runtime holding maintenance leadership, as `(pid, role, application_name)`. + async fn leader(&self) -> Option<(i32, String, String)> { + sqlx::query_as( + "SELECT l.pid, a.usename::text, a.application_name + FROM pg_locks l JOIN pg_stat_activity a ON a.pid = l.pid + WHERE l.locktype = 'advisory' AND l.granted AND a.application_name LIKE 'soak\\_%' + AND ((l.classid::bigint << 32) | l.objid::bigint) = $1 + ORDER BY l.pid LIMIT 1", + ) + .bind(MAINTENANCE_LEADER_LOCK_KEY) + .fetch_optional(&self.admin_pool) + .await + .ok() + .flatten() + } + + /// Makes one of this persona's runtimes the maintenance leader, ending other + /// runtimes' leadership if needed. Returns its role. + async fn lead_with_own_runtime(&self) -> Option { + for _ in 0..10 { + if let Some((pid, role, application)) = self.leader().await { + if application.starts_with("soak_shardops") { + self.note(format!("maintenance leader is {application} ({role})")); + return Some(role); + } + self.note(format!("ending leadership of {application} ({role})")); + let _ = sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(&self.admin_pool) + .await; + } + if self.shared.load_stop.sleep(Duration::from_secs(2)).await { + return None; + } + } + self.note("another persona's runtime kept maintenance leadership"); + None + } + + async fn lane_owners(&self, queue: &str, above: i16) -> String { + let owners: Result, _> = sqlx::query_as( + "SELECT pg_get_userbyid(c.relowner)::text, count(*)::bigint + FROM awa.queue_enqueue_heads h + JOIN pg_class c ON c.relname = h.seq_name + JOIN pg_namespace n ON n.oid = c.relnamespace AND n.nspname = 'awa' + WHERE h.queue = $1 AND h.enqueue_shard >= $2 + GROUP BY 1 ORDER BY 1", + ) + .bind(queue) + .bind(above) + .fetch_all(&self.admin_pool) + .await; + match owners { + Ok(owners) if owners.is_empty() => "none yet".to_string(), + Ok(owners) => owners + .into_iter() + .map(|(role, count)| format!("{role}={count}")) + .collect::>() + .join(","), + Err(error) => format!("lookup failed: {error}"), + } + } + + /// Returns false when the load stopped mid-change. + async fn change(&mut self, queue: &str, from: i16, to: i16) -> bool { + let raising = to > from; + let prepare = + raising && (!self.shared.config.shardops_schema_create || self.rng.chance(50)); + let leader_first = raising && !prepare && self.rng.chance(40); + self.note(format!( + "{queue}: {from} -> {to} shards{}{}", + if prepare { ", prepare_queue first" } else { "" }, + if leader_first { + ", leader creates the new lanes first" + } else { + "" + } + )); + self.shared.stats(Persona::ShardOps).add("shard_changes", 1); + if prepare { + self.prepare(queue, to).await; + } + self.shared.shardops.close_window(); + + if leader_first { + self.shared.shardops.hold_producers(true); + for _ in 0..40 { + if self + .shared + .shardops + .producers_inserting + .load(Ordering::SeqCst) + == 0 + { + break; + } + tokio::time::sleep(Duration::from_millis(250)).await; + } + self.set_shards(queue, to).await; + for (fleet, replica) in self.runtimes().await { + self.restart(fleet, replica).await; + } + for index in 0..self.stores.len() { + self.stores.renew(index); + } + let leader_role = self.lead_with_own_runtime().await; + // Scheduled jobs due shortly: the leader promotes them into new lanes before + // any producer has touched those lanes. + for _ in 0..6 { + let mut built = Vec::new(); + for _ in 0..4 { + let delay = Duration::from_millis(self.rng.range(1_500, 4_000)); + if let Ok(job) = build( + &self.shared, + &mut self.rng, + queue, + Mode::Scheduled, + None, + Some(delay), + ) { + built.push(job); + } + } + let params: Vec = + built.iter().map(|job| job.params.clone()).collect(); + match insert_many(&self.app_pool, ¶ms).await { + Ok(rows) => { + for job in &built { + self.shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&self.shared, rows.iter().map(|row| row.queue.clone())); + register_rows(&self.shared, Persona::ShardOps, &rows, |_| ANY, 0); + } + Err(error) => self.operator_error("seed scheduled jobs", error), + } + if self.shared.load_stop.sleep(Duration::from_secs(1)).await { + self.shared.shardops.hold_producers(false); + return false; + } + } + for _ in 0..15 { + let created: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.queue_enqueue_heads WHERE queue = $1 AND enqueue_shard >= $2", + ) + .bind(queue) + .bind(from) + .fetch_one(&self.admin_pool) + .await + .unwrap_or(0); + if created > 0 { + break; + } + if self.shared.load_stop.sleep(Duration::from_secs(1)).await { + self.shared.shardops.hold_producers(false); + return false; + } + } + let owners = self.lane_owners(queue, from).await; + self.note(format!( + "new lanes of {queue} created before producers resumed; sequence owners: {owners} (leader role {})", + leader_role.as_deref().unwrap_or("not ours") + )); + self.shared.shardops.hold_producers(false); + self.settle(queue); + return true; + } + + self.set_shards(queue, to).await; + if self.shared.load_stop.sleep(Duration::from_secs(4)).await { + return false; + } + // Restart some runtimes and producer stores now and the rest later, so cached + // shard counts disagree across the fleet for a while. + let runtimes = self.runtimes().await; + let mut restart_now: Vec<(usize, usize)> = runtimes + .iter() + .copied() + .filter(|_| self.rng.chance(50)) + .collect(); + if restart_now.is_empty() { + restart_now.extend(runtimes.first().copied()); + } + if restart_now.len() == runtimes.len() && runtimes.len() > 1 { + restart_now.pop(); + } + let renew_now: Vec = (0..self.stores.len()) + .filter(|_| self.rng.chance(50)) + .collect(); + self.note(format!( + "{queue}: restarting {} of {} runtimes and {} of {} producer stores; the others keep the old shard count", + restart_now.len(), + runtimes.len(), + renew_now.len(), + self.stores.len() + )); + for (fleet, replica) in &restart_now { + self.restart(*fleet, *replica).await; + } + for index in &renew_now { + self.stores.renew(*index); + } + if self.shared.load_stop.sleep(Duration::from_secs(18)).await { + return false; + } + for (fleet, replica) in runtimes + .iter() + .filter(|runtime| !restart_now.contains(runtime)) + { + self.restart(*fleet, *replica).await; + } + for index in 0..self.stores.len() { + if !renew_now.contains(&index) { + self.stores.renew(index); + } + } + self.settle(queue); + let owners = self.lane_owners(queue, from.min(to)).await; + self.note(format!( + "{queue}: settled at {to} shards; lane sequence owners above shard {}: {owners}", + from.min(to) + )); + true + } + + async fn run(&mut self, stop: &Stop, period: Duration) { + if stop.sleep(Duration::from_secs(20)).await { + return; + } + let mut step = [0usize; QUEUES.len()]; + let mut turn = 0usize; + loop { + let which = turn % QUEUES.len(); + turn += 1; + let (from, to) = STEPS[step[which] % STEPS.len()]; + step[which] += 1; + let started = Instant::now(); + if !self.change(QUEUES[which], from, to).await { + return; + } + let remaining = period + .saturating_sub(started.elapsed()) + .max(Duration::from_secs(5)); + if stop.sleep(remaining).await { + return; + } + } + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::ShardOps; + let queue_infos: Vec> = QUEUES + .iter() + .map(|name| shared.queues.register(name, persona, false)) + .collect(); + let config = &shared.config; + + let owner_pool = roles::pool_as(config, roles::OWNER, "soak_shardops_owner", 2).await?; + let migrator_pool = + roles::pool_as(config, roles::MIGRATOR, "soak_shardops_migrator", 1).await?; + let admin_pool = connect_pool(&config.database_url, "soak_shardops_admin", 2) + .await + .map_err(|error| error.to_string())?; + + // The schema owner provisions each queue's first shard before any application role runs. + let owner_store = fresh_store()?; + for name in QUEUES { + owner_store + .prepare_queue(&owner_pool, name, 1) + .await + .map_err(|error| format!("prepare_queue {name}: {error}"))?; + shared.shardops.plans.lock().expect("plans").insert( + name.to_string(), + shard_oracle::QueuePlan { + configured: 1, + changed_at_ms: 0, + settled_at_ms: Some(0), + }, + ); + } + shared.shardops.open_window(); + + let served: Vec<(Arc, i64)> = queue_infos + .iter() + .map(|info| (info.clone(), i64::from(WORKERS))) + .collect(); + let mut fleets = Vec::new(); + for (name, role) in [("shardops-a", roles::APP_A), ("shardops-b", roles::APP_B)] { + fleets.push( + Fleet::start_with_url( + &shared, + roles::url_for_role(config, role), + name, + persona, + 2, + 10, + 2, + served.clone(), + client_factory(shared.clone()), + ) + .await?, + ); + } + + let app_a = roles::pool_as( + config, + roles::APP_A, + "soak_shardops_producer_a", + shared.budget.reserve("shardops-producer-a", 6, 2), + ) + .await?; + let app_b = roles::pool_as( + config, + roles::APP_B, + "soak_shardops_producer_b", + shared.budget.reserve("shardops-producer-b", 6, 2), + ) + .await?; + let stores = Arc::new(Stores::new(2)?); + + let period = if config.load >= Duration::from_secs(600) { + Duration::from_secs(60) + } else { + Duration::from_secs(40) + }; + let queue_names: Vec = QUEUES.iter().map(|name| name.to_string()).collect(); + shared.note( + "shards", + format!( + "queues {QUEUES:?} cycle through {STEPS:?} shards every {}s per change; roles: owner {}, migrator {}, applications {} and {}", + period.as_secs(), + roles::OWNER, + roles::MIGRATOR, + roles::APP_A, + roles::APP_B + ), + ); + + let watch_stop = Stop::new(); + let watcher = tokio::spawn(shard_oracle::watch( + shared.clone(), + admin_pool.clone(), + queue_names.clone(), + watch_stop.clone(), + )); + let producers = vec![ + tokio::spawn(sql_producer( + shared.clone(), + app_a.clone(), + roles::APP_A, + false, + )), + tokio::spawn(sql_producer( + shared.clone(), + app_b.clone(), + roles::APP_B, + true, + )), + tokio::spawn(store_producer( + shared.clone(), + app_b.clone(), + roles::APP_B, + stores.clone(), + 0, + )), + tokio::spawn(keyed_producer( + shared.clone(), + app_a.clone(), + app_b.clone(), + stores.clone(), + 1, + )), + ]; + let operator_stop = shared.load_stop.clone(); + let mut operator = Operator { + shared: shared.clone(), + fleets: fleets.clone(), + stores: stores.clone(), + owner_pool: owner_pool.clone(), + migrator_pool, + admin_pool: admin_pool.clone(), + app_pool: app_a.clone(), + rng: Rng::for_task(shared.config.seed, "shardops-operator"), + }; + operator.run(&operator_stop, period).await; + shared.shardops.hold_producers(false); + for producer in producers { + let _ = producer.await; + } + shared.note("shards", "load stopped; draining"); + + shared.drain_stop.wait().await; + shard_oracle::final_check(&shared, &admin_pool, &queue_names).await; + watch_stop.trigger(); + let _ = watcher.await; + for fleet in &fleets { + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + } + for pool in [app_a, app_b, owner_pool, admin_pool] { + pool.close().await; + } + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/tom.rs b/awa/tests/persona_soak/personas/tom.rs new file mode 100644 index 000000000..af963e56c --- /dev/null +++ b/awa/tests/persona_soak/personas/tom.rs @@ -0,0 +1,279 @@ +//! Tom, multi-tenant fan-out: many queues each at a low rate, unique jobs, cron and +//! priority aging. Stresses lane and queue cardinality, per-queue maintenance cost +//! and cron fencing across runtimes. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, register_rows, tune}; +use awa::model::{insert_many, insert_with}; +use awa::{ + AwaError, Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, PeriodicJob, + QueueConfig, UniqueOpts, Worker, +}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +const PREFIX: &str = "awa_soak_tom_"; +const CRON_JOBS: usize = 4; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct TomJob { + pub seq: u64, + pub work_ms: u64, + pub unique_key: Option, +} + +struct TomWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for TomWorker { + fn kind(&self) -> &'static str { + "tom_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Tom, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(50); + tokio::time::sleep(Duration::from_millis(work_ms)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn queue_name(index: usize) -> String { + format!("{PREFIX}{index:03}") +} + +/// Four schedules firing every ten seconds into the first four tenant queues. +pub fn declare_cron(mut builder: awa::ClientBuilder) -> awa::ClientBuilder { + for cron in 0..CRON_JOBS { + let job = PeriodicJob::builder(format!("soak_tom_cron_{cron}"), "*/10 * * * * *") + .queue(queue_name(cron)) + .build(&TomJob { + seq: 0, + work_ms: 20, + unique_key: None, + }) + .expect("cron job"); + builder = builder.periodic(job); + } + builder +} + +fn client_factory(shared: Arc, queue_count: usize) -> ClientFactory { + Arc::new(move |pool| { + let mut builder = Client::builder(pool); + for index in 0..queue_count { + builder = builder.queue( + queue_name(index), + QueueConfig { + max_workers: 2, + priority_aging_interval: Duration::from_secs(5), + ..QueueConfig::default() + }, + ); + } + builder = builder.register_worker(TomWorker { + shared: shared.clone(), + }); + for cron in 0..CRON_JOBS { + let job = PeriodicJob::builder(format!("soak_tom_cron_{cron}"), "*/10 * * * * *") + .queue(queue_name(cron)) + .build(&TomJob { + seq: 0, + work_ms: 20, + unique_key: None, + }) + .expect("cron job"); + builder = builder.periodic(job); + } + attach_events(tune(builder, &shared), &shared, &["tom_job"]) + .build() + .expect("build tom client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Tom; + let queue_count = shared.config.scaled(100) as usize; + let mut served = Vec::new(); + for index in 0..queue_count { + served.push(( + shared.queues.register(&queue_name(index), persona, false), + 2, + )); + } + let fleet = Fleet::start( + &shared, + "tom", + persona, + 2, + 24, + 2, + served, + client_factory(shared.clone(), queue_count), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_tom_producer", + shared.budget.reserve("tom-producer", 4, 1), + ) + .await + .map_err(|error| error.to_string())?; + let tick = Duration::from_millis(100); + let mut rng = Rng::for_task(shared.config.seed, "tom-producer"); + let plain_per_tick = (queue_count as f64 * 0.05).max(1.0) as usize; + shared.note("tom", format!("{queue_count} queues, {plain_per_tick} jobs per 100 ms tick, {CRON_JOBS} cron schedules every 10 s")); + let unique_double_accept = Arc::new(AtomicU64::new(0)); + let mut seq = 0u64; + let mut unique_key = 0u64; + let mut unique_tasks = tokio::task::JoinSet::new(); + + while !shared.load_stop.sleep(tick).await { + let mut params = Vec::with_capacity(plain_per_tick); + for _ in 0..plain_per_tick { + seq += 1; + let scheduled = rng.chance(20); + let run_at = scheduled.then(|| { + Utc::now() + chrono::Duration::milliseconds(rng.range(1_000, 10_000) as i64) + }); + let queue_index = rng.below(queue_count as u64) as usize; + params.push( + awa::insert::params_with( + &TomJob { + seq, + work_ms: rng.range(10, 200), + unique_key: None, + }, + InsertOpts { + queue: queue_name(queue_index), + priority: rng.range(1, 4) as i16, + run_at, + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + shared.stats(persona).add( + if scheduled { + "planned_scheduled" + } else { + "planned_immediate" + }, + 1, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // One logical unique job per second: two concurrent identical inserts. + if seq.is_multiple_of(10) && unique_tasks.len() < 16 { + unique_key += 1; + let queue_index = rng.below(queue_count as u64) as usize; + let pool = producer_pool.clone(); + let shared = shared.clone(); + let double = unique_double_accept.clone(); + let key = unique_key; + let work_ms = rng.range(10, 100); + unique_tasks.spawn(async move { + let args = TomJob { + seq: key, + work_ms, + unique_key: Some(key), + }; + let opts = || InsertOpts { + queue: queue_name(queue_index), + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + by_period: None, + states: 0b0001_1111, + }), + ..Default::default() + }; + let (first, second) = tokio::join!( + insert_with(&pool, &args, opts()), + insert_with(&pool, &args, opts()) + ); + let mut accepted = 0; + for outcome in [first, second] { + match outcome { + Ok(row) => { + accepted += 1; + let queue = shared + .queues + .get(&row.queue) + .map_or(u32::MAX, |queue| queue.index); + shared.ledger.register( + row.id, + JobRecord::new( + Persona::Tom, + queue, + shared.now_ms() as u32, + expect::COMPLETED, + 0, + ), + ); + } + Err(AwaError::UniqueConflict { .. }) => { + shared.stats(Persona::Tom).add("unique_conflicts", 1) + } + Err(error) => shared.stats(Persona::Tom).add( + &format!("unique_error:{error}") + .chars() + .take(60) + .collect::(), + 1, + ), + } + } + shared.stats(Persona::Tom).add("unique_logical", 1); + if accepted > 1 { + double.fetch_add(1, Ordering::Relaxed); + shared.violation( + "unique_double_accept", + format!( + "both concurrent identical unique inserts (key {key}) were accepted" + ), + vec![format!( + "{{\"unique_key\":{key},\"queue\":\"{}\"}}", + queue_name(queue_index) + )], + ); + } + }); + } + while unique_tasks.try_join_next().is_some() {} + } + while unique_tasks.join_next().await.is_some() {} + + shared.note( + "tom", + format!("load stopped after {seq} plain jobs; draining"), + ); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/ui.rs b/awa/tests/persona_soak/personas/ui.rs new file mode 100644 index 000000000..97f5d4d4e --- /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(awa::audited_sql(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(awa::audited_sql(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(awa::audited_sql(format!( + "SELECT count(*)::bigint, count(*) FILTER (WHERE state = 'active')::bigint + FROM pg_stat_activity WHERE application_name = '{APPLICATION_NAME}'" + ))) + .fetch_one(&sampler_pool) + .await; + if let Ok((sessions, now_active)) = row { + active = now_active; + shared.run.add("ui_server_sessions_after_quiet", 0); + shared + .run + .set_max("ui_server_sessions_after_quiet", sessions); + if active == 0 { + break; + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + if active != 0 { + shared.violation( + "ui_sessions_not_released", + format!("{active} admin server database session(s) still running a query {}s after every UI client stopped", (statement_timeout + Duration::from_secs(15)).as_secs()), + vec![], + ); + } + shared.drain_stop.wait().await; + let _ = sampler.await; + sampler_pool.close().await; + Ok(()) +} + +impl Shared { + /// Sleeps unless the load phase ends first. + pub async fn stop_sleep_load(&self, duration: Duration) -> bool { + self.load_stop.sleep(duration).await + } +} + +/// After the run: latency growth and slow-endpoint findings. +pub fn analyse(shared: &Shared) { + let samples = shared.ui.samples.lock().expect("ui samples").clone(); + let run_ms = shared.config.load.as_millis() as i64; + for (endpoint, mut points) in samples { + if points.len() < 20 { + continue; + } + points.sort_by_key(|point| point.0); + let quarter = points.len() / 4; + let median = |slice: &[(i64, i64)]| { + let mut values: Vec = slice.iter().map(|point| point.1).collect(); + values.sort_unstable(); + percentile(&values, 0.5) + }; + let early = median(&points[..quarter]); + let late = median(&points[points.len() - quarter..]); + let mut all: Vec = points.iter().map(|point| point.1).collect(); + all.sort_unstable(); + let p99 = percentile(&all, 0.99); + let max = all.last().copied().unwrap_or(0); + if endpoint.starts_with("GET") && late >= 500 && late >= early.max(1) * 3 { + shared.finding( + "ui_latency_grows", + format!("{endpoint}: median latency {early} ms in the first quarter of its samples, {late} ms in the last (run {}s)", run_ms / 1000), + vec![format!("{{\"endpoint\":\"{endpoint}\",\"samples\":{},\"p99_ms\":{p99},\"max_ms\":{max}}}", points.len())], + ); + } + if p99 >= 10_000 { + shared.finding( + "ui_endpoint_slow", + format!( + "{endpoint}: p99 {p99} ms, max {max} ms over {} requests", + points.len() + ), + vec![], + ); + } + } +} diff --git a/awa/tests/persona_soak/pgserver.rs b/awa/tests/persona_soak/pgserver.rs new file mode 100644 index 000000000..10b018d99 --- /dev/null +++ b/awa/tests/persona_soak/pgserver.rs @@ -0,0 +1,99 @@ +//! Server crash detection: postmaster start time, and a scan of the server log for +//! the lines a crashed backend leaves behind. + +use super::config::SoakConfig; +use sqlx::PgPool; +use std::process::Command; + +const CRASH_PATTERNS: [&str; 6] = [ + "terminated by signal", + "TRAP:", + "server process (PID", + "PANIC:", + "all server processes terminated", + "database system was not properly shut down", +]; + +pub struct PgWatch { + pub started_at: String, + log_start_offset: u64, +} + +impl PgWatch { + pub async fn begin(config: &SoakConfig, pool: &PgPool) -> Result { + let started_at: String = sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await?; + Ok(Self { + started_at, + log_start_offset: current_log_size(config), + }) + } + + pub async fn postmaster_start(pool: &PgPool) -> Result { + sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await + } + + /// Log text written since the run began, or None when no log source is configured. + pub fn log_since_start(&self, config: &SoakConfig) -> Option { + if let Some(path) = &config.pg_log_file { + let bytes = std::fs::read(path).ok()?; + let from = (self.log_start_offset as usize).min(bytes.len()); + return Some(String::from_utf8_lossy(&bytes[from..]).into_owned()); + } + let container = config.pg_container.as_ref()?; + let output = match &config.pg_container_log { + Some(path) => Command::new("docker") + .args([ + "exec", + container, + "sh", + "-c", + &format!("tail -c +{} {path}", self.log_start_offset + 1), + ]) + .output() + .ok()?, + None => Command::new("docker") + .args(["logs", container]) + .output() + .ok()?, + }; + let mut text = String::from_utf8_lossy(&output.stdout).into_owned(); + text.push_str(&String::from_utf8_lossy(&output.stderr)); + Some(text) + } +} + +fn current_log_size(config: &SoakConfig) -> u64 { + if let Some(path) = &config.pg_log_file { + return std::fs::metadata(path).map(|meta| meta.len()).unwrap_or(0); + } + let (Some(container), Some(path)) = (&config.pg_container, &config.pg_container_log) else { + return 0; + }; + Command::new("docker") + .args(["exec", container, "sh", "-c", &format!("wc -c < {path}")]) + .output() + .ok() + .and_then(|output| String::from_utf8_lossy(&output.stdout).trim().parse().ok()) + .unwrap_or(0) +} + +/// Lines of `log` that indicate a crashed server process, each with a few lines of +/// context on either side. +pub fn crash_excerpts(log: &str) -> Vec { + let lines: Vec<&str> = log.lines().collect(); + let mut excerpts = Vec::new(); + let mut next_allowed = 0; + for (index, line) in lines.iter().enumerate() { + if index >= next_allowed && CRASH_PATTERNS.iter().any(|pattern| line.contains(pattern)) { + let from = index.saturating_sub(5); + let to = (index + 8).min(lines.len()); + excerpts.push(lines[from..to].join("\n")); + next_allowed = to; + } + } + excerpts +} diff --git a/awa/tests/persona_soak/python/soak_python.py b/awa/tests/persona_soak/python/soak_python.py new file mode 100644 index 000000000..ff6920c3e --- /dev/null +++ b/awa/tests/persona_soak/python/soak_python.py @@ -0,0 +1,359 @@ +"""Python side of the persona soak: producers, async workers and the migrate path. + +Started and stopped by the Rust harness. Every enqueue, handler start and handler end +is appended to a JSON-lines ledger that the harness folds into its own ledger, so the +oracle checks Python jobs like any others. + +Environment: + SOAK_DATABASE_URL database to use (an awa_soak* database) + SOAK_LEDGER ledger file, appended to + SOAK_SEED, SOAK_SCALE, SOAK_PACE + SOAK_ROLE "full" (default) or "migrate_probe" + SOAK_QUEUE_WORK, SOAK_QUEUE_DEADLINE, SOAK_QUEUE_SHARED queue names + +Signals: SIGUSR1 stops producing (workers keep running); SIGTERM exits at once. +""" + +import asyncio +import json +import os +import random +import signal +import sys +import time +from dataclasses import dataclass + +import awa + +URL = os.environ["SOAK_DATABASE_URL"] +ROLE = os.environ.get("SOAK_ROLE", "full") +SEED = int(os.environ.get("SOAK_SEED", "1")) +SCALE = float(os.environ.get("SOAK_SCALE", "1")) +PACE = float(os.environ.get("SOAK_PACE", "1")) +QUEUE_WORK = os.environ.get("SOAK_QUEUE_WORK", "awa_soak_py_work") +QUEUE_DEADLINE = os.environ.get("SOAK_QUEUE_DEADLINE", "awa_soak_py_deadline") +QUEUE_SHARED = os.environ.get("SOAK_QUEUE_SHARED", "awa_soak_py_shared") +WORK_WORKERS = 8 +DEADLINE_WORKERS = 16 +SHARED_WORKERS = 6 +DEADLINE_MS = 5000 + +# Outcomes a job may end in, matching the harness (1 completed, 2 failed, 4 cancelled). +COMPLETED, FAILED, CANCELLED = 1, 2, 4 + +ledger = open(os.environ["SOAK_LEDGER"], "a", buffering=1) if "SOAK_LEDGER" in os.environ else None + + +def log(**event): + if ledger is not None: + ledger.write(json.dumps(event, separators=(",", ":")) + "\n") + + +@dataclass +class PyJob: + seq: int + behavior: str + work_ms: int + born: float = 0.0 + + +@dataclass +class PySharedJob: + seq: int + behavior: str + work_ms: int + + +active = 0 +max_active = 0 +started = 0 +enqueued = 0 +errors = 0 + + +async def handle(job): + """One handler for every behaviour; see `plan` for what each one is for.""" + global active, max_active, started + args = job.args + attempt = job.attempt + active += 1 + max_active = max(max_active, active) + started += 1 + log(e="start", id=job.id, attempt=attempt, q=job.queue) + result = "completed" + try: + behavior = args.behavior + if behavior == "ok": + await asyncio.sleep(args.work_ms / 1000) + elif behavior == "progress": + for step in range(1, 5): + await asyncio.sleep(args.work_ms / 4000) + job.set_progress(step * 25, f"step {step}") + await job.flush_progress() + elif behavior == "retry": + await asyncio.sleep(0.05) + if attempt < 3: + result = "raise" + raise RuntimeError("transient failure") + elif behavior == "doomed": + result = "raise" + raise RuntimeError("permanent failure") + elif behavior == "retry_after": + if attempt == 1: + result = "retry_after" + return awa.RetryAfter(seconds=1.0) + elif behavior == "snooze": + # A snooze does not count as an attempt, so the wait is keyed on the clock. + if time.time() < args.born + 3.0: + result = "snooze" + return awa.Snooze(seconds=1.0) + elif behavior == "cancel": + result = "cancel" + return awa.Cancel("rejected") + elif behavior == "long": + # Overruns the queue deadline on the first attempt only. + await asyncio.sleep(args.work_ms / 1000 if attempt == 1 else 0.1) + elif behavior == "ignore_deadline": + # Documented behaviour: after the deadline rescue the worker slot is freed + # but this coroutine keeps running, so more handlers than max_workers run. + await asyncio.sleep(args.work_ms / 1000) + else: + await asyncio.sleep(args.work_ms / 1000) + return None + except Exception: + raise + finally: + active -= 1 + log(e="end", id=job.id, attempt=attempt, r=result) + + +def plan(rng, scale_rows=None): + """Returns (queue, behavior, work_ms, max_attempts, allowed_outcomes, kind_class).""" + roll = rng.random() * 100 + if roll < 40: + return QUEUE_WORK, "ok", rng.randint(20, 400), 5, COMPLETED + if roll < 48: + return QUEUE_WORK, "progress", rng.randint(800, 2500), 5, COMPLETED + if roll < 56: + return QUEUE_WORK, "retry", 50, 5, COMPLETED | FAILED + if roll < 60: + return QUEUE_WORK, "doomed", 50, 3, FAILED + if roll < 66: + return QUEUE_WORK, "retry_after", 50, 5, COMPLETED | FAILED + if roll < 72: + return QUEUE_WORK, "snooze", 50, 5, COMPLETED + if roll < 75: + return QUEUE_WORK, "cancel", 50, 3, CANCELLED + if roll < 80: + return QUEUE_DEADLINE, "long", rng.randint(8000, 12000), 3, COMPLETED | FAILED + if roll < 83: + return QUEUE_DEADLINE, "ignore_deadline", rng.randint(9000, 14000), 2, COMPLETED | FAILED + return QUEUE_SHARED, rng.choice(["ok", "ok", "retry", "progress"]), rng.randint(20, 300), 5, COMPLETED | FAILED + + +def kind_for(queue): + return PySharedJob if queue == QUEUE_SHARED else PyJob + + +def make_args(queue, rng, behavior, work_ms): + seq = rng.randint(1, 1 << 40) + if queue == QUEUE_SHARED: + return PySharedJob(seq=seq, behavior=behavior, work_ms=work_ms) + return PyJob(seq=seq, behavior=behavior, work_ms=work_ms, born=time.time()) + + +def kind_name(queue): + return "py_shared_job" if queue == QUEUE_SHARED else "py_job" + + +def record_enqueued(job, queue, behavior, allowed, max_attempts): + global enqueued + enqueued += 1 + log(e="enq", id=job.id, q=queue, b=behavior, x=allowed, ma=max_attempts) + + +async def producer_singles(client, rng, stop): + """Single inserts, one or two at a time.""" + period = 0.25 / max(SCALE, 0.1) + while not stop.is_set(): + for _ in range(rng.randint(1, 2)): + queue, behavior, work_ms, max_attempts, allowed = plan(rng) + args = make_args(queue, rng, behavior, work_ms) + try: + job = await client.insert(args, kind=kind_name(queue), queue=queue, max_attempts=max_attempts) + record_enqueued(job, queue, behavior, allowed, max_attempts) + except Exception as error: # connection killed by chaos, pool timeout + note_error("insert", error) + await asyncio.sleep(period) + + +async def producer_batches(client, rng, stop): + """insert_many above the old 5,461-row bind-parameter limit, and COPY.""" + rows = int(6500 * SCALE) if SCALE >= 1 else 6500 + interval = max(20.0, 75.0 * PACE) + await asyncio.sleep(min(15.0, interval)) + flip = 0 + while not stop.is_set(): + flip += 1 + args = [PyJob(seq=i, behavior="ok", work_ms=0, born=time.time()) for i in range(rows)] + try: + started_at = time.monotonic() + if flip % 3 == 0: + jobs = await client.insert_many_copy(args, kind="py_job", queue=QUEUE_WORK, max_attempts=3) + how = "copy" + else: + # One transaction, one insert_many call far above the old bind-parameter limit. + tx = await client.transaction() + try: + jobs = await tx.insert_many(args, kind="py_job", queue=QUEUE_WORK, max_attempts=3) + await tx.commit() + except Exception: + await tx.rollback() + raise + how = "insert_many" + for job in jobs: + record_enqueued(job, QUEUE_WORK, "ok", COMPLETED, 3) + print(f"BATCH how={how} rows={len(jobs)} seconds={time.monotonic() - started_at:.2f}", flush=True) + except Exception as error: + note_error(f"batch of {rows}", error) + await wait_or_stop(stop, interval) + + +async def producer_transactions(client, rng, stop): + """Enqueue inside an application transaction that sometimes rolls back.""" + while not stop.is_set(): + count = rng.randint(1, 20) + plans = [plan(rng) for _ in range(count)] + rollback = rng.random() < 0.35 + try: + tx = await client.transaction() + try: + created = [] + for queue, behavior, work_ms, max_attempts, allowed in plans: + args = make_args(queue, rng, behavior, work_ms) + job = await tx.insert(args, kind=kind_name(queue), queue=queue, max_attempts=max_attempts) + created.append((job, queue, behavior, allowed, max_attempts)) + ids = [item[0].id for item in created] + if rollback: + await tx.rollback() + log(e="rolled_back", ids=ids) + else: + await tx.commit() + for job, queue, behavior, allowed, max_attempts in created: + record_enqueued(job, queue, behavior, allowed, max_attempts) + except Exception: + try: + await tx.rollback() + except Exception: + pass + raise + except Exception as error: + note_error("transaction", error) + await asyncio.sleep(1.5 / max(SCALE, 0.1)) + + +def note_error(what, error): + global errors + errors += 1 + log(e="error", what=what, error=str(error)[:200]) + if errors <= 5: + print(f"ERROR {what}: {str(error)[:200]}", flush=True) + + +async def wait_or_stop(stop, seconds): + try: + await asyncio.wait_for(stop.wait(), timeout=seconds) + except asyncio.TimeoutError: + pass + + +async def reporter(stop): + while True: + await asyncio.sleep(2) + print(f"STAT active={active} max_active={max_active} started={started} enqueued={enqueued} errors={errors}", flush=True) + + +async def migrate_and_probe(client, label): + """What an application does at startup: migrate, then query.""" + await client.migrate() + stats = await client.queue_stats() + jobs = await client.list_jobs(limit=5) + print(f"MIGRATED {label} queues={len(stats)} sample={len(jobs)}", flush=True) + + +async def migrate_probe_main(): + client = awa.AsyncClient(URL, max_connections=4) + await migrate_and_probe(client, "probe") + # Idle with the pool open: connections stranded by migrate would show up as + # sessions beyond the pool size, or as idle-in-transaction. + await asyncio.sleep(4) + print("IDLE", flush=True) + await asyncio.sleep(4) + await client.close() + print("CLOSED", flush=True) + + +async def exit_with_parent(): + """The harness owns this process; if it dies without cleaning up, so do we.""" + parent = os.getppid() + while True: + await asyncio.sleep(3) + if os.getppid() != parent: + os._exit(1) + + +async def main(): + asyncio.create_task(exit_with_parent()) + if ROLE == "migrate_probe": + await migrate_probe_main() + return + # A chaos restart can still be migrating when the load ends. Remember the + # drain request before startup awaits, then start workers without producers. + stop = asyncio.Event() + if os.environ.get("SOAK_DRAIN") == "1": + stop.set() + loop = asyncio.get_running_loop() + loop.add_signal_handler(signal.SIGUSR1, stop.set) + rng = random.Random(SEED) + client = awa.AsyncClient(URL, max_connections=14) + await migrate_and_probe(client, "startup") + client.task(PyJob, kind="py_job", queue=QUEUE_WORK)(handle) + client.task(PySharedJob, kind="py_shared_job", queue=QUEUE_SHARED)(handle) + await client.start( + [ + {"name": QUEUE_WORK, "max_workers": WORK_WORKERS}, + {"name": QUEUE_DEADLINE, "max_workers": DEADLINE_WORKERS, "deadline_duration_ms": DEADLINE_MS}, + {"name": QUEUE_SHARED, "max_workers": SHARED_WORKERS}, + ], + heartbeat_interval_ms=2000, + heartbeat_staleness_ms=12000, + heartbeat_rescue_interval_ms=3000, + deadline_rescue_interval_ms=2000, + promote_interval_ms=250, + ) + # The deadline queue serves a second kind through the same handler. + print(f"READY pid={os.getpid()}", flush=True) + + tasks = [ + asyncio.create_task(producer_singles(client, rng, stop)), + asyncio.create_task(producer_batches(client, rng, stop)), + asyncio.create_task(producer_transactions(client, rng, stop)), + asyncio.create_task(reporter(stop)), + ] + await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.Event().wait() + + +def exit_now(signum, _frame): + os._exit(128 + signum) + + +if __name__ == "__main__": + # PyO3/tokio teardown can crash while Python unloads modules; exit without finalizers. + signal.signal(signal.SIGTERM, exit_now) + signal.signal(signal.SIGINT, exit_now) + try: + asyncio.run(main()) + finally: + sys.stdout.flush() + os._exit(0) diff --git a/awa/tests/persona_soak/python/test_soak_python.py b/awa/tests/persona_soak/python/test_soak_python.py new file mode 100644 index 000000000..337301719 --- /dev/null +++ b/awa/tests/persona_soak/python/test_soak_python.py @@ -0,0 +1,119 @@ +"""Process-boundary regressions; uses a delayed fake client, never a database.""" + +import os +from pathlib import Path +import select +import signal +import subprocess +import sys +import tempfile +import time +import unittest + + +SCRIPT = Path(os.environ.get("SOAK_PYTHON_SCRIPT", Path(__file__).with_name("soak_python.py"))) +FAKE_AWA = ''' +import asyncio +import os +from pathlib import Path +from types import SimpleNamespace + +class AsyncClient: + def __init__(self, *args, **kwargs): + pass + + async def migrate(self): + print("STARTUP_BLOCKED", flush=True) + while not Path(os.environ["SOAK_TEST_GATE"]).exists(): + await asyncio.sleep(0.01) + + async def queue_stats(self): + return [] + + async def list_jobs(self, **kwargs): + return [] + + def task(self, *args, **kwargs): + return lambda handler: handler + + async def start(self, *args, **kwargs): + Path(os.environ["SOAK_TEST_WORKERS"]).touch() + + async def insert(self, *args, **kwargs): + Path(os.environ["SOAK_TEST_PRODUCED"]).touch() + return SimpleNamespace(id=1) + + async def transaction(self): + Path(os.environ["SOAK_TEST_PRODUCED"]).touch() + raise RuntimeError("unexpected production during drain") +''' + + +class StartupDrainTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory(prefix="awa-python-drain-") + self.root = Path(self.directory.name) + (self.root / "awa.py").write_text(FAKE_AWA) + self.gate = self.root / "gate" + self.workers = self.root / "workers" + self.produced = self.root / "produced" + self.child = None + + def tearDown(self): + if self.child is not None: + if self.child.poll() is None: + self.child.terminate() + self.child.wait(timeout=10) + self.child.stdout.close() + self.child.stderr.close() + self.directory.cleanup() + + def start(self, *, draining=False): + env = os.environ | { + "PYTHONPATH": str(self.root), + "SOAK_DATABASE_URL": "unused-by-the-fake-client", + "SOAK_ROLE": "full", + "SOAK_DRAIN": "1" if draining else "0", + "SOAK_TEST_GATE": str(self.gate), + "SOAK_TEST_WORKERS": str(self.workers), + "SOAK_TEST_PRODUCED": str(self.produced), + } + self.child = subprocess.Popen( + [sys.executable, str(SCRIPT)], env=env, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, + ) + self.wait_for("STARTUP_BLOCKED") + + def wait_for(self, prefix): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + if select.select([self.child.stdout], [], [], 0.1)[0]: + line = self.child.stdout.readline() + if line.startswith(prefix): + return + if not line: + break + self.fail(f"did not see {prefix!r}; child status={self.child.poll()}") + + def assert_draining_workers(self): + self.gate.touch() + self.wait_for("READY pid=") + time.sleep(0.1) + self.assertIsNone(self.child.poll(), "workers must stay alive to drain") + self.assertTrue(self.workers.exists(), "workers must start after migration") + self.assertFalse(self.produced.exists(), "a drain restart must not enqueue new work") + + def test_stop_during_migration_is_remembered_until_workers_start(self): + self.start() + self.child.send_signal(signal.SIGUSR1) + time.sleep(0.05) + self.assertIsNone(self.child.poll(), "SIGUSR1 during startup must not terminate the worker") + self.assert_draining_workers() + + def test_restart_during_drain_starts_workers_without_producers(self): + self.start(draining=True) + self.assert_draining_workers() + + +if __name__ == "__main__": + unittest.main() diff --git a/awa/tests/persona_soak/python_probe.rs b/awa/tests/persona_soak/python_probe.rs new file mode 100644 index 000000000..c91d74247 --- /dev/null +++ b/awa/tests/persona_soak/python_probe.rs @@ -0,0 +1,451 @@ +//! A migrate probe may hold a normal transaction until its ready marker. + +use sqlx::PgPool; +use std::future::Future; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::time::Duration; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[repr(u8)] +pub enum Phase { + #[default] + Migrating, + Ready, + Idle, + Closed, +} + +impl Phase { + pub fn marker(line: &str) -> Option { + if line.starts_with("MIGRATED probe ") { + Some(Self::Ready) + } else { + match line { + "IDLE" => Some(Self::Idle), + "CLOSED" => Some(Self::Closed), + _ => None, + } + } + } + + pub fn from_byte(value: u8) -> Self { + match value { + 1 => Self::Ready, + 2 => Self::Idle, + 3 => Self::Closed, + _ => Self::Migrating, + } + } +} + +/// The unique application name isolates this process, including backends older +/// than a stalled probe's deadline. A failed query is never a zero-session fact. +pub async fn sessions(pool: &PgPool, application_name: &str) -> Result<(i64, i64), String> { + tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(2)), + sqlx::query_as::<_, (i64, i64)>( + "SELECT count(*)::bigint, + count(*) FILTER (WHERE state = 'idle in transaction')::bigint + FROM pg_stat_activity WHERE application_name = $1 AND pid <> pg_backend_pid()", + ) + .bind(application_name) + .fetch_one(pool), + ) + .await + .map_err(|error| format!("session query timed out: {error}"))? + .map_err(|error| format!("session query failed: {error}")) +} + +#[derive(Debug, Default)] +pub struct Evidence { + pub peak_sessions: i64, + pub migrating_idle_transactions: i64, + pub ready_idle_transactions: i64, + pub ready_samples: usize, + pub remaining: Option, + pub last_exit_sample_valid: bool, + pub exit_phase: Option, + pub exit_success: bool, + pub incomplete: Option, +} + +impl Evidence { + pub fn observe(&mut self, phase_before_query: Phase, result: Result<(i64, i64), String>) { + match result { + Ok((sessions, idle)) => { + self.peak_sessions = self.peak_sessions.max(sessions); + if phase_before_query == Phase::Migrating { + self.migrating_idle_transactions = self.migrating_idle_transactions.max(idle); + } else { + if phase_before_query != Phase::Closed { + self.ready_samples += 1; + } + self.ready_idle_transactions = self.ready_idle_transactions.max(idle); + } + } + Err(error) => self.incomplete(error), + } + } + + /// Process exit alone does not prove that an active server query was cancelled. + /// CLOSED is emitted only after the client acknowledged its graceful close. + pub fn record_exit(&mut self, phase: Phase, success: bool) { + self.exit_phase = Some(phase); + self.exit_success = success; + if !self.closed_exit_proven() { + self.incomplete(format!( + "Python migrate probe exited without successful CLOSED evidence (phase={phase:?}, success={success})" + )); + } + } + + fn closed_exit_proven(&self) -> bool { + self.exit_success && self.exit_phase == Some(Phase::Closed) + } + + fn post_exit_observation_complete(&self, phase: Phase, success: bool) -> bool { + self.last_exit_sample_valid + && self.remaining == Some(0) + && (!success || phase == Phase::Closed) + } + + pub fn after_exit(&mut self, result: Result<(i64, i64), String>) { + match result { + Ok((sessions, _)) => { + self.remaining = Some(sessions); + self.last_exit_sample_valid = true; + } + Err(error) => { + self.last_exit_sample_valid = false; + self.incomplete(error); + } + } + } + + /// Reaping may precede stdout consumption. Marker and session observations + /// share the existing post-exit budget, including failed/slow queries. + pub async fn finish_exit(&mut self, phase: &AtomicU8, success: bool, mut sessions: F) + where + F: FnMut() -> Fut, + Fut: Future>, + { + let deadline = + tokio::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(2)); + loop { + if tokio::time::Instant::now() >= deadline { + break; + } + let result = tokio::time::timeout_at(deadline, sessions()) + .await + .unwrap_or_else(|_| { + Err("post-exit session observation exceeded its existing deadline".into()) + }); + let failed = result.is_err(); + self.after_exit(result); + let observed_phase = Phase::from_byte(phase.load(Ordering::Acquire)); + if self.post_exit_observation_complete(observed_phase, success) || failed { + break; + } + tokio::time::sleep_until( + (tokio::time::Instant::now() + Duration::from_millis(100)).min(deadline), + ) + .await; + } + self.record_exit(Phase::from_byte(phase.load(Ordering::Acquire)), success); + } + + pub fn incomplete(&mut self, error: impl Into) { + self.incomplete.get_or_insert_with(|| error.into()); + } + + pub fn violation(&self) -> bool { + self.peak_sessions > 4 + || self.ready_idle_transactions > 0 + || (self.closed_exit_proven() + && self.last_exit_sample_valid + && self.remaining.is_some_and(|sessions| sessions > 0)) + } + + pub fn complete(&self) -> bool { + self.incomplete.is_none() + && self.ready_samples > 0 + && self.closed_exit_proven() + && self.last_exit_sample_valid + && self.remaining.is_some() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::postgres::PgPoolOptions; + + async fn pool(application_name: &str, size: u32) -> PgPool { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".into()); + PgPoolOptions::new() + .max_connections(size) + .connect(&format!( + "{url}{}application_name={application_name}", + if url.contains('?') { '&' } else { '?' } + )) + .await + .expect("probe test pool") + } + + fn name() -> String { + format!("probe_test_{}", uuid::Uuid::new_v4().simple()) + } + + #[tokio::test] + async fn healthy_migration_transaction_is_not_a_stranded_ready_connection() { + let application = name(); + let observer = pool(&name(), 1).await; + let client = pool(&application, 1).await; + let mut tx = client.begin().await.expect("migration transaction"); + sqlx::query("SELECT 1").execute(tx.as_mut()).await.unwrap(); + let mut evidence = Evidence::default(); + let migrating = sessions(&observer, &application).await.unwrap(); + assert_eq!( + migrating, + (1, 1), + "a legitimate transaction between statements" + ); + evidence.observe(Phase::Migrating, Ok(migrating)); + tx.commit().await.unwrap(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + client.close().await; + evidence.record_exit(Phase::Closed, true); + // An acknowledged close can precede the server removing its backend. + let deadline = + std::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(2)); + loop { + let count = sessions(&observer, &application).await.unwrap(); + evidence.after_exit(Ok(count)); + if count.0 == 0 || std::time::Instant::now() >= deadline { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert!(evidence.complete(), "{evidence:?}"); + assert!(!evidence.violation(), "{evidence:?}"); + observer.close().await; + } + + #[tokio::test] + async fn ready_transaction_left_open_is_a_violation() { + let application = name(); + let observer = pool(&name(), 1).await; + let client = pool(&application, 1).await; + let mut tx = client.begin().await.unwrap(); + sqlx::query("SELECT 1").execute(tx.as_mut()).await.unwrap(); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + assert!(evidence.violation(), "{evidence:?}"); + tx.rollback().await.unwrap(); + client.close().await; + observer.close().await; + } + + #[tokio::test] + async fn retained_session_after_exit_is_a_violation_and_other_probes_are_excluded() { + let application = name(); + let observer = pool(&name(), 1).await; + let leaked = pool(&application, 1).await; + let unrelated = pool(&name(), 1).await; + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(sessions(&observer, &application).await); + assert_eq!(evidence.remaining, Some(1)); + assert!(evidence.complete()); + assert!(evidence.violation(), "{evidence:?}"); + leaked.close().await; + unrelated.close().await; + observer.close().await; + } + + #[tokio::test] + async fn session_query_failure_is_inconclusive_and_not_zero() { + let observer = pool(&name(), 1).await; + observer.close().await; + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &name()).await); + evidence.after_exit(sessions(&observer, &name()).await); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, None); + } + + #[test] + fn probe_exit_forced_during_a_read_retains_incomplete_session_evidence() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Ready, false); + evidence.after_exit(Ok((1, 0))); + assert!( + !evidence.violation(), + "a forced exit does not prove graceful-close failure" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_delayed_closed_marker_uses_the_existing_session_window() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.after_exit(Ok((0, 0))); + assert!( + !evidence.post_exit_observation_complete(Phase::Ready, true), + "zero sessions must not stop the reader before its flushed CLOSED marker" + ); + // The detached reader consumes the marker during the same post-exit window. + assert!(evidence.post_exit_observation_complete(Phase::Closed, true)); + evidence.record_exit(Phase::Closed, true); + assert!(evidence.complete()); + assert!(!evidence.violation()); + } + + #[tokio::test] + async fn probe_exit_finish_waits_for_delayed_stdout_closed_marker() { + let phase = std::sync::Arc::new(AtomicU8::new(Phase::Ready as u8)); + let reader_phase = phase.clone(); + let reader = tokio::spawn(async move { + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(10)).await; + reader_phase.store(Phase::Closed as u8, Ordering::Release); + }); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + let mut calls = 0; + evidence + .finish_exit(&phase, true, || { + calls += 1; + async { Ok((0, 0)) } + }) + .await; + reader.await.unwrap(); + assert!(evidence.complete(), "{evidence:?}"); + assert!(!evidence.violation()); + assert_eq!(evidence.exit_phase, Some(Phase::Closed)); + assert!( + calls >= 2, + "the helper must continue past zero sessions before CLOSED" + ); + } + + #[tokio::test] + async fn probe_exit_finish_query_failure_does_not_prove_a_retained_session_leak() { + let phase = AtomicU8::new(Phase::Closed as u8); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + let mut calls = 0; + evidence + .finish_exit(&phase, true, || { + calls += 1; + let result = if calls == 1 { + Ok((1, 0)) + } else { + Err("injected final query failure".into()) + }; + async move { result } + }) + .await; + assert!( + !evidence.violation(), + "an earlier positive count is not a successful final observation" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!( + evidence.remaining, + Some(1), + "retain the last successful count for diagnosis" + ); + assert!(!evidence.last_exit_sample_valid); + } + + #[test] + fn probe_exit_success_without_closed_marker_is_incomplete() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Ready, true); + evidence.after_exit(Ok((1, 0))); + assert!( + !evidence.violation(), + "an exit code is not an acknowledged pool close" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_acknowledged_close_with_retained_session_is_a_violation() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(Ok((1, 0))); + assert!(evidence.violation()); + assert!( + evidence.complete(), + "complete coverage still reports the failure" + ); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_acknowledged_close_and_no_session_is_complete() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(Ok((0, 0))); + assert!(!evidence.violation()); + assert!(evidence.complete()); + } + + #[test] + fn probe_exit_unproven_close_does_not_hide_pool_or_idle_transaction_violations() { + for (phase, sample) in [(Phase::Migrating, (5, 0)), (Phase::Ready, (2, 1))] { + let mut evidence = Evidence::default(); + evidence.observe(phase, Ok(sample)); + evidence.record_exit(phase, false); + evidence.after_exit(Ok((1, 0))); + assert!(evidence.violation(), "{evidence:?}"); + assert!(!evidence.complete()); + assert_eq!(evidence.remaining, Some(1)); + } + } + + #[test] + fn ready_marker_and_missing_ready_coverage_are_explicit() { + assert_eq!( + Phase::marker("MIGRATED probe queues=1 sample=0"), + Some(Phase::Ready) + ); + assert_eq!(Phase::marker("IDLE"), Some(Phase::Idle)); + assert_eq!(Phase::marker("CLOSED"), Some(Phase::Closed)); + assert_eq!(Phase::marker("MIGRATED startup queues=1 sample=0"), None); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Migrating, Ok((4, 1))); + evidence.after_exit(Ok((0, 0))); + assert!( + !evidence.complete(), + "an exit without ready samples proves no idle invariant" + ); + evidence.observe(Phase::Migrating, Ok((5, 0))); + assert!( + evidence.violation(), + "the pool budget applies during migration too" + ); + // Even valid ready/idle/closed markers may all reach the reader before + // the next snapshot. An already closed pool is not live idle coverage. + let mut closed_only = Evidence::default(); + closed_only.observe(Phase::Closed, Ok((0, 0))); + closed_only.record_exit(Phase::Closed, true); + closed_only.after_exit(Ok((0, 0))); + assert!(!closed_only.complete()); + } +} diff --git a/awa/tests/persona_soak/registry.rs b/awa/tests/persona_soak/registry.rs new file mode 100644 index 000000000..a0cc42835 --- /dev/null +++ b/awa/tests/persona_soak/registry.rs @@ -0,0 +1,131 @@ +//! The set of queues the run owns and the live facts the oracle needs about each: +//! which persona serves it, how many workers can run it, how many are busy, and +//! whether it is intentionally not being served right now. + +use super::config::Persona; +use std::collections::HashMap; +use std::sync::atomic::{AtomicI32, AtomicI64, Ordering}; +use std::sync::{Arc, RwLock}; + +impl QueueInfo { + pub fn hold(&self) { + self.pause_holds.fetch_add(1, Ordering::Relaxed); + } + + pub fn release(&self) { + self.pause_holds.fetch_sub(1, Ordering::Relaxed); + } + + pub fn is_held(&self) -> bool { + self.pause_holds.load(Ordering::Relaxed) > 0 + } + + /// Jobs enqueued and not yet started, as far as the producers and workers know. + pub fn backlog_estimate(&self) -> i64 { + self.enqueued.load(Ordering::Relaxed) - self.started.load(Ordering::Relaxed) + } +} + +pub struct QueueInfo { + pub name: String, + pub persona: Persona, + pub index: u32, + /// Workers currently able to run this queue across all live runtimes. + pub capacity: AtomicI64, + pub in_flight: AtomicI64, + /// Jobs producers enqueued here and executions begun, for rough backlog estimates. + pub enqueued: AtomicI64, + pub started: AtomicI64, + /// An operator (or a rate limit) is deliberately holding jobs back. + /// Operators currently holding the queue paused; the stall check skips it while non-zero. + pub pause_holds: AtomicI32, + pub rate_limited: bool, + /// Milliseconds since run start of the most recent operator drain, 0 if never. + pub drained_at_ms: AtomicI64, +} + +#[derive(Default)] +pub struct QueueRegistry { + queues: RwLock>>, + by_name: RwLock>, +} + +impl QueueRegistry { + pub fn register(&self, name: &str, persona: Persona, rate_limited: bool) -> Arc { + let mut queues = self.queues.write().expect("registry lock"); + let mut by_name = self.by_name.write().expect("registry lock"); + if let Some(index) = by_name.get(name) { + return queues[*index as usize].clone(); + } + let index = queues.len() as u32; + let info = Arc::new(QueueInfo { + name: name.to_string(), + persona, + index, + capacity: AtomicI64::new(0), + in_flight: AtomicI64::new(0), + enqueued: AtomicI64::new(0), + started: AtomicI64::new(0), + pause_holds: AtomicI32::new(0), + rate_limited, + drained_at_ms: AtomicI64::new(0), + }); + by_name.insert(name.to_string(), index); + queues.push(info.clone()); + info + } + + /// Looks a queue up by its physical name; striped queues (`name#N`) resolve to + /// their logical queue. + pub fn get(&self, name: &str) -> Option> { + let by_name = self.by_name.read().expect("registry lock"); + let index = by_name.get(name).or_else(|| { + let (logical, suffix) = name.rsplit_once('#')?; + suffix.parse::().ok()?; + by_name.get(logical) + })?; + Some(self.queues.read().expect("registry lock")[*index as usize].clone()) + } + + pub fn by_index(&self, index: u32) -> Option> { + self.queues + .read() + .expect("registry lock") + .get(index as usize) + .cloned() + } + + pub fn all(&self) -> Vec> { + self.queues.read().expect("registry lock").clone() + } + + pub fn of(&self, persona: Persona) -> Vec> { + self.all() + .into_iter() + .filter(|queue| queue.persona == persona) + .collect() + } +} + +/// Counts one running handler against its queue for as long as the guard lives, +/// including when the runtime is killed and the handler future is dropped. +pub struct InFlightGuard { + queue: Option>, +} + +impl InFlightGuard { + pub fn new(queue: Option>) -> Self { + if let Some(queue) = &queue { + queue.in_flight.fetch_add(1, Ordering::Relaxed); + } + Self { queue } + } +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + if let Some(queue) = &self.queue { + queue.in_flight.fetch_sub(1, Ordering::Relaxed); + } + } +} diff --git a/awa/tests/persona_soak/report.rs b/awa/tests/persona_soak/report.rs new file mode 100644 index 000000000..d309bfc94 --- /dev/null +++ b/awa/tests/persona_soak/report.rs @@ -0,0 +1,341 @@ +//! The run report and, for a failed run, the reproducer. + +use super::config::{Persona, SoakConfig}; +use super::pgserver::crash_excerpts; +use super::shared::Shared; +use serde_json::json; +use std::fmt::Write as _; +use std::path::Path; +use std::process::Command; + +/// Printed at the top of every report; also lets a built binary be matched to this harness. +pub const HEADER: &str = "AWA PERSONA SOAK HARNESS"; + +pub fn git_revision() -> String { + Command::new("git") + .args(["rev-parse", "--short=10", "HEAD"]) + .output() + .ok() + .map(|output| String::from_utf8_lossy(&output.stdout).trim().to_string()) + .filter(|revision| !revision.is_empty()) + .unwrap_or_else(|| "unknown".into()) +} + +pub struct Outcome { + pub quiesced: bool, + pub elapsed_seconds: f64, + pub crash_excerpts: Vec, + pub log_tail: Option, +} + +pub fn render(shared: &Shared, outcome: &Outcome) -> String { + let config = &shared.config; + let info = &shared.info; + let violations = shared.violations.lock().expect("violations"); + let mut out = String::new(); + let _ = writeln!(out, "# {HEADER}"); + let _ = writeln!(out); + let _ = writeln!( + out, + "verdict: {}", + if violations + .iter() + .any(|violation| violation.kind != "oracle_incomplete") + { + "FAIL" + } else if shared + .oracle_incomplete + .load(std::sync::atomic::Ordering::Relaxed) + { + "INCONCLUSIVE" + } else { + "PASS" + } + ); + let _ = writeln!( + out, + "engine: {} (schema {})", + info.engine, + info.queue_storage_schema.as_deref().unwrap_or("-") + ); + let _ = writeln!(out, "server: PostgreSQL {}", info.version); + let _ = writeln!( + out, + "max_connections: {} (usable by personas: {})", + info.max_connections, + shared.budget.usable() + ); + let _ = writeln!(out, "connections at start: {}", info.connections_at_start); + let _ = writeln!(out, "awa: {} @ {}", info.awa_version, git_revision()); + let _ = writeln!( + out, + "seed: {} personas: {} load: {}s scale: {} elapsed: {:.0}s quiesced: {}", + config.seed, + config + .personas + .iter() + .map(|persona| persona.name()) + .collect::>() + .join(","), + config.load.as_secs(), + config.scale, + outcome.elapsed_seconds, + outcome.quiesced + ); + let _ = writeln!(out, "ledger jobs: {}", shared.ledger.len()); + let _ = writeln!(out, "NO_EVENTS durable loss detection: job_id % {} = 0; other returned NO_EVENTS jobs are excluded from terminal-commit proof.", super::ledger::TERMINAL_PROOF_STRIDE); + let _ = writeln!(out); + + let _ = writeln!(out, "## Connection budget"); + for (label, count) in shared.budget.table() { + let _ = writeln!(out, "- {label}: {count}"); + } + let _ = writeln!(out); + + for persona in Persona::ALL { + if !config.has(persona) { + continue; + } + let stats = shared.stats(persona); + let counters = stats.counters(); + let _ = writeln!(out, "## {}", persona.name()); + for (name, value) in counters { + let _ = writeln!(out, "- {name}: {value}"); + } + for (name, summary) in stats.summaries() { + let _ = writeln!( + out, + "- {name}: n={} p50={} p99={} max={}", + summary.count, summary.p50, summary.p99, summary.max + ); + } + if persona == Persona::ClaimAbandon { + let _ = writeln!(out, "{}", shared.abandon.render()); + } + let _ = writeln!(out); + } + + let _ = writeln!(out, "## Run"); + if shared.run.get("claim_ring_lock_capture_records") > 0 { + let _ = writeln!( + out, + "Lock evidence: [claim-ring-locks.jsonl](claim-ring-locks.jsonl)" + ); + } + for (name, value) in shared.run.counters() { + let _ = writeln!(out, "- {name}: {value}"); + } + for (name, summary) in shared.run.summaries() { + let _ = writeln!( + out, + "- {name}: n={} p50={} p99={} max={}", + summary.count, summary.p50, summary.p99, summary.max + ); + } + let _ = writeln!(out); + + let findings = shared.findings.lock().expect("findings"); + let _ = writeln!(out, "## Findings, not fatal ({})", findings.len()); + for finding in findings.iter() { + let _ = writeln!( + out, + "- [{:.1}s] {}: {}", + finding.at_ms as f64 / 1000.0, + finding.kind, + finding.summary + ); + for row in finding.rows.iter().take(6) { + let _ = writeln!(out, " {row}"); + } + } + let _ = writeln!(out); + + let _ = writeln!(out, "## Violations ({})", violations.len()); + for violation in violations.iter() { + let _ = writeln!( + out, + "- [{:.1}s{}] {}: {}", + violation.at_ms as f64 / 1000.0, + if violation.during_chaos { + ", during chaos" + } else { + "" + }, + violation.kind, + violation.summary + ); + for row in violation.rows.iter().take(10) { + let _ = writeln!(out, " {row}"); + } + } + if !outcome.crash_excerpts.is_empty() { + let _ = writeln!(out, "\n## Server log: crash lines"); + for excerpt in &outcome.crash_excerpts { + let _ = writeln!(out, "```\n{excerpt}\n```"); + } + } + out +} + +/// Writes `report.md` and `report.json`, and for a failed run `reproducer.md`. +pub fn write(shared: &Shared, outcome: &Outcome) -> std::io::Result<()> { + let dir = &shared.config.report_dir; + std::fs::create_dir_all(dir)?; + let markdown = render(shared, outcome); + println!("\n{markdown}"); + std::fs::write(dir.join("report.md"), &markdown)?; + + let violations = shared.violations.lock().expect("violations").clone(); + let timeline = shared.timeline.lock().expect("timeline").clone(); + let per_persona: serde_json::Map = Persona::ALL + .into_iter() + .filter(|persona| shared.config.has(*persona)) + .map(|persona| { + ( + persona.name().to_string(), + json!({ "counters": shared.stats(persona).counters(), "samples": shared.stats(persona).summaries() }), + ) + }) + .collect(); + let summary = json!({ + "verdict": if violations.iter().any(|violation| violation.kind != "oracle_incomplete") { "fail" } + else if shared.oracle_incomplete.load(std::sync::atomic::Ordering::Relaxed) { "inconclusive" } else { "pass" }, + "terminal_proof_sample_stride": super::ledger::TERMINAL_PROOF_STRIDE, + "oracle_complete": !shared.oracle_incomplete.load(std::sync::atomic::Ordering::Relaxed), + "seed": shared.config.seed, + "personas": shared.config.personas.iter().map(|p| p.name()).collect::>(), + "engine": shared.info.engine, + "server_version": shared.info.version, + "max_connections": shared.info.max_connections, + "awa_version": shared.info.awa_version, + "git": git_revision(), + "elapsed_seconds": outcome.elapsed_seconds, + "quiesced": outcome.quiesced, + "personas_stats": per_persona, + "abandon": shared.abandon.to_json(), + "run": { "counters": shared.run.counters(), "samples": shared.run.summaries() }, + "violations": violations, + "findings": shared.findings.lock().expect("findings").clone(), + "timeline": timeline, + }); + std::fs::write( + dir.join("report.json"), + serde_json::to_string_pretty(&summary).unwrap_or_default(), + )?; + + if !violations.is_empty() { + write_reproducer(dir, &shared.config, shared, outcome)?; + } + Ok(()) +} + +fn write_reproducer( + dir: &Path, + config: &SoakConfig, + shared: &Shared, + outcome: &Outcome, +) -> std::io::Result<()> { + let violations = shared.violations.lock().expect("violations").clone(); + let timeline = shared.timeline.lock().expect("timeline").clone(); + let mut out = String::new(); + let _ = writeln!(out, "# Reproducer\n"); + let _ = writeln!(out, "Re-run with the same seed, personas and duration:\n"); + let _ = writeln!( + out, + "```\nAWA_SOAK_DATABASE_URL= AWA_SOAK_SEED={} AWA_SOAK_PERSONAS={} AWA_SOAK_SECONDS={} AWA_SOAK_SCALE={} \\\n cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture\n```\n", + config.seed, + config.personas.iter().map(|persona| persona.name()).collect::>().join(","), + config.load.as_secs(), + config.scale + ); + let _ = writeln!(out, "Seed fixes the workload mix and the chaos schedule; timing of races still varies between runs.\n"); + let _ = writeln!( + out, + "Engine: {}; PostgreSQL {}; awa {} @ {}\n", + shared.info.engine, + shared.info.version, + shared.info.awa_version, + git_revision() + ); + let _ = writeln!(out, "## Violations\n"); + for violation in &violations { + let persona = violation.kind.as_str(); + let _ = writeln!( + out, + "### {persona} at {:.1}s{}\n{}\n", + violation.at_ms as f64 / 1000.0, + if violation.during_chaos { + " (during chaos)" + } else { + "" + }, + violation.summary + ); + for row in &violation.rows { + let _ = writeln!(out, " {row}"); + } + let nearby: Vec<_> = timeline + .iter() + .filter(|event| (event.at_ms - violation.at_ms).abs() <= 60_000) + .collect(); + if !nearby.is_empty() { + let _ = writeln!(out, "\nTimeline within 60s:"); + for event in nearby { + let _ = writeln!( + out, + "- {:.1}s [{}] {}", + event.at_ms as f64 / 1000.0, + event.source, + event.text + ); + } + } + let _ = writeln!(out); + } + let _ = writeln!(out, "## Full timeline\n"); + for event in &timeline { + let _ = writeln!( + out, + "- {:.1}s [{}] {}", + event.at_ms as f64 / 1000.0, + event.source, + event.text + ); + } + let _ = writeln!(out, "\n## Server log excerpt\n"); + if !outcome.crash_excerpts.is_empty() { + for excerpt in &outcome.crash_excerpts { + let _ = writeln!(out, "```\n{excerpt}\n```"); + } + } else if let Some(tail) = &outcome.log_tail { + let interesting: Vec<&str> = tail + .lines() + .filter(|line| { + line.contains("ERROR") + || line.contains("FATAL") + || line.contains("PANIC") + || line.contains("WARNING") + }) + .collect(); + let start = interesting.len().saturating_sub(80); + let _ = writeln!( + out, + "No crash lines. Last warnings and errors ({} total):\n```", + interesting.len() + ); + for line in &interesting[start..] { + let _ = writeln!(out, "{line}"); + } + let _ = writeln!(out, "```"); + } else { + let _ = writeln!( + out, + "No server log source configured (set AWA_SOAK_PG_LOG or AWA_SOAK_PG_CONTAINER)." + ); + } + std::fs::write(dir.join("reproducer.md"), out) +} + +pub fn excerpts_of(log: &str) -> Vec { + crash_excerpts(log) +} diff --git a/awa/tests/persona_soak/ring_lock_capture.sql b/awa/tests/persona_soak/ring_lock_capture.sql new file mode 100644 index 000000000..165990fe1 --- /dev/null +++ b/awa/tests/persona_soak/ring_lock_capture.sql @@ -0,0 +1,63 @@ +WITH current_db AS ( + SELECT oid FROM pg_catalog.pg_database WHERE datname = current_database() +), +relations AS MATERIALIZED ( + SELECT c.oid, c.relname::text AS relation + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname::text = ANY($2::text[]) +), +activity AS MATERIALIZED ( + SELECT pid, leader_pid, backend_start, application_name, usename, state, + wait_event_type, wait_event, xact_start, query_start, + EXTRACT(EPOCH FROM (clock_timestamp() - xact_start))::float8 AS tx_age_s, + EXTRACT(EPOCH FROM (clock_timestamp() - query_start))::float8 AS query_age_s, + left(query, 512) AS query, + query IS NULL OR query = '' AS visibility_limited + FROM pg_catalog.pg_stat_activity + WHERE datid = (SELECT oid FROM current_db) +), +locks AS MATERIALIZED ( + SELECT l.pid, r.oid AS relation_oid, r.relation, l.mode, l.granted, l.fastpath, + l.virtualtransaction, l.waitstart + FROM pg_catalog.pg_locks l + JOIN relations r ON r.oid = l.relation + LEFT JOIN activity a ON a.pid = l.pid + WHERE l.database = (SELECT oid FROM current_db) + AND l.locktype = 'relation' + AND l.pid IS DISTINCT FROM pg_backend_pid() + ORDER BY l.granted, a.xact_start NULLS LAST, a.query_start NULLS LAST, + l.pid NULLS FIRST, r.relation, l.mode + LIMIT 257 --256 retained rows plus an explicit overflow sentinel +), +interested AS MATERIALIZED ( + SELECT DISTINCT l.pid + FROM locks l + JOIN activity a ON a.pid = l.pid + WHERE a.wait_event_type = 'Lock' + ORDER BY l.pid + LIMIT 32 +), +blockers AS MATERIALIZED ( + SELECT i.pid, pg_catalog.pg_blocking_pids(i.pid) AS pids + FROM interested i +), +lock_rows AS MATERIALIZED ( +SELECT clock_timestamp() AS observed_at, l.*, + a.leader_pid, a.backend_start, a.application_name, a.usename, a.state, + a.wait_event_type, a.wait_event, a.xact_start, a.query_start, + a.tx_age_s, a.query_age_s, a.query, a.visibility_limited, + a.pid IS NULL AS activity_missing, + b.pids[1:16] AS blocking_pids, + b.pid IS NOT NULL AS blocking_pids_observed, + COALESCE(cardinality(b.pids) > 16, false) AS blocking_pids_truncated +FROM locks l +LEFT JOIN activity a ON a.pid = l.pid +LEFT JOIN blockers b ON b.pid = l.pid +) +SELECT jsonb_build_object( + 'server_statement_at', statement_timestamp(), + 'server_observed_at', clock_timestamp(), + 'resolved_relations', COALESCE((SELECT jsonb_agg(to_jsonb(r)) FROM relations r), '[]'::jsonb), + 'locks', COALESCE((SELECT jsonb_agg(to_jsonb(l) ORDER BY l.granted, l.xact_start NULLS LAST, l.pid NULLS FIRST, l.relation, l.mode) FROM lock_rows l), '[]'::jsonb) +); diff --git a/awa/tests/persona_soak/ring_locks.rs b/awa/tests/persona_soak/ring_locks.rs new file mode 100644 index 000000000..3e9a0947b --- /dev/null +++ b/awa/tests/persona_soak/ring_locks.rs @@ -0,0 +1,543 @@ +//! Supplemental lock evidence; never changes the claim-ring oracle's verdict. + +use super::abandon_oracle::RingEpisode; +use super::shared::Shared; +use serde_json::{json, Value}; +use sqlx::PgPool; +use std::collections::BTreeSet; +use std::io::Write; +use std::time::Duration; + +const SQL: &str = include_str!("ring_lock_capture.sql"); +const LOCK_ROWS: usize = 256; +const BLOCKER_PIDS: usize = 64; +const CLIENT_DEADLINE: Duration = Duration::from_millis(750); + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Stage { + Suspected, + Violation, +} + +impl Stage { + fn index(self) -> usize { + match self { + Self::Suspected => 0, + Self::Violation => 1, + } + } + + fn name(self) -> &'static str { + match self { + Self::Suspected => "suspected", + Self::Violation => "violation", + } + } +} + +#[derive(Debug, PartialEq)] +pub enum Decision { + Capture(Stage), + Exhausted(Stage), +} + +/// Four slots per stage reserve half the run budget for actual violations. +#[derive(Default)] +pub struct CaptureBudget { + episode: Option<(i64, i64)>, + seen: [bool; 2], + used: [u8; 2], + exhausted_reported: [bool; 2], +} + +impl CaptureBudget { + pub fn plan( + &mut self, + episode: &RingEpisode, + now_ms: i64, + bound_ms: i64, + zero_live: bool, + ) -> Option { + let key = (episode.generation, episode.start_ms); + if self.episode != Some(key) { + self.episode = Some(key); + self.seen = [false; 2]; + } + let stage = if episode.reported { + Stage::Violation + } else if zero_live && now_ms - episode.start_ms >= bound_ms / 2 { + Stage::Suspected + } else { + return None; + }; + let index = stage.index(); + if self.seen[index] { + return None; + } + self.seen[index] = true; + if self.used[index] < 4 { + self.used[index] += 1; + Some(Decision::Capture(stage)) + } else if !self.exhausted_reported[index] { + self.exhausted_reported[index] = true; + Some(Decision::Exhausted(stage)) + } else { + None + } + } +} + +fn names(slot: i32) -> Vec { + [ + format!("lease_claims_{slot}"), + format!("lease_claim_batches_{slot}"), + format!("lease_claim_closures_{slot}"), + format!("lease_claim_closure_batches_{slot}"), + "claim_ring_state".into(), + "claim_ring_slots".into(), + ] + .into() +} + +fn missing(error: impl std::fmt::Display) -> Value { + json!({"status": "missing", "error": error.to_string().chars().take(512).collect::()}) +} + +fn finish(mut value: Value, expected: &[String]) -> Value { + let mut gaps = Vec::new(); + let resolved: BTreeSet<_> = value["resolved_relations"] + .as_array() + .into_iter() + .flatten() + .filter_map(|row| row["relation"].as_str()) + .collect(); + let unresolved: Vec<_> = expected + .iter() + .filter(|name| !resolved.contains(name.as_str())) + .cloned() + .collect(); + if !unresolved.is_empty() { + gaps.push("unresolved_relations"); + } + if let Some(rows) = value["locks"].as_array_mut() { + if rows.len() > LOCK_ROWS { + rows.truncate(LOCK_ROWS); + gaps.push("lock_rows_truncated"); + } + if rows.iter().any(|row| { + row["blocking_pids_truncated"] == true + || (row["wait_event_type"] == "Lock" && row["blocking_pids_observed"] != true) + || row["activity_missing"] == true + || row["visibility_limited"] == true + }) { + gaps.push("incomplete_activity_or_blocker_graph"); + } + } + if value["blocker_pid_lookup_truncated"] == true { + gaps.push("blocker_pid_lookup_truncated"); + } + if value["missing_blocker_activity"] + .as_array() + .is_some_and(|rows| !rows.is_empty()) + { + gaps.push("missing_blocker_activity"); + } + if value["blocker_activity"] + .as_array() + .is_some_and(|rows| rows.iter().any(|row| row["visibility_limited"] == true)) + { + gaps.push("blocker_activity_visibility_limited"); + } + value["status"] = json!(if gaps.is_empty() { "ok" } else { "partial" }); + value["missing"] = json!(gaps); + value["unresolved_relations"] = json!(unresolved); + value +} + +async fn collect(pool: &PgPool, schema: &str, slot: i32, sql: &str) -> Value { + let expected = names(slot); + let result = tokio::time::timeout(CLIENT_DEADLINE, async { + let mut tx = pool.begin().await?; + sqlx::query("SET TRANSACTION READ ONLY") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL statement_timeout = '250ms'") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL lock_timeout = '100ms'") + .execute(tx.as_mut()) + .await?; + let mut value: Value = sqlx::query_scalar(awa::audited_sql(sql)) + .bind(schema) + .bind(&expected) + .fetch_one(tx.as_mut()) + .await?; + let pids: BTreeSet = value["locks"] + .as_array() + .into_iter() + .flatten() + .flat_map(|row| row["blocking_pids"].as_array().into_iter().flatten()) + .filter_map(|pid| pid.as_i64().and_then(|pid| i32::try_from(pid).ok())) + .filter(|pid| *pid > 0) + .collect(); + value["blocker_pid_lookup_truncated"] = json!(pids.len() > BLOCKER_PIDS); + let requested: Vec<_> = pids.into_iter().take(BLOCKER_PIDS).collect(); + let activity: Value = sqlx::query_scalar( + "SELECT COALESCE(jsonb_agg(to_jsonb(a)), '[]'::jsonb) FROM ( + SELECT pid, leader_pid, backend_start, application_name, usename, state, + wait_event_type, wait_event, xact_start, query_start, + EXTRACT(EPOCH FROM(clock_timestamp()-xact_start))::float8 AS tx_age_s, + EXTRACT(EPOCH FROM(clock_timestamp()-query_start))::float8 AS query_age_s, + left(query,512) AS query, + query IS NULL OR query = '' AS visibility_limited + FROM pg_catalog.pg_stat_activity + WHERE datid=(SELECT oid FROM pg_catalog.pg_database WHERE datname=current_database()) + AND pid=ANY($1::int4[]) ORDER BY pid LIMIT 65 + ) a", + ) + .bind(&requested) + .fetch_one(tx.as_mut()) + .await?; + let found: BTreeSet<_> = activity + .as_array() + .into_iter() + .flatten() + .filter_map(|row| row["pid"].as_i64()) + .collect(); + value["missing_blocker_activity"] = json!(requested + .iter() + .filter(|pid| !found.contains(&i64::from(**pid))) + .collect::>()); + value["blocker_activity"] = activity; + tx.rollback().await?; + Ok::<_, sqlx::Error>(value) + }) + .await; + match result { + Ok(Ok(value)) => finish(value, &expected), + Ok(Err(error)) => missing(error), + Err(_) => missing("claim-ring metadata capture exceeded its 750ms client deadline"), + } +} + +pub async fn observe( + shared: &Shared, + pool: &PgPool, + schema: &str, + slot_count: i32, + next_slot: Option<(i64, i64)>, + budget: &mut CaptureBudget, +) { + let Some(episode) = shared.abandon.current_ring_episode() else { + return; + }; + let now = shared.now_ms(); + let Some(decision) = budget.plan( + &episode, + now, + shared.config.abandon_bound.as_millis() as i64, + next_slot.is_some_and(|(live, _)| live == 0), + ) else { + return; + }; + let stage = match decision { + Decision::Capture(stage) => stage, + Decision::Exhausted(stage) => { + shared.finding( + "claim_ring_lock_capture_missing", + format!("{} lock-capture budget exhausted", stage.name()), + vec![], + ); + return; + } + }; + let incoming = (episode.slot + 1).rem_euclid(slot_count); + let mut value = collect(pool, schema, incoming, SQL).await; + value["stage"] = json!(stage.name()); + value["started_at_ms"] = json!(now); + value["finished_at_ms"] = json!(shared.now_ms()); + value["current_slot"] = json!(episode.slot); + value["next_slot"] = json!(incoming); + value["generation"] = json!(episode.generation); + value["next_slot_probe"] = json!(next_slot); + shared.run.add("claim_ring_lock_capture_attempts", 1); + if value["status"] != "ok" { + shared.finding( + "claim_ring_lock_capture_missing", + "claim-ring lock evidence is missing or partial", + vec![value.to_string()], + ); + } + let path = shared.config.report_dir.join("claim-ring-locks.jsonl"); + let written = (|| { + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path)?; + writeln!(file, "{value}") + })(); + match written { + Ok(()) => shared.run.add("claim_ring_lock_capture_records", 1), + Err(error) => shared.finding( + "claim_ring_lock_capture_missing", + format!("could not write claim-ring lock evidence: {error}"), + vec![], + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn episode(generation: i64, reported: bool) -> RingEpisode { + RingEpisode { + generation, + reported, + slot: 7, + start_ms: 0, + end_ms: None, + last_live_ms: 0, + expired_in_next_slot: 0, + } + } + + #[test] + fn ring_lock_budget_waits_half_bound_and_reserves_violation_captures() { + let mut budget = CaptureBudget::default(); + assert_eq!(budget.plan(&episode(0, false), 59_999, 120_000, true), None); + assert_eq!( + budget.plan(&episode(0, false), 60_000, 120_000, false), + None + ); + for generation in 1..=4 { + assert_eq!( + budget.plan(&episode(generation, false), 60_000, 120_000, true), + Some(Decision::Capture(Stage::Suspected)) + ); + assert_eq!( + budget.plan(&episode(generation, false), 70_000, 120_000, true), + None + ); + } + assert_eq!( + budget.plan(&episode(5, false), 60_000, 120_000, true), + Some(Decision::Exhausted(Stage::Suspected)) + ); + assert_eq!(budget.plan(&episode(6, false), 60_000, 120_000, true), None); + for generation in 7..=10 { + assert_eq!( + budget.plan(&episode(generation, true), 120_001, 120_000, true), + Some(Decision::Capture(Stage::Violation)) + ); + assert_eq!( + budget.plan(&episode(generation, true), 130_000, 120_000, true), + None + ); + } + assert_eq!( + budget.plan(&episode(11, true), 120_001, 120_000, true), + Some(Decision::Exhausted(Stage::Violation)) + ); + assert_eq!( + budget.plan(&episode(12, true), 120_001, 120_000, true), + None + ); + assert_eq!(budget.used, [4, 4]); + } + + #[test] + fn ring_lock_result_caps_and_missing_relations_stay_explicit() { + let rows = vec![json!({"blocking_pids_truncated": true}); LOCK_ROWS + 1]; + let value = finish(json!({"resolved_relations": [], "locks": rows}), &names(0)); + assert_eq!(value["status"], "partial"); + assert_eq!(value["locks"].as_array().unwrap().len(), LOCK_ROWS); + assert_eq!(value["unresolved_relations"].as_array().unwrap().len(), 6); + let gaps = value["missing"].as_array().unwrap(); + assert!(gaps.contains(&json!("lock_rows_truncated"))); + assert!(gaps.contains(&json!("incomplete_activity_or_blocker_graph"))); + } + + #[test] + fn ring_lock_hidden_second_lookup_is_partial_even_without_selected_locks() { + let expected = names(0); + let inventory: Vec<_> = expected + .iter() + .map(|name| json!({"relation": name})) + .collect(); + let value = finish( + json!({ + "resolved_relations": inventory, + "locks": [], + "blocker_activity": [{"pid": 42, "query": null, "visibility_limited": true}] + }), + &expected, + ); + assert_eq!(value["status"], "partial"); + assert!(value["missing"] + .as_array() + .unwrap() + .contains(&json!("blocker_activity_visibility_limited"))); + } + + struct Database { + config: super::super::config::SoakConfig, + pool: PgPool, + } + + impl Database { + async fn create() -> Self { + use super::super::config::{replace_database_name, SoakConfig}; + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_ring_locks_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + super::super::db::recreate_database(&config) + .await + .expect("isolated database"); + let pool = super::super::db::connect_pool(&config.database_url, "ring_capture_test", 6) + .await + .expect("fixture pool"); + Self { config, pool } + } + + async fn drop(self) { + self.pool.close().await; + let admin = + super::super::db::connect_pool(&self.config.admin_url(), "ring_capture_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(awa::audited_sql(format!( + "DROP DATABASE {} WITH (FORCE)", + self.config.database_name() + ))) + .execute(&admin) + .await + .expect("drop fixture database"); + admin.close().await; + } + } + + #[tokio::test] + async fn ring_lock_capture_identifies_holder_waiter_and_releases_cleanly() { + let db = Database::create().await; + let pool = &db.pool; + let mut holder = pool.begin().await.unwrap(); + sqlx::query("SET LOCAL application_name = 'ring_capture_holder'") + .execute(holder.as_mut()) + .await + .unwrap(); + sqlx::query("LOCK TABLE awa.lease_claims_0 IN ACCESS SHARE MODE") + .execute(holder.as_mut()) + .await + .unwrap(); + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(holder.as_mut()) + .await + .unwrap(); + let mut waiter = pool.begin().await.unwrap(); + sqlx::query("SET LOCAL application_name = 'ring_capture_waiter'") + .execute(waiter.as_mut()) + .await + .unwrap(); + let waiter_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(waiter.as_mut()) + .await + .unwrap(); + let waiting = tokio::spawn(async move { + sqlx::query("LOCK TABLE awa.lease_claims_0 IN ACCESS EXCLUSIVE MODE") + .execute(waiter.as_mut()) + .await + .unwrap(); + waiter.rollback().await.unwrap(); + }); + let deadline = + tokio::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_locks WHERE pid=$1 AND NOT granted)", + ) + .bind(waiter_pid) + .fetch_one(pool) + .await + .unwrap(); + if waiting { + break; + } + assert!( + tokio::time::Instant::now() < deadline, + "exclusive waiter did not queue" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + let value = collect(pool, "awa", 0, SQL).await; + assert_eq!(value["status"], "ok", "{value}"); + assert_eq!(value["resolved_relations"].as_array().unwrap().len(), 6); + let rows = value["locks"].as_array().unwrap(); + let held = rows + .iter() + .find(|row| row["pid"] == holder_pid && row["relation"] == "lease_claims_0") + .expect("holder evidence"); + assert_eq!(held["mode"], "AccessShareLock"); + assert_eq!(held["granted"], true); + assert_eq!(held["application_name"], "ring_capture_holder"); + assert!(held["tx_age_s"].as_f64().unwrap() >= 0.0); + let blocked = rows + .iter() + .find(|row| row["pid"] == waiter_pid && row["granted"] == false) + .expect("waiter evidence"); + assert_eq!(blocked["mode"], "AccessExclusiveLock"); + assert_eq!(blocked["application_name"], "ring_capture_waiter"); + assert_eq!(blocked["wait_event_type"], "Lock"); + assert!(blocked["blocking_pids"] + .as_array() + .unwrap() + .contains(&json!(holder_pid))); + assert!(value["blocker_activity"] + .as_array() + .unwrap() + .iter() + .any(|row| row["pid"] == holder_pid)); + holder.rollback().await.unwrap(); + tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(5)), + waiting, + ) + .await + .unwrap() + .unwrap(); + let clean = collect(pool, "awa", 0, SQL).await; + assert_eq!(clean["status"], "ok", "{clean}"); + assert!(clean["server_statement_at"].as_str().is_some()); + assert!(clean["server_observed_at"].as_str().is_some()); + assert_eq!(clean["resolved_relations"].as_array().unwrap().len(), 6); + assert!(clean["locks"] + .as_array() + .unwrap() + .iter() + .all(|row| row["pid"] != holder_pid && row["pid"] != waiter_pid)); + db.drop().await; + } + + #[tokio::test] + async fn ring_lock_capture_failure_and_unresolved_inventory_are_not_empty_success() { + let db = Database::create().await; + let value = collect(&db.pool, "missing_ring_schema", 0, SQL).await; + assert_eq!(value["status"], "partial"); + assert_eq!(value["resolved_relations"], json!([])); + assert_eq!(value["locks"], json!([])); + assert!(value["server_statement_at"].as_str().is_some()); + assert!(value["server_observed_at"].as_str().is_some()); + assert_eq!(value["unresolved_relations"].as_array().unwrap().len(), 6); + let failed = collect(&db.pool, "awa", 0, + "SELECT jsonb_build_object() FROM pg_sleep(1) WHERE $1::text IS NOT NULL AND $2::text[] IS NOT NULL").await; + assert_eq!(failed["status"], "missing", "{failed}"); + assert!( + failed.get("locks").is_none(), + "failed capture must not look unlocked" + ); + assert!(!failed["error"].as_str().unwrap().is_empty()); + db.drop().await; + } +} diff --git a/awa/tests/persona_soak/rng.rs b/awa/tests/persona_soak/rng.rs new file mode 100644 index 000000000..ca836b27b --- /dev/null +++ b/awa/tests/persona_soak/rng.rs @@ -0,0 +1,64 @@ +//! Deterministic pseudo-random numbers so a failing soak can be replayed from its seed. + +/// SplitMix64: tiny, fast, and good enough for workload shaping. +#[derive(Debug, Clone)] +pub struct Rng { + state: u64, +} + +impl Rng { + pub fn new(seed: u64) -> Self { + Self { state: seed } + } + + /// A generator for one task, derived from the run seed and a stable label so + /// that adding a persona does not shift another persona's sequence. + pub fn for_task(seed: u64, label: &str) -> Self { + let mut hash = seed ^ 0xcbf2_9ce4_8422_2325; + for byte in label.bytes() { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + Self::new(hash) + } + + pub fn next_u64(&mut self) -> u64 { + self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut value = self.state; + value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + value ^ (value >> 31) + } + + /// Uniform in `0..bound` (`bound` must be non-zero). + pub fn below(&mut self, bound: u64) -> u64 { + self.next_u64() % bound.max(1) + } + + /// Uniform in `low..=high`. + pub fn range(&mut self, low: u64, high: u64) -> u64 { + low + self.below(high.saturating_sub(low) + 1) + } + + /// True with probability `percent` / 100. + pub fn chance(&mut self, percent: u64) -> bool { + self.below(100) < percent + } + + pub fn pick<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> { + if items.is_empty() { + None + } else { + Some(&items[self.below(items.len() as u64) as usize]) + } + } +} + +/// Nearest-rank percentile of a sorted slice (`quantile` in `0.0..=1.0`). +pub fn percentile(sorted: &[i64], quantile: f64) -> i64 { + if sorted.is_empty() { + return 0; + } + let index = ((sorted.len() - 1) as f64 * quantile).round() as usize; + sorted[index] +} diff --git a/awa/tests/persona_soak/roles.rs b/awa/tests/persona_soak/roles.rs new file mode 100644 index 000000000..6bb16cc94 --- /dev/null +++ b/awa/tests/persona_soak/roles.rs @@ -0,0 +1,143 @@ +//! The database roles of the shard-ops persona: the shape of a production install that +//! runs its application processes without superuser or schema ownership. +//! +//! * `OWNER` owns the database and the `awa` schema, runs migrations and `prepare_queue`. +//! * `MIGRATOR` is a migration-like role that is a member of `OWNER` and has default +//! privileges for the objects it creates. +//! * `APP_A` and `APP_B` are members of the `RUNTIME_GROUP` role and have no default +//! privileges of their own by default. An explicit comparison profile can share +//! their newly created objects too. They run producers and the maintenance leader. +//! +//! Grants follow the restricted-runtime section of `docs/security/database-roles.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(awa::audited_sql(&statement)) + .execute(admin) + .await + .map_err(|error| format!("create roles: {error}"))?; + } + Ok(()) +} + +/// Applies the restricted-runtime grants. Run once the schema is migrated, over a +/// superuser connection to the soak database. +pub async fn grant_runtime( + superuser_pool: &PgPool, + database: &str, + schema_create: bool, + mixed_personas: bool, +) -> Result<(), String> { + // Comparison-only setup profile for older binaries whose application roles + // create their own lane sequences. Default off preserves regression coverage. + let shared_creator_grants = match std::env::var("AWA_SOAK_SHARED_CREATOR_GRANTS") { + Ok(value) if value == "1" => true, + Ok(value) if value == "0" => false, + Err(std::env::VarError::NotPresent) => false, + _ => return Err("AWA_SOAK_SHARED_CREATOR_GRANTS must be 0 or 1".into()), + }; + if shared_creator_grants && !schema_create { + return Err("shared creator grants require schema CREATE for the runtime group".into()); + } + let mut statements = vec![ + format!("GRANT USAGE ON SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("REVOKE EXECUTE ON FUNCTION awa.install_queue_storage_substrate(TEXT, INT, INT, INT, BOOLEAN) FROM {RUNTIME_GROUP}"), + format!("GRANT TEMP ON DATABASE {database} TO {RUNTIME_GROUP}"), + ]; + if schema_create { + statements.push(format!("GRANT CREATE ON SCHEMA awa TO {RUNTIME_GROUP}")); + } + let mut creators = vec![OWNER.to_string(), MIGRATOR.to_string()]; + if mixed_personas { + // The other personas keep their superuser runtimes and producers. Their objects + // must stay usable by the application roles, which now lead maintenance for every + // queue. Application-role defaults remain absent unless the explicit + // shared-creator comparison profile below enables them. + let superuser: String = sqlx::query_scalar("SELECT current_user::text") + .fetch_one(superuser_pool) + .await + .map_err(|error| format!("current_user: {error}"))?; + creators.push(superuser); + } + if shared_creator_grants { + creators.extend([APP_A.to_string(), APP_B.to_string()]); + eprintln!("SOAK_SETUP shared_creator_grants=1 schema_create=1 application_creators={APP_A},{APP_B}"); + } + for creator in &creators { + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON TABLES TO {RUNTIME_GROUP}" + )); + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO {RUNTIME_GROUP}" + )); + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT EXECUTE ON FUNCTIONS TO {RUNTIME_GROUP}" + )); + } + for statement in statements { + sqlx::query(awa::audited_sql(&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..58992d9e7 --- /dev/null +++ b/awa/tests/persona_soak/run.rs @@ -0,0 +1,303 @@ +//! 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> { + let worker_errors = super::worker_errors::install(false)?; + eprintln!( + "{}: preparing database {}", + report::HEADER, + config.database_name() + ); + db::recreate_database(&config).await?; + let oracle_pool = db::connect_pool(&config.database_url, "soak_oracle", 4) + .await + .map_err(|error| format!("oracle connect: {error}"))?; + let info = db::server_info(&oracle_pool) + .await + .map_err(|error| format!("server info: {error}"))?; + let watch = Arc::new( + PgWatch::begin(&config, &oracle_pool) + .await + .map_err(|error| format!("pg watch: {error}"))?, + ); + let budget = Budget::new(&info, HEADROOM); + eprintln!( + "{}: engine={} server={} max_connections={} connections_now={} usable={} seed={} personas={:?} load={}s scale={}", + report::HEADER, + info.engine, + info.version, + info.max_connections, + info.connections_at_start, + budget.usable(), + config.seed, + config.personas.iter().map(|persona| persona.name()).collect::>(), + config.load.as_secs(), + config.scale + ); + if info.engine != "queue_storage" { + return Err(format!( + "expected the queue_storage engine, found {}", + info.engine + )); + } + let shared = Shared::new(config, info, budget); + + // The UI personas talk to the real admin server, run as a child process. + let ui_personas = [ + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + ]; + let admin_server = if ui_personas + .iter() + .any(|persona| shared.config.has(*persona)) + { + std::fs::create_dir_all(&shared.config.report_dir).map_err(|error| error.to_string())?; + let server = AdminServer::start(&shared.config).await?; + shared.ui.set_base_url(server.base_url.clone()); + shared.note( + "run", + format!( + "admin server at {} (statement timeout {}s, client timeout {}s)", + server.base_url, + server.statement_timeout.as_secs(), + shared.config.ui_client_timeout.as_secs() + ), + ); + Some(server) + } else { + None + }; + + let terminal_observer = tokio::spawn(super::terminal_proof::run( + shared.clone(), + oracle_pool.clone(), + shared.drain_stop.clone(), + )); + + // Personas start their runtimes and then produce until the load stops. Chaos, the + // operator and the UI users reach other personas' fleets and queues, so they start + // once the rest are up. + let mut tasks: Vec<(Persona, tokio::task::JoinHandle>)> = Vec::new(); + for persona in [ + Persona::Mia, + Persona::Tom, + Persona::Priya, + Persona::Sam, + Persona::Python, + Persona::Fanout, + Persona::ShardOps, + Persona::ClaimAbandon, + ] { + if shared.config.has(persona) { + tasks.push((persona, spawn_persona(persona, shared.clone()))); + } + } + tokio::time::sleep(Duration::from_secs(8)).await; + for persona in [ + Persona::Alex, + Persona::Chaos, + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + ] { + if shared.config.has(persona) { + tasks.push((persona, spawn_persona(persona, shared.clone()))); + } + } + + let oracle_stop = Stop::new(); + let oracle_task = { + let oracle = Oracle::new(shared.clone(), oracle_pool.clone(), watch.clone()); + let stop = oracle_stop.clone(); + tokio::spawn(async move { oracle.run(stop).await }) + }; + shared.note( + "run", + format!("load phase started: {}s", shared.config.load.as_secs()), + ); + shared.started_load(Instant::now()); + + let load_deadline = Instant::now() + shared.config.load; + let mut next_progress = Instant::now() + Duration::from_secs(30); + while Instant::now() < load_deadline && !shared.load_stop.is_stopped() { + tokio::time::sleep(Duration::from_secs(1)).await; + if let Some((persona, _)) = tasks.iter().find(|(_, task)| task.is_finished()) { + // A persona only returns early on a setup failure or a panic. + let persona = *persona; + let index = tasks + .iter() + .position(|(p, _)| *p == persona) + .expect("finished task"); + let (_, task) = tasks.remove(index); + match task.await { + Ok(Ok(())) => shared.note("run", format!("{} finished early", persona.name())), + Ok(Err(error)) => shared.violation( + "persona_failed", + format!("{} stopped: {error}", persona.name()), + vec![], + ), + Err(error) => shared.violation( + "persona_panicked", + format!("{} panicked: {error}", persona.name()), + vec![], + ), + } + } + if Instant::now() >= next_progress { + next_progress = Instant::now() + Duration::from_secs(30); + progress(&shared); + } + } + + shared.note("run", "load phase over; letting the system quiesce"); + shared.load_stop.trigger(); + oracle_stop.trigger(); + let mut oracle = oracle_task + .await + .map_err(|error| format!("oracle task: {error}"))?; + let quiesced = oracle.wait_quiescent(shared.config.drain_grace).await; + shared + .quiesced + .store(quiesced, std::sync::atomic::Ordering::Relaxed); + shared.note("run", format!("quiesced: {quiesced}")); + super::terminal_proof::observe(&shared, &oracle_pool).await; + oracle.final_check(quiesced).await; + let excerpts = oracle.crash_scan(); + for excerpt in &excerpts { + shared.violation( + "server_crash_log", + "the Postgres log contains crash lines", + vec![excerpt.clone()], + ); + } + let log_tail = watch.log_since_start(&shared.config); + if shared.config.has(Persona::ShardOps) { + match &log_tail { + Some(log) => super::shard_oracle::scan_server_log(&shared, &oracle_pool, log).await, + None => shared.note( + "run", + "no Postgres log source is configured, so permission errors raised inside the runtimes are not scanned", + ), + } + } + + shared.drain_stop.trigger(); + for (persona, task) in tasks { + match tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(180)), + task, + ) + .await + { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => shared.violation( + "persona_failed", + format!("{} stopped: {error}", persona.name()), + vec![], + ), + Ok(Err(error)) => shared.violation( + "persona_panicked", + format!("{} panicked: {error}", persona.name()), + vec![], + ), + Err(_) => shared.violation( + "shutdown_hang", + format!("{} did not shut down within 180s", persona.name()), + vec![], + ), + } + } + + terminal_observer + .await + .map_err(|error| format!("terminal observer task: {error}"))?; + personas::ui::analyse(&shared); + if let Some(mut server) = admin_server { + if !server.is_running() { + shared.violation( + "admin_server_died", + "the awa serve process exited during the run; see awa-serve.log", + vec![], + ); + } + server.stop(); + } + super::worker_errors::reconcile(&shared, &worker_errors).await; + let outcome = Outcome { + quiesced, + elapsed_seconds: shared.started.elapsed().as_secs_f64(), + crash_excerpts: excerpts, + log_tail, + }; + if let Err(error) = report::write(&shared, &outcome) { + eprintln!("could not write the report: {error}"); + } + oracle_pool.close().await; + Ok(shared) +} + +fn spawn_persona( + persona: Persona, + shared: Arc, +) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + match persona { + Persona::Mia => personas::mia::run(shared).await, + Persona::Tom => personas::tom::run(shared).await, + Persona::Priya => personas::priya::run(shared).await, + Persona::Sam => personas::sam::run(shared).await, + Persona::Alex => personas::alex::run(shared).await, + Persona::Chaos => personas::chaos::run(shared).await, + Persona::Watchers => personas::ui::run_watchers(shared).await, + Persona::Support => personas::ui::run_support(shared).await, + Persona::UiOps => personas::ui::run_ui_operator(shared).await, + Persona::Impatient => personas::ui::run_impatient(shared).await, + Persona::Python => personas::python::run(shared).await, + Persona::Fanout => personas::fanout::run(shared).await, + Persona::ShardOps => personas::shardops::run(shared).await, + Persona::ClaimAbandon => personas::abandon::run(shared).await, + } + }) +} + +fn progress(shared: &Arc) { + let mut line = format!("t={:.0}s", shared.now_ms() as f64 / 1000.0); + for persona in Persona::ALL { + if !shared.config.has(persona) { + continue; + } + let counters = shared.stats(persona).counters(); + let enqueued: i64 = counters + .iter() + .filter(|(name, _)| name.starts_with("enqueued")) + .map(|(_, value)| *value) + .sum(); + line.push_str(&format!( + " {}[enq={} exec={}]", + persona.name(), + enqueued, + counters.get("executions").copied().unwrap_or(0) + )); + } + let violations = shared.violations.lock().expect("violations").len(); + line.push_str(&format!(" violations={violations}")); + shared.note("run", line); +} diff --git a/awa/tests/persona_soak/shard_oracle.rs b/awa/tests/persona_soak/shard_oracle.rs new file mode 100644 index 000000000..33768ac2a --- /dev/null +++ b/awa/tests/persona_soak/shard_oracle.rs @@ -0,0 +1,638 @@ +//! Checks specific to live `enqueue_shards` changes (ADR-025): every shard drains, +//! claim cursors keep moving while their shard has backlog, jobs run exactly once, +//! `ordering_key` jobs stay FIFO within a settled window, and no role ever hits a +//! permission or missing-lane error. + +use super::config::Persona; +use super::ledger::flag; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +const MAX_ROWS: usize = 25; +/// Same-key starts further behind the newest start than this are an order inversion. +const INVERSION_WINDOW: u64 = 8; +/// Seconds a shard above the configured count may keep backlog once every runtime has +/// observed the new count. +const ABOVE_COUNT_DRAIN_SECONDS: i64 = 45; + +#[derive(Debug, Clone, Default)] +pub struct QueuePlan { + pub configured: i16, + pub changed_at_ms: i64, + /// When every runtime and producer store had observed `configured`; None while the + /// fleet disagrees. + pub settled_at_ms: Option, +} + +/// What the shard-ops persona and the oracle share. +#[derive(Default)] +pub struct ShardOpsState { + pub plans: Mutex>, + /// Id of the current settled window; 0 while a change is propagating. + window: AtomicU64, + window_counter: AtomicU64, + producers_held: AtomicBool, + pub producers_inserting: AtomicI64, + seq: AtomicU64, + start_tick: AtomicU64, + /// Executions each enqueued job should have: 1, or 2 for a job that retries once. + expected: Mutex>, + /// Executions per job sequence number, and the job id that ran. + executions: Mutex>, + keyed_window: Mutex>, + keyed_starts: Mutex>, + reported: Mutex>, +} + +impl ShardOpsState { + pub fn next_seq(&self) -> u64 { + self.seq.fetch_add(1, Ordering::Relaxed) + 1 + } + + pub fn window(&self) -> u64 { + self.window.load(Ordering::SeqCst) + } + + pub fn open_window(&self) -> u64 { + let id = self.window_counter.fetch_add(1, Ordering::SeqCst) + 1; + self.window.store(id, Ordering::SeqCst); + id + } + + pub fn close_window(&self) { + self.window.store(0, Ordering::SeqCst); + } + + pub fn hold_producers(&self, held: bool) { + self.producers_held.store(held, Ordering::SeqCst); + } + + pub fn producers_held(&self) -> bool { + self.producers_held.load(Ordering::SeqCst) + } + + pub fn note_enqueued(&self, seq: u64, executions: u8) { + self.expected + .lock() + .expect("expected") + .insert(seq, executions); + } + + pub fn note_keyed(&self, key: u64, key_seq: u64, window: u64) { + self.keyed_window + .lock() + .expect("keyed window") + .insert((key, key_seq), window); + } + + /// Called by the worker at the start of every execution. + pub fn note_execution(&self, seq: u64, job_id: i64, key: Option<(u64, u64)>) { + let mut executions = self.executions.lock().expect("executions"); + let entry = executions.entry(seq).or_insert((0, job_id)); + entry.0 += 1; + drop(executions); + if let Some((key, key_seq)) = key { + let tick = self.start_tick.fetch_add(1, Ordering::SeqCst); + self.keyed_starts + .lock() + .expect("keyed starts") + .push((key, key_seq, tick)); + } + } + + pub fn configured(&self, queue: &str) -> i16 { + self.plans + .lock() + .expect("plans") + .get(queue) + .map_or(1, |plan| plan.configured.max(1)) + } + + /// Reports `summary` as a violation the first time `key` is seen and only counts + /// repeats, so a persistent failure does not bury the report. + pub fn violation_once( + &self, + shared: &Shared, + kind: &str, + key: &str, + summary: String, + rows: Vec, + ) { + let first = { + let mut reported = self.reported.lock().expect("reported"); + let count = reported.entry(format!("{kind}:{key}")).or_insert(0); + *count += 1; + *count == 1 + }; + if first { + shared.violation(kind, summary, rows); + } + shared + .stats(Persona::ShardOps) + .add(&format!("{kind}_total"), 1); + } +} + +/// One lane's cursors and backlog, as read from the store. +#[derive(Debug, Clone)] +struct ShardRow { + queue: String, + priority: i16, + shard: i16, + claim_seq: i64, + next_seq: i64, + available: i64, +} + +impl ShardRow { + fn json(&self, extra: &str) -> String { + format!( + "{{\"queue\":\"{}\",\"priority\":{},\"enqueue_shard\":{},\"claim_seq\":{},\"next_seq\":{},\"lag\":{},\"available\":{}{extra}}}", + self.queue, + self.priority, + self.shard, + self.claim_seq, + self.next_seq, + self.next_seq - self.claim_seq, + self.available + ) + } +} + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Per-shard claim cursor, enqueue cursor and the rows the cursor has not yet claimed +/// (the lag query from ADR-025, with the unclaimed rows counted exactly because aborted +/// transactions leave gaps in the sequences). +async fn shard_rows( + pool: &PgPool, + schema: &str, + queues: &[String], +) -> Result, sqlx::Error> { + let rows: Vec<(String, i16, i16, i64, i64, i64)> = sqlx::query_as(awa::audited_sql(format!( + "WITH claims AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS claim_seq + FROM {schema}.queue_claim_heads WHERE queue = ANY($1) + ), + enqueues AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS next_seq + FROM {schema}.queue_enqueue_heads WHERE queue = ANY($1) + ), + available AS ( + SELECT r.queue, r.priority, r.enqueue_shard, count(*)::bigint AS n + FROM {schema}.ready_entries r + JOIN claims c ON c.queue = r.queue AND c.priority = r.priority AND c.enqueue_shard = r.enqueue_shard + WHERE r.queue = ANY($1) AND r.lane_seq >= c.claim_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones t + WHERE t.ready_slot = r.ready_slot AND t.ready_generation = r.ready_generation + AND t.queue = r.queue AND t.priority = r.priority + AND t.enqueue_shard = r.enqueue_shard AND t.lane_seq = r.lane_seq + ) + GROUP BY r.queue, r.priority, r.enqueue_shard + ) + SELECT c.queue, c.priority, c.enqueue_shard, c.claim_seq::bigint, e.next_seq::bigint, + COALESCE(a.n, 0)::bigint + FROM claims c + JOIN enqueues e USING (queue, priority, enqueue_shard) + LEFT JOIN available a USING (queue, priority, enqueue_shard) + ORDER BY c.queue, c.priority, c.enqueue_shard" + ))) + .bind(queues) + .fetch_all(pool) + .await?; + Ok(rows + .into_iter() + .map( + |(queue, priority, shard, claim_seq, next_seq, available)| ShardRow { + queue, + priority, + shard, + claim_seq, + next_seq, + available, + }, + ) + .collect()) +} + +/// Periodic per-shard checks until `stop`: claim cursors advance while their shard has +/// backlog and idle workers, and shards above a lowered count drain. +pub async fn watch(shared: Arc, pool: PgPool, queues: Vec, stop: Stop) { + let schema = schema(&shared); + // (queue, priority, shard) -> (claim cursor, when it last moved, already reported) + let mut cursors: HashMap<(String, i16, i16), (i64, i64, bool)> = HashMap::new(); + // (queue, priority, shard) -> when backlog above the configured count was first seen + let mut above_since: HashMap<(String, i16, i16), (i64, bool)> = HashMap::new(); + while !stop.sleep(Duration::from_secs(5)).await { + let rows = match shard_rows(&pool, &schema, &queues).await { + Ok(rows) => rows, + Err(error) => { + shared.incomplete(format!("shard cursor query failed: {error}"), vec![]); + continue; + } + }; + let now = shared.now_ms(); + let mut window = shared.config.stall_window; + if shared.chaos_active() { + window += shared.config.chaos_slack; + } + let mut shards_seen: BTreeMap = BTreeMap::new(); + for row in rows { + let key = (row.queue.clone(), row.priority, row.shard); + let seen = shards_seen.entry(row.queue.clone()).or_insert(0); + *seen = (*seen).max(i64::from(row.shard) + 1); + + let entry = cursors + .entry(key.clone()) + .or_insert((row.claim_seq, now, false)); + if entry.0 != row.claim_seq { + *entry = (row.claim_seq, now, false); + } else if let Some(queue) = shared.queues.get(&row.queue) { + let capacity = queue.capacity.load(Ordering::Relaxed); + let in_flight = queue.in_flight.load(Ordering::Relaxed); + let idle_workers = capacity > 0 && in_flight < capacity; + if row.available <= 0 || !idle_workers || queue.is_held() { + entry.1 = now; + } else if !entry.2 && (now - entry.1) as u64 >= window.as_millis() as u64 { + entry.2 = true; + shared.violation( + "shard_cursor_stalled", + format!( + "queue {} priority {} shard {}: {} unclaimed job(s), claim cursor frozen at {} for {}s with {in_flight}/{capacity} workers busy", + row.queue, + row.priority, + row.shard, + row.available, + row.claim_seq, + (now - entry.1) / 1000 + ), + vec![row.json(&format!( + ",\"configured_shards\":{}", + shared.shardops.configured(&row.queue) + ))], + ); + } + } + + let plan = shared + .shardops + .plans + .lock() + .expect("plans") + .get(&row.queue) + .cloned(); + let above = plan.as_ref().is_some_and(|plan| { + i32::from(row.shard) >= i32::from(plan.configured) + && plan + .settled_at_ms + .is_some_and(|at| now - at >= ABOVE_COUNT_DRAIN_SECONDS * 1000) + }); + if above && row.available > 0 { + let since = above_since.entry(key).or_insert((now, false)); + if !since.1 && now - since.0 >= ABOVE_COUNT_DRAIN_SECONDS * 1000 { + since.1 = true; + shared.violation( + "shard_above_count_not_draining", + format!( + "queue {} shard {} is above the configured count {} and still holds {} unclaimed job(s) {}s after every runtime observed the change", + row.queue, + row.shard, + plan.as_ref().map_or(0, |plan| plan.configured), + row.available, + (now - since.0) / 1000 + ), + vec![row.json("")], + ); + } + } else { + above_since.remove(&key); + } + } + for (queue, shards) in shards_seen { + shared + .stats(Persona::ShardOps) + .set_max(&format!("highest_shard_count_{queue}"), shards); + } + } +} + +/// After the system quiesced: every shard of every shard-ops queue is empty, every +/// enqueued job ran exactly the expected number of times, and keyed jobs kept their +/// order within each settled window. +pub async fn final_check(shared: &Arc, pool: &PgPool, queues: &[String]) { + let quiesced = shared.quiesced.load(Ordering::Relaxed); + let schema = schema(shared); + match shard_rows(pool, &schema, queues).await { + Ok(rows) => { + let stuck: Vec = rows + .iter() + .filter(|row| row.available > 0) + .map(|row| { + row.json(&format!( + ",\"configured_shards\":{}", + shared.shardops.configured(&row.queue) + )) + }) + .take(MAX_ROWS) + .collect(); + shared.run.add("shardops_lanes_final", rows.len() as i64); + if !stuck.is_empty() { + let message = format!( + "{} shard(s) still hold unclaimed jobs after the load stopped", + stuck.len() + ); + if quiesced { + shared.violation("shard_not_drained", message, stuck); + } else { + shared.finding("shard_backlog_at_end", message, stuck); + } + } + } + Err(error) => shared.incomplete(format!("final shard query failed: {error}"), vec![]), + } + + let state = &shared.shardops; + let expected = state.expected.lock().expect("expected").clone(); + let executions = state.executions.lock().expect("executions").clone(); + let mut missing = Vec::new(); + let mut repeated = Vec::new(); + let mut rescued = Vec::new(); + let mut short = Vec::new(); + for (seq, want) in &expected { + let (ran, job_id) = executions.get(seq).copied().unwrap_or((0, 0)); + let record = shared.ledger.get(job_id); + let row = format!( + "{{\"seq\":{seq},\"job_id\":{job_id},\"expected_executions\":{want},\"executions\":{ran},\"max_attempt\":{},\"ledger_flags\":{}}}", + record.as_ref().map_or(0, |record| record.max_attempt), + record.as_ref().map_or(0, |record| record.flags) + ); + if ran == 0 { + missing.push(row); + } else if ran > u16::from(*want) { + // A rescue re-runs a job whose attempt lost its lease, which is at-least-once + // delivery working as designed; any other repeat is a duplicate. + let was_rescued = record + .as_ref() + .is_some_and(|record| record.has(flag::RESCUED)); + if was_rescued { + rescued.push(row); + } else { + repeated.push(row); + } + } else if ran < u16::from(*want) { + short.push(row); + } + } + shared + .stats(Persona::ShardOps) + .add("tracked_jobs", expected.len() as i64); + let report = |kind: &str, what: &str, rows: Vec, fatal: bool| { + if rows.is_empty() { + return; + } + let summary = format!("{} job(s) {what}", rows.len()); + let rows = rows.into_iter().take(MAX_ROWS).collect(); + if fatal { + shared.violation(kind, summary, rows); + } else { + shared.finding(kind, summary, rows); + } + }; + report( + "lost_job", + "were enqueued across a shard change and never ran", + missing, + quiesced, + ); + report( + "rescued_reexecution", + "ran again after a stale-heartbeat rescue", + rescued, + false, + ); + report( + "duplicate_execution", + "ran more often than their handler asked for", + repeated, + true, + ); + report( + "retry_not_executed", + "ran fewer times than their handler asked for", + short, + quiesced, + ); + + let windows = state.keyed_window.lock().expect("keyed window").clone(); + let mut starts = state.keyed_starts.lock().expect("keyed starts").clone(); + starts.sort_by_key(|(_, _, tick)| *tick); + let mut newest: HashMap<(u64, u64), u64> = HashMap::new(); + let mut inversions = Vec::new(); + let mut checked = 0u64; + let mut worst = 0u64; + let mut seen_per_key: HashSet<(u64, u64)> = HashSet::new(); + for (key, key_seq, _) in starts { + let window = windows.get(&(key, key_seq)).copied().unwrap_or(0); + if window == 0 { + continue; + } + if !seen_per_key.insert((key, key_seq)) { + continue; + } + checked += 1; + let top = newest.entry((key, window)).or_insert(0); + if key_seq + INVERSION_WINDOW < *top { + worst = worst.max(*top - key_seq); + if inversions.len() < MAX_ROWS { + inversions.push(format!( + "{{\"ordering_key\":{key},\"window\":{window},\"started_seq\":{key_seq},\"newest_started_seq\":{top}}}" + )); + } + } else { + *top = (*top).max(key_seq); + } + } + shared + .stats(Persona::ShardOps) + .add("keyed_starts_checked", checked as i64); + shared + .stats(Persona::ShardOps) + .set_max("keyed_worst_inversion", worst as i64); + if !inversions.is_empty() { + shared.violation( + "ordering_key_inversion", + format!( + "keyed jobs started out of order within a settled shard-count window (worst distance {worst})" + ), + inversions, + ); + } +} + +struct LoggedError { + role: String, + application: String, + message: String, + statement: String, +} + +fn is_lane_error(message: &str) -> bool { + message.contains("permission denied") + || message.contains("must be owner") + || message.contains("has not been provisioned") + || message.contains("provision") + || is_lane_race(message) +} + +/// Two sessions creating the same lane sequence at once: `CREATE SEQUENCE IF NOT EXISTS` +/// loses the race with a unique violation on the catalog. +fn is_lane_race(message: &str) -> bool { + message.contains("pg_class_relname_nsp_index") +} + +/// Permission-denied and missing-lane errors in the server log, one violation per role +/// and error message with the lane, the sequence's owner and the statement. +pub async fn scan_server_log(shared: &Arc, pool: &PgPool, log: &str) { + let mut errors: Vec = Vec::new(); + let lines: Vec<&str> = log.lines().collect(); + let is_entry = |line: &str| { + let bytes = line.as_bytes(); + bytes.len() > 20 && bytes[4] == b'-' && bytes[7] == b'-' && bytes[10] == b' ' + }; + let mut index = 0; + while index < lines.len() { + let line = lines[index]; + index += 1; + let Some((head, message)) = line.split_once(" ERROR: ") else { + continue; + }; + if !is_lane_error(message) { + continue; + } + let identity = head + .rsplit_once("] ") + .map_or("?", |(_, identity)| identity.trim()); + let (role, application) = identity.split_once('@').unwrap_or((identity, "")); + let (role, application) = (role.to_string(), application.to_string()); + let mut statement = String::new(); + while index < lines.len() && !is_entry(lines[index]) { + if let Some(text) = lines[index].trim_start().strip_prefix("STATEMENT:") { + statement = text.trim().chars().take(400).collect(); + } else if !statement.is_empty() { + statement.push(' '); + statement.push_str(lines[index].trim()); + } + index += 1; + } + errors.push(LoggedError { + role, + application, + message: message.trim().to_string(), + statement, + }); + } + if errors.is_empty() { + return; + } + let schema = schema(shared); + let mut groups: BTreeMap<(String, String), (u64, String, BTreeSet)> = BTreeMap::new(); + for error in &errors { + let entry = groups + .entry((error.role.clone(), error.message.clone())) + .or_insert((0, error.statement.clone(), BTreeSet::new())); + entry.0 += 1; + entry.2.insert(error.application.clone()); + } + shared + .run + .add("shard_lane_errors_logged", errors.len() as i64); + for ((role, message), (count, statement, applications)) in groups.into_iter().take(MAX_ROWS) { + let (lane, owner) = describe_sequence(pool, &schema, &message).await; + let applications = applications.into_iter().collect::>().join(","); + let kind = if is_lane_race(&message) { + "lane_creation_race" + } else { + "permission_denied" + }; + shared.violation( + kind, + format!("role {role}: {message} ({count}x; lane {lane}; sequence owner {owner})"), + vec![format!( + "{{\"role\":\"{role}\",\"applications\":{},\"error\":{},\"lane\":{},\"sequence_owner\":{},\"occurrences\":{count},\"statement\":{}}}", + json_string(&applications), + json_string(&message), + json_string(&lane), + json_string(&owner), + json_string(&statement) + )], + ); + } +} + +fn json_string(text: &str) -> String { + serde_json::Value::String(text.to_string()).to_string() +} + +/// The lane and owning role of the sequence an error message names, when it names one. +async fn describe_sequence(pool: &PgPool, schema: &str, message: &str) -> (String, String) { + let Some(name) = message + .split("sequence ") + .nth(1) + .map(|rest| { + rest.split_whitespace() + .next() + .unwrap_or("") + .trim_matches('"') + }) + .map(|name| { + name.rsplit('.') + .next() + .unwrap_or(name) + .trim_matches('"') + .to_string() + }) + .filter(|name| !name.is_empty()) + else { + return ("unknown".into(), "unknown".into()); + }; + let owner: Option = sqlx::query_scalar( + "SELECT pg_get_userbyid(c.relowner) FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname = $2", + ) + .bind(schema) + .bind(&name) + .fetch_optional(pool) + .await + .ok() + .flatten(); + let lane: Option = sqlx::query_scalar(awa::audited_sql(format!( + "SELECT queue || ' priority ' || priority || ' shard ' || enqueue_shard + FROM {schema}.queue_enqueue_heads WHERE seq_name = $1 + UNION ALL + SELECT queue || ' priority ' || priority || ' shard ' || enqueue_shard + FROM {schema}.queue_claim_heads WHERE seq_name = $1 + LIMIT 1" + ))) + .bind(&name) + .fetch_optional(pool) + .await + .ok() + .flatten(); + ( + lane.unwrap_or_else(|| format!("sequence {name}")), + owner.unwrap_or_else(|| "unknown".into()), + ) +} diff --git a/awa/tests/persona_soak/shared.rs b/awa/tests/persona_soak/shared.rs new file mode 100644 index 000000000..2bd038081 --- /dev/null +++ b/awa/tests/persona_soak/shared.rs @@ -0,0 +1,468 @@ +//! State shared by every persona, the oracle and the report. + +use super::config::{Persona, SoakConfig}; +use super::db::{Budget, ServerInfo}; +use super::fleet::Fleet; +use super::ledger::{expect, flag, JobRecord, Ledger}; +use super::registry::{InFlightGuard, QueueRegistry}; +use super::rng::percentile; +use super::stop::Stop; +use awa::{JobRow, UntypedJobEvent}; +use serde::Serialize; +use std::collections::{BTreeMap, HashMap}; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +#[derive(Debug, Clone, Serialize)] +pub struct TimelineEvent { + pub at_ms: i64, + pub source: String, + pub text: String, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Violation { + pub kind: String, + pub at_ms: i64, + pub during_chaos: bool, + pub summary: String, + /// The violating rows, as JSON text. + pub rows: Vec, +} + +#[derive(Default)] +pub struct Stats { + counters: Mutex>, + samples: Mutex>>, +} + +impl Stats { + pub fn add(&self, name: &str, amount: i64) { + *self + .counters + .lock() + .expect("stats") + .entry(name.to_string()) + .or_insert(0) += amount; + } + + pub fn set_max(&self, name: &str, value: i64) { + let mut counters = self.counters.lock().expect("stats"); + let slot = counters.entry(name.to_string()).or_insert(i64::MIN); + *slot = (*slot).max(value); + } + + pub fn get(&self, name: &str) -> i64 { + self.counters + .lock() + .expect("stats") + .get(name) + .copied() + .unwrap_or(0) + } + + /// Records a latency-style sample; kept bounded so long runs stay small. + pub fn observe(&self, name: &str, value: i64) { + let mut samples = self.samples.lock().expect("stats"); + let list = samples.entry(name.to_string()).or_default(); + if list.len() < 400_000 { + list.push(value); + } + } + + pub fn counters(&self) -> BTreeMap { + self.counters.lock().expect("stats").clone() + } + + pub fn summaries(&self) -> BTreeMap { + self.samples + .lock() + .expect("stats") + .iter() + .map(|(name, values)| { + let mut sorted = values.clone(); + sorted.sort_unstable(); + ( + name.clone(), + SampleSummary { + count: sorted.len() as i64, + p50: percentile(&sorted, 0.5), + p99: percentile(&sorted, 0.99), + max: sorted.last().copied().unwrap_or(0), + }, + ) + }) + .collect() + } +} + +#[derive(Debug, Clone, Serialize)] +pub struct SampleSummary { + pub count: i64, + pub p50: i64, + pub p99: i64, + pub max: i64, +} + +pub struct Shared { + pub config: SoakConfig, + pub info: ServerInfo, + pub started: Instant, + pub ledger: Ledger, + pub queues: QueueRegistry, + pub budget: Budget, + pub load_stop: Stop, + pub drain_stop: Stop, + pub timeline: Mutex>, + pub violations: Mutex>, + /// Required oracle evidence was unavailable; the run cannot establish a pass. + pub oracle_incomplete: AtomicBool, + /// Problems that are not invariant violations (slow admin queries); reported, not fatal. + pub findings: Mutex>, + pub stats: HashMap, + pub ui: super::personas::ui::UiMetrics, + /// Handlers currently running in the Python process, from its own reports. + pub python_active: AtomicI64, + /// Set once the post-load drain finished with no live work left. + pub quiesced: std::sync::atomic::AtomicBool, + /// The running Python process, so the chaos persona can kill it. + pub python: Mutex>>, + /// Run-wide measurements (resource use, oracle probes). + pub run: Stats, + pub fleets: Mutex>>, + /// State of the shard-ops persona and the checks that go with it. + pub shardops: super::shard_oracle::ShardOpsState, + /// State of the claim-and-abandon persona and the checks that go with it. + pub abandon: super::abandon_oracle::AbandonState, + chaos_until_ms: AtomicI64, +} + +impl Shared { + pub fn new(config: SoakConfig, info: ServerInfo, budget: Budget) -> Arc { + let stats = Persona::ALL + .into_iter() + .map(|persona| (persona, Stats::default())) + .collect(); + Arc::new(Self { + config, + info, + started: Instant::now(), + ledger: Ledger::new(), + queues: QueueRegistry::default(), + budget, + load_stop: Stop::new(), + drain_stop: Stop::new(), + timeline: Mutex::new(Vec::new()), + violations: Mutex::new(Vec::new()), + oracle_incomplete: AtomicBool::new(false), + findings: Mutex::new(Vec::new()), + stats, + run: Stats::default(), + ui: Default::default(), + python_active: AtomicI64::new(0), + quiesced: std::sync::atomic::AtomicBool::new(false), + python: Mutex::new(None), + fleets: Mutex::new(HashMap::new()), + shardops: Default::default(), + abandon: Default::default(), + chaos_until_ms: AtomicI64::new(0), + }) + } + + pub fn started_load(&self, _at: Instant) {} + + pub fn now_ms(&self) -> i64 { + self.started.elapsed().as_millis() as i64 + } + + pub fn stats(&self, persona: Persona) -> &Stats { + &self.stats[&persona] + } + + pub fn note(&self, source: &str, text: impl Into) { + let text = text.into(); + let at_ms = self.now_ms(); + eprintln!("[{:>7.1}s] {source:<6} {text}", at_ms as f64 / 1000.0); + self.timeline.lock().expect("timeline").push(TimelineEvent { + at_ms, + source: source.to_string(), + text, + }); + } + + /// Marks a disruptive window (kill, failover, connection termination) so that the + /// time-based checks allow the system its recovery time. + pub fn chaos_window(&self, duration: std::time::Duration) { + let until = self.now_ms() + duration.as_millis() as i64; + self.chaos_until_ms.fetch_max(until, Ordering::Relaxed); + } + + pub fn chaos_active(&self) -> bool { + self.now_ms() < self.chaos_until_ms.load(Ordering::Relaxed) + } + + pub fn violation(&self, kind: &str, summary: impl Into, rows: Vec) { + let violation = Violation { + kind: kind.to_string(), + at_ms: self.now_ms(), + during_chaos: self.chaos_active(), + summary: summary.into(), + rows, + }; + eprintln!( + "[{:>7.1}s] VIOLATION {}: {}", + violation.at_ms as f64 / 1000.0, + violation.kind, + violation.summary + ); + self.violations.lock().expect("violations").push(violation); + } + + pub fn incomplete(&self, summary: impl Into, rows: Vec) { + self.run.add("oracle_incomplete_checks", 1); + if !self.oracle_incomplete.swap(true, Ordering::Relaxed) { + self.violation("oracle_incomplete", summary, rows); + } + } + + /// Records a non-fatal finding: something worth a bug report that is not one of the + /// invariants the soak fails on. + pub fn finding(&self, kind: &str, summary: impl Into, rows: Vec) { + let finding = Violation { + kind: kind.to_string(), + at_ms: self.now_ms(), + during_chaos: self.chaos_active(), + summary: summary.into(), + rows, + }; + eprintln!( + "[{:>7.1}s] FINDING {}: {}", + finding.at_ms as f64 / 1000.0, + finding.kind, + finding.summary + ); + self.findings.lock().expect("findings").push(finding); + } + + /// Records a successful operator cancel: the job may legitimately end cancelled, a + /// handler must not start after it, and it must not overwrite a completion. + pub fn record_operator_cancel( + &self, + persona: Persona, + job_id: i64, + completed_before: bool, + call_started_ms: u32, + ) { + let now = self.now_ms() as u32; + self.stats(persona).add("cancelled", 1); + self.ledger.with( + job_id, + || JobRecord::new(persona, u32::MAX, now, expect::ANY_TERMINAL, 0), + |record| { + record.flags |= flag::OPERATOR_REMOVED | flag::OPERATOR_CANCEL_OK; + record.operator_cancel_ms = now; + }, + ); + if completed_before { + self.violation( + "cancel_overwrote_completed", + format!("operator cancel of job {job_id} succeeded after it was reported completed (call began at {call_started_ms} ms)"), + vec![format!("{{\"job_id\":{job_id}}}")], + ); + } + } + + pub fn fleet(&self, name: &str) -> Option> { + self.fleets.lock().expect("fleets").get(name).cloned() + } + + pub fn fleets_of(&self, persona: Persona) -> Vec> { + self.fleets + .lock() + .expect("fleets") + .values() + .filter(|fleet| fleet.persona == persona) + .cloned() + .collect() + } + + /// Records the start of a handler. Hold the returned guard for the duration of + /// `perform`; dropping it, including when the runtime is killed, ends the execution. + pub fn begin_execution( + self: &Arc, + persona: Persona, + job: &JobRow, + no_events: bool, + ) -> ExecGuard { + let queue = self.queues.get(&job.queue); + let queue_index = queue.as_ref().map_or(u32::MAX, |queue| queue.index); + let now = self.now_ms() as u32; + let default_flags = if no_events { flag::NO_EVENTS } else { 0 }; + let (ran_after_cancel, overlapped) = self.ledger.with( + job.id, + || { + JobRecord::new( + persona, + queue_index, + now, + expect::ANY_TERMINAL, + default_flags, + ) + }, + |record| { + let overlapped = record.in_flight > 0; + record.overlapped |= overlapped; + record.in_flight = record.in_flight.saturating_add(1); + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(job.attempt); + record.latest_run_lease = record.latest_run_lease.max(job.run_lease); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + ( + record.operator_cancel_ms != 0 && now > record.operator_cancel_ms, + overlapped, + ) + }, + ); + if ran_after_cancel { + self.violation( + "ran_after_cancel", + format!("job {} started after its operator cancel committed", job.id), + vec![format!( + "{{\"job_id\":{},\"queue\":\"{}\",\"attempt\":{}}}", + job.id, job.queue, job.attempt + )], + ); + } + if overlapped { + self.stats(persona).add("overlapping_executions", 1); + } + if let Some(name) = job + .metadata + .get("cron_name") + .and_then(|value| value.as_str()) + { + let fire = job + .metadata + .get("cron_fire_time") + .and_then(|value| value.as_str()) + .unwrap_or_default(); + // Manually triggered jobs carry the schedule's name but no fire time. + if !fire.is_empty() + && !self + .ledger + .note_cron_fire(name.to_string(), fire.to_string(), job.id) + { + self.violation( + "cron_double_fire", + format!( + "cron {name} fire {fire} produced more than one job (latest {})", + job.id + ), + vec![], + ); + } + } + if let Some(queue) = &queue { + queue.started.fetch_add(1, Ordering::Relaxed); + } + self.stats(persona).add("executions", 1); + ExecGuard { + _flight: InFlightGuard::new(queue), + shared: self.clone(), + job_id: job.id, + } + } + + /// Records a lifecycle event emitted by a runtime after the transition committed. + pub fn on_event(&self, event: &UntypedJobEvent) { + let job = event.job(); + let now = self.now_ms() as u32; + let queue_index = self + .queues + .get(&job.queue) + .map_or(u32::MAX, |queue| queue.index); + let persona = self + .queues + .get(&job.queue) + .map_or(Persona::Chaos, |queue| queue.persona); + let mut double_complete = false; + self.ledger.with( + job.id, + || JobRecord::new(persona, queue_index, now, expect::ANY_TERMINAL, 0), + |record| match event { + UntypedJobEvent::Completed { .. } => { + record.note_terminal(job.run_lease); + record.completed_events = record.completed_events.saturating_add(1); + record.completed_event_ms = now; + record.flags |= flag::COMPLETED_EVENT; + double_complete = record.completed_events > 1 && !record.has(flag::DLQ_RETRIED); + } + UntypedJobEvent::Exhausted { .. } => { + record.note_terminal(job.run_lease); + record.exhausted_events = record.exhausted_events.saturating_add(1); + record.flags |= flag::EXHAUSTED; + } + UntypedJobEvent::Cancelled { .. } => { + record.note_terminal(job.run_lease); + record.cancel_events = record.cancel_events.saturating_add(1); + record.flags |= flag::HANDLER_CANCELLED; + } + UntypedJobEvent::Rescued { .. } => record.flags |= flag::RESCUED, + _ => {} + }, + ); + match event { + UntypedJobEvent::Rescued { reason, .. } => { + if persona == Persona::ShardOps { + self.note( + "events", + format!( + "job {} on {} rescued ({}) at attempt {}", + job.id, + job.queue, + reason.as_str(), + job.attempt + ), + ); + } + self.stats(persona) + .add(&format!("rescued_{}", reason.as_str()), 1); + } + UntypedJobEvent::Exhausted { .. } => self.stats(persona).add("exhausted", 1), + UntypedJobEvent::Completed { .. } => self.stats(persona).add("completed_events", 1), + UntypedJobEvent::Cancelled { .. } => self.stats(persona).add("handler_cancelled", 1), + UntypedJobEvent::Retried { .. } => self.stats(persona).add("retried", 1), + _ => {} + } + if double_complete { + self.violation( + "duplicate_terminal_event", + format!("job {} reported completed more than once", job.id), + vec![format!( + "{{\"job_id\":{},\"queue\":\"{}\"}}", + job.id, job.queue + )], + ); + } + } +} + +pub struct ExecGuard { + _flight: InFlightGuard, + shared: Arc, + job_id: i64, +} + +impl Drop for ExecGuard { + fn drop(&mut self) { + self.shared.ledger.with( + self.job_id, + || JobRecord::new(Persona::Chaos, u32::MAX, 0, expect::ANY_TERMINAL, 0), + |record| record.in_flight = record.in_flight.saturating_sub(1), + ); + } +} diff --git a/awa/tests/persona_soak/stop.rs b/awa/tests/persona_soak/stop.rs new file mode 100644 index 000000000..6264c06e2 --- /dev/null +++ b/awa/tests/persona_soak/stop.rs @@ -0,0 +1,50 @@ +//! A cheap cooperative stop signal shared by every task in the run. + +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::watch; + +#[derive(Clone)] +pub struct Stop { + sender: Arc>, +} + +impl Stop { + pub fn new() -> Self { + Self { + sender: Arc::new(watch::channel(false).0), + } + } + + pub fn trigger(&self) { + self.sender.send_replace(true); + } + + pub fn is_stopped(&self) -> bool { + *self.sender.borrow() + } + + /// Sleeps for `duration` or until stopped. Returns true when stopped. + pub async fn sleep(&self, duration: Duration) -> bool { + if self.is_stopped() { + return true; + } + let mut receiver = self.sender.subscribe(); + tokio::select! { + _ = tokio::time::sleep(duration) => false, + _ = receiver.wait_for(|stopped| *stopped) => true, + } + } + + /// Completes once stopped. + pub async fn wait(&self) { + let mut receiver = self.sender.subscribe(); + let _ = receiver.wait_for(|stopped| *stopped).await; + } +} + +impl Default for Stop { + fn default() -> Self { + Self::new() + } +} diff --git a/awa/tests/persona_soak/terminal_proof.rs b/awa/tests/persona_soak/terminal_proof.rs new file mode 100644 index 000000000..2ecec2b20 --- /dev/null +++ b/awa/tests/persona_soak/terminal_proof.rs @@ -0,0 +1,122 @@ +//! Durable evidence for returned eventful jobs and the NO_EVENTS sample. +//! Handler returns are not commits; lifecycle hooks are best-effort notifications. + +use super::config::Persona; +use super::ledger::{expect, flag, JobRecord, TERMINAL_PROOF_STRIDE}; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::sync::Arc; +use std::time::Duration; + +const MAX_ROWS: usize = 100_000; + +pub async fn run(shared: Arc, pool: PgPool, stop: Stop) { + let mut tick = tokio::time::interval(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = stop.wait() => return, + _ = tick.tick() => observe(&shared, &pool).await, + } + } +} + +/// One statement gives one snapshot. Explicit eventful attempts have bounded +/// input/result counts and a two-second statement timeout. NO_EVENTS keeps its +/// immutable sample; neither a returned handler nor a missed query proves a commit. +pub async fn observe(shared: &Shared, pool: &PgPool) { + let schema = shared.info.queue_storage_schema.as_deref().unwrap_or("awa"); + let mut pending = Vec::new(); + shared.ledger.for_each(|id, record| { + if record.has(flag::RETURNED_COMPLETED) + && !record.has(flag::NO_EVENTS) + && !record.has_durable_terminal() + && !record.has(flag::OPERATOR_REMOVED) + && !record.has(flag::DLQ_PURGED) + && pending.len() <= MAX_ROWS + { + pending.push((id, record.latest_run_lease)); + } + }); + if pending.len() > MAX_ROWS { + shared.incomplete( + "durable terminal observer exceeded its pending job cap", + vec![], + ); + return; + } + let (pending_ids, pending_leases): (Vec<_>, Vec<_>) = pending.into_iter().unzip(); + let result = async { + let mut tx = pool.begin().await?; + sqlx::query("SET LOCAL statement_timeout = '2s'") + .execute(tx.as_mut()) + .await?; + let rows: Vec<(i64, i64, String)> = sqlx::query_as(awa::audited_sql(format!( + "WITH pending AS ( + SELECT * FROM unnest($3::bigint[], $4::bigint[]) AS p(job_id, run_lease) + ) + SELECT done.job_id, done.run_lease, done.state::text + FROM {schema}.done_entries AS done + JOIN pending USING (job_id, run_lease) + UNION ALL + SELECT dlq.job_id, dlq.run_lease, 'failed' + FROM {schema}.dlq_entries AS dlq + JOIN pending USING (job_id, run_lease) + UNION ALL + SELECT item.job_id, item.run_lease, 'completed' + FROM {schema}.receipt_completion_batches AS batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + JOIN pending USING (job_id, run_lease) + WHERE batch.job_ids && $3 + UNION ALL + SELECT job_id, run_lease, state::text FROM {schema}.done_entries + WHERE job_id % $1 = 0 AND finalized_at >= now() - interval '30 seconds' + UNION ALL + SELECT job_id, run_lease, 'failed' FROM {schema}.dlq_entries + WHERE job_id % $1 = 0 AND dlq_at >= now() - interval '30 seconds' + UNION ALL + SELECT item.job_id, item.run_lease, 'completed' + FROM {schema}.receipt_completion_batches AS batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + WHERE item.job_id % $1 = 0 AND batch.finalized_at >= now() - interval '30 seconds' + LIMIT $2" + ))) + .bind(TERMINAL_PROOF_STRIDE) + .bind(MAX_ROWS as i64 + 1) + .bind(pending_ids) + .bind(pending_leases) + .fetch_all(tx.as_mut()) + .await?; + tx.commit().await?; + Ok::<_, sqlx::Error>(rows) + } + .await; + let rows = match result { + Ok(rows) if rows.len() <= MAX_ROWS => rows, + Ok(_) => { + shared.incomplete("durable terminal observer exceeded its row cap", vec![]); + return; + } + Err(error) => { + shared.incomplete( + format!("durable terminal observer query failed: {error}"), + vec![], + ); + return; + } + }; + shared + .run + .add("terminal_proof_observer_rows", rows.len() as i64); + let now = shared.now_ms() as u32; + for (id, run_lease, state) in rows { + shared.ledger.with( + id, + || JobRecord::new(Persona::Chaos, u32::MAX, now, expect::ANY_TERMINAL, 0), + |record| { + record.note_observed_terminal(run_lease, &state); + }, + ); + } +} diff --git a/awa/tests/persona_soak/worker_errors.rs b/awa/tests/persona_soak/worker_errors.rs new file mode 100644 index 000000000..89dfb5070 --- /dev/null +++ b/awa/tests/persona_soak/worker_errors.rs @@ -0,0 +1,241 @@ +//! A structured gate for a known worker validation failure, independent of log text. +//! +//! This subscriber covers Rust runtimes in this process. The separately spawned +//! Rust claiming child forwards only a fixed marker over its existing output pipe. +//! Python-extension internals have no tracing bridge and are not covered by it. + +use super::shared::Shared; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; +use tracing::field::{Field, Visit}; +use tracing::{Event, Subscriber}; +use tracing_subscriber::filter::LevelFilter; +use tracing_subscriber::layer::{Context, SubscriberExt}; +use tracing_subscriber::{Layer, Registry}; + +const MISSING_READY: &str = "ready row missing for deleted lease"; +const AMBIGUOUS_READY: &str = "ready body ambiguous for deleted lease"; +const CHILD_MARKER: &str = "AWA_SOAK_WORKER_ERROR ready_body_unavailable"; +static EVIDENCE: OnceLock, String>> = OnceLock::new(); + +#[derive(Default)] +pub struct WorkerErrors { + invalid_ready_body: AtomicU64, + child_readers: AtomicUsize, +} + +#[derive(Default)] +struct ErrorFields { + invalid_ready_body: bool, +} + +impl Visit for ErrorFields { + fn record_str(&mut self, field: &Field, value: &str) { + if matches!(field.name(), "error" | "message") { + self.invalid_ready_body |= + value.contains(MISSING_READY) || value.contains(AMBIGUOUS_READY); + } + } + + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + if matches!(field.name(), "error" | "message") { + self.invalid_ready_body |= { + let formatted = format!("{value:?}"); + formatted.contains(MISSING_READY) || formatted.contains(AMBIGUOUS_READY) + }; + } + } +} + +#[derive(Clone)] +struct ErrorLayer { + evidence: Arc, + child: bool, +} + +impl Layer for ErrorLayer { + fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) { + if *event.metadata().level() > tracing::Level::WARN { + return; + } + let mut fields = ErrorFields::default(); + event.record(&mut fields); + if fields.invalid_ready_body { + self.evidence + .invalid_ready_body + .fetch_add(1, Ordering::Relaxed); + if self.child { + // One short synchronous line only for this invariant failure. + // The parent's reader is tracked through EOF before reconciliation. + eprintln!("{CHILD_MARKER}"); + } + } + } +} + +/// Install once, before spawning any worker. An existing unrelated subscriber +/// makes setup fail rather than silently dropping the error gate. +pub fn install(child: bool) -> Result, String> { + EVIDENCE + .get_or_init(|| { + let evidence = Arc::new(WorkerErrors::default()); + let subscriber = Registry::default() + .with( + ErrorLayer { + evidence: evidence.clone(), + child, + } + .with_filter(LevelFilter::WARN), + ) + .with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(std::io::stderr) + .with_filter(LevelFilter::WARN), + ); + tracing::subscriber::set_global_default(subscriber) + .map_err(|error| format!("worker error tracing unavailable: {error}"))?; + Ok(evidence) + }) + .clone() +} + +pub struct ChildOutputGuard(Arc); + +impl Drop for ChildOutputGuard { + fn drop(&mut self) { + self.0.child_readers.fetch_sub(1, Ordering::Release); + } +} + +/// Register before spawning the existing output reader, so a fast child exit +/// cannot race the final pending-reader check. +pub fn child_output_guard() -> ChildOutputGuard { + let evidence = EVIDENCE + .get() + .expect("worker error tracing installed before child spawn") + .as_ref() + .expect("worker error tracing installation succeeded") + .clone(); + evidence.child_readers.fetch_add(1, Ordering::Relaxed); + ChildOutputGuard(evidence) +} + +pub fn absorb_child_marker(line: &str) -> bool { + if line != CHILD_MARKER { + return false; + } + EVIDENCE + .get() + .expect("worker error tracing installed before child spawn") + .as_ref() + .expect("worker error tracing installation succeeded") + .invalid_ready_body + .fetch_add(1, Ordering::Relaxed); + true +} + +/// Call after actor shutdown and the admin child stop, before final report. +pub async fn reconcile(shared: &Shared, evidence: &WorkerErrors) { + let drained = tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(5)), + async { + while evidence.child_readers.load(Ordering::Acquire) != 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }, + ) + .await + .is_ok(); + if !drained { + shared.incomplete("claiming-child error output did not drain", vec![]); + } + let missing = evidence.invalid_ready_body.load(Ordering::Relaxed); + if missing != 0 { + shared + .run + .add("worker_invalid_ready_body_events", missing as i64); + shared.violation( + "worker_ready_body_invalid", + format!("{missing} worker event(s) reported missing or ambiguous ready bodies; see WARN/ERROR output"), + vec![], + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct DisplayOnly; + + impl std::fmt::Display for DisplayOnly { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "queue storage {MISSING_READY} job 17 run_lease 2" + ) + } + } + + impl std::fmt::Debug for DisplayOnly { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("debug form omits the error") + } + } + + fn count_events(emit: impl FnOnce()) -> u64 { + let evidence = Arc::new(WorkerErrors::default()); + let subscriber = Registry::default().with(ErrorLayer { + evidence: evidence.clone(), + child: false, + }); + tracing::subscriber::with_default(subscriber, emit); + evidence.invalid_ready_body.load(Ordering::Relaxed) + } + + #[test] + fn display_error_field_is_counted_once_even_when_message_also_matches() { + assert_eq!( + count_events(|| { + tracing::warn!(error = %DisplayOnly, "{MISSING_READY}"); + }), + 1 + ); + } + + #[test] + fn string_error_and_message_only_forms_are_counted() { + assert_eq!( + count_events(|| { + tracing::error!(error = MISSING_READY, "rescue failed"); + tracing::warn!("queue storage {MISSING_READY} job 18 run_lease 3"); + }), + 2 + ); + } + + #[test] + fn ambiguous_ready_body_failure_is_also_gated() { + assert_eq!( + count_events(|| { + tracing::error!(error = AMBIGUOUS_READY, "hydrate failed"); + }), + 1 + ); + } + + #[test] + fn unrelated_chaos_warnings_and_payload_text_are_not_gated() { + assert_eq!( + count_events(|| { + tracing::warn!(error = "canceling statement due to user request", "chaos"); + tracing::error!(error = "lock timeout", "expected contention"); + tracing::warn!(payload = MISSING_READY, "a job argument"); + tracing::info!(error = MISSING_READY, "not a worker failure event"); + }), + 0 + ); + } +} diff --git a/awa/tests/persona_soak_test.rs b/awa/tests/persona_soak_test.rs new file mode 100644 index 000000000..1c33ba709 --- /dev/null +++ b/awa/tests/persona_soak_test.rs @@ -0,0 +1,43 @@ +//! Persona soak: several application-shaped workloads (Mia, Tom, Priya, Sam, Alex and +//! a chaos agent) run concurrently against one disposable database on the queue-storage +//! engine while an invariant oracle checks the system. Ignored by default. +//! +//! ```text +//! AWA_SOAK_DATABASE_URL=postgres://postgres:pw@localhost:5432/awa_soak_run \ +//! AWA_SOAK_SECONDS=900 \ +//! cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture +//! ``` +//! +//! See `docs/test-plan.md` for the environment variables and `scripts/persona-soak.sh` +//! for the engine matrix. + +#![allow(dead_code)] +mod ci_timing; +mod persona_soak; + +use persona_soak::config::SoakConfig; + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "drops and recreates an awa_soak* database; see the module docs"] +async fn persona_soak() { + let config = SoakConfig::from_env().unwrap_or_else(|error| panic!("{error}")); + let shared = persona_soak::run::run_soak(config) + .await + .unwrap_or_else(|error| panic!("soak setup failed: {error}")); + let violations = shared.violations.lock().expect("violations").len(); + assert_eq!( + violations, + 0, + "{violations} invariant violation(s); see {}", + shared.config.report_dir.display() + ); +} + +/// The claiming process of the claim-and-abandon persona. The harness runs this test binary +/// again with only this test selected and kills the process with SIGKILL; started any +/// other way it returns at once. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "run by the claim-and-abandon persona as a separate process"] +async fn claim_abandon_child() { + persona_soak::personas::abandon::run_child().await; +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index c693f4e30..36f622613 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -6,12 +6,13 @@ mod ci_timing; #[path = "../../awa-model/src/storage_footprint_sql.rs"] mod storage_footprint_sql; +use ci_timing::scaled_timeout; use awa::audited_sql; use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, - BatchOperationSpec, PruneOutcome, QueueStorage, QueueStorageConfig, RingRotationCadence, - RotateOutcome, SkipReason, SubmitBatchOperation, + BatchOperationSpec, PromotionLane, PruneOutcome, QueueStorage, QueueStorageConfig, + RingRotationCadence, RotateOutcome, SkipReason, SubmitBatchOperation, }; use awa::{ Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, JobRow, JobState, QueueConfig, @@ -2048,6 +2049,231 @@ impl Worker for MultiClientTrackingWorker { } } +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct GracefulRestartJob { + id: i64, +} + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct GracefulRestartFollowUp { + id: i64, +} + +/// The first execution reports its attempt and returns `Completed` once +/// shutdown signals cancellation; later executions complete at once. +#[derive(Clone, Default)] +struct GracefulRestartWorker { + executions: Arc, + first_attempt: Arc>>, + entered: Arc, +} + +#[async_trait::async_trait] +impl Worker for GracefulRestartWorker { + fn kind(&self) -> &'static str { + "graceful_restart_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + if self.executions.fetch_add(1, Ordering::SeqCst) > 0 { + return Ok(JobResult::Completed); + } + *self.first_attempt.lock().expect("first attempt") = Some((ctx.job.id, ctx.job.run_lease)); + self.entered.notify_one(); + while !ctx.is_cancelled() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + Ok(JobResult::Completed) + } +} + +/// A runtime on its own tokio runtime and pool, so stopping it after a +/// graceful shutdown drops whatever the client left running, as a process +/// exit does. +struct ProcessHost { + runtime: tokio::runtime::Runtime, + client: Arc, +} + +impl ProcessHost { + async fn start( + options: sqlx::postgres::PgConnectOptions, + build: impl FnOnce(sqlx::PgPool) -> Client + Send + 'static, + ) -> Self { + tokio::task::spawn_blocking(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("build process runtime"); + let client = runtime.block_on(async move { + let pool = PgPoolOptions::new() + .max_connections(6) + .connect_with(options) + .await + .expect("connect process pool"); + let client = build(pool); + client.start().await.expect("start process client"); + client + }); + Self { + runtime, + client: Arc::new(client), + } + }) + .await + .expect("start process host") + } + + async fn exit_gracefully(self, timeout: Duration) { + let client = self.client.clone(); + self.runtime + .spawn(async move { client.shutdown(timeout).await }) + .await + .expect("graceful shutdown"); + tokio::task::spawn_blocking(move || self.runtime.shutdown_background()) + .await + .expect("stop process runtime"); + } +} + +/// A graceful shutdown must finalize every attempt whose handler returned +/// during the drain. Finalization used to run detached from the drain, so a +/// process that exited after `shutdown` returned could drop it: the attempt +/// stayed open with no heartbeat and stale-heartbeat rescue ran the completed +/// job again. A held receipt-attempt lock keeps this attempt's finalize in +/// flight while shutdown drains, which the detached version did not wait for. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_graceful_shutdown_finalizes_returned_attempts() { + // Each host below opens its own pool of 6, and only one runs at a time. + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_graceful_restart"; + let store_config = QueueStorageConfig { + schema: "awa_qs_graceful_restart".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + let job_id = enqueue_job( + &pool, + &store, + &GracefulRestartJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + + let worker = GracefulRestartWorker::default(); + let options = (*pool.connect_options()).clone(); + let build = |worker: GracefulRestartWorker, store_config: QueueStorageConfig| { + move |pool: sqlx::PgPool| { + Client::builder(pool) + .queue( + queue, + QueueConfig { + max_workers: 2, + poll_interval: Duration::from_millis(25), + deadline_duration: Duration::ZERO, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(worker) + // The follow-up's queue is not served; it only routes the + // completion through a transaction of its own. + .on_completed_enqueue::( + |args, _job| GracefulRestartFollowUp { id: args.id }, + ) + .heartbeat_interval(Duration::from_millis(100)) + .heartbeat_staleness(ci_timing::scaled_staleness(Duration::from_secs(1))) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .promote_interval(Duration::from_millis(25)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build graceful restart client") + } + }; + + let host = + ProcessHost::start(options.clone(), build(worker.clone(), store_config.clone())).await; + tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(10)), + worker.entered.notified(), + ) + .await + .expect("the job should start"); + let (attempt_job_id, run_lease) = worker + .first_attempt + .lock() + .expect("first attempt") + .expect("first attempt recorded"); + assert_eq!(attempt_job_id, job_id); + + let mut lock_holder = pool.acquire().await.expect("lock connection"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_lock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(job_id) + .bind(run_lease) + .execute(&mut *lock_holder) + .await + .expect("hold the attempt's finalize lock"); + + let exit = + tokio::spawn(host.exit_gracefully(ci_timing::scaled_timeout(Duration::from_secs(10)))); + tokio::time::sleep(Duration::from_millis(500)).await; + sqlx::query( + "SELECT pg_catalog.pg_advisory_unlock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(job_id) + .bind(run_lease) + .execute(&mut *lock_holder) + .await + .expect("release the attempt's finalize lock"); + drop(lock_holder); + exit.await.expect("graceful exit"); + + let after_exit = store + .load_job(&pool, job_id) + .await + .expect("load job") + .expect("job row"); + + // A later runtime rescues anything the exit left open. + let rescuer = ProcessHost::start(options, build(worker.clone(), store_config)).await; + let finished = wait_for_job_state( + &store, + &pool, + job_id, + &[JobState::Completed], + ci_timing::scaled_timeout(Duration::from_secs(15)), + ) + .await; + rescuer + .exit_gracefully(ci_timing::scaled_timeout(Duration::from_secs(10))) + .await; + + assert_eq!( + (after_exit.state, after_exit.run_lease), + (JobState::Completed, run_lease), + "graceful shutdown returned before the completed attempt was finalized" + ); + assert_eq!( + (worker.executions.load(Ordering::SeqCst), finished.run_lease), + (1, run_lease), + "the job ran again after a graceful shutdown" + ); +} + /// ADR-023 claim-ring control-plane smoke test. Exercises /// `rotate_claims`, the #371 idle gate, and `prune_oldest_claims` on an /// empty schema: rotate is idle-skipped when nothing has been written @@ -11108,9 +11334,10 @@ fn branch_overruns() -> Arc>> { /// snoozed and retry-after jobs runnable, so overrunning promotion bodies must /// never arm a cooldown on it. /// -/// The first three promotion bodies (two `DELETE`s each) take 1.2s against a -/// 25ms promote interval, so the maintenance branch tracker flags promotion as -/// overrunning. The test asserts on that overrun warning rather than on job +/// Every promotion body runs the canonical-plane promotion `DELETE` once per +/// state, even with nothing due. A statement trigger slows the first six, so +/// the first three bodies take 1.2s against a 25ms promote interval and the +/// maintenance branch tracker flags promotion as overrunning. The test asserts on that overrun warning rather than on job /// latency: the promotion warning must report no cooldown, while another /// branch in the same run (cleanup, whose bodies exceed its 1ms interval) /// does report one, so the check cannot pass vacuously. Jobs that fall due @@ -11137,7 +11364,7 @@ async fn test_queue_storage_runtime_promotion_not_suppressed_after_overrunning_b END $$" ), format!( - "CREATE TRIGGER slow_first_promote_deletes BEFORE DELETE ON {schema}.deferred_jobs \ + "CREATE TRIGGER slow_first_promote_deletes BEFORE DELETE ON awa.scheduled_jobs \ FOR EACH STATEMENT EXECUTE FUNCTION {schema}.slow_first_promote_deletes()" ), ] { @@ -11272,86 +11499,670 @@ async fn test_queue_storage_runtime_promotion_not_suppressed_after_overrunning_b client.shutdown(Duration::from_secs(5)).await; } +/// Advisory lock key of a lane's enqueue reservation, as +/// `reserve_enqueue_seq` takes it. The lane must already exist. +async fn enqueue_lane_lock_key( + pool: &sqlx::PgPool, + store: &QueueStorage, + queue: &str, + priority: i16, +) -> i64 { + sqlx::query_scalar(audited_sql(format!( + "SELECT hashtextextended(format('%I.%I', $1::text, seq_name), 0) \ + FROM {}.queue_enqueue_heads \ + WHERE queue = $2 AND priority = $3 AND enqueue_shard = 0", + store.schema() + ))) + .bind(store.schema()) + .bind(queue) + .bind(priority) + .fetch_one(pool) + .await + .expect("derive enqueue lane lock key") +} + +/// A promotion batch that meets a lane whose enqueue lock is held promotes +/// the other lanes, reports the held lane as contended after its lock timeout +/// and leaves its rows deferred. Promoted later, the lane's rows take lane +/// sequence numbers in due order and are all claimable. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn test_queue_storage_runtime_stale_heartbeat_rescue() { - let (_db_guard, pool) = setup_pool(10).await; - let queue = "qs_heartbeat_rescue"; - let schema = "awa_qs_runtime_heartbeat_rescue"; - let store = create_store(&pool, schema).await; - let job_id = enqueue_job( +async fn test_queue_storage_promotion_leaves_a_contended_lane_deferred_and_in_order() { + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(6).await; + let hot = "qs_promote_contended_hot"; + let other = "qs_promote_contended_other"; + let store = create_store(&pool, "awa_qs_promote_contended_lane").await; + let opts = |queue: &str, run_at: Option>| InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at, + ..Default::default() + }; + + let warm_job_id = enqueue_job(&pool, &store, &CompleteJob { id: 0 }, opts(hot, None)).await; + // Later job ids fall due earlier, so due order differs from id order. + let due_base = Utc::now() - chrono::Duration::seconds(60); + let mut hot_due = Vec::new(); + for n in 0..6 { + let run_at = due_base - chrono::Duration::seconds(n); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: n }, + opts(hot, Some(run_at)), + ) + .await; + hot_due.push((run_at, job_id)); + } + hot_due.sort(); + let hot_in_due_order: Vec = hot_due.iter().map(|(_, job_id)| *job_id).collect(); + let other_job_id = enqueue_job( &pool, &store, - &HeartbeatRescueJob { id: 3 }, - InsertOpts { - queue: queue.to_string(), - ..Default::default() - }, + &CompleteJob { id: 100 }, + opts(other, Some(due_base)), ) .await; - let client = Client::builder(pool.clone()) - .queue( - queue, - QueueConfig { - max_workers: 4, - poll_interval: Duration::from_millis(25), - deadline_duration: Duration::from_secs(30), - ..QueueConfig::default() - }, - ) - .queue_storage( - QueueStorageConfig { - schema: schema.to_string(), - queue_slot_count: 4, - lease_slot_count: 2, - lease_claim_receipts: false, - ..Default::default() - }, - Duration::from_millis(1_000), - Duration::from_millis(50), + let hot_lane = PromotionLane { + queue: hot.to_string(), + priority: PRIORITY, + }; + let lock_key = enqueue_lane_lock_key(&pool, &store, hot, PRIORITY).await; + let mut holder = pool.begin().await.expect("begin lane holder"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(lock_key) + .execute(holder.as_mut()) + .await + .expect("hold the hot lane's enqueue lock"); + + let lock_timeout = Duration::from_millis(100); + let started = Instant::now(); + let outcome = store + .promote_due_lanes(&pool, JobState::Scheduled, 100, Some(lock_timeout), &[]) + .await + .expect("promote with the hot lane held"); + let elapsed = started.elapsed(); + assert_eq!(outcome.promoted, 1, "{outcome:?}"); + assert_eq!(outcome.selected, 7, "{outcome:?}"); + assert_eq!(outcome.contended, vec![hot_lane.clone()]); + assert!( + elapsed >= lock_timeout && elapsed < scaled_timeout(Duration::from_secs(1)), + "the held lane should cost one lock timeout, took {elapsed:?}" + ); + let other_job = store + .load_job(&pool, other_job_id) + .await + .expect("load other job") + .expect("other job exists"); + assert_eq!(other_job.state, JobState::Available); + for job_id in &hot_in_due_order { + let job = store + .load_job(&pool, *job_id) + .await + .expect("load hot job") + .expect("hot job exists"); + assert_eq!(job.state, JobState::Scheduled, "job {job_id}"); + } + + let skipped = store + .promote_due_lanes( + &pool, + JobState::Scheduled, + 100, + Some(lock_timeout), + std::slice::from_ref(&hot_lane), ) - .register_worker(StaleHeartbeatWorker) - .heartbeat_interval(Duration::from_secs(5)) - .promote_interval(Duration::from_millis(25)) - .leader_election_interval(Duration::from_millis(100)) - .leader_check_interval(Duration::from_millis(50)) - .heartbeat_rescue_interval(Duration::from_millis(100)) - .heartbeat_staleness(Duration::from_secs(2)) - .deadline_rescue_interval(Duration::from_secs(10)) - .callback_rescue_interval(Duration::from_secs(10)) - .build() - .expect("Failed to build heartbeat rescue client"); - client - .start() .await - .expect("Failed to start heartbeat rescue client"); + .expect("promote with the hot lane skipped"); + assert_eq!(skipped, Default::default()); - let completed = wait_for_job_state( - &store, - &pool, - job_id, - &[JobState::Completed], - Duration::from_secs(15), - ) - .await; - assert_eq!(completed.state, JobState::Completed); - assert_eq!(completed.attempt, 2); - assert_eq!(attempt_state_count(&pool, &store).await, 0); + holder.commit().await.expect("release the hot lane"); + let mut batches = 0; + loop { + let outcome = store + .promote_due_lane(&pool, JobState::Scheduled, &hot_lane, 4, None) + .await + .expect("promote the hot lane"); + assert!(outcome.contended.is_empty(), "{outcome:?}"); + if outcome.promoted == 0 { + break; + } + batches += 1; + } + assert_eq!(batches, 2); - client.shutdown(Duration::from_secs(5)).await; + let in_lane_order: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT job_id FROM {}.ready_entries WHERE job_id = ANY($1) ORDER BY lane_seq", + store.schema() + ))) + .bind(&hot_in_due_order) + .fetch_all(&pool) + .await + .expect("read hot lane order"); + assert_eq!(in_lane_order, hot_in_due_order); + + let mut claimed: Vec = store + .claim_runtime_batch(&pool, hot, 100, Duration::from_secs(60)) + .await + .expect("claim the hot lane") + .into_iter() + .map(|claimed| claimed.job.id) + .collect(); + claimed.sort_unstable(); + let mut expected = hot_in_due_order.clone(); + expected.push(warm_job_id); + expected.sort_unstable(); + assert_eq!(claimed, expected); +} + +/// Holds PostgreSQL's cluster-wide notify-queue lock for about `commits` × +/// 100 ms. A transaction that issued NOTIFY takes that lock at COMMIT and +/// keeps it until its commit record is flushed, which `commit_delay` stretches. +/// Returns once the lock is held. +async fn hold_notify_commit_lock( + pool: &sqlx::PgPool, + commits: usize, +) -> tokio::task::JoinHandle<()> { + sqlx::query("CREATE TABLE IF NOT EXISTS notify_lock_holder (id bigint)") + .execute(pool) + .await + .expect("create notify holder table"); + let holder = PgPoolOptions::new() + .max_connections(1) + .connect_with( + pool.connect_options() + .as_ref() + .clone() + .options([("commit_delay", "100000"), ("commit_siblings", "0")]), + ) + .await + .expect("connect the notify holder"); + let holder_task = tokio::spawn(async move { + for _ in 0..commits { + sqlx::raw_sql( + "BEGIN; INSERT INTO notify_lock_holder VALUES (1); \ + SELECT pg_notify('awa_test_notify_holder', ''); COMMIT;", + ) + .execute(&holder) + .await + .expect("notify holder commit"); + } + holder.close().await; + }); + let deadline = Instant::now() + scaled_timeout(Duration::from_secs(5)); + loop { + let held: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'object' \ + AND classid = 'pg_database'::regclass AND objid = 0 AND granted)", + ) + .fetch_one(pool) + .await + .expect("read notify lock"); + if held { + return holder_task; + } + assert!( + Instant::now() < deadline, + "notify holder never took the lock" + ); + tokio::time::sleep(Duration::from_millis(1)).await; + } } +/// Promotion's lane transactions do not wait for the notify-queue lock: a lane +/// promotes under its short lock timeout while another notifying transaction +/// holds that lock, and a wakeup that cannot get it, under the connection's own +/// `lock_timeout`, does not undo or fail the promotion. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn test_queue_storage_admin_queries_cover_running_and_failed_rows() { - let (_db_guard, pool) = setup_pool(10).await; - let queue = "qs_admin_runtime"; - let schema = "awa_qs_admin_runtime"; - let store = create_store(&pool, schema).await; - let gate = BlockingCompleteWorkerGate::new(); +async fn test_queue_storage_promotion_does_not_wait_for_the_notify_lock_in_its_lanes() { + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_promote_notify_commit"; + let store = create_store(&pool, "awa_qs_promote_notify_commit").await; + let lane = PromotionLane { + queue: queue.to_string(), + priority: PRIORITY, + }; + let enqueue_due = |id: i64| { + let pool = pool.clone(); + let store = &store; + async move { + enqueue_job( + &pool, + store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at: Some(Utc::now() - chrono::Duration::seconds(60)), + ..Default::default() + }, + ) + .await + } + }; + for id in 0..3 { + enqueue_due(id).await; + } - let running_job_id = enqueue_job( - &pool, - &store, + let holder = hold_notify_commit_lock(&pool, 3).await; + let outcome = store + .promote_due_lanes( + &pool, + JobState::Scheduled, + 100, + Some(Duration::from_millis(50)), + &[], + ) + .await + .expect("promote while the notify lock is held"); + holder.await.expect("notify holder"); + assert_eq!(outcome.promoted, 3, "{outcome:?}"); + assert!(outcome.contended.is_empty(), "{outcome:?}"); + + let caller_lock_timeout_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + pool.connect_options() + .as_ref() + .clone() + .options([("lock_timeout", "20ms")]), + ) + .await + .expect("connect with a session lock_timeout"); + enqueue_due(3).await; + let holder = hold_notify_commit_lock(&pool, 3).await; + let result = store + .promote_due_lanes( + &caller_lock_timeout_pool, + JobState::Scheduled, + 100, + Some(Duration::from_millis(50)), + &[], + ) + .await; + holder.await.expect("notify holder"); + caller_lock_timeout_pool.close().await; + let outcome = result.expect("a failed wakeup does not fail the promotion"); + assert_eq!(outcome.promoted, 1, "{outcome:?}"); + assert!(outcome.contended.is_empty(), "{outcome:?}"); + let deferred: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {}.deferred_jobs WHERE queue = $1", + store.schema() + ))) + .bind(queue) + .fetch_one(&pool) + .await + .expect("count deferred rows"); + assert_eq!(deferred, 0); + + let claimed = store + .claim_runtime_batch(&pool, queue, 10, Duration::from_secs(60)) + .await + .expect("claim promoted rows"); + assert_eq!(claimed.len(), 4); + assert!(store + .promote_due_lane(&pool, JobState::Scheduled, &lane, 100, None) + .await + .expect("nothing left to promote") + .contended + .is_empty()); +} + +/// The maintenance leader must not wait on a lane whose enqueue lock never +/// comes free. Two sessions take turns holding the hot lane's lock for +/// `HOLD` each, so a lock request queued behind them waits up to two holds. +/// A due job on another queue must still run, and an abandoned attempt must +/// still be rescued, well within one hold; the hot queue's due jobs must +/// still be promoted while the lane stays busy. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_busy_lane_holds_back_neither_promotion_nor_rescue() { + const HOLD: Duration = Duration::from_secs(5); + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(12).await; + let hot = "qs_busy_lane_hot"; + let other = "qs_busy_lane_other"; + let abandoned = "qs_busy_lane_abandoned"; + let schema = "awa_qs_busy_lane_leader"; + let store = create_store(&pool, schema).await; + let opts = |queue: &str, run_at: Option>| InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at, + ..Default::default() + }; + + enqueue_job(&pool, &store, &CompleteJob { id: 0 }, opts(hot, None)).await; + let lock_key = enqueue_lane_lock_key(&pool, &store, hot, PRIORITY).await; + + let stop = Arc::new(AtomicBool::new(false)); + let mut holders = Vec::new(); + for _ in 0..2 { + let (held_tx, held_rx) = tokio::sync::oneshot::channel::<()>(); + let pool = pool.clone(); + let stop = Arc::clone(&stop); + holders.push(tokio::spawn(async move { + let mut held_tx = Some(held_tx); + let mut holds = 0u32; + while !stop.load(Ordering::SeqCst) { + let mut tx = pool.begin().await.expect("begin lane holder"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(lock_key) + .execute(tx.as_mut()) + .await + .expect("hold the hot lane's enqueue lock"); + holds += 1; + if let Some(held_tx) = held_tx.take() { + let _ = held_tx.send(()); + } + tokio::time::sleep(HOLD).await; + tx.commit().await.expect("release the hot lane"); + } + holds + })); + // The first holder takes the lock; wait for it so the second queues + // behind it rather than racing it. + if holders.len() == 1 { + held_rx.await.expect("first holder took the lane"); + } + } + + let due = Utc::now() - chrono::Duration::seconds(1); + let mut hot_jobs = Vec::new(); + for n in 1..=3 { + hot_jobs + .push(enqueue_job(&pool, &store, &CompleteJob { id: n }, opts(hot, Some(due))).await); + } + let other_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: 10 }, + opts(other, Some(due)), + ) + .await; + let abandoned_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: 20 }, + opts(abandoned, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, abandoned, 1, Duration::from_secs(3_600)) + .await + .expect("claim the abandoned attempt"); + assert_eq!(claimed.len(), 1); + + let heartbeat_staleness = Duration::from_millis(500); + let worker_queue = || QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + ..QueueConfig::default() + }; + let client = Client::builder(pool.clone()) + .queue(hot, worker_queue()) + .queue(other, worker_queue()) + .queue_storage( + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }, + Duration::from_secs(60), + Duration::from_secs(60), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(CompleteWorker) + .promote_interval(Duration::from_millis(50)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .heartbeat_staleness(heartbeat_staleness) + .deadline_rescue_interval(Duration::from_secs(10)) + .callback_rescue_interval(Duration::from_secs(10)) + .build() + .expect("Failed to build busy-lane client"); + client + .start() + .await + .expect("Failed to start busy-lane client"); + + let bound = scaled_timeout(Duration::from_secs(2)); + wait_for_job_state(&store, &pool, other_job, &[JobState::Completed], bound).await; + wait_for_job_state( + &store, + &pool, + abandoned_job, + &[JobState::Retryable, JobState::Available], + bound + heartbeat_staleness, + ) + .await; + + for job_id in &hot_jobs { + wait_for_job_state( + &store, + &pool, + *job_id, + &[JobState::Completed], + scaled_timeout(HOLD * 3), + ) + .await; + } + assert!( + holders.iter().all(|holder| !holder.is_finished()), + "the hot lane must stay busy until its jobs are promoted" + ); + + stop.store(true, Ordering::SeqCst); + for holder in holders { + assert!(holder.await.expect("lane holder") >= 1); + } + client.shutdown(Duration::from_secs(5)).await; +} + +/// Rows another transaction holds locked, as a lane promotion still waiting +/// under a former leader does, must not fill every batch of a promotion tick. +/// With more of the oldest due rows locked than one batch selects, a due job +/// on another queue must still be promoted and run promptly, and the locked +/// rows must be promoted once released. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_locked_due_rows_do_not_starve_promotion() { + const LOCKED: usize = 4_096 + 64; + let (_db_guard, pool) = setup_pool(10).await; + let locked_queue = "qs_locked_due_rows"; + let other = "qs_locked_due_other"; + let schema = "awa_qs_locked_due_rows"; + let store = create_store(&pool, schema).await; + + let oldest = Utc::now() - chrono::Duration::seconds(120); + let params: Vec<_> = (0..LOCKED) + .map(|n| { + insert::params_with( + &CompleteJob { id: n as i64 }, + InsertOpts { + queue: locked_queue.to_string(), + run_at: Some(oldest), + ..Default::default() + }, + ) + .expect("build locked-row params") + }) + .collect(); + for chunk in params.chunks(2_000) { + store + .enqueue_params_batch(&pool, chunk) + .await + .expect("enqueue rows to lock"); + } + let mut holder = pool.begin().await.expect("begin row holder"); + let locked: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT job_id FROM {schema}.deferred_jobs WHERE queue = $1 FOR UPDATE" + ))) + .bind(locked_queue) + .fetch_all(holder.as_mut()) + .await + .expect("lock the oldest due rows"); + assert_eq!(locked.len(), LOCKED); + + let other_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: -1 }, + InsertOpts { + queue: other.to_string(), + run_at: Some(Utc::now() - chrono::Duration::seconds(1)), + ..Default::default() + }, + ) + .await; + + let client = Client::builder(pool.clone()) + .queue( + other, + QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + ..QueueConfig::default() + }, + ) + .queue_storage( + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }, + Duration::from_secs(60), + Duration::from_secs(60), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(CompleteWorker) + .promote_interval(Duration::from_millis(50)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("Failed to build locked-rows client"); + client + .start() + .await + .expect("Failed to start locked-rows client"); + + wait_for_job_state( + &store, + &pool, + other_job, + &[JobState::Completed], + scaled_timeout(Duration::from_secs(2)), + ) + .await; + + holder.rollback().await.expect("release the locked rows"); + let deadline = Instant::now() + scaled_timeout(Duration::from_secs(10)); + loop { + let remaining: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.deferred_jobs WHERE queue = $1" + ))) + .bind(locked_queue) + .fetch_one(&pool) + .await + .expect("count deferred rows"); + if remaining == 0 { + break; + } + assert!( + Instant::now() < deadline, + "{remaining} released rows were never promoted" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + + client.shutdown(Duration::from_secs(5)).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_stale_heartbeat_rescue() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_heartbeat_rescue"; + let schema = "awa_qs_runtime_heartbeat_rescue"; + let store = create_store(&pool, schema).await; + let job_id = enqueue_job( + &pool, + &store, + &HeartbeatRescueJob { id: 3 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + deadline_duration: Duration::from_secs(30), + ..QueueConfig::default() + }, + ) + .queue_storage( + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }, + Duration::from_millis(1_000), + Duration::from_millis(50), + ) + .register_worker(StaleHeartbeatWorker) + .heartbeat_interval(Duration::from_secs(5)) + .promote_interval(Duration::from_millis(25)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .heartbeat_staleness(Duration::from_secs(2)) + .deadline_rescue_interval(Duration::from_secs(10)) + .callback_rescue_interval(Duration::from_secs(10)) + .build() + .expect("Failed to build heartbeat rescue client"); + client + .start() + .await + .expect("Failed to start heartbeat rescue client"); + + let completed = wait_for_job_state( + &store, + &pool, + job_id, + &[JobState::Completed], + Duration::from_secs(15), + ) + .await; + assert_eq!(completed.state, JobState::Completed); + assert_eq!(completed.attempt, 2); + assert_eq!(attempt_state_count(&pool, &store).await, 0); + + client.shutdown(Duration::from_secs(5)).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_admin_queries_cover_running_and_failed_rows() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_admin_runtime"; + let schema = "awa_qs_admin_runtime"; + let store = create_store(&pool, schema).await; + let gate = BlockingCompleteWorkerGate::new(); + + let running_job_id = enqueue_job( + &pool, + &store, &CompleteJob { id: 91 }, InsertOpts { queue: queue.to_string(), @@ -11474,14 +12285,130 @@ async fn test_queue_storage_admin_queries_cover_running_and_failed_rows() { ..Default::default() }, ) - .await - .expect("Failed to list failed queue_storage jobs"); - assert_eq!(failed_jobs.len(), 1); - assert_eq!(failed_jobs[0].id, failed_job_id); + .await + .expect("Failed to list failed queue_storage jobs"); + assert_eq!(failed_jobs.len(), 1); + assert_eq!(failed_jobs[0].id, failed_job_id); + + gate.wait_until_entered(Duration::from_secs(5)).await; + gate.release(); + client.shutdown(Duration::from_secs(5)).await; +} + +/// Open receipt claims count as running until durable evidence closes them: +/// a completion batch, a retryable transition, or materialization into a +/// lease, which must be counted once. +async fn receipt_admin_counts_follow_durable_dispositions(row_local: bool) { + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_receipt_dispositions"; + let schema = if row_local { + "awa_qs_row_receipt_dispositions" + } else { + "awa_qs_receipt_dispositions" + }; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + lease_claim_receipts: true, + ..Default::default() + }, + ) + .await; - gate.wait_until_entered(Duration::from_secs(5)).await; - gate.release(); - client.shutdown(Duration::from_secs(5)).await; + async fn counts(pool: &sqlx::PgPool, queue: &str) -> (i64, i64, i64) { + let overview = admin::queue_overview(pool, queue) + .await + .expect("read receipt disposition counts") + .expect("receipt queue exists"); + ( + overview.running, + overview.retryable, + overview.completed_last_hour, + ) + } + + store.enqueue_batch(&pool, queue, 2, 3).await.unwrap(); + let deadline = Duration::from_secs(300); + let mut claimed = store + .claim_runtime_batch(&pool, queue, 3, deadline) + .await + .expect("claim three deadline-backed receipts"); + assert_eq!(claimed.len(), 3); + if row_local { + rewrite_as_row_local_claims(&pool, schema, &mut claimed, deadline).await; + assert_eq!(lease_claim_count(&pool, &store).await, 3); + } + assert_eq!(lease_count(&pool, &store).await, 0); + assert_eq!(counts(&pool, queue).await, (3, 0, 0)); + + let completed = store + .complete_runtime_batch(&pool, &claimed[..1]) + .await + .unwrap(); + assert_eq!( + completed, + vec![(claimed[0].job.id, claimed[0].job.run_lease)] + ); + assert_eq!(lease_claim_closure_batch_count(&pool, &store).await, 1); + assert_eq!(counts(&pool, queue).await, (2, 0, 1)); + + let retried = store + .fail_retryable( + &pool, + claimed[1].job.id, + claimed[1].job.run_lease, + "retry", + None, + ) + .await + .unwrap() + .expect("second attempt moves to retryable"); + assert_eq!(retried.state, JobState::Retryable); + assert_eq!(counts(&pool, queue).await, (1, 1, 1)); + + let callback = store + .register_callback( + &pool, + claimed[2].job.id, + claimed[2].job.run_lease, + Duration::from_secs(300), + ) + .await + .unwrap(); + assert_eq!(lease_count(&pool, &store).await, 1); + assert_eq!( + counts(&pool, queue).await, + (1, 1, 1), + "materialization counts the attempt once" + ); + let finished = store + .complete_external(&pool, callback, None, Some(claimed[2].job.run_lease), false) + .await + .unwrap(); + assert_eq!(finished.state, JobState::Completed); + assert_eq!(counts(&pool, queue).await, (0, 1, 2)); + for (claim, expected) in claimed.iter().zip([ + JobState::Completed, + JobState::Retryable, + JobState::Completed, + ]) { + let job = store.load_job(&pool, claim.job.id).await.unwrap().unwrap(); + assert_eq!( + job.state, expected, + "point reads and aggregate counts agree" + ); + } +} + +#[tokio::test] +async fn test_receipt_admin_counts_follow_durable_dispositions() { + receipt_admin_counts_follow_durable_dispositions(false).await; +} + +#[tokio::test] +async fn test_row_local_receipt_admin_counts_follow_durable_dispositions() { + receipt_admin_counts_follow_durable_dispositions(true).await; } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] @@ -21849,119 +22776,597 @@ async fn moved_ready_row_is_joined_once_by_attempt_reads(receipts: bool) { } else { "awa_qs_moved_ready_lease_reads" }; - let store = create_store_with_config( + let store = create_store_with_config( + &pool, + QueueStorageConfig { + lease_claim_receipts: receipts, + ..unique_cancel_store_config(schema) + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &unique_cancel_job("moved"), + InsertOpts { + priority: 2, + ..unique_cancel_opts(queue, None) + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + let aged = store + .age_waiting_priorities(&pool, Duration::from_secs(60), 10) + .await + .expect("age waiting priorities"); + assert_eq!(aged, vec![job_id]); + + let same_position: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(DISTINCT priority)::bigint + FROM {schema}.ready_entries + WHERE job_id = $1 + GROUP BY ready_slot, ready_generation, enqueue_shard, lane_seq" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("ready rows for the moved job"); + assert_eq!( + same_position, 2, + "the source and destination rows must share a position for this test to mean anything" + ); + + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim moved job"); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].job.id, job_id); + assert_eq!(claimed[0].claim.claim_batch_id.is_some(), receipts); + + // This is the production CTE used by current-job projections, not only + // the point-read candidate list. A compact item must hydrate exactly once. + let projected: Vec = sqlx::query_scalar(audited_sql(format!( + "{} SELECT job_id FROM current_jobs WHERE queue = $1 AND state = 'running'", + admin::queue_storage_current_jobs_cte(schema) + ))) + .bind(queue) + .fetch_all(&pool) + .await + .expect("project current running jobs"); + assert_eq!(projected, vec![job_id], "hydrate only the live ready row"); + + let running = admin::list_jobs( + &pool, + &admin::ListJobsFilter { + state: Some(JobState::Running), + queue: Some(queue.to_string()), + ..Default::default() + }, + ) + .await + .expect("list running jobs"); + assert_eq!( + running.iter().map(|job| job.id).collect::>(), + vec![job_id], + "the moved running job is listed once" + ); + + let overview = admin::job_kind_overviews(&pool) + .await + .expect("job kind overviews"); + let kind = overview + .iter() + .find(|row| row.kind == UniqueCancelJob::kind()) + .expect("the moved job's kind is listed"); + assert_eq!(kind.job_count, 1, "the moved running job is counted once"); + + let loaded = store + .load_job(&pool, job_id) + .await + .expect("load moved job") + .expect("moved job exists"); + assert_eq!(loaded.state, JobState::Running); + assert_eq!(loaded.priority, 1, "reads the live destination row"); + assert_eq!( + loaded.metadata.get("_awa_original_priority"), + Some(&serde_json::json!(2)), + "payload comes from the aged destination, not its tombstoned source" + ); +} + +#[tokio::test] +async fn test_moved_ready_row_is_joined_once_by_lease_reads() { + moved_ready_row_is_joined_once_by_attempt_reads(false).await; +} + +#[tokio::test] +async fn test_moved_ready_row_is_joined_once_by_compact_receipt_reads() { + moved_ready_row_is_joined_once_by_attempt_reads(true).await; +} + +/// Rewrites compact receipt claims as the row-local `lease_claims` rows a 0.6 +/// runtime writes for deadline-backed claims, which this runtime must still +/// read and close during a rolling upgrade, and clears the compact identity on +/// the claims so completion takes the row-local path. +async fn rewrite_as_row_local_claims( + pool: &sqlx::PgPool, + schema: &str, + claims: &mut [awa::model::ClaimedRuntimeJob], + deadline: Duration, +) { + let batch_ids: Vec = claims + .iter() + .map(|claim| claim.claim.claim_batch_id.expect("a compact claim")) + .collect(); + sqlx::query(audited_sql(format!( + "DELETE FROM {schema}.lease_claim_batches WHERE batch_id = ANY($1)" + ))) + .bind(&batch_ids) + .execute(pool) + .await + .expect("remove the compact claims"); + for claim in claims.iter_mut() { + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.lease_claims ( + claim_slot, receipt_id, job_id, run_lease, ready_slot, + ready_generation, queue, priority, attempt, max_attempts, + lane_seq, enqueue_shard, deadline_at + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, + clock_timestamp() + make_interval(secs => $13))" + ))) + .bind(claim.claim.claim_slot) + .bind(claim.claim.receipt_id) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .bind(claim.claim.ready_slot) + .bind(claim.claim.ready_generation) + .bind(&claim.claim.queue) + .bind(claim.claim.priority) + .bind(claim.job.attempt) + .bind(claim.job.max_attempts) + .bind(claim.claim.lane_seq) + .bind(claim.claim.enqueue_shard) + .bind(deadline.as_secs_f64()) + .execute(pool) + .await + .expect("write the row-local claim"); + claim.claim.claim_batch_id = None; + claim.claim.claim_batch_index = None; + } +} + +async fn receipt_admin_ids(pool: &sqlx::PgPool, filter: admin::ListJobsFilter) -> Vec { + admin::list_jobs(pool, &filter) + .await + .expect("list receipt-held jobs") + .into_iter() + .map(|job| job.id) + .collect() +} + +async fn receipt_admin_reads_preserve_scope_and_disposition(receipts: bool, row_local: bool) { + let (_db_guard, pool) = setup_pool(6).await; + let schema = match (receipts, row_local) { + (false, _) => "awa_qs_admin_scope_leases", + (true, false) => "awa_qs_admin_scope_compact", + (true, true) => "awa_qs_admin_scope_row_claims", + }; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + lease_claim_receipts: receipts, + ..unique_cancel_store_config(schema) + }, + ) + .await; + let queue = "receipt_admin_scope"; + let other_queue = "receipt_admin_other"; + let first = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let second = enqueue_job( + &pool, + &store, + &CompleteJob { id: 2 }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let other_kind = enqueue_job( + &pool, + &store, + &DlqJob { id: 3 }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let other = enqueue_job( + &pool, + &store, + &CompleteJob { id: 4 }, + InsertOpts { + queue: other_queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let deadline = if receipts && !row_local { + Duration::ZERO + } else { + Duration::from_secs(30) + }; + let mut claims = store + .claim_runtime_batch(&pool, queue, 3, deadline) + .await + .unwrap(); + if receipts { + assert!(matches!( + store + .rotate_claims(&pool) + .await + .expect("rotate between queues"), + RotateOutcome::Rotated { slot: 1, .. } + )); + } + claims.extend( + store + .claim_runtime_batch(&pool, other_queue, 1, deadline) + .await + .unwrap(), + ); + assert_eq!(claims.len(), 4); + if row_local { + rewrite_as_row_local_claims(&pool, schema, &mut claims, deadline).await; + } + assert!(claims + .iter() + .all(|claim| claim.claim.lease_claim_receipt == receipts)); + assert!(claims + .iter() + .all(|claim| claim.claim.claim_batch_id.is_some() == (receipts && !row_local))); + let row_claims: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.lease_claims" + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row_claims, if row_local { 4 } else { 0 }); + + let scoped = admin::ListJobsFilter { + state: Some(JobState::Running), + kind: Some(CompleteJob::kind().into()), + queue: Some(queue.into()), + tag: Some("selected".into()), + limit: Some(1), + ..Default::default() + }; + assert_eq!(receipt_admin_ids(&pool, scoped.clone()).await, vec![second]); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + before_id: Some(second), + ..scoped.clone() + } + ) + .await, + vec![first] + ); + assert!(receipt_admin_ids( &pool, - QueueStorageConfig { - lease_claim_receipts: receipts, - ..unique_cancel_store_config(schema) - }, + admin::ListJobsFilter { + before_id: Some(first), + ..scoped.clone() + } ) - .await; - let job_id = enqueue_job( + .await + .is_empty()); + assert!(receipt_admin_ids( &pool, - &store, - &unique_cancel_job("moved"), - InsertOpts { - priority: 2, - ..unique_cancel_opts(queue, None) - }, + admin::ListJobsFilter { + tag: Some("absent".into()), + ..scoped.clone() + } ) - .await; - backdate_one_lane_ready_job_run_at( + .await + .is_empty()); + assert!(receipt_admin_ids( &pool, - schema, - job_id, - Utc::now() - chrono::Duration::minutes(5), + admin::ListJobsFilter { + state: Some(JobState::WaitingExternal), + ..scoped.clone() + } ) - .await; - let aged = store - .age_waiting_priorities(&pool, Duration::from_secs(60), 10) - .await - .expect("age waiting priorities"); - assert_eq!(aged, vec![job_id]); - - let same_position: i64 = sqlx::query_scalar(audited_sql(format!( - "SELECT count(DISTINCT priority)::bigint - FROM {schema}.ready_entries - WHERE job_id = $1 - GROUP BY ready_slot, ready_generation, enqueue_shard, lane_seq" - ))) - .bind(job_id) - .fetch_one(&pool) .await - .expect("ready rows for the moved job"); + .is_empty()); assert_eq!( - same_position, 2, - "the source and destination rows must share a position for this test to mean anything" + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + state: None, + limit: Some(10), + ..scoped.clone() + } + ) + .await, + vec![second, first] ); - let claimed = store - .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) - .await - .expect("claim moved job"); - assert_eq!(claimed.len(), 1); - assert_eq!(claimed[0].job.id, job_id); - assert_eq!(claimed[0].claim.claim_batch_id.is_some(), receipts); - - // This is the production CTE used by current-job projections, not only - // the point-read candidate list. A compact item must hydrate exactly once. let projected: Vec = sqlx::query_scalar(audited_sql(format!( - "{} SELECT job_id FROM current_jobs WHERE queue = $1 AND state = 'running'", + "{} SELECT job_id FROM current_jobs WHERE state='running' ORDER BY job_id DESC", admin::queue_storage_current_jobs_cte(schema) ))) - .bind(queue) .fetch_all(&pool) .await - .expect("project current running jobs"); - assert_eq!(projected, vec![job_id], "hydrate only the live ready row"); + .unwrap(); + assert_eq!(projected, vec![other, other_kind, second, first]); + let kinds = admin::job_kind_overviews(&pool).await.unwrap(); + let complete = kinds + .iter() + .find(|kind| kind.kind == CompleteJob::kind()) + .unwrap(); + assert_eq!((complete.job_count, complete.queue_count), (3, 2)); + let alternate = kinds + .iter() + .find(|kind| kind.kind == DlqJob::kind()) + .unwrap(); + assert_eq!((alternate.job_count, alternate.queue_count), (1, 1)); - let running = admin::list_jobs( - &pool, - &admin::ListJobsFilter { - state: Some(JobState::Running), - queue: Some(queue.to_string()), - ..Default::default() - }, - ) - .await - .expect("list running jobs"); + // Materialization keeps the same attempt running. The receipt leg must + // yield to the lease leg rather than report both representations. + let first_claim = claims.iter().find(|claim| claim.job.id == first).unwrap(); + store + .register_callback( + &pool, + first, + first_claim.job.run_lease, + Duration::from_secs(60), + ) + .await + .expect("materialize one running attempt"); assert_eq!( - running.iter().map(|job| job.id).collect::>(), - vec![job_id], - "the moved running job is listed once" + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + limit: Some(10), + ..scoped.clone() + } + ) + .await, + vec![second, first] ); + let running: i64 = sqlx::query_scalar(audited_sql(format!( + "{} SELECT count(*) FROM current_jobs WHERE state='running'", + admin::queue_storage_current_jobs_cte(schema) + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(running, 4, "materialization does not duplicate the claim"); - let overview = admin::job_kind_overviews(&pool) - .await - .expect("job kind overviews"); - let kind = overview + let closed = claims .iter() - .find(|row| row.kind == UniqueCancelJob::kind()) - .expect("the moved job's kind is listed"); - assert_eq!(kind.job_count, 1, "the moved running job is counted once"); - - let loaded = store - .load_job(&pool, job_id) - .await - .expect("load moved job") - .expect("moved job exists"); - assert_eq!(loaded.state, JobState::Running); - assert_eq!(loaded.priority, 1, "reads the live destination row"); + .find(|claim| claim.job.id == second) + .unwrap() + .clone(); assert_eq!( - loaded.metadata.get("_awa_original_priority"), - Some(&serde_json::json!(2)), - "payload comes from the aged destination, not its tombstoned source" + store + .complete_runtime_batch(&pool, &[closed]) + .await + .unwrap() + .len(), + 1 + ); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + limit: Some(10), + ..scoped + } + ) + .await, + vec![first], + "closed receipts leave running enumeration" + ); + let running: i64 = sqlx::query_scalar(audited_sql(format!( + "{} SELECT count(*) FROM current_jobs WHERE state='running'", + admin::queue_storage_current_jobs_cte(schema) + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(running, 3); + let kinds = admin::job_kind_overviews(&pool).await.unwrap(); + let complete = kinds + .iter() + .find(|kind| kind.kind == CompleteJob::kind()) + .unwrap(); + assert_eq!( + (complete.job_count, complete.queue_count), + (3, 2), + "the completed job is counted once, not as both running and terminal" ); } #[tokio::test] -async fn test_moved_ready_row_is_joined_once_by_lease_reads() { - moved_ready_row_is_joined_once_by_attempt_reads(false).await; +async fn test_receipt_admin_scope_compact() { + receipt_admin_reads_preserve_scope_and_disposition(true, false).await; } #[tokio::test] -async fn test_moved_ready_row_is_joined_once_by_compact_receipt_reads() { - moved_ready_row_is_joined_once_by_attempt_reads(true).await; +async fn test_receipt_admin_scope_row_local() { + receipt_admin_reads_preserve_scope_and_disposition(true, true).await; +} + +#[tokio::test] +async fn test_receipt_admin_scope_materialized_control() { + receipt_admin_reads_preserve_scope_and_disposition(false, false).await; +} + +#[tokio::test] +async fn test_receipt_list_refills_pages_past_closed_and_other_kind_claims() { + let (_guard, pool) = setup_pool(6).await; + let queue = "receipt_candidate_refill"; + let store = create_store_with_config( + &pool, + unique_cancel_store_config("awa_qs_receipt_candidate_refill"), + ) + .await; + let mut targets = Vec::new(); + for id in 0..3 { + targets.push( + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await, + ); + } + let targets_claimed = store + .claim_runtime_batch(&pool, queue, 3, Duration::ZERO) + .await + .expect("claim older targets"); + assert_eq!(targets_claimed.len(), 3); + let first_target = targets_claimed + .iter() + .find(|claim| claim.job.id == targets[0]) + .expect("first target was claimed"); + // The same job appears in both the lease and raw-receipt candidate sources. + store + .register_callback( + &pool, + targets[0], + first_target.job.run_lease, + Duration::from_secs(60), + ) + .await + .expect("materialize one target"); + + let mut all_ids = targets.clone(); + for id in 10..22 { + all_ids.push( + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await, + ); + } + let finished = store + .claim_runtime_batch(&pool, queue, 12, Duration::ZERO) + .await + .expect("claim newer completed candidates"); + assert_eq!(finished.len(), 12); + assert_eq!( + store + .complete_runtime_batch(&pool, &finished) + .await + .unwrap() + .len(), + 12 + ); + for id in 30..42 { + all_ids.push( + enqueue_job( + &pool, + &store, + &DlqJob { id }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await, + ); + } + assert_eq!( + store + .claim_runtime_batch(&pool, queue, 12, Duration::ZERO) + .await + .unwrap() + .len(), + 12 + ); + + let filter = admin::ListJobsFilter { + state: Some(JobState::Running), + kind: Some(CompleteJob::kind().into()), + queue: Some(queue.into()), + limit: Some(2), + ..Default::default() + }; + targets.sort_unstable_by(|a, b| b.cmp(a)); + let first = receipt_admin_ids(&pool, filter.clone()).await; + assert_eq!(first, targets[..2]); + let second = receipt_admin_ids( + &pool, + admin::ListJobsFilter { + before_id: first.last().copied(), + ..filter.clone() + }, + ) + .await; + assert_eq!(second, targets[2..]); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + tag: Some("selected".into()), + limit: Some(100), + ..filter + }, + ) + .await, + targets + ); + all_ids.sort_unstable_by(|a, b| b.cmp(a)); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + queue: Some(queue.into()), + limit: Some(100), + ..Default::default() + }, + ) + .await, + all_ids, + "closed receipts and materialized claims must not duplicate results" + ); } /// A point read's snapshot can outlive completion. TRUNCATE is not MVCC-safe: @@ -23816,15 +25221,36 @@ async fn moved_tombstoned_claim_lane( assert_eq!(loaded.max_attempts, 1); } if !receipts { - let err = store + let other_queue = "qs_moved_tomb_claim_other"; + let other_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9044 }, + InsertOpts { + queue: other_queue.into(), + ..Default::default() + }, + ) + .await; + let other_claims = store + .claim_runtime_batch(&pool, other_queue, 1, Duration::from_secs(300)) + .await + .expect("claim an unrelated job"); + assert_eq!(other_claims.len(), 1); + assert_eq!(lease_count(&pool, &store).await, 2); + + let rescued = store .rescue_stale_heartbeats(&pool, Duration::ZERO) .await - .expect_err("legacy ambiguous physical lanes must not be selected arbitrarily"); - assert!(err.to_string().contains("ready body ambiguous"), "{err}"); + .expect("an ambiguous attempt does not stop rescue of the others"); + assert_eq!( + rescued.iter().map(|job| job.id).collect::>(), + vec![other_id] + ); assert_eq!( lease_count(&pool, &store).await, 1, - "ambiguous mutation rolls back" + "the ambiguous attempt's mutation rolls back and it stays running" ); return; } @@ -23865,7 +25291,7 @@ async fn test_historical_tombstoned_claim_uses_emitted_lane_high_to_low() { moved_tombstoned_claim_lane(4, 2, true, false, false).await; } #[tokio::test] -async fn test_historical_tombstoned_legacy_ambiguous_lanes_roll_back() { +async fn test_historical_tombstoned_legacy_ambiguous_lane_is_skipped_by_rescue() { moved_tombstoned_claim_lane(2, 4, false, false, false).await; } diff --git a/correctness/README.md b/correctness/README.md index 799f80c45..8a2059bf0 100644 --- a/correctness/README.md +++ b/correctness/README.md @@ -34,6 +34,7 @@ What is intentionally not modeled: - Python bridge details - real-time token bucket math for rate limiting - unbounded retry histories; `AwaExtended` caps attempts so TLC can close the state graph +- ready-body identity. Job bodies are abstract, so the rescue liveness property `RecoverableAbandoned(j) ~> jobState[j] # "running"` does not cover an attempt that an older release's aging race left with two equally plausible bodies; rescue skips such an attempt and logs it on every pass. - full advisory-lock mechanics; leadership is an abstract exclusive token - exact permit identity per task attempt; the model is job-centric and approximates task-held capacity closely enough to check the protocol invariants - **general priority fairness and starvation.** The focused `AwaClaimLaneScheduling` check covers one high/low lane-selection trace, not general scheduling fairness. Priority is a scheduling heuristic, not a safety invariant — all formal invariants (NoDuplicateClaim, TaskLeaseBounded, RunningHasPermit, etc.) are independent of dispatch order. Cross-priority fairness is enforced by the maintenance leader's priority aging task (see ADR-005) diff --git a/correctness/storage/MAPPING.md b/correctness/storage/MAPPING.md index d0a6b5236..55fb7b178 100644 --- a/correctness/storage/MAPPING.md +++ b/correctness/storage/MAPPING.md @@ -39,7 +39,7 @@ The TLA+ lifecycle model does not represent the completed-history rollup cache, | --- | --- | --- | | `EnqueueReady(j)` | `QueueStorage::insert_ready_rows_tx`, `QueueStorage::insert_ready_rows_copy_tx`, and SQL compatibility `awa.insert_job_compat`; producer entry points include `insert` / `insert_with`, `enqueue_batch`, `enqueue_runtime_rows`, `enqueue_params_batch`, and `enqueue_params_copy` | reserve a lane sequence range through `{schema}.reserve_enqueue_seq` under a transaction-scoped lane advisory lock, sync enqueue-time uniqueness claims, append to `{schema}.ready_entries` via INSERT or COPY, append one `{schema}.ready_segments` row per contiguous committed lane range, and notify logical queues in one tx | | `EnqueueDeferred(j)` | `QueueStorage::insert_deferred_rows_tx` and `QueueStorage::insert_deferred_rows_copy_tx`; producer entry points include `enqueue_params_batch` and `enqueue_params_copy` | allocate job ids, sync enqueue-time uniqueness claims, append to `{schema}.deferred_jobs` via INSERT or COPY in one tx | -| `PromoteDeferred(j)` | maintenance promote loop in `awa-worker/src/maintenance.rs::promote_due_state` | `DELETE FROM deferred_jobs ... INSERT INTO ready_entries ...` in one tx | +| `PromoteDeferred(j)` | `QueueStoragePromoter::promote_due` in `awa-worker/src/maintenance.rs` -> `QueueStorage::promote_due_lanes` / `promote_due_lane` -> `promote_lane_tx` | Per `(queue, priority)` lane, `DELETE FROM deferred_jobs ... INSERT INTO ready_entries ...` in one tx under a short transaction-local `lock_timeout`. A timeout rolls the lane back with its rows still deferred (a stutter step in the model) and hands the lane to a bounded waiter. Claimers are notified once per batch after the lanes commit. | | `AdvanceClaimCursor` | claim path spent-prefix advancement and post-commit `QueueStorage::advance_claim_cursors` | `claim_ready_runtime()` may advance the sequence-backed claim cursor only across a contiguous prefix of committed spent/tombstoned lanes. In receipts mode, already-emitted attempts are proved from queue-slot-local `{schema}.ready_claim_attempt_batches` range evidence so recovery does not scan claim-ring parent tables while old claim partitions are locked for prune. Newly emitted claims advance the cursor after the claim transaction commits. | | `Claim(w, j)` | `QueueStorage::claim_runtime_batch` (`queue_storage.rs`) -> `claim_runtime_batch_with_aging_for_instance` -> dispatcher (`awa-worker/src/dispatcher.rs`) | `claim_ready_runtime()`: lane selection locks `queue_claim_heads` with `FOR UPDATE SKIP LOCKED`, consults `{schema}.ready_segments` (ordered by `next_lane_seq` for an index short-circuit) to identify the committed ready slot/generation for the lane head. It then validates the target ready row, reads the `lease_ring_rotations` and `claim_ring_rotations` cursors (max-generation rows), treats `ready_tombstones` and durable attempt evidence as spent lane evidence for stale-cursor advancement, and in receipts mode INSERTs one `{schema}.ready_claim_attempt_batches_` row for the claimed lane range before inserting compact `lease_claim_batches_` evidence for the whole claim batch — with the batch's shared `deadline_at` populated for deadline-backed claims and NULL for zero-deadline claims (#246), so both go through the compact representation. Legacy mode inserts into `leases_`. The common live-head claim plan filters ready tombstones only; already-emitted attempt ranges are consulted after the head row is proved spent, where they advance a stale claim cursor instead of making every successful claim anti-join the attempt ledger. The broad legacy spent-evidence battery is retained only for non-receipts custom schemas. From v051 the claim CTE itself runs in `{schema}.claim_ready_slot_(...)`: `claim_ready_runtime` keeps lane selection, cursor catch-up and the head/spent-prefix checks, then dispatches to the slot function with the resolved lane state; the slot function addresses `ready_entries_`, `ready_tombstones_` and `ready_claim_attempt_batches_` directly and writes the same evidence as before. | | `MaterializeAttemptState(j)` | `QueueStorage::upsert_attempt_state_from_receipts_tx` and `QueueStorage::upsert_attempt_state_progress_from_receipts_tx` | `INSERT INTO attempt_state ... ON CONFLICT (job_id, run_lease) DO NOTHING` | @@ -55,12 +55,12 @@ The TLA+ lifecycle model does not represent the completed-history rollup cache, | `StatefulComplete(w, j)` | `QueueStorage::complete_runtime_batch` + `DELETE FROM attempt_state` | same as above with an observed materialized `attempt_state`; terminal reads hydrate the immutable body through `{schema}.terminal_jobs` and the retained ready row | | `FailToDlq(w, j)` | `QueueStorage::fail_to_dlq` / `QueueStorage::fail_terminal` via executor terminal failure path | `DELETE FROM leases`, `DELETE FROM attempt_state`, `INSERT INTO dlq_entries` in one tx | | `RetryToDeferred(w, j)` | `QueueStorage::retry_after` / `QueueStorage::snooze` on `JobError::RetryAfter` / `Snooze` | `DELETE FROM leases`, `INSERT INTO deferred_jobs` | -| `RescueToReady(j)` | lease-side `rescue_stale_heartbeats` / `rescue_expired_deadlines` in maintenance, plus receipt-side deadline rescue once it has force-closed an open receipt, when the rescued attempt was not the job's last | lease-side paths `DELETE FROM leases ... RETURNING ...`, then `hydrate_deleted_leases_tx` (drop `attempt_state`, close the receipt) and `settle_rescued_attempt_tx` inserts `deferred_jobs` (retryable, with backoff); the spec collapses the deferred hop into the ready lane. Receipt-side deadline rescue first appends a `deadline_expired` closure row in `lease_claim_closures` and then takes the same path | +| `RescueToReady(j)` | lease-side `rescue_stale_heartbeats` / `rescue_expired_deadlines` in maintenance, plus receipt-side deadline rescue once it has force-closed an open receipt, when the rescued attempt was not the job's last | lease-side paths `DELETE FROM leases ... RETURNING ...`, then `hydrate_deleted_leases_tx` (drop `attempt_state`, close the receipt) and `settle_rescued_attempt_tx` inserts `deferred_jobs` (retryable, with backoff); the spec collapses the deferred hop into the ready lane. Receipt-side deadline rescue first appends a `deadline_expired` closure row in `lease_claim_closures` and then takes the same path. Each candidate is hydrated inside a savepoint: an attempt whose claim heads are busy is rolled back and deferred to a later pass, and an attempt whose retained ready bodies are ambiguous (legacy aging data) is skipped and logged on every pass. | | `RescueToDlq(j)` | the same rescue paths via `rescue_stale_heartbeats_with_dlq` / `rescue_expired_deadlines_with_dlq`, when the rescued attempt was the job's last (`attempt >= max_attempts`) and the queue's DLQ is enabled | same lease delete, hydrate and receipt closure, then `settle_rescued_attempt_tx` inserts the `dlq_entries` row (reason `heartbeat_stale` / `deadline_exceeded`) in that same transaction, as `FailToDlq` does | | `RescueToTerminal(j)` | the same exhausted rescue on a queue without a DLQ | same lease delete, hydrate and receipt closure, then `settle_rescued_attempt_tx` inserts a `failed` `done_entries` row and its terminal-count delta | | `CancelWaitingToTerminal(j)` | waiting branch in `QueueStorage::cancel_job_tx` | `DELETE FROM leases WHERE state IN ('running', 'waiting_external')`, hydrate from `ready_entries`, insert `done_entries`, close the matching receipt | | `CancelReadyToTerminal(j)` | available branch in `cancel_job_tx` | `SELECT ready_entries ... FOR UPDATE SKIP LOCKED`, append `ready_tombstones`, insert `done_entries(cancelled)`, and post-commit advance the claim cursor only if the cancelled lane was the current head. The ready row remains until queue prune. | -| `ReprioritizeReady(j)` | `QueueStorage::age_waiting_priorities` and batch `set_priority` via `QueueStorage::set_priority_tx` / `move_ready_batch_fields_tx` | `SELECT ready_entries ... FOR UPDATE SKIP LOCKED`, append `ready_tombstones` for the old lane, append a new `ready_entries` row at the target priority, and notify the queue. Batch operations snapshot job ids in `awa.batch_operation_items`; item status and row mutation commit in one transaction so resume does not replay a completed item. | +| `ReprioritizeReady(j)` | `QueueStorage::age_waiting_priorities` and batch `set_priority` via `QueueStorage::set_priority_tx` / `move_ready_batch_fields_tx` | Lock the affected `queue_claim_heads` rows in sorted order, recheck that no claim has emitted the row, `SELECT ready_entries ... FOR UPDATE SKIP LOCKED`, append `ready_tombstones` for the old lane, append a new `ready_entries` row at the target priority, and notify the queue. Batch operations snapshot job ids in `awa.batch_operation_items`; item status and row mutation commit in one transaction so resume does not replay a completed item. | | `MoveQueueReady(j)` | batch `move_queue` via `QueueStorage::move_queue_tx` / `move_ready_batch_fields_tx` | Same abstract storage transition as `ReprioritizeReady`: tombstone source ready lane, append destination ready row with rewritten queue/optional priority, and notify the destination queue. Queue labels and cross-queue routing are outside `AwaSegmentedStorage`'s one-lane abstraction; lock-order implications are covered by `AwaStorageLockOrder`. | | `StaleCompleteRejected(w, j)` | `complete_runtime_batch` returning `CompletionOutcome::IgnoredStale` | `UPDATE leases ... WHERE run_lease = $2` matching 0 rows | | `MoveFailedToDlq(j)` | `QueueStorage::move_failed_to_dlq`; admin entry in `awa-model/src/dlq.rs` | hydrate from retained ready body, `DELETE FROM done_entries`, append a negative `queue_terminal_count_deltas` row, and `INSERT INTO dlq_entries` guarded by state=failed. The retained ready row remains until queue prune. | diff --git a/docs/architecture.md b/docs/architecture.md index d3c527f61..9fc093412 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -147,7 +147,7 @@ sequenceDiagram Enqueue is transactional: if the producer's outer transaction rolls back, the job never becomes visible. Immediate jobs reserve a lane sequence range under a transaction-scoped lane lock, append to `ready_entries_*`, and append a compact `ready_segments` range for the committed lane sequence range; future scheduled or retryable jobs append to `deferred_jobs`. The lane lock is held until commit or rollback so a later producer cannot make lane `N+1` visible while lane `N` can still commit. -`deferred_jobs` is deliberately outside the claim path. The maintenance leader promotes due `scheduled` and `retryable` rows into `ready_entries_*` in batches and notifies the target queues. Cron schedules are just producers for ordinary jobs: when a schedule fires, the cron transaction records the fire and enqueues the job atomically, using `ready_entries_*` for immediate fires or `deferred_jobs` when the enqueue carries a future `run_at`. +`deferred_jobs` is deliberately outside the claim path. The maintenance leader promotes due `scheduled` and `retryable` rows into `ready_entries_*` in batches and notifies the target queues. Re-enqueueing a row takes its lane lock, which producers hold until they commit, so each `(queue, priority)` lane's share of a batch is promoted in its own transaction under a short `lock_timeout`. A lane that times out is handed to a dedicated waiter task, which queues for the lock behind the producers already waiting, so a busy lane is still promoted in turn without holding up other lanes or the leader's rescue, rotation and prune. Cron schedules are just producers for ordinary jobs: when a schedule fires, the cron transaction records the fire and enqueues the job atomically, using `ready_entries_*` for immediate fires or `deferred_jobs` when the enqueue carries a future `run_at`. There are two COPY-shaped producer paths: diff --git a/docs/configuration.md b/docs/configuration.md index f4e36389f..cc9b035fd 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -104,7 +104,7 @@ The key `QueueConfig` fields: | --- | --- | --- | | `max_workers` | `50` | Always — this is your concurrency cap per queue | | `rate_limit` | `None` | External API rate limits, backpressure | -| `deadline_duration` | `5m` | Hard upper bound on a single attempt. Set to `Duration::ZERO` to skip the deadline rescue path; receipts mode (the 0.6 default storage) supports both shapes — the deadline lands on `lease_claims.deadline_at` and the maintenance rescue path force-closes expired claims. | +| `deadline_duration` | `5m` | Hard upper bound on a single attempt. Set to `Duration::ZERO` to skip the deadline rescue path; receipts mode (the 0.6 default storage) supports both shapes — the deadline lands on the claim batch's shared `lease_claim_batches.deadline_at` and the maintenance rescue path force-closes expired claims. | | `poll_interval` | `200ms` | Tune if NOTIFY latency matters (rare) | | `min_workers` / `weight` | `0` / `1` | Only in weighted mode | | `claimers` | `1` | Hot queue-storage queues that need more than one dispatcher/claimer loop inside a single runtime. Claimers share the queue's worker permits. | @@ -245,6 +245,7 @@ Retry backoff and `Snooze` use the same mechanism: the current attempt closes, a Operational knobs: - `promote_interval` controls how often maintenance checks for due deferred rows. +- A queue whose enqueue lane is kept busy by producers does not delay promotion of other queues or the leader's rescue: promotion waits 50 ms for a lane's lock, then hands the lane to a dedicated waiter that is served in turn. Each waiter holds one connection from the worker's pool while it waits, on top of the leader's own. The leader runs at most one waiter per eight pool connections, never fewer than one or more than four, so a 10-connection pool gives up at most one extra connection to a busy lane. `awa_maintenance_promote_lane_contention_total` counts those hand-offs by queue. A steady `waiter_limit` rate means more lanes are saturated than the waiters cover, and a larger pool allows more waiters. `waiter_timed_out` means a transaction held a lane lock for longer than 30 s. - Per-job `run_at` controls the due time. - Periodic schedules are declared in worker code with `periodic()` and managed through the CLI/UI `cron` surface. @@ -474,7 +475,7 @@ A handler that ignores cancellation and is still running two `deadline_rescue_in | `QueueConfig.deadline_duration` (Rust) / `deadline_duration_ms` (Python dict form) | `5m` | Per-queue hard upper bound on one attempt. `Duration::ZERO` / `0` skips deadline rescue for that queue. | | `ClientBuilder::deadline_rescue_interval` (Rust) / `deadline_rescue_interval_ms` (Python kwarg) | `30s` | How often the maintenance leader scans for expired deadlines. | -Receipts mode (the 0.6 default storage) supports both shapes: the deadline lands on `lease_claims.deadline_at` and is rescued there for short claims, or carried onto `leases.deadline_at` if the claim materializes for a long-running attempt. See [Queue storage tuning](#queue-storage-tuning) and ADR-023. +Receipts mode (the 0.6 default storage) supports both shapes: the deadline lands on the claim batch's shared `lease_claim_batches.deadline_at` and is rescued there for short claims, or carried onto `leases.deadline_at` if the claim materializes for a long-running attempt. See [Queue storage tuning](#queue-storage-tuning) and ADR-023. ### Callback timeout — bounding `wait_for_callback` { #callback-timeout } @@ -577,7 +578,7 @@ await client.start( | `lease_slot_count` | `8` | Number of rotating lease partitions | | `claim_slot_count` | `8` | Number of rotating ADR-023 claim-ring partitions (`lease_claims`, `lease_claim_closures`, and `lease_claim_closure_batches` children). All three tables share the same `claim_slot` so each partition's claims and closure evidence are reclaimed together by `TRUNCATE`; compact closure batches also carry ready segment metadata so prune can prove sealed segments by counts and skip conservatively when counts do not prove closure. Each slot has tiny stale-rescue and deadline-rescue cursors in `claim_ring_slots`. | | `queue_stripe_count` / `queue_storage_queue_stripe_count` | `1` | Number of physical stripes behind each logical queue. `1` is the normal unstriped path. For a single very hot queue on many small replicas, `2` is the current tuned recommendation; higher values should be benchmarked before use. | -| `lease_claim_receipts` | `true` | Use the receipt-plane short path (claim writes compact attempt-emission range evidence into `ready_claim_attempt_batches` and receipt evidence into row-local `lease_claims` or compact `lease_claim_batches`; successful compact completion validates the returned receipt identity, writes `receipt_completion_batches` terminal history, and writes `lease_claim_closure_batches` closure evidence; non-success exits write explicit closures into `lease_claim_closures`; receipt tables are reclaimed by claim-ring prune, and attempt-emission evidence is reclaimed by queue-ring prune). Receipts mode supports per-claim deadlines: when `QueueConfig.deadline_duration > 0`, the claim writes `clock_timestamp() + interval` onto `lease_claims.deadline_at` and the maintenance rescue path force-closes expired claims with a `'deadline_expired'` closure. Set to `false` to force every claim through the legacy `leases` materialization path. See ADR-023. | +| `lease_claim_receipts` | `true` | Use the receipt-plane short path (claim writes compact attempt-emission range evidence into `ready_claim_attempt_batches` and receipt evidence into row-local `lease_claims` or compact `lease_claim_batches`; successful compact completion validates the returned receipt identity, writes `receipt_completion_batches` terminal history, and writes `lease_claim_closure_batches` closure evidence; non-success exits write explicit closures into `lease_claim_closures`; receipt tables are reclaimed by claim-ring prune, and attempt-emission evidence is reclaimed by queue-ring prune). Receipts mode supports per-claim deadlines: when `QueueConfig.deadline_duration > 0`, the claim writes `clock_timestamp() + interval` onto the batch's `lease_claim_batches.deadline_at` and the maintenance rescue path force-closes expired claims with a `'deadline_expired'` closure. Set to `false` to force every claim through the legacy `leases` materialization path. See ADR-023. | | `queue_rotate_interval` | `1000ms` | How often ready/terminal segments rotate | | `lease_rotate_interval` | `1000ms` | How often lease segments rotate | | `claim_rotate_interval` | matches `queue_rotate_interval` | How often the ADR-023 claim-ring rotates. Set with `ClientBuilder::claim_rotate_interval` (Rust) or `queue_storage_claim_rotate_interval_ms` (Python). Test harnesses sometimes set this to a long interval to pin claim-ring layout for deterministic count assertions. | @@ -597,9 +598,13 @@ Use the defaults unless you have a reason not to: - Increase `claim_slot_count` if the rotation cadence (`claim_rotate_interval`) plus the slot count combine to a partition retention window shorter than your longest in-flight zero-deadline short job; running out of empty slots forces `rotate_claims` to return `SkippedBusy` and the receipt-plane churn falls back onto a smaller working set of partitions. - Increase `queue_stripe_count` only for measured workloads where many small replicas contend on the same logical queue. This is queue-storage configuration, not a per-queue `QueueConfig` field: it gives each logical queue the same number of physical stripes. Striping spreads a hot logical queue over `queue#N` physical coordination paths, but it weakens perfect global ordering and can regress calmer shapes if overused. For single hot-queue workloads, benchmark `2` stripes first; higher values should be treated as workload-specific tuning rather than a general default. - Increase `prune_lock_timeout` when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}[15m])` stays above zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Among its initial partition locks, queue ring prune waits under the timeout only on `ready_segments` and takes the others with `NOWAIT`, so a longer timeout does not help when readers hold those. Later reads and writes also run under the timeout. Raise it in small steps (for example `50ms`, then `100ms`, then `250ms`) and keep it well below PostgreSQL's `deadlock_timeout` (default `1s`) and below any `statement_timeout` on the pool's connections. A single lock wait that reaches `deadlock_timeout` makes PostgreSQL run its deadlock detector, and if claims or enqueues are part of a lock cycle with it, the detector can cancel them rather than the prune; a short `lock_timeout` abandons the prune first, so the prune is always the one that gives way. A `statement_timeout` at or below the lock timeout cancels the prune statement before the lock timeout fires, so it fails with an error instead of reporting `blocked`. Startup logs a warning when the configured timeout is at or above either setting. While prune waits, new accesses to the same partition also queue behind its lock request, so a longer timeout trades claim latency for reclamation. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for a few ticks after repeated `skipped_active`, and `error` is a database failure. -- Increase rotation intervals to reduce partition churn and metadata activity. +- Consider increasing rotation intervals before raising `prune_lock_timeout`: longer intervals can reduce how often prune requests exclusive locks, while a longer timeout lets each request wait longer. Confirm the effect with workload measurements. - Decrease rotation intervals to tighten dead-tuple bounds at the cost of more frequent rotate/prune work. +Longer rotation intervals accumulate more work in each slot. `queue_rotate_interval` is how often the leader decides whether to rotate, not how long a slot stays open: a non-empty slot rotates at the first decision after a `ring_rotation_cadence` gate passes, so with the default cadence it stays open about 5 s once it holds 256 rows, less if it reaches 2 MiB first, and up to 60 s under a trickle of writes. Only with `RingRotationCadence::EAGER` does the slot rotate on every decision. At a steady enqueue rate, ready rows per queue slot are roughly the enqueue rate multiplied by that rotation period; retries and blocked recycling can increase the retained footprint. Budget for larger partitions and indexes. Queue-ring history becomes coarser and generally spans about `queue_slot_count` rotation periods; active or blocked slots can retain it longer. The independent `claim_rotate_interval` changes claim-evidence retention, not completed-job history in the ready/terminal ring. + +A longer interval does not remove the locks required by each prune. `prune_lock_timeout` still applies separately to each acquisition; queue prune can lock up to eight partitions before its later reads and writes. + ### Sharding the enqueue head per queue Each queue's `awa.queue_meta.enqueue_shards` (default `1`, range `1..=64`) governs how many independent enqueue head rows the queue has. With the default value, the queue uses a single head row and the ordering contract is strict FIFO per `(queue, priority)` — identical to v016. diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 68c0e0ae6..f1d031b28 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -254,6 +254,7 @@ In Prometheus they appear under the OpenTelemetry [OTLP-to-Prometheus translatio | `awa.maintenance.promote_batches` | Counter | state | Promotion batches | | `awa.maintenance.promote_batch_size` | Histogram | state | Jobs promoted | | `awa.maintenance.promote_duration` | Histogram (s) | state | Promotion time | +| `awa.maintenance.promote.lane_contention` | Counter | state, queue, outcome | Promotion lanes whose lock was busy: `handed_off` to a dedicated waiter, `waiter_limit` when every waiter was busy, `waiter_timed_out` when a waiter's 30 s wait ran out | | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | | `awa.maintenance.prune.attempts` | Counter | ring, outcome, reason | Ring prune outcomes, including destructive `pruned` and no-DDL `already_pruned` | | `awa.maintenance.prune.duration` | Histogram (s) | ring, phase | Successful destructive prune time split into `lock`, `truncate`, and `commit` | diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index 2803a20b5..8863fc1a5 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -65,7 +65,7 @@ If an existing installation created objects as `awa_migrator` without first sett ## Lane sequences -Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. From schema v053 they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after checking that the queue, priority and shard are set and the shard is 0–63. A runtime therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. From schema v053 they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after checking that the queue, priority and shard are set and the shard is 0–63. A runtime therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. Their inserts into `queue_lanes`, `queue_enqueue_heads` and `queue_claim_heads` fire those tables' triggers as the schema owner, so don't grant `TRIGGER` on them to any other role. A new lane sequence gets the privileges of the head table it backs. A role that may `INSERT` or `UPDATE` `queue_enqueue_heads` gets `USAGE, SELECT, UPDATE` on the enqueue sequences, and one that may only `SELECT` it gets `SELECT`; `queue_claim_heads` does the same for the claim sequences. Granting the runtime the table privileges above is enough, and no sequence default privileges are needed for lanes. diff --git a/docs/test-plan.md b/docs/test-plan.md index b35c0fb7f..da0a6bc6f 100644 --- a/docs/test-plan.md +++ b/docs/test-plan.md @@ -31,6 +31,7 @@ Tests run against real Postgres 15+ (not managed services) using a dedicated tes | B5 | Deadline rescue signals ctx.is_cancelled() | ✓ | | | B6 | UniqueConflict.constraint field | ✓ | | | B7 | Admin cancel signals in-flight handler cancellation | ✓ | ✓ | +| B8 | Shutdown finalizes attempts whose handler returned during the drain, so a process exit cannot leave them for rescue | ✓ | | | T75 | Priority aging maintenance task promotes long-waiting low-priority jobs | ✓ | | ### Uniqueness (Python) @@ -331,6 +332,75 @@ V25 adds a focused `AwaKeyedExecution` model for grant safety, partial-legacy co | V30 | Partitioned-queue preset ≥9k jobs/s e2e on the reference 24-CPU harness | Not a 0.7 gate (moved to 0.8) | | V31 | Gate A evidence: allocator bake-off (E1) + WAL/job decomposition (E3) run ids recorded on #295 | #295 | +## Persona soak + +Durable loss detection covers every lifecycle-observed job and a deterministic `job_id % 64 = 0` sample of `NO_EVENTS` workloads. The other 63/64 returned `NO_EVENTS` jobs remain in enqueue/execution accounting, but their terminal commits are not proven. A bulk observer polls sampled terminal rows and eventful attempts whose handlers returned without terminal proof once per second (2 s query timeout, 100,000-row cap) and keeps the run lease in the client ledger, so captured proof survives pruning. A handler return is never terminal proof. Missing sampled proof after pruning or a failed invariant query makes the report `INCONCLUSIVE` and fails the test; it does not assert job loss or establish a pass. Reports include sampled, confirmed and excluded counts. + +On CI, assertion waits and heartbeat staleness use `ci_timing.rs` scaling; load durations, kill schedules and measurement windows retain their configured values. + +`awa/tests/persona_soak_test.rs` (harness in `awa/tests/persona_soak/`, Python side in `awa/tests/persona_soak/python/`) runs several application-shaped workloads concurrently against one database, the way different applications share an `awa` schema, and fails the run on any invariant violation. + +| Persona | Workload | Stresses | +| --- | --- | --- | +| Mia, ML pipeline | jobs of 0.1-60 s with per-job deadlines, progress updates, retries into the DLQ, handler cancels, latency spikes | deadline and stale-heartbeat rescues racing heartbeats and completions, DLQ | +| Tom, multi-tenant | 100 queues at a low rate each, unique jobs, scheduled jobs, cron every 10 s on two runtimes, mixed priorities with aging | lane and queue cardinality, per-queue maintenance, cron fencing | +| Priya, firehose | one hot queue, 100k-row `insert_many` bursts, COPY batches, steady small batches, tiny jobs, three runtimes | bind-parameter limit, ring rotation and prune, claim contention | +| Sam, ordered streams | a queue with 4 enqueue shards and 64 ordering keys, a 4-stripe queue (`queue_stripe_count`), scheduled jobs, a rate-limited queue | routing, promotion, fairness | +| Alex, operator | pause/resume, drain, cancel and bulk cancel, move and reprioritise batch operations, DLQ retry and purge, the admin read queries, rolling restarts of other personas' runtimes | admin paths under live load | +| Chaos | kills a runtime or the Python process without a shutdown handshake, terminates the maintenance leader's connection and random runtime connections, exhausts a runtime's pool, holds an idle-in-transaction session | recovery paths | +| Watchers | six dashboard tabs polling overview, queues, counts, timeseries, runtime and storage at the UI's refresh interval, some also the kinds, queue-runtime and cron pages | admin read queries under load, through the real `awa serve` HTTP server | +| Support engineer | searches jobs by kind, queue, state and tag, pages through large result sets with `before_id`, opens job detail and follows a job at the detail page's poll rate, pages the DLQ, opens runtime, cron and kind pages | search and paging cost against retained history and backlog | +| UI operator | the UI's mutations over HTTP: pause and resume queues and cron schedules, cancel and bulk-cancel, DLQ retry and purge (single and by filter), batch operation preview and submit | mutation endpoints racing live work | +| Impatient user | requests abandoned after a few milliseconds, bursts of refreshes, while sampling the server's database sessions | an abandoned request must not keep a pool connection past the statement timeout | +| Fanout, replay fan-out | parent jobs whose handlers batch-insert 100-185 unique child jobs (`insert_many`, unique by queue and args) into their own queue from many concurrent handlers; some parents replay a batch with an already-claimed key, hit the unique violation and fall back to row-by-row inserts | concurrent batch enqueue into one lane, unique claims, inserting from inside a handler | +| Shard operator | an operator changes `enqueue_shards` on two hot queues while producers, promotion and claimers run: 1 → 4 → 8, down to 2 and back to 1, sometimes after `prepare_queue` and sometimes not, restarting some runtimes and producer `QueueStorage` handles before others so cached shard counts disagree. Producers use `insert_many`, `insert_many_copy` and `QueueStorage::enqueue_params_batch`, with scheduled jobs, retries (promotion into shards) and `ordering_key` jobs. The actors run as non-superuser roles: a schema owner that migrates and runs `prepare_queue`, a migrator role with default privileges, and two application roles in a runtime group role with no default privileges, granted as in [Database roles](security/database-roles.md). For some raises the producers are held until the maintenance leader's runtime has created the new lanes, so a different role is the next to use them | live shard-count changes (ADR-025), lane provisioning by different roles, stale per-process shard counts, promotion into new shards | +| Claim and abandon | a separate process (the harness binary started again with one test selected) that claims as much as it can: `max_workers` equal to the target on two queues, one with zero deadlines (compact claim batches) and one with a deadline, jobs held until the process dies, heartbeats running. It is killed with SIGKILL, repeatedly, and replaced after a few seconds by a process that claims the rescued jobs again; some jobs have `max_attempts` 2 so repeated kills dead-letter them. A surviving in-process runtime leads maintenance and serves a canary queue of short jobs, and the other personas keep working; after the last kill an in-process runtime finishes the rescued jobs | rescue of a large abandoned-claim backlog, claim-ring rotation while slots hold abandoned claims, false rescues of live owners, `max_attempts` and the DLQ under repeated rescue | +| Python users | a separate process using `awa-python`: single inserts, a 6,500-row `insert_many` in one transaction and a COPY batch, enqueues in transactions that sometimes roll back; async workers that complete, raise, `RetryAfter`, `Snooze`, report progress, overrun a queue deadline and ignore their deadline; a queue shared with a Rust runtime; a periodic migrate-then-query probe | Python bindings against live load, mixed Rust and Python runtimes on one lane, the documented slot behaviour of an ignored deadline, connections stranded after `migrate()` | + +The oracle checks, while the load runs and again after it stops: + +- a job attempt with more than one terminal row (`done_entries`, `dlq_entries`, compact receipt completions), a job reported completed twice, and a job that is terminal and still available, deferred or leased +- a cancelled job nobody asked to cancel, an operator cancel that overwrote a completion, and a handler that started after its operator cancel committed +- a queue with claimable jobs and idle workers whose claim cursor stops advancing (stall), and scheduled jobs overdue beyond a threshold (promotion stall) +- a lost job: enqueued, never terminal after the system quiesces, and not removed by an operator action; a cron fire that produced more than one job +- shard operations: every job enqueued before, during and after a shard-count change runs and completes exactly once (a repeat is allowed only after a stale-heartbeat rescue and is reported as a finding), every shard drains including those above a lowered count, claim cursors keep advancing on every shard that has backlog while workers are idle (the per-shard lag query from ADR-025), `ordering_key` jobs start in order per key within a window in which every runtime had observed the shard count, and no role raises a permission-denied or missing-lane error (found in the server log, with the role, lane, sequence owner and statement; needs `AWA_SOAK_PG_CONTAINER` or `AWA_SOAK_PG_LOG`, and `log_line_prefix` containing `%u@%a`, as `scripts/persona-soak.sh` sets) +- claim and abandon: after each SIGKILL the rescue backlog (open claims whose last heartbeat or claim time is older than the staleness window, or whose deadline passed, counted every few seconds from the claim tables and leases) clears within `AWA_SOAK_ABANDON_BOUND_SECONDS`, with the age and count of the oldest unrescued expired claim sampled into `abandon-backlog.csv` and the time to clear reported per kill; the claim ring does not stay still for longer than the bound once the slot it must enter next holds no claim whose owner is alive (a ring pinned by a live process is expected); no job is run again while the process that held it is alive (every start is recorded with the process that ran it, and a repeat is allowed only after that process was killed); no attempt exceeds `max_attempts`, no job is dead-lettered early, and a job abandoned on its last attempt ends in the DLQ +- a transaction the Python application rolled back that left a job in the store, and database sessions left stranded by the Python `migrate()` path +- an admin-server query that outlives its statement timeout, and sessions still running a query after every UI client has stopped +- a server crash: `pg_postmaster_start_time()` changing, or `terminated by signal`, `TRAP:`, `server process (PID` and similar lines in the Postgres log + +Slow or failing admin endpoints, latency that grows with retained history or backlog, and admin queries cancelled by the statement timeout are reported as non-fatal findings in the report, next to the violations. + +Terminal rows are pruned by ring rotation within seconds, so the oracle keeps a client-side ledger (job ids per persona, executions per job, lifecycle events, operator actions) and reconciles it with the database. Time-based checks are widened while a chaos action is recovering. + +Environment: + +| Variable | Meaning | +| --- | --- | +| `AWA_SOAK_DATABASE_URL` | required; database name must match `awa_soak[A-Za-z0-9_]*` and is dropped and recreated | +| `AWA_SOAK_PERSONAS` | comma list of `mia,tom,priya,sam,alex,chaos,watchers,support,uiops,impatient,python,fanout,shardops,abandon`, default all | +| `AWA_SOAK_ABANDON_CLAIMS`, `AWA_SOAK_ABANDON_KILL_SECONDS`, `AWA_SOAK_ABANDON_KILLS` | the claim-and-abandon persona's scale, independent of `AWA_SOAK_SCALE`: claims each killed process holds (default 4000; 25,000 to 70,000 reproduces a large backlog), seconds between kills (default 180, long enough for the process to claim everything and heartbeat for a while) and how many processes are killed (default 4; the last kill is the end of the load phase if that comes first) | +| `AWA_SOAK_ABANDON_BOUND_SECONDS`, `AWA_SOAK_ABANDON_RESTART_SECONDS` | the bound after a kill within which the rescue backlog must clear and the claim ring must be free to rotate (default 120), and the pause before the replacement process starts (default 3) | +| `AWA_SOAK_HEARTBEAT_SECONDS`, `AWA_SOAK_STALENESS_SECONDS` | heartbeat interval (default 2) and heartbeat staleness after which a claim may be rescued (default 12) of every runtime, including the claim-and-abandon process; the library's own defaults are 30 and 90. The claim-and-abandon persona also records the duration of each heartbeat pass of its process against the number of handlers running, read from the timestamps the heartbeats write (`abandon-heartbeat-passes.csv`, and percentiles by handler count in the report) | +| `AWA_SOAK_RESCUE_INTERVAL_SECONDS` | heartbeat and deadline rescue interval of every runtime; unset keeps the harness's 3 s and 2 s, the library default for heartbeat rescue is 30 s | +| `AWA_SOAK_SHARDOPS_SCALE`, `AWA_SOAK_SHARDOPS_SCHEMA_CREATE` | the shard operator's rate multiplier, independent of `AWA_SOAK_SCALE` (default 1), and whether its application roles may create objects in the schema (default 1: lanes are created on the fly; 0 withholds schema `CREATE` and the operator provisions every lane before raising a count) | +| `AWA_SOAK_SHARED_CREATOR_GRANTS` | default `0`; `1` shares application-created objects for comparisons with older binaries and requires schema `CREATE`. Keep it off when checking application-role privileges. | +| `AWA_SOAK_UI_POOL_MAX`, `AWA_SOAK_ADMIN_READ_SECONDS` | the admin server's database pool size (default 10) and the seconds between the operator persona's admin read queries (default 2); lower the load on a slow server by shrinking the pool and stretching the interval | +| `AWA_SOAK_AWA_BIN` | the `awa` binary for the UI personas (`cargo build --release -p awa-cli`); defaults to the one next to the test binary | +| `AWA_SOAK_PYTHON` | interpreter of a virtualenv with `awa-python` installed (`uv venv awa-python/.venv`, `uv pip install maturin`, `maturin develop --release` in `awa-python/`); required by the `python` persona | +| `AWA_SOAK_UI_STATEMENT_TIMEOUT_SECONDS`, `AWA_SOAK_UI_CLIENT_TIMEOUT_SECONDS` | the admin server's statement timeout (default 30) and the UI personas' request timeout (default 45); a request that hits either is reported as a finding | +| `AWA_SOAK_SECONDS` | load phase length, default 900 | +| `AWA_SOAK_SEED` | fixes the workload mix and chaos schedule; printed at start | +| `AWA_SOAK_SCALE` | multiplies rates, burst sizes and queue counts, default 1 | +| `AWA_SOAK_DRAIN_GRACE_SECONDS` | how long to wait for the system to quiesce after load stops, default 120 | +| `AWA_SOAK_PG_LOG` / `AWA_SOAK_PG_CONTAINER` / `AWA_SOAK_PG_CONTAINER_LOG` | where to read the server log for the crash scan: a file on this host, or a Docker container (with an optional log path inside it, otherwise `docker logs`) | +| `AWA_SOAK_REPORT_DIR` | where `report.md`, `report.json` and, for a failed run, `reproducer.md` are written | +| `AWA_SOAK_STALL_SECONDS`, `AWA_SOAK_PROMOTION_OVERDUE_SECONDS`, `AWA_SOAK_CHAOS_SLACK_SECONDS` | oracle thresholds | + +The report records the engine, server version, `max_connections`, connections at start and the per-persona connection budget. A failed run leaves `reproducer.md` with the seed, personas and duration, the violating rows, the timeline around each violation and the relevant server log excerpt. `scripts/persona-soak.sh` runs the engine matrix (`postgres:17`, `postgres:18`, and an assert-enabled PostgreSQL 18 image) in private containers on separate ports and removes them afterwards. + +After-exit retained-session findings require a successful `CLOSED` acknowledgement and a valid final session observation. Forced or unproven exits remain incomplete evidence. Pool-budget and idle-transaction checks after readiness remain independent. + ## Running Tests CI runs the Rust suite as a sharded matrix (`scripts/ci-test-shard.sh`): the diff --git a/scripts/persona-soak-against.sh b/scripts/persona-soak-against.sh new file mode 100755 index 000000000..ef6838238 --- /dev/null +++ b/scripts/persona-soak-against.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# Runs the persona soak against another revision of awa. +# +# scripts/persona-soak-against.sh [engine ...] +# +# The harness lives entirely in awa/tests, so any revision that builds with the same public +# API can be soaked: a detached worktree of is created, this checkout's harness +# files (including uncommitted changes) are copied over it, scripts/persona-soak.sh runs +# there with the arguments after the ref, and the worktree is removed afterwards. Results go +# to this checkout's soak-results/- unless SOAK_RESULTS_DIR is set. +# +# Environment is passed through to scripts/persona-soak.sh. Build outputs go to +# CARGO_TARGET_DIR when it is set, so builds of several revisions share one cache; the +# harness files are copied with a fresh modification time so cargo rebuilds the test binary +# rather than reusing one built from another revision. +set -euo pipefail + +if [ $# -lt 1 ]; then + echo "usage: $0 [engine ...]" >&2 + exit 2 +fi +ref="$1" +shift + +here="$(cd "$(dirname "$0")/.." && pwd)" +label="$(printf '%s' "$ref" | tr -c 'A-Za-z0-9._-' '_')" +export SOAK_RESULTS_DIR="${SOAK_RESULTS_DIR:-$here/soak-results/$(date +%Y%m%d-%H%M%S)-$label}" +worktree="$(mktemp -d "${TMPDIR:-/tmp}/awa-soak-against.XXXXXX")" +cleanup() { git -C "$here" worktree remove --force "$worktree" >/dev/null 2>&1 || rm -rf "$worktree"; } +trap cleanup EXIT + +git -C "$here" worktree add --detach "$worktree" "$ref" >/dev/null +tar -C "$here" -c awa/tests/persona_soak awa/tests/persona_soak_test.rs awa/tests/ci_timing.rs \ + scripts/persona-soak.sh \ + | tar -C "$worktree" -x -m +echo "== persona soak against $ref ($(git -C "$worktree" rev-parse --short=10 HEAD))" +(cd "$worktree" && scripts/persona-soak.sh "$@") diff --git a/scripts/persona-soak.sh b/scripts/persona-soak.sh new file mode 100755 index 000000000..b1f6a079c --- /dev/null +++ b/scripts/persona-soak.sh @@ -0,0 +1,213 @@ +#!/usr/bin/env bash +# Runs the persona soak (awa/tests/persona_soak_test.rs) against a matrix of +# throwaway PostgreSQL containers, one engine after another. +# +# scripts/persona-soak.sh [engine ...] +# +# Engines: pg17 (postgres:17), pg18 (postgres:18) and pg18-cassert (an +# assert-enabled, -O0 build of PostgreSQL 18; set SOAK_CASSERT_IMAGE). With no +# arguments all three run in that order. +# +# Environment: +# AWA_SOAK_SECONDS load phase length per engine (default 900) +# AWA_SOAK_PERSONAS comma list, default all of mia,tom,priya,sam,alex,chaos, +# watchers,support,uiops,impatient,python,fanout,shardops,abandon +# AWA_SOAK_SEED fixes the workload mix and chaos schedule +# AWA_SOAK_SCALE multiplies rates and queue counts (default 1) +# AWA_SOAK_SHARDOPS_SCALE multiplies the shardops persona's rates on its own (default 1) +# AWA_SOAK_SHARDOPS_SCHEMA_CREATE +# 1 (default) lets the shardops application roles create objects in +# the schema, so lanes are created on the fly; 0 withholds schema +# CREATE and has the operator provision lanes before raising a count +# AWA_SOAK_ABANDON_CLAIMS claims held by each killed process of the abandon persona (default 4000; +# 25000 or more reproduces a large abandoned-claim backlog) +# AWA_SOAK_ABANDON_KILLS, AWA_SOAK_ABANDON_KILL_SECONDS +# how many processes are SIGKILLed (default 4) and the seconds between +# kills (default 180) +# AWA_SOAK_ABANDON_BOUND_SECONDS +# seconds after a kill within which the rescue backlog must clear and the +# claim ring must be free to rotate (default 120) +# AWA_SOAK_RESCUE_INTERVAL_SECONDS +# heartbeat and deadline rescue interval of every runtime (default: the +# harness's 3 s and 2 s; the library default for heartbeat rescue is 30 s) +# AWA_SOAK_HEARTBEAT_SECONDS, AWA_SOAK_STALENESS_SECONDS +# heartbeat interval and heartbeat staleness of every runtime (default 2 s +# and 12 s, so rescues happen often; the library's are 30 s and 90 s) +# SOAK_RESULTS_DIR output directory (default ./soak-results/) +# SOAK_BASE_PORT first host port; engines use consecutive ports (default 15471) +# SOAK_MAX_CONNECTIONS server max_connections (default 400) +# SOAK_PG_CPUS CPU limit for each database container (default: none) +# SOAK_CASSERT_IMAGE image for pg18-cassert (default awa-pg18-cassert:18.6-O0; the +# -epqlog5 build, which logs every partition-pruning setup with its +# query text, is opt-in for diagnosing pruning crashes) +# SOAK_SKIP_PREPARE set to 1 to skip building the awa binary and the Python bindings +# (neither is needed without the UI and python personas) +# SOAK_TEST_BIN a prebuilt harness binary to run instead of `cargo test` +# CARGO_TARGET_DIR cargo target directory, as for any cargo build +# +# The admin UI personas drive the real `awa serve` binary and the python persona runs +# the awa-python bindings in a separate process; both are built once, before the first +# engine, with `cargo build --release -p awa-cli` and `uv` + `maturin develop --release` +# into awa-python/.venv. +# +# Every container is private to the run (name awa-soak--) and is +# removed on exit, including on failure or interruption. Containers that already +# exist are never touched. +set -uo pipefail + +cd "$(dirname "$0")/.." + +SECONDS_PER_ENGINE="${AWA_SOAK_SECONDS:-900}" +BASE_PORT="${SOAK_BASE_PORT:-15471}" +MAX_CONNECTIONS="${SOAK_MAX_CONNECTIONS:-400}" +CASSERT_IMAGE="${SOAK_CASSERT_IMAGE:-awa-pg18-cassert:18.6-O0}" +RESULTS_DIR="${SOAK_RESULTS_DIR:-$PWD/soak-results/$(date +%Y%m%d-%H%M%S)}" +ENGINES=("$@") +if [ ${#ENGINES[@]} -eq 0 ]; then + ENGINES=(pg17 pg18 pg18-cassert) +fi + +mkdir -p "$RESULTS_DIR" + +prepare() { + echo "== building the awa binary and the Python bindings" + cargo build --release -p awa-cli || return 1 + if [ ! -x awa-python/.venv/bin/python ]; then + uv venv --quiet awa-python/.venv || return 1 + uv pip install --quiet --python awa-python/.venv maturin || return 1 + fi + (cd awa-python && VIRTUAL_ENV="$PWD/.venv" PATH="$PWD/.venv/bin:$PATH" maturin develop --release) || return 1 +} +if [ "${SOAK_SKIP_PREPARE:-0}" != 1 ]; then + prepare >"$RESULTS_DIR/prepare.log" 2>&1 || { echo "prepare failed; see $RESULTS_DIR/prepare.log" >&2; exit 2; } +fi +export AWA_SOAK_PYTHON="${AWA_SOAK_PYTHON:-$PWD/awa-python/.venv/bin/python}" +export AWA_SOAK_AWA_BIN="${AWA_SOAK_AWA_BIN:-${CARGO_TARGET_DIR:-$PWD/target}/release/awa}" +CONTAINERS=() +STATS_PID="" + +cleanup() { + [ -n "$STATS_PID" ] && kill "$STATS_PID" 2>/dev/null + for container in "${CONTAINERS[@]:-}"; do + [ -n "$container" ] && docker rm -f -v "$container" >/dev/null 2>&1 + done +} +trap cleanup EXIT +trap 'exit 130' INT TERM + +start_container() { + local engine="$1" name="$2" port="$3" + local cpu_args=() + [ -n "${SOAK_PG_CPUS:-}" ] && cpu_args=(--cpus "$SOAK_PG_CPUS") + local common=(-c "max_connections=$MAX_CONNECTIONS" -c max_locks_per_transaction=256 -c shared_buffers=512MB -c "log_line_prefix=%m [%p] %u@%a ") + case "$engine" in + pg17) + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -e POSTGRES_PASSWORD=test -p "$port:5432" postgres:17 "${common[@]}" >/dev/null ;; + pg18) + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -e POSTGRES_PASSWORD=test -p "$port:5432" postgres:18 "${common[@]}" >/dev/null ;; + pg18-cassert) + # The image ships an initialised, trust-authenticated cluster in /data. Parallel query is + # off: at -O0 a handful of admin reads would otherwise take every core. + local pgargs='-c listen_addresses=0.0.0.0 -c max_connections='"$MAX_CONNECTIONS"' -c max_locks_per_transaction=256 -c shared_buffers=512MB -c track_activity_query_size=65536 -c max_parallel_workers_per_gather=0 -c logging_collector=on -c log_directory=/data/log -c log_filename=pg.log -c log_rotation_size=0' + case "$CASSERT_IMAGE" in + *epqlog*) + # The instrumented build logs an AWAEPQ line, with the full query text, for every + # partition-pruning setup: tens of gigabytes per hour under this workload. The log + # collector writes into a FIFO whose reader drops those messages and keeps the rest + # (assertion failures, crash notices, errors) in /data/log/pg-filtered.log. + SOAK_PG_LOG_FILE=/data/log/pg-filtered.log + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -p "$port:5432" "$CASSERT_IMAGE" sh -c ' + mkdir -p /data/log + rm -f /data/log/pg.log /data/log/pg-filtered.log + mkfifo /data/log/pg.log + awk -W interactive " + /^[0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9] /{ + pid=\"\"; if (match(\$0,/\\[[0-9]+\\]/)) pid=substr(\$0,RSTART,RLENGTH) + if (\$0 ~ /AWAEPQ/) {skip=1; epq[pid]=1} + else if (\$0 ~ /STATEMENT:/ && (pid in epq)) {skip=1; delete epq[pid]} + else {skip=0; delete epq[pid]} + } + !skip{print; fflush()}" \ + >/data/log/pg-filtered.log & + exec postgres -D /data '"$pgargs"'' >/dev/null ;; + *) + SOAK_PG_LOG_FILE=/data/log/pg.log + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -p "$port:5432" "$CASSERT_IMAGE" \ + sh -c 'mkdir -p /data/log && exec postgres -D /data '"$pgargs" >/dev/null ;; + esac ;; + *) + echo "unknown engine: $engine (expected pg17, pg18 or pg18-cassert)" >&2 + return 1 ;; + esac +} + +wait_ready() { + local name="$1" deadline=$((SECONDS + 120)) + until docker exec "$name" psql -U postgres -d postgres -Atc 'select 1' >/dev/null 2>&1; do + if [ "$SECONDS" -ge "$deadline" ]; then + echo "container $name did not become ready" >&2 + docker logs --tail 40 "$name" >&2 + return 1 + fi + sleep 1 + done + sleep 2 +} + +sample_stats() { + local name="$1" out="$2" + echo "time,cpu_percent,mem" >"$out" + while true; do + docker stats --no-stream --format '{{.CPUPerc}},{{.MemUsage}}' "$name" 2>/dev/null | + sed "s/^/$(date +%H:%M:%S),/" >>"$out" + sleep 15 + done +} + +if [ -n "${SOAK_TEST_BIN:-}" ]; then + harness_cmd=("$SOAK_TEST_BIN" persona_soak --ignored --nocapture --exact) +else + harness_cmd=(cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture) +fi + +overall=0 +for index in "${!ENGINES[@]}"; do + engine="${ENGINES[$index]}" + port=$((BASE_PORT + index)) + name="awa-soak-$engine-$$" + out="$RESULTS_DIR/$engine" + mkdir -p "$out" + echo "== $engine on port $port (container $name)" + start_container "$engine" "$name" "$port" || exit 2 + CONTAINERS+=("$name") + wait_ready "$name" || exit 2 + + log_env=(AWA_SOAK_PG_CONTAINER="$name") + [ "$engine" = pg18-cassert ] && log_env+=(AWA_SOAK_PG_CONTAINER_LOG="$SOAK_PG_LOG_FILE") + + sample_stats "$name" "$out/container-stats.csv" & + STATS_PID=$! + + env AWA_SOAK_DATABASE_URL="postgres://postgres:test@localhost:$port/awa_soak_${engine//-/_}" \ + AWA_SOAK_SECONDS="$SECONDS_PER_ENGINE" \ + AWA_SOAK_REPORT_DIR="$out" \ + "${log_env[@]}" \ + "${harness_cmd[@]}" >"$out/run.log" 2>&1 + status=$? + + kill "$STATS_PID" 2>/dev/null + STATS_PID="" + if [ $status -ne 0 ]; then + overall=1 + echo " FAILED (exit $status); see $out/report.md and $out/reproducer.md" + # Keep the server log with the reproducer when the run failed. + docker logs "$name" >"$out/server.log" 2>&1 || true + [ "$engine" = pg18-cassert ] && docker exec "$name" cat "$SOAK_PG_LOG_FILE" >"$out/server.log" 2>/dev/null + else + echo " passed; report in $out/report.md" + fi + docker rm -f -v "$name" >/dev/null 2>&1 +done + +echo "results: $RESULTS_DIR" +exit $overall diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index a3f2cea37..1b2639f8d 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -140,7 +140,11 @@ tick). The leader alone runs cluster-wide scans: deferred→ready promotion, the three rescue paths (stale-heartbeat, hard-deadline, callback-timeout), ring rotation and prune, DLQ and descriptor cleanup, and cron evaluation. During an unfinalized transition it promotes both planes' deferred backlogs, so drain and -post-flip execution do not depend on which runtime holds leadership. Every +post-flip execution do not depend on which runtime holds leadership. +Promotion gives each queue lane its own short transaction and hands a lane +whose enqueue lock stays busy to a dedicated waiter, so a hot queue delays +neither other queues' promotion nor rescue +(`awa_maintenance_promote_lane_contention_total` counts the hand-offs). Every worker dispatches and heartbeats its own attempts. There is no `pg_cron` or external scheduler.