diff --git a/engine/packages/depot/tests/compaction_byte_volume_txn_window.rs b/engine/packages/depot/tests/compaction_byte_volume_txn_window.rs index cac1fd2c55..8e7b4ec913 100644 --- a/engine/packages/depot/tests/compaction_byte_volume_txn_window.rs +++ b/engine/packages/depot/tests/compaction_byte_volume_txn_window.rs @@ -127,7 +127,7 @@ async fn unbounded_compaction_read_ages_out_bounded_survives() -> Result<()> { let aged_out = err.chain().any(|cause| { matches!( cause.downcast_ref::(), - Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached) + Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached(_)) ) }); assert!( diff --git a/engine/packages/depot/tests/compaction_hot_read_byte_volume_txn_window.rs b/engine/packages/depot/tests/compaction_hot_read_byte_volume_txn_window.rs index 4f46704884..9a76db62aa 100644 --- a/engine/packages/depot/tests/compaction_hot_read_byte_volume_txn_window.rs +++ b/engine/packages/depot/tests/compaction_hot_read_byte_volume_txn_window.rs @@ -196,7 +196,7 @@ async fn hot_input_read_stays_bounded_at_byte_scale() -> Result<()> { let aged_out = err.chain().any(|cause| { matches!( cause.downcast_ref::(), - Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached) + Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached(_)) ) }); assert!( diff --git a/engine/packages/depot/tests/sqlite_byte_volume_txn_window.rs b/engine/packages/depot/tests/sqlite_byte_volume_txn_window.rs index f948356e44..74670c70fa 100644 --- a/engine/packages/depot/tests/sqlite_byte_volume_txn_window.rs +++ b/engine/packages/depot/tests/sqlite_byte_volume_txn_window.rs @@ -131,7 +131,7 @@ async fn get_pages_ages_out_unbounded_bounded_survives() -> Result<()> { let aged_out = err.chain().any(|cause| { matches!( cause.downcast_ref::(), - Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached) + Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached(_)) ) }); assert!( diff --git a/engine/packages/universaldb/src/driver/postgres/commit.rs b/engine/packages/universaldb/src/driver/postgres/commit.rs index 778c122a52..3310ca5b43 100644 --- a/engine/packages/universaldb/src/driver/postgres/commit.rs +++ b/engine/packages/universaldb/src/driver/postgres/commit.rs @@ -74,14 +74,20 @@ async fn submit_local( if commit_tx.send(job).await.is_err() { // The leader drain loop is gone (driver shutting down). Retryable. - return Err(DatabaseError::NotCommitted.into()); + return Err( + anyhow::Error::from(DatabaseError::NotCommitted).context("leader drain loop is gone") + ); } match response_rx.await { Ok(CommitOutcome::Committed { .. }) => Ok(()), - Ok(CommitOutcome::Conflict) => Err(DatabaseError::NotCommitted.into()), + // The leader resolved this commit as a loser. A cold-window rejection during leader recovery + // arrives as the same outcome, so the leader's batch log is what separates the two. + Ok(CommitOutcome::Conflict) => Err(anyhow::Error::from(DatabaseError::NotCommitted) + .context("leader resolved the commit as a conflict")), // The leader dropped the job without responding; it was not applied. - Err(_) => Err(DatabaseError::NotCommitted.into()), + Err(_) => Err(anyhow::Error::from(DatabaseError::NotCommitted) + .context("leader dropped the commit without responding")), } } @@ -128,7 +134,9 @@ async fn submit_nats( return Ok(()); } Ok(CommitOutcome::Conflict) => { - return Err(DatabaseError::NotCommitted.into()); + // As in the single-node path, a cold-window rejection is reported as a conflict. + return Err(anyhow::Error::from(DatabaseError::NotCommitted) + .context("leader resolved the commit as a conflict")); } Err(err) => { tracing::warn!(?err, client_seq, "malformed udb commit reply; resending"); @@ -161,7 +169,11 @@ async fn submit_nats( wait_ms = submit_start.elapsed().as_millis() as u64, "udb commit exhausted resend attempts; treating as not committed" ); - Err(DatabaseError::NotCommitted.into()) + Err( + anyhow::Error::from(DatabaseError::NotCommitted).context(format!( + "exhausted {MAX_SUBMIT_ATTEMPTS} commit resend attempts without a determinate reply" + )), + ) } /// Wait for a known leader, returning a retryable error if none is elected in time. @@ -172,7 +184,12 @@ async fn wait_for_leader(shared: &Arc) -> Result { return Ok(lease); } if Instant::now() >= deadline { - return Err(DatabaseError::NotCommitted.into()); + return Err( + anyhow::Error::from(DatabaseError::NotCommitted).context(format!( + "no leader elected within {}s", + LEADER_WAIT_TIMEOUT.as_secs() + )), + ); } tokio::time::sleep(LEADER_POLL_INTERVAL).await; } diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index d78d18a18b..28963b6d2c 100644 --- a/engine/packages/universaldb/src/driver/postgres/database.rs +++ b/engine/packages/universaldb/src/driver/postgres/database.rs @@ -47,6 +47,35 @@ const POOL_METRICS_INTERVAL: Duration = Duration::from_secs(1); /// deleted while a resend that needs it could still arrive. const DEDUP_ROW_MAX_AGE_SECS: i64 = 120; +/// The schema every node applies on startup. `kv` is the durable latest-value store; the rest is the +/// leader lease, commit version allocation, and failover dedup. +pub(super) const SCHEMA: &str = "CREATE TABLE IF NOT EXISTS kv ( + key BYTEA PRIMARY KEY, + value BYTEA NOT NULL + ); + + CREATE TABLE IF NOT EXISTS udb_lease ( + id INT PRIMARY KEY DEFAULT 1 CHECK (id = 1), + epoch BIGINT NOT NULL, + leader_addr TEXT NOT NULL, + durable_version BIGINT NOT NULL DEFAULT 0, + expires_at TIMESTAMPTZ NOT NULL + ); + + CREATE SEQUENCE IF NOT EXISTS udb_version_seq AS BIGINT + START WITH 1 INCREMENT BY 1 MINVALUE 1; + + CREATE TABLE IF NOT EXISTS udb_applied ( + client_node_id BYTEA NOT NULL, + client_seq BIGINT NOT NULL, + commit_version BIGINT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + PRIMARY KEY (client_node_id, client_seq) + ); + + CREATE INDEX IF NOT EXISTS udb_applied_created_at_idx + ON udb_applied (created_at);"; + #[derive(Clone, Debug)] pub struct PostgresConfig { pub connection_string: String, @@ -224,37 +253,9 @@ impl PostgresDatabaseDriver { } async fn init_schema(conn: &deadpool_postgres::Client) -> Result<()> { - // Durable latest-value store. - conn.batch_execute( - "CREATE TABLE IF NOT EXISTS kv ( - key BYTEA PRIMARY KEY, - value BYTEA NOT NULL - ); - - CREATE TABLE IF NOT EXISTS udb_lease ( - id INT PRIMARY KEY DEFAULT 1 CHECK (id = 1), - epoch BIGINT NOT NULL, - leader_addr TEXT NOT NULL, - durable_version BIGINT NOT NULL DEFAULT 0, - expires_at TIMESTAMPTZ NOT NULL - ); - - CREATE SEQUENCE IF NOT EXISTS udb_version_seq AS BIGINT - START WITH 1 INCREMENT BY 1 MINVALUE 1; - - CREATE TABLE IF NOT EXISTS udb_applied ( - client_node_id BYTEA NOT NULL, - client_seq BIGINT NOT NULL, - commit_version BIGINT NOT NULL, - created_at TIMESTAMPTZ NOT NULL DEFAULT now(), - PRIMARY KEY (client_node_id, client_seq) - ); - - CREATE INDEX IF NOT EXISTS udb_applied_created_at_idx - ON udb_applied (created_at);", - ) - .await - .context("failed to initialize postgres schema")?; + conn.batch_execute(SCHEMA) + .await + .context("failed to initialize postgres schema")?; Ok(()) } @@ -339,18 +340,6 @@ impl DatabaseDriver for PostgresDatabaseDriver { let mut attempt = 0; loop { - // Re-read every iteration. The first attempt always runs, because nothing has called - // `retry_limit` yet; from then on the closure's limit wins over the database-wide one. - let limit = retry_limit.load(Ordering::SeqCst); - let max_attempts = if limit == RETRY_LIMIT_UNSET { - max_retries - } else { - limit.saturating_add(1) - }; - if attempt >= max_attempts { - break; - } - let tx = Transaction::new(Arc::new(PostgresTransactionDriver::with_retry_limit( self.shared.clone(), retry_limit.clone(), @@ -380,17 +369,31 @@ impl DatabaseDriver for PostgresDatabaseDriver { maybe_committed = MaybeCommitted(true); } + // Re-read every iteration. Nothing has called `retry_limit` before the first + // attempt; from then on the closure's limit wins over the database-wide one. + // The check runs after an attempt failed, so both values bound retries rather + // than total attempts. + let limit = retry_limit.load(Ordering::SeqCst); + let retry_budget = if limit == RETRY_LIMIT_UNSET { + max_retries + } else { + limit + }; + if attempt >= retry_budget { + return Err(DatabaseError::MaxRetriesReached(error).into()); + } + + attempt += 1; + let backoff_ms = calculate_tx_retry_backoff(attempt as usize); tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await; - attempt += 1; - continue; + } else { + return Err(error); } + } else { + return Err(error); } - - return Err(error); } - - Err(DatabaseError::MaxRetriesReached.into()) }) } diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index 303b293b55..eaff651ef2 100644 --- a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs +++ b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs @@ -8,7 +8,7 @@ use std::{ }; use anyhow::{Context, Result, bail}; -use futures_util::StreamExt; +use futures_util::{StreamExt, future::try_join_all}; use tokio::sync::mpsc; use tokio_util::task::AbortOnDropHandle; @@ -389,6 +389,30 @@ enum BatchOutcome { LostLease, } +/// Clears each key range with its own statement. +/// +/// Passing every range to one statement as `unnest` arrays turns the bounds into join columns, so the +/// planner cannot estimate a range's width and prices each one as a fixed fraction of `kv`. Once `kv` +/// outgrows the page cache that estimate makes a full table scan per range look cheaper than the +/// primary key, and the batch transaction runs for minutes while holding its locks. As plain +/// parameters the bounds are planned with their real values, so each range walks the primary key. +/// +/// The statements are sent concurrently so tokio-postgres pipelines them, and each is prepared fresh +/// so a cached generic plan never replaces the planner's per-range estimate. +async fn clear_ranges( + txn: &tokio_postgres::Transaction<'_>, + ranges: &[(Vec, Vec)], +) -> Result<()> { + try_join_all(ranges.iter().map(|(begin, end)| async move { + txn.execute("DELETE FROM kv WHERE key >= $1 AND key < $2", &[begin, end]) + .await + })) + .await + .context("failed to clear ranges")?; + + Ok(()) +} + async fn drain_batch( shared: &Arc, epoch: i64, @@ -418,6 +442,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 +500,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 +523,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 +573,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,28 +591,31 @@ 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>) = - range_deletes.into_iter().unzip(); - - // 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. - if !range_begins.is_empty() { - txn.execute( - "DELETE FROM kv USING unnest($1::bytea[], $2::bytea[]) AS r(b, e) - WHERE key >= r.b AND key < r.e", - &[&range_begins, &range_ends], - ) - .await - .context("failed to clear ranges")?; - } + + // Range deletes run before the apply CTE: a range delete and an in-range upsert in one CTE would + // have unspecified ordering, so the clear must take effect first and the upsert then re-inserts the + // key. + let range_delete_start = Instant::now(); + clear_ranges(&txn, &range_deletes).await?; + 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 +659,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,8 +704,29 @@ 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" ); Ok(BatchOutcome::Processed) } + +#[cfg(test)] +#[path = "../../../../tests/unit/postgres_resolver.rs"] +mod tests; diff --git a/engine/packages/universaldb/src/driver/rocksdb/database.rs b/engine/packages/universaldb/src/driver/rocksdb/database.rs index 6a68df8abc..9c4a6bf758 100644 --- a/engine/packages/universaldb/src/driver/rocksdb/database.rs +++ b/engine/packages/universaldb/src/driver/rocksdb/database.rs @@ -79,18 +79,6 @@ impl DatabaseDriver for RocksDbDatabaseDriver { let mut attempt = 0; loop { - // Re-read every iteration. The first attempt always runs, because nothing has called - // `retry_limit` yet; from then on the closure's limit wins over the database-wide one. - let limit = retry_limit.load(Ordering::SeqCst); - let max_attempts = if limit == RETRY_LIMIT_UNSET { - max_retries - } else { - limit.saturating_add(1) - }; - if attempt >= max_attempts { - break; - } - let tx = Transaction::new(Arc::new(RocksDbTransactionDriver::with_retry_limit( self.db.clone(), self.txn_conflict_tracker.clone(), @@ -121,17 +109,31 @@ impl DatabaseDriver for RocksDbDatabaseDriver { maybe_committed = MaybeCommitted(true); } + // Re-read every iteration. Nothing has called `retry_limit` before the first + // attempt; from then on the closure's limit wins over the database-wide one. + // The check runs after an attempt failed, so both values bound retries rather + // than total attempts. + let limit = retry_limit.load(Ordering::SeqCst); + let retry_budget = if limit == RETRY_LIMIT_UNSET { + max_retries + } else { + limit + }; + if attempt >= retry_budget { + return Err(DatabaseError::MaxRetriesReached(error).into()); + } + + attempt += 1; + let backoff_ms = calculate_tx_retry_backoff(attempt as usize); tokio::time::sleep(tokio::time::Duration::from_millis(backoff_ms)).await; - attempt += 1; - continue; + } else { + return Err(error); } + } else { + return Err(error); } - - return Err(error); } - - Err(DatabaseError::MaxRetriesReached.into()) }) } diff --git a/engine/packages/universaldb/src/error.rs b/engine/packages/universaldb/src/error.rs index ef98a05039..b6ddcbccf5 100644 --- a/engine/packages/universaldb/src/error.rs +++ b/engine/packages/universaldb/src/error.rs @@ -7,8 +7,10 @@ pub enum DatabaseError { #[error("transaction is too old to perform reads or be committed")] TransactionTooOld, - #[error("max number of transaction retries reached")] - MaxRetriesReached, + // Stores the last error. The alternate format prints the whole context chain, so the cause the + // context names is reported alongside the underlying variant. + #[error("max number of transaction retries reached, last error: {0:#}")] + MaxRetriesReached(anyhow::Error), #[error("operation issued while a commit was outstanding")] UsedDuringCommit, @@ -22,7 +24,7 @@ impl DatabaseError { use DatabaseError::*; match self { - NotCommitted | TransactionTooOld | MaxRetriesReached => true, + NotCommitted | TransactionTooOld | MaxRetriesReached(_) => true, _ => false, } } diff --git a/engine/packages/universaldb/tests/conflict_parity.rs b/engine/packages/universaldb/tests/conflict_parity.rs index bde30d1733..f10403c62e 100644 --- a/engine/packages/universaldb/tests/conflict_parity.rs +++ b/engine/packages/universaldb/tests/conflict_parity.rs @@ -403,7 +403,7 @@ async fn writer_conflicts_when_a_key_it_read_was_written(db: Database) { assert!( err.chain().any(|x| matches!( x.downcast_ref::(), - Some(universaldb::error::DatabaseError::MaxRetriesReached) + Some(universaldb::error::DatabaseError::MaxRetriesReached(_)) )), "expected the conflict to exhaust retries, got {err:?}" ); 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; +} diff --git a/engine/packages/universaldb/tests/unit/postgres_resolver.rs b/engine/packages/universaldb/tests/unit/postgres_resolver.rs new file mode 100644 index 0000000000..145c9afec9 --- /dev/null +++ b/engine/packages/universaldb/tests/unit/postgres_resolver.rs @@ -0,0 +1,124 @@ +//! Query-plan checks for the leader's batch apply statements. +//! +//! These run against a real Postgres and read back the plan each statement actually executed through +//! `auto_explain`, which reports plans to the session as notices. + +use futures_util::future::poll_fn; +use rivet_test_deps_docker::TestDatabase; +use tokio::sync::mpsc; +use tokio_postgres::{AsyncMessage, NoTls}; +use uuid::Uuid; + +use super::{super::database::SCHEMA, clear_ranges}; + +const WORKFLOWS: i64 = 12_500; +const CHUNKS_PER_WORKFLOW: i64 = 8; +const CLEARED_WORKFLOWS: i64 = 8; + +fn state_range(workflow: i64) -> (Vec, Vec) { + let begin = format!("wf/{workflow:08}/state/").into_bytes(); + let mut end = begin.clone(); + end.push(0xff); + (begin, end) +} + +/// Clearing ranges must walk the primary key even when the planner prices a full scan of `kv` as +/// competitive. +/// +/// In production the planner makes that call once `kv` outgrows the page cache and random reads get +/// expensive, and a range clear that falls back to a full scan holds the leader's batch transaction +/// for minutes. Raising `random_page_cost` prices random reads the same way on a table small enough +/// to build here. +#[tokio::test] +async fn clear_ranges_uses_primary_key_when_scans_look_cheap() { + 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!(); + }; + let url = postgres_config.url.read().clone(); + + let (mut client, mut connection) = tokio_postgres::connect(&url, NoTls).await.unwrap(); + let (notice_tx, mut notice_rx) = mpsc::unbounded_channel(); + tokio::spawn(async move { + // The connection yields a notice before it routes the response that follows it, so every plan + // is in the channel by the time the statement that produced it resolves. + while let Some(message) = poll_fn(|cx| connection.poll_message(cx)).await { + match message { + Ok(AsyncMessage::Notice(notice)) => { + let _ = notice_tx.send(notice.message().to_string()); + } + // `AsyncMessage` is non-exhaustive, so other messages need a catch-all. + Ok(_) => {} + Err(_) => break, + } + } + }); + + client.batch_execute(SCHEMA).await.unwrap(); + // Random insertion order leaves no correlation between key order and heap order, as in + // production, so the planner cannot count on range reads touching adjacent pages. + client + .execute( + "INSERT INTO kv (key, value) + SELECT convert_to(format('wf/%s/state/%s', lpad(w::text, 8, '0'), lpad(c::text, 4, '0')), 'UTF8'), + repeat('x', 64)::bytea + FROM generate_series(1, $1::bigint) w, generate_series(1, $2::bigint) c + ORDER BY random()", + &[&WORKFLOWS, &CHUNKS_PER_WORKFLOW], + ) + .await + .unwrap(); + client + .batch_execute( + "ANALYZE kv; + LOAD 'auto_explain'; + SET auto_explain.log_min_duration = 0; + SET auto_explain.log_level = notice; + SET random_page_cost = 40;", + ) + .await + .unwrap(); + + let ranges: Vec<_> = (1..=CLEARED_WORKFLOWS).map(state_range).collect(); + let txn = client.transaction().await.unwrap(); + clear_ranges(&txn, &ranges).await.unwrap(); + + let plans: Vec = std::iter::from_fn(|| notice_rx.try_recv().ok()) + .filter(|notice| notice.contains("plan:")) + .collect(); + for plan in &plans { + assert!( + !plan.contains("Seq Scan on kv"), + "range clear fell back to a full scan of kv:\n{plan}" + ); + assert!( + plan.contains("kv_pkey"), + "range clear did not use the kv primary key:\n{plan}" + ); + } + assert_eq!( + plans.len(), + ranges.len(), + "expected one executed plan per cleared range: {plans:#?}" + ); + + let remaining: i64 = txn + .query_one("SELECT count(*) FROM kv", &[]) + .await + .unwrap() + .get(0); + assert_eq!( + remaining, + (WORKFLOWS - CLEARED_WORKFLOWS) * CHUNKS_PER_WORKFLOW, + "range clear removed the wrong rows" + ); +}