fix(universaldb): clear postgres key ranges one statement per range so the planner keeps index access

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(())
 	}

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<u8>, Vec<u8>)],
+) -> 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<PostgresShared>,
 	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<u8>>, Vec<Vec<u8>>) = upserts.into_iter().unzip();
-	let (range_begins, range_ends): (Vec<Vec<u8>>, Vec<Vec<u8>>) =
-		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;

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<u8>, Vec<u8>) {
+	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<String> = 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"
+	);
+}