test(gasoline): reproduce slow workflow pulls on postgres

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::<DatabaseError>(),
-			Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached)
+			Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached(_))
 		)
 	});
 	assert!(

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::<DatabaseError>(),
-				Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached)
+				Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached(_))
 			)
 		});
 		assert!(

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::<DatabaseError>(),
-			Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached)
+			Some(DatabaseError::TransactionTooOld | DatabaseError::MaxRetriesReached(_))
 		)
 	});
 	assert!(

engine/packages/gasoline/tests/pull_backlog.rs

@@ -0,0 +1,175 @@
+//! Investigation harness for workflow pulls against a large wake backlog.
+//!
+//! A worker gives up when `pull_workflows` takes longer than its fixed pull timeout, and giving up
+//! stops every workflow it is running. This harness builds a backlog of awake workflows, adds network
+//! round-trip latency to the database the way a managed Postgres has it, and times one pull phase by
+//! phase.
+//!
+//! Run with:
+//!
+//! ```text
+//! RIVET_TEST_DATABASE=postgres cargo test -p gasoline --test pull_backlog -- --ignored --nocapture
+//! ```
+//!
+//! `PULL_BACKLOG_WORKFLOWS` sets the backlog size and `PULL_BACKLOG_RTT_MS` the added latency.
+
+use std::{
+	collections::HashMap,
+	process::Command,
+	time::{Duration, Instant},
+};
+
+use anyhow::{Context, Result, ensure};
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::Id;
+use gasoline as gas;
+use gasoline::db::{Database, DatabaseKv};
+use serde_json::json;
+use tracing::{Subscriber, span};
+use tracing_subscriber::{
+	Layer, filter::filter_fn, layer::SubscriberExt, registry::LookupSpan, util::SubscriberInitExt,
+};
+use uuid::Uuid;
+
+const WORKFLOW_NAME: &str = "pull_backlog_test";
+/// The worker's pull timeout in `gasoline::worker`.
+const PULL_WORKFLOWS_TIMEOUT: Duration = Duration::from_secs(10);
+/// Spans that bound the phases of `pull_workflows`.
+const PHASE_SPANS: &[&str] = &[
+	"pull_workflows",
+	"read_wake_conditions",
+	"map_to_leased_workflows",
+	"pull_workflows_tx",
+	"clear_workflow_secondary_idx_tx",
+	"pull_workflow_history_tx",
+];
+
+fn env_or<T: std::str::FromStr>(name: &str, default: T) -> T {
+	std::env::var(name)
+		.ok()
+		.and_then(|value| value.parse().ok())
+		.unwrap_or(default)
+}
+
+/// Prints how long each pull phase span was open.
+struct PhaseTimer;
+
+impl<S> Layer<S> for PhaseTimer
+where
+	S: Subscriber + for<'a> LookupSpan<'a>,
+{
+	fn on_new_span(
+		&self,
+		_attrs: &span::Attributes<'_>,
+		id: &span::Id,
+		ctx: tracing_subscriber::layer::Context<'_, S>,
+	) {
+		if let Some(span) = ctx.span(id) {
+			span.extensions_mut().insert(Instant::now());
+		}
+	}
+
+	fn on_close(&self, id: span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
+		if let Some(span) = ctx.span(&id) {
+			if let Some(opened_at) = span.extensions().get::<Instant>() {
+				println!("phase {:<32} {:?}", span.name(), opened_at.elapsed());
+			}
+		}
+	}
+}
+
+/// Adds `rtt_ms` of egress delay to a container's network, which every response to the engine then
+/// pays. The sidecar shares the container's network namespace, so the database image needs no tools.
+fn add_network_latency(container_name: &str, rtt_ms: u64) -> Result<()> {
+	let output = Command::new("docker")
+		.args([
+			"run",
+			"--rm",
+			"--net",
+			&format!("container:{container_name}"),
+			"--cap-add",
+			"NET_ADMIN",
+			"alpine:3",
+			"sh",
+			"-c",
+			&format!(
+				"apk add --no-cache iproute2-tc >/dev/null && tc qdisc add dev eth0 root netem delay {rtt_ms}ms"
+			),
+		])
+		.output()
+		.context("failed to run the netem sidecar")?;
+	ensure!(
+		output.status.success(),
+		"netem sidecar failed: {}",
+		String::from_utf8_lossy(&output.stderr)
+	);
+	Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+#[ignore = "investigation harness; run explicitly with --ignored --nocapture"]
+async fn pull_workflows_against_wake_backlog() -> Result<()> {
+	tracing_subscriber::registry()
+		.with(PhaseTimer.with_filter(filter_fn(|metadata| {
+			metadata.is_span() && PHASE_SPANS.contains(&metadata.name())
+		})))
+		.init();
+
+	let backlog: usize = env_or("PULL_BACKLOG_WORKFLOWS", 5_000);
+	let rtt_ms: u64 = env_or("PULL_BACKLOG_RTT_MS", 1);
+	let on_postgres = std::env::var("RIVET_TEST_DATABASE").as_deref() == Ok("postgres");
+
+	let test_id = Uuid::new_v4();
+	let test_deps = rivet_test_deps::TestDeps::new_with_test_id(test_id).await?;
+	let config = test_deps.config().clone();
+	let db = <DatabaseKv as Database>::new(config.clone(), test_deps.pools().clone()).await?;
+
+	let dispatch_start = Instant::now();
+	let input = serde_json::value::to_raw_value(&json!({}))?;
+	futures_util::stream::iter(0..backlog)
+		.map(|_| {
+			db.dispatch_workflow(
+				Id::new_v1(config.dc_label()),
+				Id::new_v1(config.dc_label()),
+				WORKFLOW_NAME,
+				None,
+				input.as_ref(),
+				false,
+			)
+		})
+		.buffer_unordered(128)
+		.try_collect::<Vec<_>>()
+		.await?;
+	println!(
+		"dispatched {backlog} workflows in {:?}",
+		dispatch_start.elapsed()
+	);
+
+	// Latency is added after the backlog exists so building it stays fast.
+	if on_postgres {
+		add_network_latency(&format!("test-postgres-{test_id}-1"), rtt_ms)?;
+	} else {
+		println!("not on postgres, so no network latency was added");
+	}
+
+	let worker_id = Id::new_v1(config.dc_label());
+	db.update_worker_ping(worker_id, 1, true).await?;
+
+	let pull_start = Instant::now();
+	let pulled = db
+		.pull_workflows(worker_id, 1, &[WORKFLOW_NAME], &HashMap::new())
+		.await?;
+	let pull_elapsed = pull_start.elapsed();
+
+	println!(
+		"pulled {} of {backlog} workflows in {pull_elapsed:?} (rtt {}ms)",
+		pulled.len(),
+		if on_postgres { rtt_ms } else { 0 }
+	);
+	ensure!(
+		pull_elapsed < PULL_WORKFLOWS_TIMEOUT,
+		"pull took {pull_elapsed:?}, past the worker's {PULL_WORKFLOWS_TIMEOUT:?} pull timeout"
+	);
+
+	Ok(())
+}

engine/packages/universaldb/src/driver/postgres/chunks.rs

@@ -0,0 +1,102 @@
+use std::{
+	collections::HashMap,
+	time::{Duration, Instant},
+};
+
+/// A partially received commit request is dropped once no piece of it has arrived for this long. The
+/// follower gives up on an attempt after its request timeout and resends from the first piece, so
+/// anything older can never complete.
+pub const PENDING_CHUNK_MAX_IDLE: Duration = Duration::from_secs(30);
+
+/// One piece of a commit request that a follower split across several NATS messages.
+pub struct CommitChunk {
+	pub client_node_id: Vec<u8>,
+	pub client_seq: u64,
+	pub attempt: u32,
+	pub index: u32,
+	pub count: u32,
+	pub data: Vec<u8>,
+}
+
+struct PendingRequest {
+	attempt: u32,
+	count: u32,
+	next_index: u32,
+	bytes: Vec<u8>,
+	updated_at: Instant,
+}
+
+/// Reassembles chunked commit requests on the leader.
+///
+/// NATS delivers one publisher's messages on a subject in order, so pieces of an attempt arrive in
+/// sequence unless a reconnect loses some. Any gap abandons the attempt instead of waiting for the
+/// missing piece: the follower's request times out and it resends every piece under a new attempt.
+/// Owned by the single commit subscriber task.
+#[derive(Default)]
+pub struct ChunkAssembler {
+	pending: HashMap<(Vec<u8>, u64), PendingRequest>,
+}
+
+impl ChunkAssembler {
+	/// Accepts one piece and returns the complete encoded request once its last piece arrives.
+	pub fn push(&mut self, chunk: CommitChunk, now: Instant) -> Option<Vec<u8>> {
+		let key = (chunk.client_node_id, chunk.client_seq);
+
+		if chunk.index == 0 {
+			if let Some(existing) = self.pending.get(&key) {
+				if chunk.attempt < existing.attempt {
+					return None;
+				}
+			}
+			if chunk.count <= 1 {
+				self.pending.remove(&key);
+				return Some(chunk.data);
+			}
+			self.pending.insert(
+				key,
+				PendingRequest {
+					attempt: chunk.attempt,
+					count: chunk.count,
+					next_index: 1,
+					bytes: chunk.data,
+					updated_at: now,
+				},
+			);
+			return None;
+		}
+
+		let Some(pending) = self.pending.get_mut(&key) else {
+			return None;
+		};
+		if chunk.attempt < pending.attempt {
+			return None;
+		}
+		if chunk.attempt != pending.attempt
+			|| chunk.count != pending.count
+			|| chunk.index != pending.next_index
+		{
+			self.pending.remove(&key);
+			return None;
+		}
+
+		pending.bytes.extend_from_slice(&chunk.data);
+		pending.next_index += 1;
+		pending.updated_at = now;
+
+		if pending.next_index == pending.count {
+			return self.pending.remove(&key).map(|pending| pending.bytes);
+		}
+		None
+	}
+
+	/// Drops requests that stopped receiving pieces, so a follower that died mid-send does not pin its
+	/// partial request in memory.
+	pub fn evict_idle(&mut self, now: Instant) {
+		self.pending
+			.retain(|_, pending| now.duration_since(pending.updated_at) < PENDING_CHUNK_MAX_IDLE);
+	}
+}
+
+#[cfg(test)]
+#[path = "../../../tests/unit/postgres_chunks.rs"]
+mod tests;

engine/packages/universaldb/src/driver/postgres/codec.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{Context, Result};
 use rivet_universaldb_commit::{self as proto, versioned};
 use vbare::OwnedVersionedData;
 
@@ -7,7 +7,11 @@ use crate::{
 	tx_ops::Operation,
 };
 
-use super::transport::CommitOutcome;
+use super::{chunks::CommitChunk, transport::CommitOutcome};
+
+/// Protocol version that introduced [`proto::CommitRequestChunk`]. A fleet negotiated below it has
+/// leaders that cannot reassemble a chunked request.
+pub const CHUNKED_COMMIT_PROTOCOL_VERSION: u16 = 2;
 
 /// Decoded form of a commit request payload sent from a follower to the leader over NATS.
 pub struct DecodedCommit {
@@ -80,6 +84,60 @@ pub fn decode_commit_request(payload: &[u8]) -> Result<DecodedCommit> {
 	})
 }
 
+/// Split an encoded commit request into chunk messages that each fit within `max_payload` bytes,
+/// encoded at `protocol_version`, which must be at least [`CHUNKED_COMMIT_PROTOCOL_VERSION`].
+pub fn encode_commit_request_chunks(
+	request: &[u8],
+	client_node_id: &[u8],
+	client_seq: u64,
+	attempt: u32,
+	max_payload: usize,
+	protocol_version: u16,
+) -> Result<Vec<Vec<u8>>> {
+	let encode = |index: u32, count: u32, data: Vec<u8>| {
+		versioned::CommitRequestChunk::wrap_latest(proto::CommitRequestChunk {
+			client_node_id: client_node_id.to_vec(),
+			client_seq,
+			attempt,
+			index,
+			count,
+			data,
+		})
+		.serialize_with_embedded_version(protocol_version)
+	};
+
+	// Every field except `data` has the same encoded width in every chunk, so an empty chunk measures
+	// the envelope. The length prefix of `data` grows from one byte to at most five as a piece grows.
+	let overhead = encode(0, 0, Vec::new())?.len() + 4;
+	let piece_len = max_payload
+		.checked_sub(overhead)
+		.filter(|len| *len > 0)
+		.with_context(|| {
+			format!("nats max_payload of {max_payload} bytes cannot fit a commit request chunk")
+		})?;
+	let count = u32::try_from(request.len().div_ceil(piece_len))
+		.context("commit request needs too many chunks")?;
+
+	request
+		.chunks(piece_len)
+		.zip(0..)
+		.map(|(piece, index)| encode(index, count, piece.to_vec()))
+		.collect()
+}
+
+/// Decode one chunk produced by [`encode_commit_request_chunks`].
+pub fn decode_commit_request_chunk(payload: &[u8]) -> Result<CommitChunk> {
+	let chunk = versioned::CommitRequestChunk::deserialize_with_embedded_version(payload)?;
+	Ok(CommitChunk {
+		client_node_id: chunk.client_node_id,
+		client_seq: chunk.client_seq,
+		attempt: chunk.attempt,
+		index: chunk.index,
+		count: chunk.count,
+		data: chunk.data,
+	})
+}
+
 /// Encode a leader's commit reply at `protocol_version`, the version negotiated across the fleet,
 /// with an embedded version header.
 pub fn encode_commit_reply(outcome: CommitOutcome, protocol_version: u16) -> Result<Vec<u8>> {