chore(pegboard): consolidate to single subscriber per gateway

Author: NathanFlurry

Cargo.lock

@@ -12,7 +12,6 @@ async-trait.workspace = true
 base64.workspace = true
 deadpool-postgres.workspace = true
 futures-util.workspace = true
-moka.workspace = true
 rivet-error.workspace = true
 rivet-ups-protocol.workspace = true
 serde_json.workspace = true
@@ -21,6 +20,7 @@ serde.workspace = true
 sha2.workspace = true
 tokio-postgres.workspace = true
 tokio.workspace = true
+tokio-util.workspace = true
 tracing.workspace = true
 uuid.workspace = true
 

packages/common/universalpubsub/Cargo.toml

@@ -1,13 +1,13 @@
+use std::collections::HashMap;
 use std::hash::{DefaultHasher, Hash, Hasher};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use anyhow::*;
 use async_trait::async_trait;
 use base64::Engine;
 use base64::engine::general_purpose::STANDARD_NO_PAD as BASE64;
 use deadpool_postgres::{Config, ManagerConfig, Pool, PoolConfig, RecyclingMethod, Runtime};
 use futures_util::future::poll_fn;
-use moka::future::Cache;
 use tokio_postgres::{AsyncMessage, NoTls};
 use tracing::Instrument;
 
@@ -18,6 +18,15 @@ use crate::pubsub::DriverOutput;
 struct Subscription {
 	// Channel to send messages to this subscription
 	tx: tokio::sync::broadcast::Sender<Vec<u8>>,
+	// Cancellation token shared by all subscribers of this subject
+	token: tokio_util::sync::CancellationToken,
+}
+
+impl Subscription {
+	fn new(tx: tokio::sync::broadcast::Sender<Vec<u8>>) -> Self {
+		let token = tokio_util::sync::CancellationToken::new();
+		Self { tx, token }
+	}
 }
 
 /// > In the default configuration it must be shorter than 8000 bytes
@@ -40,7 +49,7 @@ pub const POSTGRES_MAX_MESSAGE_SIZE: usize =
 pub struct PostgresDriver {
 	pool: Arc<Pool>,
 	client: Arc<tokio_postgres::Client>,
-	subscriptions: Cache<String, Subscription>,
+	subscriptions: Arc<Mutex<HashMap<String, Subscription>>>,
 }
 
 impl PostgresDriver {
@@ -65,8 +74,8 @@ impl PostgresDriver {
 			.context("failed to create postgres pool")?;
 		tracing::debug!("postgres pool created successfully");
 
-		let subscriptions: Cache<String, Subscription> =
-			Cache::builder().initial_capacity(5).build();
+		let subscriptions: Arc<Mutex<HashMap<String, Subscription>>> =
+			Arc::new(Mutex::new(HashMap::new()));
 		let subscriptions2 = subscriptions.clone();
 
 		let (client, mut conn) = tokio_postgres::connect(&conn_str, tokio_postgres::NoTls).await?;
@@ -75,7 +84,9 @@ impl PostgresDriver {
 			loop {
 				match poll_fn(|cx| conn.poll_message(cx)).await {
 					Some(std::result::Result::Ok(AsyncMessage::Notification(note))) => {
-						if let Some(sub) = subscriptions2.get(note.channel()).await {
+						if let Some(sub) =
+							subscriptions2.lock().unwrap().get(note.channel()).cloned()
+						{
 							let bytes = match BASE64.decode(note.payload()) {
 								std::result::Result::Ok(b) => b,
 								std::result::Result::Err(err) => {
@@ -121,7 +132,7 @@ impl PostgresDriver {
 #[async_trait]
 impl PubSubDriver for PostgresDriver {
 	async fn subscribe(&self, subject: &str) -> Result<SubscriberDriverHandle> {
-		// TODO: To match NATS implementation, LIST must be pipelined (i.e. wait for the command
+		// TODO: To match NATS implementation, LISTEN must be pipelined (i.e. wait for the command
 		// to reach the server, but not wait for it to respond). However, this has to ensure that
 		// NOTIFY & LISTEN are called on the same connection (not diff connections in a pool) or
 		// else there will be race conditions where messages might be published before
@@ -135,33 +146,57 @@ impl PubSubDriver for PostgresDriver {
 		let hashed = self.hash_subject(subject);
 
 		// Check if we already have a subscription for this channel
-		let rx = if let Some(existing_sub) = self.subscriptions.get(&hashed).await {
-			// Reuse the existing broadcast channel
-			existing_sub.tx.subscribe()
-		} else {
-			// Create a new broadcast channel for this subject
-			let (tx, rx) = tokio::sync::broadcast::channel(1024);
-			let subscription = Subscription { tx: tx.clone() };
-
-			// Register subscription
-			self.subscriptions
-				.insert(hashed.clone(), subscription)
-				.await;
-
-			// Execute LISTEN command on the async client (for receiving notifications)
-			// This only needs to be done once per channel
-			let span = tracing::trace_span!("pg_listen");
-			self.client
-				.execute(&format!("LISTEN \"{hashed}\""), &[])
-				.instrument(span)
-				.await?;
-
-			rx
-		};
+		let (rx, drop_guard) =
+			if let Some(existing_sub) = self.subscriptions.lock().unwrap().get(&hashed).cloned() {
+				// Reuse the existing broadcast channel
+				let rx = existing_sub.tx.subscribe();
+				let drop_guard = existing_sub.token.clone().drop_guard();
+				(rx, drop_guard)
+			} else {
+				// Create a new broadcast channel for this subject
+				let (tx, rx) = tokio::sync::broadcast::channel(1024);
+				let subscription = Subscription::new(tx.clone());
+
+				// Register subscription
+				self.subscriptions
+					.lock()
+					.unwrap()
+					.insert(hashed.clone(), subscription.clone());
+
+				// Execute LISTEN command on the async client (for receiving notifications)
+				// This only needs to be done once per channel
+				let span = tracing::trace_span!("pg_listen");
+				self.client
+					.execute(&format!("LISTEN \"{hashed}\""), &[])
+					.instrument(span)
+					.await?;
+
+				// Spawn a single cleanup task for this subscription waiting on its token
+				let driver = self.clone();
+				let hashed_clone = hashed.clone();
+				let tx_clone = tx.clone();
+				let token_clone = subscription.token.clone();
+				tokio::spawn(async move {
+					token_clone.cancelled().await;
+					if tx_clone.receiver_count() == 0 {
+						let sql = format!("UNLISTEN \"{}\"", hashed_clone);
+						if let Err(err) = driver.client.execute(sql.as_str(), &[]).await {
+							tracing::warn!(?err, %hashed_clone, "failed to UNLISTEN channel");
+						} else {
+							tracing::trace!(%hashed_clone, "unlistened channel");
+						}
+						driver.subscriptions.lock().unwrap().remove(&hashed_clone);
+					}
+				});
+
+				let drop_guard = subscription.token.clone().drop_guard();
+				(rx, drop_guard)
+			};
 
 		Ok(Box::new(PostgresSubscriber {
 			subject: subject.to_string(),
-			rx,
+			rx: Some(rx),
+			_drop_guard: drop_guard,
 		}))
 	}
 
@@ -191,13 +226,18 @@ impl PubSubDriver for PostgresDriver {
 
 pub struct PostgresSubscriber {
 	subject: String,
-	rx: tokio::sync::broadcast::Receiver<Vec<u8>>,
+	rx: Option<tokio::sync::broadcast::Receiver<Vec<u8>>>,
+	_drop_guard: tokio_util::sync::DropGuard,
 }
 
 #[async_trait]
 impl SubscriberDriver for PostgresSubscriber {
 	async fn next(&mut self) -> Result<DriverOutput> {
-		match self.rx.recv().await {
+		let rx = match self.rx.as_mut() {
+			Some(rx) => rx,
+			None => return Ok(DriverOutput::Unsubscribed),
+		};
+		match rx.recv().await {
 			std::result::Result::Ok(payload) => Ok(DriverOutput::Message {
 				subject: self.subject.clone(),
 				payload,

packages/common/universalpubsub/src/driver/postgres/mod.rs

@@ -20,6 +20,7 @@ hyper-tungstenite.workspace = true
 tower.workspace = true
 udb-util.workspace = true
 universaldb.workspace = true
+universalpubsub.workspace = true
 futures.workspace = true
 # TODO: Make this use workspace version
 hyper = "1.6.0"

packages/core/guard/server/Cargo.toml

@@ -5,6 +5,7 @@ pub mod cache;
 pub mod errors;
 pub mod middleware;
 pub mod routing;
+pub mod shared_state;
 pub mod tls;
 
 #[tracing::instrument(skip_all)]
@@ -26,8 +27,12 @@ pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> R
 		tracing::warn!("crypto provider already installed in this process");
 	}
 
+	// Share shared context
+	let shared_state = shared_state::SharedState::new(ctx.ups()?);
+	shared_state.start().await?;
+
 	// Create handlers
-	let routing_fn = routing::create_routing_function(ctx.clone());
+	let routing_fn = routing::create_routing_function(ctx.clone(), shared_state.clone());
 	let cache_key_fn = cache::create_cache_key_function(ctx.clone());
 	let middleware_fn = middleware::create_middleware_function(ctx.clone());
 	let cert_resolver = tls::create_cert_resolver(&ctx).await?;

packages/core/guard/server/src/lib.rs

@@ -5,7 +5,7 @@ use gas::prelude::*;
 use hyper::header::HeaderName;
 use rivet_guard_core::RoutingFn;
 
-use crate::errors;
+use crate::{errors, shared_state::SharedState};
 
 mod api_peer;
 mod api_public;
@@ -17,13 +17,14 @@ pub(crate) const X_RIVET_TARGET: HeaderName = HeaderName::from_static("x-rivet-t
 
 /// Creates the main routing function that handles all incoming requests
 #[tracing::instrument(skip_all)]
-pub fn create_routing_function(ctx: StandaloneCtx) -> RoutingFn {
+pub fn create_routing_function(ctx: StandaloneCtx, shared_state: SharedState) -> RoutingFn {
 	Arc::new(
 		move |hostname: &str,
 		      path: &str,
 		      port_type: rivet_guard_core::proxy_service::PortType,
 		      headers: &hyper::HeaderMap| {
 			let ctx = ctx.clone();
+			let shared_state = shared_state.clone();
 
 			Box::pin(
 				async move {
@@ -41,9 +42,15 @@ pub fn create_routing_function(ctx: StandaloneCtx) -> RoutingFn {
 							return Ok(routing_output);
 						}
 
-						if let Some(routing_output) =
-							pegboard_gateway::route_request(&ctx, target, host, path, headers)
-								.await?
+						if let Some(routing_output) = pegboard_gateway::route_request(
+							&ctx,
+							&shared_state,
+							target,
+							host,
+							path,
+							headers,
+						)
+						.await?
 						{
 							return Ok(routing_output);
 						}

packages/core/guard/server/src/routing/mod.rs

@@ -6,7 +6,7 @@ use hyper::header::HeaderName;
 use rivet_guard_core::proxy_service::{RouteConfig, RouteTarget, RoutingOutput, RoutingTimeout};
 use udb_util::{SERIALIZABLE, TxnExt};
 
-use crate::errors;
+use crate::{errors, shared_state::SharedState};
 
 const ACTOR_READY_TIMEOUT: Duration = Duration::from_secs(10);
 pub const X_RIVET_ACTOR: HeaderName = HeaderName::from_static("x-rivet-actor");
@@ -16,6 +16,7 @@ pub const X_RIVET_PORT: HeaderName = HeaderName::from_static("x-rivet-port");
 #[tracing::instrument(skip_all)]
 pub async fn route_request(
 	ctx: &StandaloneCtx,
+	shared_state: &SharedState,
 	target: &str,
 	_host: &str,
 	path: &str,
@@ -73,7 +74,7 @@ pub async fn route_request(
 	let port_name = port_name.to_str()?;
 
 	// Lookup actor
-	find_actor(ctx, actor_id, port_name, path).await
+	find_actor(ctx, shared_state, actor_id, port_name, path).await
 }
 
 struct FoundActor {
@@ -86,6 +87,7 @@ struct FoundActor {
 #[tracing::instrument(skip_all, fields(%actor_id, %port_name, %path))]
 async fn find_actor(
 	ctx: &StandaloneCtx,
+	shared_state: &SharedState,
 	actor_id: Id,
 	port_name: &str,
 	path: &str,
@@ -158,10 +160,10 @@ async fn find_actor(
 		actor_ids: vec![actor_id],
 	});
 	let res = tokio::time::timeout(Duration::from_secs(5), get_runner_fut).await??;
-	let runner_info = res.actors.into_iter().next().filter(|x| x.is_connectable);
+	let actor = res.actors.into_iter().next().filter(|x| x.is_connectable);
 
-	let runner_id = if let Some(runner_info) = runner_info {
-		runner_info.runner_id
+	let runner_id = if let Some(actor) = actor {
+		actor.runner_id
 	} else {
 		tracing::info!(?actor_id, "waiting for actor to become ready");
 
@@ -185,11 +187,23 @@ async fn find_actor(
 
 	tracing::debug!(?actor_id, ?runner_id, "actor ready");
 
+	// Get runner key from runner_id
+	let runner_key = ctx
+		.udb()?
+		.run(|tx, _mc| async move {
+			let txs = tx.subspace(pegboard::keys::subspace());
+			let key_key = pegboard::keys::runner::KeyKey::new(runner_id);
+			txs.read_opt(&key_key, SERIALIZABLE).await
+		})
+		.await?
+		.context("runner key not found")?;
+
 	// Return pegboard-gateway instance
 	let gateway = pegboard_gateway::PegboardGateway::new(
 		ctx.clone(),
+		shared_state.pegboard_gateway.clone(),
 		actor_id,
-		runner_id,
+		runner_key,
 		port_name.to_string(),
 	);
 	Ok(Some(RoutingOutput::CustomServe(std::sync::Arc::new(

packages/core/guard/server/src/routing/pegboard_gateway.rs

@@ -12,13 +12,10 @@ pub async fn route_request(
 	_host: &str,
 	_path: &str,
 ) -> Result<Option<RoutingOutput>> {
-	// Check target
 	if target != "tunnel" {
 		return Ok(None);
 	}
 
-	// Create pegboard-tunnel service instance
-	let tunnel = pegboard_tunnel::PegboardTunnelCustomServe::new(ctx.clone()).await?;
-
+	let tunnel = pegboard_tunnel::PegboardTunnelCustomServe::new(ctx.clone());
 	Ok(Some(RoutingOutput::CustomServe(Arc::new(tunnel))))
 }