chore(pegboard-gateway): add new message id format & add deprecated tunnel ack

Author: NathanFlurry

Cargo.lock

@@ -4,12 +4,13 @@ use bytes::Bytes;
 use http_body_util::Full;
 use hyper::{Request, Response};
 use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame;
-use uuid::Uuid;
 
 use crate::WebSocketHandle;
 use crate::proxy_service::ResponseBody;
 use crate::request_context::RequestContext;
 
+use pegboard::tunnel::id::RequestId;
+
 pub enum HibernationResult {
 	Continue,
 	Close,
@@ -23,6 +24,7 @@ pub trait CustomServeTrait: Send + Sync {
 		&self,
 		req: Request<Full<Bytes>>,
 		request_context: &mut RequestContext,
+		request_id: RequestId,
 	) -> Result<Response<ResponseBody>>;
 
 	/// Handle a WebSocket connection after upgrade. Supports connection retries.
@@ -33,7 +35,7 @@ pub trait CustomServeTrait: Send + Sync {
 		_path: &str,
 		_request_context: &mut RequestContext,
 		// Identifies the websocket across retries.
-		_unique_request_id: Uuid,
+		_unique_request_id: RequestId,
 		// True if this websocket is reconnecting after hibernation.
 		_after_hibernation: bool,
 	) -> Result<Option<CloseFrame>> {
@@ -44,7 +46,7 @@ pub trait CustomServeTrait: Send + Sync {
 	async fn handle_websocket_hibernation(
 		&self,
 		_websocket: WebSocketHandle,
-		_unique_request_id: Uuid,
+		_unique_request_id: RequestId,
 	) -> Result<HibernationResult> {
 		bail!("service does not support websocket hibernation");
 	}

engine/packages/guard-core/src/custom_serve.rs

@@ -12,6 +12,7 @@ pub mod websocket_handle;
 
 pub use cert_resolver::CertResolverFn;
 pub use custom_serve::CustomServeTrait;
+pub use pegboard::tunnel::id::{RequestId, generate_request_id};
 pub use proxy_service::{
 	CacheKeyFn, MiddlewareFn, ProxyService, ProxyState, RouteTarget, RoutingFn, RoutingOutput,
 };

engine/packages/guard-core/src/lib.rs

@@ -13,9 +13,11 @@ use rivet_error::{INTERNAL_ERROR, RivetError};
 use rivet_metrics::KeyValue;
 use rivet_util::Id;
 use serde_json;
+
+use pegboard::tunnel::id::{RequestId, generate_request_id};
 use std::{
 	borrow::Cow,
-	collections::HashMap as StdHashMap,
+	collections::{HashMap as StdHashMap, HashSet},
 	net::SocketAddr,
 	sync::Arc,
 	time::{Duration, Instant},
@@ -28,7 +30,6 @@ use tokio_tungstenite::tungstenite::{
 };
 use tracing::Instrument;
 use url::Url;
-use uuid::Uuid;
 
 use crate::{
 	WebSocketHandle,
@@ -349,6 +350,7 @@ pub struct ProxyState {
 	route_cache: RouteCache,
 	rate_limiters: Cache<(Id, std::net::IpAddr), Arc<Mutex<RateLimiter>>>,
 	in_flight_counters: Cache<(Id, std::net::IpAddr), Arc<Mutex<InFlightCounter>>>,
+	inflight_requests: Arc<Mutex<HashSet<RequestId>>>,
 	port_type: PortType,
 	clickhouse_inserter: Option<clickhouse_inserter::ClickHouseInserterHandle>,
 	tasks: Arc<TaskGroup>,
@@ -377,6 +379,7 @@ impl ProxyState {
 				.max_capacity(10_000)
 				.time_to_live(PROXY_STATE_CACHE_TTL)
 				.build(),
+			inflight_requests: Arc::new(Mutex::new(HashSet::new())),
 			port_type,
 			clickhouse_inserter,
 			tasks: TaskGroup::new(),
@@ -649,6 +652,42 @@ impl ProxyState {
 			counter.release();
 		}
 	}
+
+	/// Generate a unique request ID that is not currently in flight
+	pub async fn generate_unique_request_id(&self) -> anyhow::Result<RequestId> {
+		const MAX_TRIES: u32 = 100;
+		let mut requests = self.inflight_requests.lock().await;
+
+		for attempt in 0..MAX_TRIES {
+			let request_id = generate_request_id();
+
+			// Check if this ID is already in use
+			if !requests.contains(&request_id) {
+				// Insert the ID and return it
+				requests.insert(request_id);
+				return Ok(request_id);
+			}
+
+			// Collision occurred (extremely rare with 4 bytes = 4 billion possibilities)
+			// Generate a new ID and try again
+			tracing::warn!(
+				?request_id,
+				attempt,
+				"request id collision, generating new id"
+			);
+		}
+
+		anyhow::bail!(
+			"failed to generate unique request id after {} attempts",
+			MAX_TRIES
+		);
+	}
+
+	/// Release a request ID when the request is complete
+	pub async fn release_request_id(&self, request_id: RequestId) {
+		let mut requests = self.inflight_requests.lock().await;
+		requests.remove(&request_id);
+	}
 }
 
 // Helper function to choose a random target from a list of targets
@@ -1044,6 +1083,9 @@ impl ProxyService {
 			ResolveRouteOutput::CustomServe(mut handler) => {
 				let req_headers = req.headers().clone();
 
+				// Generate unique request ID
+				let request_id = self.state.generate_unique_request_id().await?;
+
 				// Collect request body
 				let (req_parts, body) = req.into_parts();
 				let collected_body = match http_body_util::BodyExt::collect(body).await {
@@ -1064,7 +1106,7 @@ impl ProxyService {
 					attempts += 1;
 
 					let res = handler
-						.handle_request(req_collected.clone(), request_context)
+						.handle_request(req_collected.clone(), request_context, request_id)
 						.await;
 					if should_retry_request(&res) {
 						// Request connect error, might retry
@@ -1838,12 +1880,13 @@ impl ProxyService {
 				tracing::debug!(%req_path, "Spawning task to handle WebSocket communication");
 				let mut request_context = request_context.clone();
 				let req_headers = req_headers.clone();
+				let state = self.state.clone();
 				let req_path = req_path.clone();
 				let req_host = req_host.clone();
 
 				self.state.tasks.spawn(
 					async move {
-						let request_id = Uuid::new_v4();
+						let request_id = state.generate_unique_request_id().await?;
 						let mut ws_hibernation_close = false;
 						let mut after_hibernation = false;
 						let mut attempts = 0u32;

engine/packages/guard-core/src/proxy_service.rs

@@ -5,29 +5,30 @@ use futures_util::TryStreamExt;
 use gas::prelude::*;
 use http_body_util::{BodyExt, Full};
 use hyper::{Request, Response, StatusCode};
+use pegboard::tunnel::id::{self as tunnel_id, RequestId};
 use rand::Rng;
 use rivet_error::*;
 use rivet_guard_core::{
-	WebSocketHandle,
 	custom_serve::{CustomServeTrait, HibernationResult},
 	errors::{
 		ServiceUnavailable, WebSocketServiceHibernate, WebSocketServiceTimeout,
 		WebSocketServiceUnavailable,
 	},
-	proxy_service::{ResponseBody, is_ws_hibernate},
+	proxy_service::{is_ws_hibernate, ResponseBody},
 	request_context::RequestContext,
 	websocket_handle::WebSocketReceiver,
+	WebSocketHandle,
 };
 use rivet_runner_protocol as protocol;
 use rivet_util::serde::HashableMap;
 use std::{sync::Arc, time::Duration};
 use tokio::{
-	sync::{Mutex, watch},
+	sync::{watch, Mutex},
 	task::JoinHandle,
 };
 use tokio_tungstenite::tungstenite::{
+	protocol::frame::{coding::CloseCode, CloseFrame},
 	Message,
-	protocol::frame::{CloseFrame, coding::CloseCode},
 };
 
 use crate::shared_state::{InFlightRequestHandle, SharedState};
@@ -86,6 +87,7 @@ impl CustomServeTrait for PegboardGateway {
 		&self,
 		req: Request<Full<Bytes>>,
 		_request_context: &mut RequestContext,
+		request_id: RequestId,
 	) -> Result<Response<ResponseBody>> {
 		// Use the actor ID from the gateway instance
 		let actor_id = self.actor_id.to_string();
@@ -163,7 +165,6 @@ impl CustomServeTrait for PegboardGateway {
 			pegboard::pubsub_subjects::RunnerReceiverSubject::new(self.runner_id).to_string();
 
 		// Start listening for request responses
-		let request_id = Uuid::new_v4().into_bytes();
 		let InFlightRequestHandle {
 			mut msg_rx,
 			mut drop_rx,
@@ -212,7 +213,7 @@ impl CustomServeTrait for PegboardGateway {
 							}
 						} else {
 							tracing::warn!(
-								request_id=?Uuid::from_bytes(request_id),
+								request_id=?tunnel_id::request_id_to_string(&request_id),
 								"received no message response during request init",
 							);
 							break;
@@ -274,7 +275,7 @@ impl CustomServeTrait for PegboardGateway {
 		headers: &hyper::HeaderMap,
 		_path: &str,
 		_request_context: &mut RequestContext,
-		unique_request_id: Uuid,
+		unique_request_id: RequestId,
 		after_hibernation: bool,
 	) -> Result<Option<CloseFrame>> {
 		// Use the actor ID from the gateway instance
@@ -298,7 +299,7 @@ impl CustomServeTrait for PegboardGateway {
 			pegboard::pubsub_subjects::RunnerReceiverSubject::new(self.runner_id).to_string();
 
 		// Start listening for WebSocket messages
-		let request_id = unique_request_id.into_bytes();
+		let request_id = unique_request_id;
 		let InFlightRequestHandle {
 			mut msg_rx,
 			mut drop_rx,
@@ -348,7 +349,7 @@ impl CustomServeTrait for PegboardGateway {
 								}
 							} else {
 								tracing::warn!(
-									request_id=?Uuid::from_bytes(request_id),
+									request_id=?tunnel_id::request_id_to_string(&request_id),
 									"received no message response during ws init",
 								);
 								break;
@@ -414,7 +415,7 @@ impl CustomServeTrait for PegboardGateway {
 									}
 										protocol::ToServerTunnelMessageKind::ToServerWebSocketMessageAck(ack) => {
 										tracing::debug!(
-											request_id=?Uuid::from_bytes(request_id),
+											request_id=?tunnel_id::request_id_to_string(&request_id),
 											ack_index=?ack.index,
 											"received WebSocketMessageAck from runner"
 										);
@@ -477,8 +478,6 @@ impl CustomServeTrait for PegboardGateway {
 										let ws_message =
 											protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage(
 												protocol::ToClientWebSocketMessage {
-													// NOTE: This gets set in shared_state.ts
-													index: 0,
 													data: data.into(),
 													binary: true,
 												},
@@ -491,8 +490,6 @@ impl CustomServeTrait for PegboardGateway {
 										let ws_message =
 											protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage(
 												protocol::ToClientWebSocketMessage {
-													// NOTE: This gets set in shared_state.ts
-													index: 0,
 													data: text.as_bytes().to_vec(),
 													binary: false,
 												},
@@ -615,12 +612,14 @@ impl CustomServeTrait for PegboardGateway {
 	async fn handle_websocket_hibernation(
 		&self,
 		client_ws: WebSocketHandle,
-		unique_request_id: Uuid,
+		unique_request_id: RequestId,
 	) -> Result<HibernationResult> {
+		let request_id = unique_request_id;
+
 		// Immediately rewake if we have pending messages
 		if self
 			.shared_state
-			.has_pending_websocket_messages(unique_request_id.into_bytes())
+			.has_pending_websocket_messages(request_id)
 			.await?
 		{
 			tracing::debug!(
@@ -633,6 +632,8 @@ impl CustomServeTrait for PegboardGateway {
 		// Start keepalive task
 		let ctx = self.ctx.clone();
 		let actor_id = self.actor_id;
+		let gateway_id = self.shared_state.gateway_id();
+		let request_id = unique_request_id;
 		let keepalive_handle: JoinHandle<Result<()>> = tokio::spawn(async move {
 			let mut ping_interval = tokio::time::interval(Duration::from_millis(
 				(ctx.config()
@@ -652,7 +653,8 @@ impl CustomServeTrait for PegboardGateway {
 
 				ctx.op(pegboard::ops::actor::hibernating_request::upsert::Input {
 					actor_id,
-					request_id: unique_request_id,
+					gateway_id,
+					request_id,
 				})
 				.await?;
 			}
@@ -669,6 +671,7 @@ impl CustomServeTrait for PegboardGateway {
 				self.ctx
 					.op(pegboard::ops::actor::hibernating_request::delete::Input {
 						actor_id: self.actor_id,
+						gateway_id: self.shared_state.gateway_id(),
 						request_id: unique_request_id,
 					})
 					.await?;