chore(pegboard): consolidate to single subscriber per gateway

Author: NathanFlurry

Cargo.lock

@@ -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))))
 }

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

@@ -0,0 +1,32 @@
+use anyhow::*;
+use gas::prelude::*;
+use std::{ops::Deref, sync::Arc};
+use universalpubsub::PubSub;
+
+#[derive(Clone)]
+pub struct SharedState(Arc<SharedStateInner>);
+
+impl SharedState {
+	pub fn new(pubsub: PubSub) -> SharedState {
+		SharedState(Arc::new(SharedStateInner {
+			pegboard_gateway: pegboard_gateway::shared_state::SharedState::new(pubsub),
+		}))
+	}
+
+	pub async fn start(&self) -> Result<()> {
+		self.pegboard_gateway.start().await?;
+		Ok(())
+	}
+}
+
+impl Deref for SharedState {
+	type Target = SharedStateInner;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
+}
+
+pub struct SharedStateInner {
+	pub pegboard_gateway: pegboard_gateway::shared_state::SharedState,
+}

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

@@ -14,12 +14,14 @@ gas.workspace = true
 http-body-util.workspace = true
 hyper = "1.6"
 hyper-tungstenite.workspace = true
-pegboard = { path = "../../services/pegboard" }
+pegboard.workspace = true
 rand.workspace = true
 rivet-error.workspace = true
-rivet-guard-core = { path = "../guard/core" }
+rivet-guard-core.workspace = true
 rivet-tunnel-protocol.workspace = true
 rivet-util.workspace = true
 tokio-tungstenite.workspace = true
 tokio.workspace = true
 universalpubsub.workspace = true
+versioned-data-util.workspace = true
+thiserror.workspace = true