diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index 303b293b55..376ce3df57 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -418,6 +418,7 @@ async fn drain_batch( .start() .await .context("failed to start drain batch txn")?; + let begin_ms = batch_start.elapsed().as_millis() as u64 - pool_wait.as_millis() as u64; // Build the failover dedup keys: a job whose (client_node_id, client_seq) is already recorded in // udb_applied was committed by a prior leader; respond with the recorded version and do not @@ -475,7 +476,9 @@ async fn drain_batch( .collect::>(); anyhow::Ok(versions) }; + let prepare_start = Instant::now(); let (applied, mut versions) = tokio::try_join!(dedup_fut, versions_fut)?; + let prepare_ms = prepare_start.elapsed().as_millis() as u64; // Postgres does not guarantee nextval is evaluated in row order, so the versions are sorted and // assigned to to-resolve jobs in arrival order to keep versionstamps monotonic with commit order @@ -496,6 +499,7 @@ async fn drain_batch( resolve_indices.push(i); } + let resolve_start = Instant::now(); let cold_window = Instant::now() < recovery_deadline; let mut winners: Vec = Vec::new(); let mut winner_dedup_nids: Vec> = Vec::new(); @@ -545,6 +549,9 @@ async fn drain_batch( // Bulk-read the pre-batch value of every key a winner's atomic op reads, then fold all winners // into one materialized write-set in memory. + let resolve_ms = resolve_start.elapsed().as_millis() as u64; + + let atomic_start = Instant::now(); let atomic_keys = apply::atomic_read_keys(&winners); let base = if atomic_keys.is_empty() { HashMap::new() @@ -560,11 +567,21 @@ async fn drain_batch( .collect() }; + let atomic_read_count = atomic_keys.len(); + let atomic_ms = atomic_start.elapsed().as_millis() as u64; + + let fold_start = Instant::now(); let apply::WriteSet { upserts, point_deletes, range_deletes, } = apply::fold_winners(winners, &base).context("failed to fold batch winners")?; + let fold_ms = fold_start.elapsed().as_millis() as u64; + + let upsert_count = upserts.len(); + let point_delete_count = point_deletes.len(); + let range_delete_count = range_deletes.len(); + let upsert_bytes: usize = upserts.iter().map(|(k, v)| k.len() + v.len()).sum(); let (upsert_keys, upsert_values): (Vec>, Vec>) = upserts.into_iter().unzip(); let (range_begins, range_ends): (Vec>, Vec>) = @@ -573,6 +590,7 @@ async fn drain_batch( // Range deletes run in their own statement before the apply CTE: a range delete and an in-range // upsert in one CTE would have unspecified ordering, so the clear must commit its effect first and // the upsert then re-inserts the key. + let range_delete_start = Instant::now(); if !range_begins.is_empty() { txn.execute( "DELETE FROM kv USING unnest($1::bytea[], $2::bytea[]) AS r(b, e) @@ -582,6 +600,8 @@ async fn drain_batch( .await .context("failed to clear ranges")?; } + let range_delete_ms = range_delete_start.elapsed().as_millis() as u64; + let apply_start = Instant::now(); // Apply the rest of the batch in one CTE: point deletes, the kv upsert, the dedup records for // multi-node winners, and the epoch-fenced watermark advance. A zombie old leader whose epoch was @@ -625,7 +645,11 @@ async fn drain_batch( } }; + let apply_ms = apply_start.elapsed().as_millis() as u64; + + let commit_start = Instant::now(); txn.commit().await.context("failed to commit drain batch")?; + let commit_ms = commit_start.elapsed().as_millis() as u64; // The watermark advances strictly after the apply txn is durably committed and visible, so a // reader handed this read_version can never miss a write with commit_version <= read_version. @@ -666,6 +690,23 @@ async fn drain_batch( cold_window, new_durable, batch_ms = batch_start.elapsed().as_millis() as u64, + // Phase breakdown, so a slow batch says which statement was slow instead of only that the + // apply was slow overall. Every phase is milliseconds and they sum to roughly `batch_ms`. + pool_wait_ms = pool_wait.as_millis() as u64, + begin_ms, + prepare_ms, + resolve_ms, + atomic_ms, + fold_ms, + range_delete_ms, + apply_ms, + commit_ms, + // Work volume, to separate a large batch from a slow one. + upserts = upsert_count, + point_deletes = point_delete_count, + range_deletes = range_delete_count, + atomic_reads = atomic_read_count, + upsert_bytes, "udb leader processed commit batch" ); diff --git a/engine/packages/universaldb/tests/leader_apply_stall.rs b/engine/packages/universaldb/tests/leader_apply_stall.rs new file mode 100644 index 0000000000..b2c7ca52ab --- /dev/null +++ b/engine/packages/universaldb/tests/leader_apply_stall.rs @@ -0,0 +1,309 @@ +//! Investigation harness for leader drain-batch stalls. +//! +//! Production traces show the leader's `drain_batch` blocking for roughly seven seconds at a time on +//! batches of ten to twenty-five jobs, with no pool wait, no conflicts, and no leadership change. The +//! batch log reports only the total, so this harness drives sustained multi-node commit load in the +//! shape gasoline produces and reads the per-phase timings back out of the tracing output. It exists +//! to localize which statement in the apply is slow, not to assert a latency bound. +//! +//! Run with `--ignored --nocapture`; these boot containers and push real load, so they are not part +//! of the default suite. + +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex, OnceLock}, + time::Duration, +}; + +use futures_util::future::join_all; +use rivet_test_deps_docker::{TestDatabase, TestPubSub}; +use tracing_subscriber::{Layer, layer::SubscriberExt, util::SubscriberInitExt}; +use universaldb::{ + Database, + driver::postgres::{NatsConfig, PostgresConfig}, +}; +use uuid::Uuid; + +/// Workflow-state values gasoline writes are chunked; this is a representative chunk size. +const CHUNK_BYTES: usize = 8 * 1024; + +fn test_config() -> rivet_config::Config { + rivet_config::Config::from_root_with_build_meta( + rivet_config::config::Root::default(), + rivet_config::BuildMeta::default(), + rivet_config::RuntimeProtocols { + universaldb_commit: rivet_config::RuntimeProtocol::new( + rivet_config::RuntimeProtocolKind::UniversaldbCommit, + rivet_universaldb_commit::PROTOCOL_VERSION, + ), + ..Default::default() + }, + ) +} + +/// Collects the fields of every `udb leader processed commit batch` event so the test can report the +/// phase breakdown itself. The batch log is `debug`, and capturing it is the whole point of the +/// harness, so it is parsed rather than eyeballed. +#[derive(Default)] +struct BatchCollector { + batches: Mutex>>, +} + +impl BatchCollector { + fn clear(&self) { + self.batches.lock().unwrap().clear(); + } +} + +/// The drain loop runs on its own spawned task, so a thread-local `set_default` subscriber never +/// sees its events. The collector is installed globally once and shared by every test here, and the +/// tests serialize on [`RUN_LOCK`] so one test's batches cannot land in another's report. +static COLLECTOR: OnceLock> = OnceLock::new(); +static RUN_LOCK: Mutex<()> = Mutex::new(()); + +fn collector() -> Arc { + COLLECTOR + .get_or_init(|| { + let collector = Arc::new(BatchCollector::default()); + tracing_subscriber::registry() + .with(BatchLayer(collector.clone())) + .init(); + collector + }) + .clone() +} + +struct BatchLayer(Arc); + +impl Layer for BatchLayer { + fn on_event( + &self, + event: &tracing::Event<'_>, + _ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + let mut fields = BTreeMap::new(); + event.record(&mut FieldVisitor(&mut fields)); + if fields.get("message").map(String::as_str) == Some("udb leader processed commit batch") { + self.0.batches.lock().unwrap().push(fields); + } + } +} + +struct FieldVisitor<'a>(&'a mut BTreeMap); + +impl tracing::field::Visit for FieldVisitor<'_> { + fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { + self.0 + .insert(field.name().to_string(), format!("{value:?}")); + } + fn record_u64(&mut self, field: &tracing::field::Field, value: u64) { + self.0.insert(field.name().to_string(), value.to_string()); + } + fn record_i64(&mut self, field: &tracing::field::Field, value: i64) { + self.0.insert(field.name().to_string(), value.to_string()); + } + fn record_str(&mut self, field: &tracing::field::Field, value: &str) { + self.0.insert(field.name().to_string(), value.to_string()); + } +} + +async fn setup_postgres() -> (String, rivet_test_deps_docker::DockerRunConfig) { + let (db_config, docker_config) = TestDatabase::Postgres + .config(Uuid::new_v4(), 1) + .await + .unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + TestDatabase::Postgres + .wait_for_ready(&docker_config) + .await + .unwrap(); + let rivet_config::config::Database::Postgres(postgres_config) = db_config else { + unreachable!(); + }; + (postgres_config.url.read().clone(), docker_config) +} + +async fn setup_nats() -> (NatsConfig, rivet_test_deps_docker::DockerRunConfig) { + let (pubsub_config, docker_config) = TestPubSub::Nats.config(Uuid::new_v4(), 1).await.unwrap(); + let mut docker_config = docker_config.unwrap(); + docker_config.start().await.unwrap(); + tokio::time::sleep(Duration::from_secs(1)).await; + let rivet_config::config::PubSub::Nats(nats) = pubsub_config else { + unreachable!(); + }; + ( + NatsConfig { + addresses: nats.addresses.clone(), + username: nats.username.clone(), + password: nats.password.as_ref().map(|p| p.read().clone()), + client_capacity: nats.client_capacity, + subscription_capacity: nats.subscription_capacity, + }, + docker_config, + ) +} + +async fn make_db(connection_string: &str, nats: Option<&NatsConfig>) -> Database { + let mut config = PostgresConfig::new(connection_string.to_string()); + config.nats = nats.cloned(); + let driver = + universaldb::driver::PostgresDatabaseDriver::new_with_config(test_config(), config) + .await + .unwrap(); + Database::new(Arc::new(driver)) +} + +fn state_prefix(workflow: usize) -> Vec { + format!("wf/{workflow:08}/state/").into_bytes() +} + +fn state_chunk_key(workflow: usize, chunk: usize) -> Vec { + format!("wf/{workflow:08}/state/{chunk:04}").into_bytes() +} + +fn range_end(prefix: &[u8]) -> Vec { + let mut end = prefix.to_vec(); + end.push(0xff); + end +} + +/// One `update_workflow_state`-shaped transaction: clear the whole state subspace, then write the +/// state back as chunks. This is the write pattern that dominates gasoline's commit volume. +async fn write_state(db: &Database, workflow: usize, chunks: usize) { + db.txn("test_update_workflow_state", move |tx| async move { + let prefix = state_prefix(workflow); + tx.clear_range(&prefix, &range_end(&prefix)); + for chunk in 0..chunks { + tx.set(&state_chunk_key(workflow, chunk), &vec![b'x'; CHUNK_BYTES]); + } + Ok(()) + }) + .await + .unwrap(); +} + +/// Report the phase breakdown of collected batches, sorted by total time. +fn report(collector: &BatchCollector, label: &str) { + let batches = collector.batches.lock().unwrap(); + let num = |b: &BTreeMap, k: &str| -> u64 { + b.get(k).and_then(|v| v.parse().ok()).unwrap_or(0) + }; + let mut sorted: Vec<_> = batches.iter().collect(); + sorted.sort_by_key(|b| std::cmp::Reverse(num(b, "batch_ms"))); + + let total = batches.len(); + let slow = batches + .iter() + .filter(|b| num(b, "batch_ms") >= 1000) + .count(); + println!("\n===== {label} ====="); + println!("batches={total} slow(>=1s)={slow}"); + println!( + "{:>8} {:>5} {:>5} {:>5} {:>5} {:>5} {:>5} {:>5} {:>7} {:>6} {:>6} {:>8} {:>7}", + "batch_ms", + "pool", + "begin", + "prep", + "resol", + "atomi", + "fold", + "rdel", + "apply", + "commit", + "len", + "upserts", + "bytes" + ); + for b in sorted.iter().take(10) { + println!( + "{:>8} {:>5} {:>5} {:>5} {:>5} {:>5} {:>5} {:>5} {:>7} {:>6} {:>6} {:>8} {:>7}", + num(b, "batch_ms"), + num(b, "pool_wait_ms"), + num(b, "begin_ms"), + num(b, "prepare_ms"), + num(b, "resolve_ms"), + num(b, "atomic_ms"), + num(b, "fold_ms"), + num(b, "range_delete_ms"), + num(b, "apply_ms"), + num(b, "commit_ms"), + num(b, "batch_len"), + num(b, "upserts"), + num(b, "upsert_bytes"), + ); + } + let sum = |k: &str| -> u64 { batches.iter().map(|b| num(b, k)).sum() }; + println!( + "totals: batch={} pool={} prepare={} resolve={} atomic={} fold={} range_delete={} apply={} commit={}", + sum("batch_ms"), + sum("pool_wait_ms"), + sum("prepare_ms"), + sum("resolve_ms"), + sum("atomic_ms"), + sum("fold_ms"), + sum("range_delete_ms"), + sum("apply_ms"), + sum("commit_ms"), + ); +} + +/// Drive sustained multi-node write load and report where the leader's apply time goes. +/// +/// `workflows` sets how many distinct state subspaces churn, `chunks` how many chunks each state +/// carries, and `rounds` how many times every workflow rewrites its state. +async fn run_load(workflows: usize, chunks: usize, rounds: usize, concurrency: usize, label: &str) { + let _serialized = RUN_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + let collector = collector(); + collector.clear(); + + let (connection_string, _pg) = setup_postgres().await; + let (nats, _nats_docker) = setup_nats().await; + + let leader = Arc::new(make_db(&connection_string, Some(&nats)).await); + let follower = Arc::new(make_db(&connection_string, Some(&nats)).await); + // Let one node win the lease before load starts, so the run measures steady-state apply cost + // rather than election. + tokio::time::sleep(Duration::from_secs(2)).await; + + for round in 0..rounds { + let mut handles = Vec::new(); + for batch in 0..concurrency { + let db = if batch % 2 == 0 { + leader.clone() + } else { + follower.clone() + }; + let start = (round * concurrency + batch) % workflows; + handles.push(tokio::spawn(async move { + write_state(&db, start, chunks).await; + })); + } + join_all(handles).await; + } + + report(&collector, label); +} + +/// Baseline: small states, high transaction rate. Establishes what a healthy apply costs. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn leader_apply_small_states() { + run_load(512, 1, 40, 64, "small states (1 chunk)").await; +} + +/// Large states: the same transaction shape, but each one clears and rewrites a much bigger +/// subspace. Tests whether apply cost tracks byte volume rather than job count. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn leader_apply_large_states() { + run_load(128, 64, 20, 64, "large states (64 chunks)").await; +} + +/// Accumulated table: many distinct workflows churn so `kv` grows and range deletes scan more, which +/// is closer to a long-lived production table than a freshly created one. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn leader_apply_wide_table() { + run_load(8192, 8, 12, 96, "wide table (8192 workflows)").await; +}