fix(universaldb): chunk postgres commit requests that exceed the nats max_payload

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>> {

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

@@ -3,7 +3,8 @@ use std::{
 	time::{Duration, Instant},
 };
 
-use anyhow::{Context, Result};
+use anyhow::{Context, Result, bail};
+use futures_util::FutureExt;
 use tokio::sync::oneshot;
 
 use crate::{
@@ -30,6 +31,8 @@ const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
 const MAX_SUBMIT_ATTEMPTS: usize = 8;
 /// Backoff between multi-node resends.
 const RESEND_BACKOFF: Duration = Duration::from_millis(100);
+/// The NATS server's `max_payload` when it is not configured.
+const NATS_DEFAULT_MAX_PAYLOAD: usize = 1024 * 1024;
 
 /// Submit a follower transaction's commit to the leader and await the result.
 ///
@@ -106,22 +109,51 @@ async fn submit_nats(
 	// One dedup key for this logical commit, reused across every resend so the leader applies it at
 	// most once even if an earlier attempt was applied but its reply was lost to a failover.
 	let client_seq = shared.next_commit_seq();
+	let protocol_version = shared.commit_protocol_version();
 	let payload = codec::encode_commit_request(
 		read_version.max(0) as u64,
 		&conflict_ranges,
 		&operations,
 		shared.node_id.as_bytes(),
 		client_seq as u64,
-		shared.commit_protocol_version(),
+		protocol_version,
 	)
 	.context("failed to encode commit request")?;
 
 	let submit_start = Instant::now();
 	for attempt in 0..MAX_SUBMIT_ATTEMPTS {
 		let lease = wait_for_leader(shared).await?;
-		let subject = nats.subjects.commit(&lease.leader_addr);
 
-		let request = nats.client.request(subject, payload.clone().into());
+		// async-nats does not check a request against the server's max_payload, and the server
+		// answers an oversized message by closing the whole connection. A request that would not fit
+		// is split into chunks instead, which only a leader at the chunked protocol version accepts.
+		let max_payload = nats_max_payload(&nats.client);
+		let request = if payload.len() <= max_payload {
+			let subject = nats.subjects.commit(&lease.leader_addr);
+			let request = nats.client.request(subject, payload.clone().into());
+			async { request.await.map_err(anyhow::Error::from) }.boxed()
+		} else {
+			if protocol_version < codec::CHUNKED_COMMIT_PROTOCOL_VERSION {
+				bail!(
+					"commit request is {} bytes, over the nats server max_payload of {max_payload} bytes, \
+					 and the fleet has not negotiated chunked commit requests (protocol version \
+					 {protocol_version}); raise the nats max_payload or finish upgrading every node",
+					payload.len()
+				);
+			}
+			let chunks = codec::encode_commit_request_chunks(
+				&payload,
+				shared.node_id.as_bytes(),
+				client_seq as u64,
+				attempt as u32,
+				max_payload,
+				protocol_version,
+			)
+			.context("failed to chunk commit request")?;
+			let subject = nats.subjects.commit_chunk(&lease.leader_addr);
+			send_chunks(&nats.client, subject, chunks).boxed()
+		};
+
 		match tokio::time::timeout(REQUEST_TIMEOUT, request).await {
 			Ok(Ok(msg)) => match codec::decode_commit_reply(&msg.payload) {
 				Ok(CommitOutcome::Committed { .. }) => {
@@ -176,6 +208,38 @@ async fn submit_nats(
 	)
 }
 
+/// The largest message the connected NATS server accepts.
+fn nats_max_payload(client: &async_nats::Client) -> usize {
+	match client.server_info().max_payload {
+		// A client that has not received the server's INFO yet reports zero, so assume the server
+		// default rather than refusing to send.
+		0 => NATS_DEFAULT_MAX_PAYLOAD,
+		max_payload => max_payload,
+	}
+}
+
+/// Send every chunk except the last as a plain publish and the last as the request, so the reply
+/// arrives once the leader holds the whole commit.
+async fn send_chunks(
+	client: &async_nats::Client,
+	subject: String,
+	mut chunks: Vec<Vec<u8>>,
+) -> Result<async_nats::Message> {
+	let last = chunks
+		.pop()
+		.context("chunked commit request has no chunks")?;
+	for chunk in chunks {
+		client
+			.publish(subject.clone(), chunk.into())
+			.await
+			.context("failed to publish commit request chunk")?;
+	}
+	client
+		.request(subject, last.into())
+		.await
+		.context("commit request chunk got no reply")
+}
+
 /// Wait for a known leader, returning a retryable error if none is elected in time.
 async fn wait_for_leader(shared: &Arc<PostgresShared>) -> Result<LeaseInfo> {
 	let deadline = Instant::now() + LEADER_WAIT_TIMEOUT;

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

@@ -1,3 +1,4 @@
+mod chunks;
 mod codec;
 mod commit;
 mod database;