diff --git a/engine/packages/universaldb/src/driver/postgres/database.rs b/engine/packages/universaldb/src/driver/postgres/database.rs index 25e31b47bf..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(()) } diff --git a/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs b/engine/packages/universaldb/src/driver/postgres/resolver/mod.rs index 376ce3df57..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, @@ -584,22 +608,12 @@ async fn drain_batch( 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. + // 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(); - 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")?; - } + clear_ranges(&txn, &range_deletes).await?; let range_delete_ms = range_delete_start.elapsed().as_millis() as u64; let apply_start = Instant::now(); @@ -712,3 +726,7 @@ async fn drain_batch( Ok(BatchOutcome::Processed) } + +#[cfg(test)] +#[path = "../../../../tests/unit/postgres_resolver.rs"] +mod tests; 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" + ); +}