@rivetkit/engine 1.0.0

package/packages/universalpubsub/src/chunking.rs

@@ -0,0 +1,226 @@
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+use anyhow::*;
+use rivet_ups_protocol::versioned::UpsMessage;
+use rivet_ups_protocol::{MessageBody, MessageChunk, MessageStart, PROTOCOL_VERSION};
+use vbare::OwnedVersionedData;
+
+pub const CHUNK_BUFFER_GC_INTERVAL: Duration = Duration::from_secs(60);
+pub const CHUNK_BUFFER_MAX_AGE: Duration = Duration::from_secs(300);
+
+#[derive(Debug)]
+pub struct ChunkBuffer {
+	pub message_id: [u8; 16],
+	pub received_chunks: u32,
+	pub last_chunk_ts: Instant,
+	pub buffer: Vec<u8>,
+	pub chunk_count: u32,
+	pub reply_subject: Option<String>,
+}
+
+pub struct ChunkTracker {
+	chunks_in_process: HashMap<[u8; 16], ChunkBuffer>,
+}
+
+impl ChunkTracker {
+	pub fn new() -> Self {
+		Self {
+			chunks_in_process: HashMap::new(),
+		}
+	}
+
+	pub fn process_chunk(
+		&mut self,
+		raw_message: &[u8],
+	) -> Result<Option<(Vec<u8>, Option<String>)>> {
+		let message = UpsMessage::deserialize_with_embedded_version(raw_message)?;
+
+		match message.body {
+			MessageBody::MessageStart(msg) => {
+				// If only one chunk, return immediately
+				if msg.chunk_count == 1 {
+					return Ok(Some((msg.payload, msg.reply_subject)));
+				}
+
+				// Start of a multi-chunk message
+				let buffer = ChunkBuffer {
+					message_id: msg.message_id,
+					received_chunks: 1,
+					last_chunk_ts: Instant::now(),
+					buffer: msg.payload,
+					chunk_count: msg.chunk_count,
+					reply_subject: msg.reply_subject,
+				};
+				self.chunks_in_process.insert(msg.message_id, buffer);
+				Ok(None)
+			}
+			MessageBody::MessageChunk(msg) => {
+				// Find the matching buffer using message_id
+				let buffer = self.chunks_in_process.get_mut(&msg.message_id);
+
+				let Some(buffer) = buffer else {
+					bail!(
+						"received chunk {} for message {:?} but no matching buffer found",
+						msg.chunk_index,
+						msg.message_id
+					);
+				};
+
+				// Validate chunk order
+				if buffer.received_chunks != msg.chunk_index {
+					bail!(
+						"received chunk {} but expected chunk {} for message {:?}",
+						msg.chunk_index,
+						buffer.received_chunks,
+						msg.message_id
+					);
+				}
+
+				// Update buffer
+				buffer.buffer.extend_from_slice(&msg.payload);
+				buffer.received_chunks += 1;
+				buffer.last_chunk_ts = Instant::now();
+				let is_complete = buffer.received_chunks == buffer.chunk_count;
+
+				if is_complete {
+					let completed_buffer = self.chunks_in_process.remove(&msg.message_id).unwrap();
+					Ok(Some((
+						completed_buffer.buffer,
+						completed_buffer.reply_subject,
+					)))
+				} else {
+					Ok(None)
+				}
+			}
+		}
+	}
+
+	pub fn gc(&mut self) {
+		let now = Instant::now();
+		let size_before = self.chunks_in_process.len();
+		self.chunks_in_process
+			.retain(|_, buffer| now.duration_since(buffer.last_chunk_ts) < CHUNK_BUFFER_MAX_AGE);
+		let size_after = self.chunks_in_process.len();
+
+		tracing::debug!(
+			?size_before,
+			?size_after,
+			"performed chunk buffer garbage collection"
+		);
+	}
+}
+
+/// Splits a payload into chunks that fit within message size limits.
+///
+/// This function handles chunking by accounting for different overhead
+/// between the first chunk (MessageStart) and subsequent chunks (MessageChunk).
+///
+/// The first chunk carries additional metadata like the reply_subject and chunk_count,
+/// which means it has more protocol overhead and less room for payload data.
+/// Subsequent chunks only carry a chunk_index, allowing them to fit more payload.
+///
+/// This optimization ensures:
+/// - Reply subject is only transmitted once (in MessageStart)
+/// - Maximum payload utilization in each chunk
+/// - Efficient bandwidth usage for multi-chunk messages
+///
+/// # Returns
+/// A vector of payload chunks, where each chunk is sized to fit within the message limit
+/// after accounting for protocol overhead.
+pub fn split_payload_into_chunks(
+	payload: &[u8],
+	max_message_size: usize,
+	message_id: [u8; 16],
+	reply_subject: Option<&str>,
+) -> Result<Vec<Vec<u8>>> {
+	// Calculate overhead for MessageStart (first chunk)
+	let start_message = MessageStart {
+		message_id,
+		chunk_count: 1,
+		reply_subject: reply_subject.map(|s| s.to_string()),
+		payload: vec![],
+	};
+	let start_ups_message = rivet_ups_protocol::UpsMessage {
+		body: MessageBody::MessageStart(start_message),
+	};
+	let start_overhead = UpsMessage::wrap_latest(start_ups_message)
+		.serialize_with_embedded_version(PROTOCOL_VERSION)?
+		.len();
+
+	// Calculate overhead for MessageChunk (subsequent chunks)
+	let chunk_message = MessageChunk {
+		message_id,
+		chunk_index: 0,
+		payload: vec![],
+	};
+	let chunk_ups_message = rivet_ups_protocol::UpsMessage {
+		body: MessageBody::MessageChunk(chunk_message),
+	};
+	let chunk_overhead = UpsMessage::wrap_latest(chunk_ups_message)
+		.serialize_with_embedded_version(PROTOCOL_VERSION)?
+		.len();
+
+	// Calculate max payload sizes
+	let first_chunk_max_payload = max_message_size.saturating_sub(start_overhead);
+	let other_chunk_max_payload = max_message_size.saturating_sub(chunk_overhead);
+
+	if first_chunk_max_payload == 0 || other_chunk_max_payload == 0 {
+		bail!("message overhead exceeds max message size");
+	}
+
+	// Calculate how many chunks we need
+	if payload.len() <= first_chunk_max_payload {
+		// Single chunk - all data fits in first message
+		return Ok(vec![payload.to_vec()]);
+	}
+
+	// Multi-chunk: first chunk + remaining chunks
+	let remaining_after_first = payload.len() - first_chunk_max_payload;
+	let additional_chunks =
+		(remaining_after_first + other_chunk_max_payload - 1) / other_chunk_max_payload;
+
+	let mut chunks = Vec::new();
+
+	// First chunk (smaller due to reply_subject overhead)
+	chunks.push(payload[..first_chunk_max_payload].to_vec());
+
+	// Subsequent chunks
+	let mut offset = first_chunk_max_payload;
+	for _ in 0..additional_chunks {
+		let end = std::cmp::min(offset + other_chunk_max_payload, payload.len());
+		chunks.push(payload[offset..end].to_vec());
+		offset = end;
+	}
+
+	Ok(chunks)
+}
+
+/// Encodes a chunk to the resulting BARE message.
+pub fn encode_chunk(
+	payload: Vec<u8>,
+	chunk_idx: u32,
+	chunk_count: u32,
+	message_id: [u8; 16],
+	reply_subject: Option<String>,
+) -> Result<Vec<u8>> {
+	let body = if chunk_idx == 0 {
+		// First chunk - MessageStart
+		MessageBody::MessageStart(MessageStart {
+			message_id,
+			chunk_count,
+			reply_subject,
+			payload,
+		})
+	} else {
+		// Subsequent chunks - MessageChunk
+		MessageBody::MessageChunk(MessageChunk {
+			message_id,
+			chunk_index: chunk_idx,
+			payload,
+		})
+	};
+
+	let ups_message = rivet_ups_protocol::UpsMessage { body };
+	UpsMessage::wrap_latest(ups_message).serialize_with_embedded_version(PROTOCOL_VERSION)
+}

package/packages/universalpubsub/src/driver/memory/mod.rs

@@ -0,0 +1,91 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use anyhow::*;
+use async_trait::async_trait;
+use tokio::sync::{RwLock, mpsc};
+
+use crate::driver::{PubSubDriver, SubscriberDriver, SubscriberDriverHandle};
+use crate::pubsub::DriverOutput;
+
+type Subscribers = Arc<RwLock<HashMap<String, Vec<mpsc::UnboundedSender<Vec<u8>>>>>>;
+
+/// This is arbitrary.
+const MEMORY_MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024; // 10MiB
+
+#[derive(Clone)]
+pub struct MemoryDriver {
+	channel: String,
+	subscribers: Subscribers,
+}
+
+impl MemoryDriver {
+	pub fn new(channel: String) -> Self {
+		Self {
+			channel,
+			subscribers: Arc::new(RwLock::new(HashMap::new())),
+		}
+	}
+
+	fn subject_with_channel(&self, subject: &str) -> String {
+		format!("{}::{}", self.channel, subject)
+	}
+}
+
+#[async_trait]
+impl PubSubDriver for MemoryDriver {
+	async fn subscribe(&self, subject: &str) -> Result<SubscriberDriverHandle> {
+		let (tx, rx) = mpsc::unbounded_channel();
+		let subject_with_channel = self.subject_with_channel(subject);
+
+		let mut subscribers = self.subscribers.write().await;
+		subscribers
+			.entry(subject_with_channel.clone())
+			.or_default()
+			.push(tx);
+
+		Ok(Box::new(MemorySubscriber {
+			subject: subject_with_channel,
+			rx,
+		}))
+	}
+
+	async fn publish(&self, subject: &str, payload: &[u8]) -> Result<()> {
+		let subject_with_channel = self.subject_with_channel(subject);
+		let subscribers = self.subscribers.read().await;
+
+		if let Some(subs) = subscribers.get(&subject_with_channel) {
+			for tx in subs {
+				let _ = tx.send(payload.to_vec());
+			}
+		}
+
+		Ok(())
+	}
+
+	async fn flush(&self) -> Result<()> {
+		Ok(())
+	}
+
+	fn max_message_size(&self) -> usize {
+		MEMORY_MAX_MESSAGE_SIZE
+	}
+}
+
+pub struct MemorySubscriber {
+	subject: String,
+	rx: mpsc::UnboundedReceiver<Vec<u8>>,
+}
+
+#[async_trait]
+impl SubscriberDriver for MemorySubscriber {
+	async fn next(&mut self) -> Result<DriverOutput> {
+		match self.rx.recv().await {
+			Some(payload) => Ok(DriverOutput::Message {
+				subject: self.subject.clone(),
+				payload,
+			}),
+			None => Ok(DriverOutput::Unsubscribed),
+		}
+	}
+}

package/packages/universalpubsub/src/driver/mod.rs

@@ -0,0 +1,58 @@
+use std::sync::Arc;
+
+use anyhow::*;
+use async_trait::async_trait;
+
+pub mod memory;
+pub mod nats;
+pub mod postgres;
+
+pub type PubSubDriverHandle = Arc<dyn PubSubDriver>;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum PublishBehavior {
+	/// Publishes a message to a single subscriber.
+	///
+	/// This should not be used if there will ever be more than one subscription at a time to the given topic
+	/// on a global scale. Its intended to enable in-memory optimizations where a subscription that exists on
+	/// the same machine as the published message will not have to communicate with the driver but can instead
+	/// be delivered entirely in-memory.
+	OneSubscriber,
+
+	/// Publishes a message to multiple subscribers.
+	Broadcast,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct PublishOpts {
+	pub behavior: PublishBehavior,
+}
+
+impl PublishOpts {
+	pub const fn one() -> Self {
+		Self {
+			behavior: PublishBehavior::OneSubscriber,
+		}
+	}
+
+	pub const fn broadcast() -> Self {
+		Self {
+			behavior: PublishBehavior::Broadcast,
+		}
+	}
+}
+
+#[async_trait]
+pub trait PubSubDriver: Send + Sync {
+	async fn subscribe(&self, subject: &str) -> Result<Box<dyn SubscriberDriver>>;
+	async fn publish(&self, subject: &str, message: &[u8]) -> Result<()>;
+	async fn flush(&self) -> Result<()>;
+	fn max_message_size(&self) -> usize;
+}
+
+pub type SubscriberDriverHandle = Box<dyn SubscriberDriver>;
+
+#[async_trait]
+pub trait SubscriberDriver: Send + Sync {
+	async fn next(&mut self) -> Result<crate::pubsub::DriverOutput>;
+}

package/packages/universalpubsub/src/driver/nats/mod.rs

@@ -0,0 +1,79 @@
+use anyhow::*;
+use async_nats::Client;
+use async_trait::async_trait;
+use futures_util::StreamExt;
+
+use crate::driver::{PubSubDriver, SubscriberDriver, SubscriberDriverHandle};
+use crate::pubsub::DriverOutput;
+
+/// > The size is set to 1 MB by default, but can be increased up to 64 MB if needed (though we recommend keeping the max message size to something more reasonable like 8 MB).
+///
+/// https://docs.nats.io/reference/faq#is-there-a-message-size-limitation-in-nats
+///
+/// When they say "MB" they mean "MiB." Ignorance strikes again.
+pub const NATS_MAX_MESSAGE_SIZE: usize = 1024 * 1024;
+
+#[derive(Clone)]
+pub struct NatsDriver {
+	client: Client,
+}
+
+impl NatsDriver {
+	pub async fn connect(
+		options: async_nats::ConnectOptions,
+		server_addrs: impl async_nats::ToServerAddrs,
+	) -> Result<Self> {
+		tracing::debug!("nats connecting");
+		let client = options.connect(server_addrs).await?;
+		tracing::debug!("nats connected");
+
+		Ok(Self { client })
+	}
+}
+
+#[async_trait]
+impl PubSubDriver for NatsDriver {
+	async fn subscribe(&self, subject: &str) -> Result<SubscriberDriverHandle> {
+		let subscriber = self.client.subscribe(subject.to_string()).await?;
+		Ok(Box::new(NatsSubscriber { subscriber }))
+	}
+
+	async fn publish(&self, subject: &str, payload: &[u8]) -> Result<()> {
+		self.client
+			.publish(subject.to_string(), payload.to_vec().into())
+			.await?;
+		Ok(())
+	}
+
+	async fn flush(&self) -> Result<()> {
+		self.client.flush().await?;
+		Ok(())
+	}
+
+	fn max_message_size(&self) -> usize {
+		NATS_MAX_MESSAGE_SIZE
+	}
+}
+
+pub struct NatsSubscriber {
+	subscriber: async_nats::Subscriber,
+}
+
+#[async_trait]
+impl SubscriberDriver for NatsSubscriber {
+	async fn next(&mut self) -> Result<DriverOutput> {
+		match self.subscriber.next().await {
+			Some(msg) => Ok(DriverOutput::Message {
+				subject: msg.subject.to_string(),
+				payload: msg.payload.to_vec(),
+			}),
+			None => Ok(DriverOutput::Unsubscribed),
+		}
+	}
+}
+
+impl Drop for NatsSubscriber {
+	fn drop(&mut self) {
+		let _ = self.subscriber.unsubscribe();
+	}
+}