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

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

@@ -1,30 +1,95 @@
-use std::collections::HashMap;
+use std::ops::Deref;
 use std::sync::Arc;
+use std::time::Duration;
 
-use anyhow::*;
+use anyhow::Result;
 use async_trait::async_trait;
-use tokio::sync::{RwLock, mpsc};
+use rand::seq::SliceRandom;
+use scc::HashMap;
+use tokio::sync::mpsc;
 
 use crate::driver::{PubSubDriver, SubscriberDriver, SubscriberDriverHandle};
+use crate::metrics;
 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
+const GC_INTERVAL: Duration = Duration::from_secs(60);
 
-#[derive(Clone)]
-pub struct MemoryDriver {
+pub struct MemoryDriverInner {
 	channel: String,
-	subscribers: Subscribers,
+	// Map<topic, Vec<sub>>
+	subscribers: HashMap<String, Vec<mpsc::UnboundedSender<Vec<u8>>>>,
+	// Map<topic, Map<queue, Vec<sub>>>
+	queue_subscribers: HashMap<String, HashMap<String, Vec<mpsc::UnboundedSender<Vec<u8>>>>>,
+}
+
+#[derive(Clone)]
+pub struct MemoryDriver(Arc<MemoryDriverInner>);
+
+impl Deref for MemoryDriver {
+	type Target = MemoryDriverInner;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
 }
 
 impl MemoryDriver {
 	pub fn new(channel: String) -> Self {
-		Self {
+		let inner = Arc::new(MemoryDriverInner {
 			channel,
-			subscribers: Arc::new(RwLock::new(HashMap::new())),
-		}
+			subscribers: HashMap::new(),
+			queue_subscribers: HashMap::new(),
+		});
+
+		// TODO: Why not use drop impl?
+		// Spawn GC task to clean up closed subscribers
+		let gc_inner = Arc::downgrade(&inner);
+		tokio::spawn(async move {
+			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 closed senders
+					inner
+						.subscribers
+						.retain_async(|_, senders| {
+							// Retain only senders that are not closed
+							senders.retain(|sender| !sender.is_closed());
+							// Remove the entire subject entry if no senders remain
+							!senders.is_empty()
+						})
+						.await;
+
+					inner
+						.queue_subscribers
+						.retain_async(|_, queues| {
+							// TODO: Is retain_sync bad here?
+							// Retain only queues with senders
+							queues.retain_sync(|_, senders| {
+								// Retain only senders that are not closed
+								senders.retain(|sender| !sender.is_closed());
+								// Remove the entire subject entry if no senders remain
+								!senders.is_empty()
+							});
+
+							// Remove the entire subject entry if no queues remain
+							!queues.is_empty()
+						})
+						.await;
+
+					metrics::MEMORY_SUBSCRIBER_COUNT
+						.set((inner.subscribers.len() + inner.queue_subscribers.len()) as i64);
+				} else {
+					break;
+				}
+			}
+		});
+
+		Self(inner)
 	}
 
 	fn subject_with_channel(&self, subject: &str) -> String {
@@ -38,11 +103,34 @@ impl PubSubDriver for MemoryDriver {
 		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())
+		self.subscribers
+			.entry_async(subject_with_channel.clone())
+			.await
+			.or_default()
+			.push(tx);
+		metrics::MEMORY_SUBSCRIBER_COUNT
+			.set((self.subscribers.len() + self.queue_subscribers.len()) as i64);
+
+		Ok(Box::new(MemorySubscriber {
+			subject: subject_with_channel,
+			rx,
+		}))
+	}
+
+	async fn queue_subscribe(&self, subject: &str, queue: &str) -> Result<SubscriberDriverHandle> {
+		let (tx, rx) = mpsc::unbounded_channel();
+		let subject_with_channel = self.subject_with_channel(subject);
+
+		self.queue_subscribers
+			.entry_async(subject_with_channel.clone())
+			.await
+			.or_default()
+			.entry_async(queue.to_string())
+			.await
 			.or_default()
 			.push(tx);
+		metrics::MEMORY_SUBSCRIBER_COUNT
+			.set((self.subscribers.len() + self.queue_subscribers.len()) as i64);
 
 		Ok(Box::new(MemorySubscriber {
 			subject: subject_with_channel,
@@ -52,13 +140,36 @@ impl PubSubDriver for MemoryDriver {
 
 	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());
-			}
-		}
+		tokio::join!(
+			// Send to subs
+			async {
+				if let Some(subs) = self.subscribers.get_async(&subject_with_channel).await {
+					for tx in &*subs {
+						let _ = tx.send(payload.to_vec());
+					}
+				}
+			},
+			// Send to queue subs
+			async {
+				if let Some(queues) = self
+					.queue_subscribers
+					.get_async(&subject_with_channel)
+					.await
+				{
+					queues
+						.iter_async(|_, subs| {
+							// Choose random sub to receive message
+							if let Some(tx) = subs.choose(&mut rand::thread_rng()) {
+								let _ = tx.send(payload.to_vec());
+							}
+
+							true
+						})
+						.await;
+				}
+			},
+		);
 
 		Ok(())
 	}

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

@@ -45,6 +45,11 @@ impl PublishOpts {
 #[async_trait]
 pub trait PubSubDriver: Send + Sync {
 	async fn subscribe(&self, subject: &str) -> Result<Box<dyn SubscriberDriver>>;
+	async fn queue_subscribe(
+		&self,
+		subject: &str,
+		queue: &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;

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

@@ -38,6 +38,14 @@ impl PubSubDriver for NatsDriver {
 		Ok(Box::new(NatsSubscriber { subscriber }))
 	}
 
+	async fn queue_subscribe(&self, subject: &str, queue: &str) -> Result<SubscriberDriverHandle> {
+		let subscriber = self
+			.client
+			.queue_subscribe(subject.to_string(), queue.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())