@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/packages/universalpubsub/src/lib.rs

@@ -1,7 +1,10 @@
 pub mod chunking;
 pub mod driver;
 pub mod errors;
+pub mod metrics;
 pub mod pubsub;
+pub mod subject;
 
 pub use driver::*;
 pub use pubsub::{Message, NextOutput, PubSub, Response, Subscriber};
+pub use subject::Subject;

package/packages/universalpubsub/src/metrics.rs

@@ -0,0 +1,60 @@
+use rivet_metrics::{REGISTRY, prometheus::*};
+
+lazy_static::lazy_static! {
+	pub static ref SUBSCRIBER_COUNT: IntCounterVec = register_int_counter_vec_with_registry!(
+		"ups_subscriber_count",
+		"Total number of subscribers ever created by subject.",
+		&["subject"],
+		*REGISTRY
+	).unwrap();
+	pub static ref ACTIVE_SUBSCRIBER_COUNT: IntGaugeVec = register_int_gauge_vec_with_registry!(
+		"ups_active_subscriber_count",
+		"Number of active subscribers by subject.",
+		&["subject"],
+		*REGISTRY
+	).unwrap();
+	// PubSub layer metrics
+	pub static ref LOCAL_SUBSCRIBER_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"ups_local_subscriber_count",
+		"Number of local in-memory subscriber entries.",
+		*REGISTRY
+	).unwrap();
+	pub static ref REPLY_SUBSCRIBER_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"ups_reply_subscriber_count",
+		"Number of pending reply subscriber entries.",
+		*REGISTRY
+	).unwrap();
+	// Memory driver metrics
+	pub static ref MEMORY_SUBSCRIBER_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"ups_memory_subscriber_count",
+		"Number of subject entries in the memory driver.",
+		*REGISTRY
+	).unwrap();
+	// Postgres driver metrics
+	pub static ref POSTGRES_SUBSCRIPTION_COUNT: IntGauge = register_int_gauge_with_registry!(
+		"ups_postgres_subscription_count",
+		"Number of subscription entries in the postgres driver.",
+		*REGISTRY
+	).unwrap();
+
+	// Message metrics
+	pub static ref MESSAGE_RECV_COUNT: IntCounterVec = register_int_counter_vec_with_registry!(
+		"ups_message_recv_count",
+		"Total number of messages ever received by subject.",
+		&["subject"],
+		*REGISTRY
+	).unwrap();
+	pub static ref MESSAGE_SEND_COUNT: IntCounterVec = register_int_counter_vec_with_registry!(
+		"ups_message_send_count",
+		"Total number of messages ever sent by subject.",
+		&["kind", "subject"],
+		*REGISTRY
+	).unwrap();
+	pub static ref BYTES_PER_MESSAGE: HistogramVec = register_histogram_vec_with_registry!(
+		"ups_bytes_per_message",
+		"Amount of bytes per message received.",
+		&["subject"],
+		vec![16.0, 32.0, 64.0, 128.0, 256.0, 1024.0, 4096.0, 16384.0, 65536.0, 262144.0, 1048576.0, 4194304.0, 16777216.0],
+		*REGISTRY
+	).unwrap();
+}

package/packages/universalpubsub/src/pubsub.rs

@@ -1,24 +1,28 @@
-use std::collections::HashMap;
 use std::ops::Deref;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
-use anyhow::*;
-use tokio::sync::broadcast;
-use tokio::sync::{RwLock, oneshot};
+use anyhow::Result;
+use scc::HashMap;
+use tokio::sync::{broadcast, oneshot};
+use tracing::Instrument;
 use uuid::Uuid;
 
 use rivet_util::backoff::Backoff;
 
 use crate::chunking::{ChunkTracker, encode_chunk, split_payload_into_chunks};
 use crate::driver::{PubSubDriverHandle, PublishOpts, SubscriberDriverHandle};
+use crate::metrics;
+use crate::subject::Subject;
+
+const GC_INTERVAL: Duration = Duration::from_secs(60);
 
 pub struct PubSubInner {
 	driver: PubSubDriverHandle,
 	chunk_tracker: Mutex<ChunkTracker>,
-	reply_subscribers: RwLock<HashMap<String, oneshot::Sender<Vec<u8>>>>,
+	reply_subscribers: HashMap<String, oneshot::Sender<Vec<u8>>>,
 	// Local in-memory subscribers by subject (shared across all drivers)
-	local_subscribers: RwLock<HashMap<String, broadcast::Sender<Vec<u8>>>>,
+	local_subscribers: HashMap<String, broadcast::Sender<Vec<u8>>>,
 	// Enables/disables local fast-path across all drivers
 	memory_optimization: bool,
 }
@@ -46,19 +50,29 @@ impl PubSub {
 		let inner = Arc::new(PubSubInner {
 			driver,
 			chunk_tracker: Mutex::new(ChunkTracker::new()),
-			reply_subscribers: RwLock::new(HashMap::new()),
-			local_subscribers: RwLock::new(HashMap::new()),
+			reply_subscribers: HashMap::new(),
+			local_subscribers: HashMap::new(),
 			memory_optimization,
 		});
 
-		// Spawn GC task for chunk buffers
+		// Spawn GC task for chunk buffers and local subscribers
 		let gc_inner = Arc::downgrade(&inner);
 		tokio::spawn(async move {
-			let mut interval = tokio::time::interval(crate::chunking::CHUNK_BUFFER_GC_INTERVAL);
+			let mut interval = tokio::time::interval(GC_INTERVAL);
+			interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+
 			loop {
 				interval.tick().await;
 				if let Some(inner) = gc_inner.upgrade() {
+					// Clean up chunk buffers
 					inner.chunk_tracker.lock().unwrap().gc();
+
+					// Clean up local subscribers with no receivers
+					inner
+						.local_subscribers
+						.retain_async(|_, sender| sender.receiver_count() > 0)
+						.await;
+					metrics::LOCAL_SUBSCRIBER_COUNT.set(inner.local_subscribers.len() as i64);
 				} else {
 					break;
 				}
@@ -68,83 +82,94 @@ impl PubSub {
 		Self(inner)
 	}
 
-	pub async fn subscribe(&self, subject: &str) -> Result<Subscriber> {
+	#[tracing::instrument(skip_all, fields(%subject))]
+	pub async fn subscribe<T: Subject>(&self, subject: T) -> Result<Subscriber> {
 		// Underlying driver subscription
-		let driver = self.driver.subscribe(subject).await?;
+		let driver = self.driver.subscribe(&subject.as_cow()).await?;
 
 		if !self.memory_optimization {
-			return Ok(Subscriber::new(driver, self.clone()));
+			return Ok(Subscriber::new(
+				driver,
+				self.clone(),
+				false,
+				subject.to_string(),
+				T::root().map(|x| x.to_string()),
+			));
 		}
 
 		// Ensure a local broadcast channel exists for this subject
-		let rx = {
-			// Try read first for fast path
-			if let Some(sender) = self.local_subscribers.read().await.get(subject) {
-				sender.subscribe()
-			} else {
-				// Create and insert
-				let (tx, rx) = broadcast::channel(1024);
-				let mut map = self.local_subscribers.write().await;
-				// Double-check in case of race
-				let rx = if let Some(existing) = map.get(subject) {
-					existing.subscribe()
-				} else {
-					map.insert(subject.to_string(), tx);
-					rx
-				};
-				rx
-			}
+		let local_rx = {
+			let rx = self
+				.local_subscribers
+				.entry_async(subject.to_string())
+				.await
+				.or_insert_with(|| broadcast::channel(1024).0)
+				.subscribe();
+			metrics::LOCAL_SUBSCRIBER_COUNT.set(self.local_subscribers.len() as i64);
+			rx
 		};
 
 		// Wrap the driver
 		let optimized_driver: SubscriberDriverHandle = Box::new(LocalOptimizedSubscriberDriver {
 			subject: subject.to_string(),
 			driver,
-			local_rx: rx,
+			local_rx,
 		});
 
-		Ok(Subscriber::new(optimized_driver, self.clone()))
+		Ok(Subscriber::new(
+			optimized_driver,
+			self.clone(),
+			true,
+			subject.to_string(),
+			T::root().map(|x| x.to_string()),
+		))
 	}
 
-	pub async fn publish(&self, subject: &str, payload: &[u8], opts: PublishOpts) -> Result<()> {
-		let message_id = *Uuid::new_v4().as_bytes();
-		let chunks =
-			split_payload_into_chunks(payload, self.driver.max_message_size(), message_id, None)?;
-		let chunk_count = chunks.len() as u32;
-
-		let use_local = self
-			.should_use_local_subscriber(subject, opts.behavior)
-			.await;
+	#[tracing::instrument(skip_all, fields(%subject))]
+	pub async fn queue_subscribe<T: Subject>(&self, subject: T, queue: &str) -> Result<Subscriber> {
+		// Underlying driver subscription
+		let driver = self
+			.driver
+			.queue_subscribe(&subject.as_cow(), queue)
+			.await?;
 
-		for (chunk_idx, chunk_payload) in chunks.into_iter().enumerate() {
-			let encoded = encode_chunk(
-				chunk_payload,
-				chunk_idx as u32,
-				chunk_count,
-				message_id,
-				None,
-			)?;
+		return Ok(Subscriber::new(
+			driver,
+			self.clone(),
+			false,
+			subject.to_string(),
+			T::root().map(|x| x.to_string()),
+		));
+	}
 
-			if use_local {
-				if let Some(sender) = self.local_subscribers.read().await.get(subject) {
-					let _ = sender.send(encoded);
-				} else {
-					tracing::warn!(%subject, "local subscriber disappeared");
-					break;
-				}
-			} else {
-				// Use backoff when publishing through the driver
-				self.publish_with_backoff(subject, &encoded).await?;
-			}
-		}
-		Ok(())
+	#[tracing::instrument(skip_all, fields(%subject))]
+	pub async fn publish(
+		&self,
+		subject: impl Subject,
+		payload: &[u8],
+		opts: PublishOpts,
+	) -> Result<()> {
+		self.publish_inner(subject, payload, None::<&str>, opts)
+			.await
 	}
 
+	#[tracing::instrument(skip_all, fields(%subject, %reply_subject))]
 	pub async fn publish_with_reply(
 		&self,
-		subject: &str,
+		subject: impl Subject,
 		payload: &[u8],
-		reply_subject: &str,
+		reply_subject: impl Subject,
+		opts: PublishOpts,
+	) -> Result<()> {
+		self.publish_inner(subject, payload, Some(reply_subject), opts)
+			.await
+	}
+
+	async fn publish_inner<T: Subject>(
+		&self,
+		subject: T,
+		payload: &[u8],
+		reply_subject: Option<impl Subject>,
 		opts: PublishOpts,
 	) -> Result<()> {
 		let message_id = *Uuid::new_v4().as_bytes();
@@ -152,25 +177,27 @@ impl PubSub {
 			payload,
 			self.driver.max_message_size(),
 			message_id,
-			Some(reply_subject),
+			reply_subject.as_ref().map(|x| x.as_cow()).as_deref(),
 		)?;
 		let chunk_count = chunks.len() as u32;
 
 		let use_local = self
-			.should_use_local_subscriber(subject, opts.behavior)
+			.should_use_local_subscriber(&subject, opts.behavior)
 			.await;
 
+		let subject_cow = subject.as_cow();
+
 		for (chunk_idx, chunk_payload) in chunks.into_iter().enumerate() {
 			let encoded = encode_chunk(
 				chunk_payload,
 				chunk_idx as u32,
 				chunk_count,
 				message_id,
-				Some(reply_subject.to_string()),
+				reply_subject.as_ref().map(|x| x.to_string()),
 			)?;
 
 			if use_local {
-				if let Some(sender) = self.local_subscribers.read().await.get(subject) {
+				if let Some(sender) = self.local_subscribers.get_async(&*subject_cow).await {
 					let _ = sender.send(encoded);
 				} else {
 					tracing::warn!(%subject, "local subscriber disappeared");
@@ -178,17 +205,35 @@ impl PubSub {
 				}
 			} else {
 				// Use backoff when publishing through the driver
-				self.publish_with_backoff(subject, &encoded).await?;
+				self.publish_with_backoff(&subject, &encoded).await?;
 			}
 		}
+
+		let subject_str = T::root();
+		let subject_str = subject_str.as_deref().unwrap_or("unknown");
+		if use_local {
+			metrics::MESSAGE_SEND_COUNT
+				.with_label_values(&["local", subject_str])
+				.inc();
+		} else {
+			metrics::MESSAGE_SEND_COUNT
+				.with_label_values(&["driver", subject_str])
+				.inc();
+		}
+
 		Ok(())
 	}
 
-	async fn publish_with_backoff(&self, subject: &str, encoded: &[u8]) -> Result<()> {
+	#[tracing::instrument(skip_all, fields(%subject))]
+	async fn publish_with_backoff(&self, subject: &impl Subject, encoded: &[u8]) -> Result<()> {
+		let subject = subject.as_cow();
+
 		let mut backoff = Backoff::default();
 		loop {
-			match self.driver.publish(subject, encoded).await {
-				Result::Ok(_) => break,
+			match self.driver.publish(&subject, encoded).await {
+				Result::Ok(_) => {
+					break;
+				}
 				Err(err) if !backoff.tick().await => {
 					tracing::warn!(?err, "error publishing, cannot retry again");
 					return Err(crate::errors::Ups::PublishFailed.build().into());
@@ -202,18 +247,21 @@ impl PubSub {
 		Ok(())
 	}
 
+	#[tracing::instrument(skip_all)]
 	pub async fn flush(&self) -> Result<()> {
 		self.driver.flush().await
 	}
 
-	pub async fn request(&self, subject: &str, payload: &[u8]) -> Result<Response> {
+	#[tracing::instrument(skip_all, fields(%subject))]
+	pub async fn request(&self, subject: impl Subject, payload: &[u8]) -> Result<Response> {
 		self.request_with_timeout(subject, payload, Duration::from_secs(30))
 			.await
 	}
 
+	#[tracing::instrument(skip_all, fields(%subject))]
 	pub async fn request_with_timeout(
 		&self,
-		subject: &str,
+		subject: impl Subject,
 		payload: &[u8],
 		timeout: Duration,
 	) -> Result<Response> {
@@ -224,10 +272,10 @@ impl PubSub {
 		let (tx, rx) = oneshot::channel();
 
 		// Register the reply handler
-		{
-			let mut subscribers = self.reply_subscribers.write().await;
-			subscribers.insert(reply_subject.clone(), tx);
-		}
+		self.reply_subscribers
+			.upsert_async(reply_subject.clone(), tx)
+			.await;
+		metrics::REPLY_SUBSCRIBER_COUNT.set(self.reply_subscribers.len() as i64);
 
 		// Subscribe to the reply subject (use local-aware subscribe)
 		let mut reply_subscriber = self.subscribe(&reply_subject).await?;
@@ -244,10 +292,14 @@ impl PubSub {
 				match reply_subscriber.next().await {
 					std::result::Result::Ok(NextOutput::Message(msg)) => {
 						// Already decoded; forward payload
-						let mut subscribers = inner.reply_subscribers.write().await;
-						if let Some(tx) = subscribers.remove(&reply_subject_clone) {
+						if let Some((_, tx)) = inner
+							.reply_subscribers
+							.remove_async(&reply_subject_clone)
+							.await
+						{
 							let _ = tx.send(msg.payload);
 						}
+						metrics::REPLY_SUBSCRIBER_COUNT.set(inner.reply_subscribers.len() as i64);
 						break;
 					}
 					std::result::Result::Ok(NextOutput::Unsubscribed)
@@ -261,12 +313,14 @@ impl PubSub {
 			std::result::Result::Ok(std::result::Result::Ok(payload)) => Response { payload },
 			std::result::Result::Ok(std::result::Result::Err(_)) => {
 				// Clean up the reply subscription
-				self.reply_subscribers.write().await.remove(&reply_subject);
+				self.reply_subscribers.remove_async(&reply_subject).await;
+				metrics::REPLY_SUBSCRIBER_COUNT.set(self.reply_subscribers.len() as i64);
 				return Err(crate::errors::Ups::RequestTimeout.build().into());
 			}
 			std::result::Result::Err(_) => {
 				// Timeout elapsed
-				self.reply_subscribers.write().await.remove(&reply_subject);
+				self.reply_subscribers.remove_async(&reply_subject).await;
+				metrics::REPLY_SUBSCRIBER_COUNT.set(self.reply_subscribers.len() as i64);
 				return Err(crate::errors::Ups::RequestTimeout.build().into());
 			}
 		};
@@ -274,9 +328,10 @@ impl PubSub {
 		Ok(response)
 	}
 
+	#[tracing::instrument(skip_all, fields(%subject))]
 	async fn should_use_local_subscriber(
 		&self,
-		subject: &str,
+		subject: &impl Subject,
 		behavior: crate::driver::PublishBehavior,
 	) -> bool {
 		// Local fast-path for one-subscriber behavior:
@@ -294,7 +349,7 @@ impl PubSub {
 		if !matches!(behavior, crate::driver::PublishBehavior::OneSubscriber) {
 			return false;
 		}
-		if let Some(sender) = self.local_subscribers.read().await.get(subject) {
+		if let Some(sender) = self.local_subscribers.get_async(&*subject.as_cow()).await {
 			sender.receiver_count() > 0
 		} else {
 			false
@@ -305,13 +360,41 @@ impl PubSub {
 pub struct Subscriber {
 	driver: SubscriberDriverHandle,
 	pubsub: PubSub,
+	memory_optimization: bool,
+	subject: String,
+	root_subject: Option<String>,
 }
 
 impl Subscriber {
-	pub fn new(driver: SubscriberDriverHandle, pubsub: PubSub) -> Self {
-		Self { driver, pubsub }
+	fn new(
+		driver: SubscriberDriverHandle,
+		pubsub: PubSub,
+		memory_optimization: bool,
+		subject: String,
+		root_subject: Option<String>,
+	) -> Self {
+		let subject_str = if let Some(root_subject) = &root_subject {
+			root_subject.as_str()
+		} else {
+			"unknown"
+		};
+		metrics::SUBSCRIBER_COUNT
+			.with_label_values(&[subject_str])
+			.inc();
+		metrics::ACTIVE_SUBSCRIBER_COUNT
+			.with_label_values(&[subject_str])
+			.inc();
+
+		Self {
+			driver,
+			pubsub,
+			memory_optimization,
+			subject,
+			root_subject,
+		}
 	}
 
+	#[tracing::instrument(skip_all)]
 	pub async fn next(&mut self) -> Result<NextOutput> {
 		loop {
 			match self.driver.next().await? {
@@ -323,6 +406,26 @@ impl Subscriber {
 					let mut tracker = self.pubsub.chunk_tracker.lock().unwrap();
 					match tracker.process_chunk(&payload) {
 						std::result::Result::Ok(Some((payload, reply_subject))) => {
+							metrics::MESSAGE_RECV_COUNT
+								.with_label_values(&[
+									if let Some(root_subject) = &self.root_subject {
+										root_subject.as_str()
+									} else {
+										"unknown"
+									},
+								])
+								.inc();
+
+							metrics::BYTES_PER_MESSAGE
+								.with_label_values(&[
+									if let Some(root_subject) = &self.root_subject {
+										root_subject.as_str()
+									} else {
+										"unknown"
+									},
+								])
+								.observe(payload.len() as f64);
+
 							return Ok(NextOutput::Message(Message {
 								pubsub: self.pubsub.clone(),
 								payload,
@@ -342,6 +445,32 @@ impl Subscriber {
 	}
 }
 
+impl Drop for Subscriber {
+	fn drop(&mut self) {
+		metrics::ACTIVE_SUBSCRIBER_COUNT
+			.with_label_values(&[if let Some(root_subject) = &self.root_subject {
+				root_subject.as_str()
+			} else {
+				"unknown"
+			}])
+			.dec();
+
+		// Clean up local subscriber entry immediately if memory_optimization was enabled
+		if self.memory_optimization {
+			let pubsub = self.pubsub.clone();
+			let subject = self.subject.clone();
+			tokio::spawn(async move {
+				if let Some(sender) = pubsub.local_subscribers.get_async(&subject).await {
+					if sender.receiver_count() == 0 {
+						let _ = sender.remove();
+						metrics::LOCAL_SUBSCRIBER_COUNT.set(pubsub.local_subscribers.len() as i64);
+					}
+				}
+			});
+		}
+	}
+}
+
 // Output from drivers (raw binary messages)
 pub enum DriverOutput {
 	Message { subject: String, payload: Vec<u8> },
@@ -354,6 +483,15 @@ pub enum NextOutput {
 	Unsubscribed,
 }
 
+impl From<NextOutput> for Option<Message> {
+	fn from(value: NextOutput) -> Self {
+		match value {
+			NextOutput::Message(msg) => Some(msg),
+			NextOutput::Unsubscribed => None,
+		}
+	}
+}
+
 pub struct Message {
 	pub pubsub: PubSub,
 	pub payload: Vec<u8>,
@@ -361,6 +499,7 @@ pub struct Message {
 }
 
 impl Message {
+	#[tracing::instrument(skip_all)]
 	pub async fn reply(&self, payload: &[u8]) -> Result<()> {
 		if let Some(ref reply_subject) = self.reply {
 			// Replies expect exactly one subscriber and should use local fast-path
@@ -385,6 +524,7 @@ struct LocalOptimizedSubscriberDriver {
 
 #[async_trait::async_trait]
 impl crate::driver::SubscriberDriver for LocalOptimizedSubscriberDriver {
+	#[tracing::instrument(skip_all)]
 	async fn next(&mut self) -> Result<DriverOutput> {
 		loop {
 			tokio::select! {
@@ -406,7 +546,7 @@ impl crate::driver::SubscriberDriver for LocalOptimizedSubscriberDriver {
 						}
 					}
 				}
-				res = self.driver.next() => {
+				res = self.driver.next().in_current_span() => {
 					return res;
 				}
 			}

package/packages/universalpubsub/src/subject.rs

@@ -0,0 +1,32 @@
+use std::{borrow::Cow, fmt::Display};
+
+pub trait Subject: Display {
+	/// Used for metrics.
+	fn root<'a>() -> Option<Cow<'a, str>> {
+		None
+	}
+
+	fn as_str(&self) -> Option<&str> {
+		None
+	}
+
+	fn as_cow<'a>(&'a self) -> Cow<'a, str> {
+		if let Some(subject) = self.as_str() {
+			Cow::Borrowed(subject)
+		} else {
+			Cow::Owned(self.to_string())
+		}
+	}
+}
+
+impl Subject for &str {
+	fn as_str(&self) -> Option<&str> {
+		Some(self)
+	}
+}
+
+impl Subject for &String {
+	fn as_str(&self) -> Option<&str> {
+		Some(self)
+	}
+}