fix(pb): fix tunnel race condition handling (#3043)

Fixes RVT-5193

Author: MasterPtato

Cargo.lock

@@ -15,6 +15,7 @@ import {
 	type DefaultDataProvider,
 	RECORDS_PER_PAGE,
 } from "./default-data-provider";
+import { getConfig } from "@/components/lib/config";
 
 export type CreateNamespace = {
 	displayName: string;
@@ -29,7 +30,7 @@ export type Namespace = {
 };
 
 export function createClient(
-	baseUrl = engineEnv().VITE_APP_API_URL,
+	baseUrl = getConfig().apiUrl,
 	opts: { token: (() => string) | string },
 ) {
 	return new RivetClient({
@@ -42,7 +43,7 @@ export function createClient(
 export const createGlobalContext = (opts: {
 	engineToken: (() => string) | string;
 }) => {
-	const client = createClient(engineEnv().VITE_APP_API_URL, {
+	const client = createClient(getConfig().apiUrl, {
 		token: opts.engineToken,
 	});
 	return {

frontend/src/app/data-providers/engine-data-provider.tsx

@@ -1798,7 +1798,7 @@ impl ProxyService {
 								.await
 							{
 								Result::Ok(()) => {
-									tracing::debug!("websocket closed");
+									tracing::debug!("websocket handler complete, closing");
 
 									// Send graceful close
 									ws_handle

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

@@ -23,5 +23,6 @@ rivet-util.workspace = true
 thiserror.workspace = true
 tokio-tungstenite.workspace = true
 tokio.workspace = true
+tracing.workspace = true
 universalpubsub.workspace = true
 vbare.workspace = true

packages/core/pegboard-gateway/Cargo.toml

@@ -21,7 +21,7 @@ use crate::shared_state::{SharedState, TunnelMessageData};
 
 pub mod shared_state;
 
-const UPS_REQ_TIMEOUT: Duration = Duration::from_secs(2);
+const TUNNEL_ACK_TIMEOUT: Duration = Duration::from_secs(2);
 const SEC_WEBSOCKET_PROTOCOL: HeaderName = HeaderName::from_static("sec-websocket-protocol");
 const WS_PROTOCOL_ACTOR: &str = "rivet_actor.";
 
@@ -33,6 +33,7 @@ pub struct PegboardGateway {
 }
 
 impl PegboardGateway {
+	#[tracing::instrument(skip_all, fields(?actor_id, ?runner_id))]
 	pub fn new(ctx: StandaloneCtx, shared_state: SharedState, runner_id: Id, actor_id: Id) -> Self {
 		Self {
 			ctx,
@@ -45,6 +46,7 @@ impl PegboardGateway {
 
 #[async_trait]
 impl CustomServeTrait for PegboardGateway {
+	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id, runner_id=?self.runner_id))]
 	async fn handle_request(
 		&self,
 		req: Request<Full<Bytes>>,
@@ -67,6 +69,7 @@ impl CustomServeTrait for PegboardGateway {
 		}
 	}
 
+	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id, runner_id=?self.runner_id))]
 	async fn handle_websocket(
 		&self,
 		client_ws: WebSocketHandle,
@@ -91,6 +94,7 @@ impl CustomServeTrait for PegboardGateway {
 }
 
 impl PegboardGateway {
+	#[tracing::instrument(skip_all)]
 	async fn handle_request_inner(
 		&self,
 		req: Request<Full<Bytes>>,
@@ -155,9 +159,16 @@ impl PegboardGateway {
 		self.shared_state.send_message(request_id, message).await?;
 
 		// Wait for response
-		tracing::debug!("starting response handler task");
+		tracing::debug!("gateway waiting for response from tunnel");
 		let response_start = loop {
-			let Some(msg) = msg_rx.recv().await else {
+			let Some(msg) = tokio::time::timeout(TUNNEL_ACK_TIMEOUT, msg_rx.recv())
+				.await
+				.map_err(|_| {
+					tracing::warn!("timed out waiting for tunnel ack");
+
+					RequestError::ServiceUnavailable
+				})?
+			else {
 				tracing::warn!("received no message response");
 				return Err(RequestError::ServiceUnavailable.into());
 			};
@@ -195,6 +206,7 @@ impl PegboardGateway {
 		Ok(response)
 	}
 
+	#[tracing::instrument(skip_all)]
 	async fn handle_websocket_inner(
 		&self,
 		client_ws: WebSocketHandle,
@@ -247,34 +259,42 @@ impl PegboardGateway {
 			.send_message(request_id, open_message)
 			.await?;
 
+		tracing::debug!("gateway waiting for websocket open from tunnel");
+
 		// Wait for WebSocket open acknowledgment
-		let open_ack_received = loop {
-			let Some(msg) = msg_rx.recv().await else {
-				bail!("received no websocket open response");
+		loop {
+			let Some(msg) = tokio::time::timeout(TUNNEL_ACK_TIMEOUT, msg_rx.recv())
+				.await
+				.map_err(|_| {
+					tracing::warn!("timed out waiting for tunnel ack");
+
+					RequestError::ServiceUnavailable
+				})?
+			else {
+				tracing::warn!("received no message response");
+				return Err(RequestError::ServiceUnavailable.into());
 			};
 
 			match msg {
 				TunnelMessageData::Message(
 					protocol::ToServerTunnelMessageKind::ToServerWebSocketOpen,
 				) => {
-					break true;
+					break;
 				}
 				TunnelMessageData::Message(
 					protocol::ToServerTunnelMessageKind::ToServerWebSocketClose(close),
 				) => {
-					bail!("websocket closed before opening: {close:?}");
+					tracing::warn!(?close, "websocket closed before opening");
+					return Err(RequestError::ServiceUnavailable.into());
 				}
 				TunnelMessageData::Timeout => {
-					bail!("websocket open timeout");
+					tracing::warn!("websocket open timeout");
+					return Err(RequestError::ServiceUnavailable.into());
 				}
 				_ => {
 					tracing::warn!("received unexpected message while waiting for websocket open");
 				}
 			}
-		};
-
-		if !open_ack_received {
-			bail!("failed to open websocket");
 		}
 
 		// Accept the WebSocket

packages/core/pegboard-gateway/src/lib.rs

@@ -16,6 +16,7 @@ use crate::{
 	utils::{self},
 };
 
+#[tracing::instrument(skip_all, fields(runner_id=?conn.runner_id, workflow_id=?conn.workflow_id, protocol_version=%conn.protocol_version))]
 pub async fn task(ctx: StandaloneCtx, conn: Arc<Conn>, ws_rx: WebSocketReceiver) {
 	match task_inner(ctx, conn, ws_rx).await {
 		Ok(_) => {}
@@ -25,6 +26,7 @@ pub async fn task(ctx: StandaloneCtx, conn: Arc<Conn>, ws_rx: WebSocketReceiver)
 	}
 }
 
+#[tracing::instrument(skip_all)]
 async fn task_inner(
 	ctx: StandaloneCtx,
 	conn: Arc<Conn>,
@@ -77,6 +79,7 @@ async fn task_inner(
 	Ok(())
 }
 
+#[tracing::instrument(skip_all)]
 async fn handle_message(
 	ctx: &StandaloneCtx,
 	conn: &Arc<Conn>,
@@ -354,6 +357,7 @@ async fn handle_message(
 	Ok(())
 }
 
+#[tracing::instrument(skip_all)]
 async fn handle_tunnel_message(
 	ctx: &StandaloneCtx,
 	conn: &Arc<Conn>,

packages/core/pegboard-runner/src/client_to_pubsub_task.rs

@@ -12,6 +12,7 @@ use crate::{
 	utils,
 };
 
+#[tracing::instrument(skip_all, fields(runner_id=?conn.runner_id, workflow_id=?conn.workflow_id, protocol_version=%conn.protocol_version))]
 pub async fn task(ctx: StandaloneCtx, conn: Arc<Conn>, sub: Subscriber) {
 	match task_inner(ctx, conn, sub).await {
 		Ok(_) => {}
@@ -21,6 +22,7 @@ pub async fn task(ctx: StandaloneCtx, conn: Arc<Conn>, sub: Subscriber) {
 	}
 }
 
+#[tracing::instrument(skip_all)]
 async fn task_inner(ctx: StandaloneCtx, conn: Arc<Conn>, mut sub: Subscriber) -> Result<()> {
 	while let Result::Ok(NextOutput::Message(ups_msg)) = sub.next().await {
 		tracing::debug!(
@@ -37,6 +39,7 @@ async fn task_inner(ctx: StandaloneCtx, conn: Arc<Conn>, mut sub: Subscriber) ->
 				continue;
 			}
 		};
+		let is_close = utils::is_to_client_close(&msg);
 
 		// Handle tunnel messages
 		if let protocol::ToClient::ToClientTunnelMessage(tunnel_msg) = &mut msg {
@@ -57,11 +60,17 @@ async fn task_inner(ctx: StandaloneCtx, conn: Arc<Conn>, mut sub: Subscriber) ->
 			tracing::error!(?e, "failed to send message to WebSocket");
 			break;
 		}
+
+		if is_close {
+			tracing::debug!("manually closing websocket");
+			break;
+		}
 	}
 
 	Ok(())
 }
 
+#[tracing::instrument(skip_all)]
 async fn handle_tunnel_message(conn: &Arc<Conn>, msg: &mut protocol::ToClientTunnelMessage) {
 	// Save active request
 	//

packages/core/pegboard-runner/src/pubsub_to_client_task.rs

@@ -4,10 +4,6 @@ use hyper_tungstenite::tungstenite::Message as WsMessage;
 use hyper_util::rt::TokioIo;
 use rivet_error::*;
 use rivet_runner_protocol as protocol;
-use tokio_tungstenite::{
-	WebSocketStream,
-	tungstenite::protocol::frame::{CloseFrame, coding::CloseCode},
-};
 
 #[derive(Clone)]
 pub struct UrlData {
@@ -73,3 +69,10 @@ pub fn is_to_client_tunnel_message_kind_request_close(
 		_ => false,
 	}
 }
+
+pub fn is_to_client_close(kind: &protocol::ToClient) -> bool {
+	match kind {
+		protocol::ToClient::ToClientClose => true,
+		_ => false,
+	}
+}