feat(gateway): websocket hibernation

Author: MasterPtato

CLAUDE.md

@@ -125,7 +125,7 @@ Key points:
 - Add `Serialize, Deserialize` derives for errors with metadata fields
 - Always return anyhow errors from failable functions
 	- For example: `fn foo() -> Result<i64> { /* ... */ }`
-- Import anyhow using `use anyhow::*` instead of importing individual types
+- Do not glob import (`::*`) from anyhow. Instead, import individual types and traits
 
 **Dependency Management**
 - When adding a dependency, check for a workspace dependency in Cargo.toml

engine/artifacts/errors/guard.websocket_service_hibernate.json

@@ -30,4 +30,7 @@ pub trait CustomServeTrait: Send + Sync {
 		// Identifies the websocket across retries.
 		unique_request_id: Uuid,
 	) -> Result<Option<CloseFrame>>;
+
+	/// Returns true if the websocket should close.
+	async fn handle_websocket_hibernation(&self, websocket: WebSocketHandle) -> Result<bool>;
 }

engine/artifacts/errors/guard.websocket_service_retry.json

@@ -93,8 +93,12 @@ pub struct ServiceUnavailable;
 pub struct WebSocketServiceUnavailable;
 
 #[derive(RivetError, Serialize, Deserialize)]
-#[error("guard", "websocket_service_retry", "WebSocket service retry.")]
-pub struct WebSocketServiceRetry;
+#[error(
+	"guard",
+	"websocket_service_hibernate",
+	"Initiate WebSocket service hibernation."
+)]
+pub struct WebSocketServiceHibernate;
 
 #[derive(RivetError, Serialize, Deserialize)]
 #[error("guard", "websocket_service_timeout", "WebSocket service timed out.")]

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

@@ -1828,7 +1828,7 @@ impl ProxyService {
 				);
 			}
 			ResolveRouteOutput::Response(_) => unreachable!(),
-			ResolveRouteOutput::CustomServe(mut handlers) => {
+			ResolveRouteOutput::CustomServe(mut handler) => {
 				tracing::debug!(%req_path, "Spawning task to handle WebSocket communication");
 				let mut request_context = request_context.clone();
 				let req_headers = req_headers.clone();
@@ -1845,7 +1845,7 @@ impl ProxyService {
 							.context("failed initiating websocket handle")?;
 
 						loop {
-							match handlers
+							match handler
 								.handle_websocket(
 									ws_handle.clone(),
 									&req_headers,
@@ -1895,18 +1895,33 @@ impl ProxyService {
 								Err(err) => {
 									tracing::debug!(?err, "websocket handler error");
 
-									// Denotes that the connection did not fail, but needs to be retried to
-									// resole a new target
-									let ws_retry = is_ws_retry(&err);
+									// Denotes that the connection did not fail, but the downstream has closed
+									let ws_hibernate = is_ws_hibernate(&err);
 
-									if ws_retry {
+									if ws_hibernate {
 										attempts = 0;
 									} else {
 										attempts += 1;
 									}
 
-									if attempts > max_attempts
-										|| (!is_retryable_ws_error(&err) && !ws_retry)
+									if ws_hibernate {
+										// After this function returns false:
+										// - the route will be resolved again
+										// - the websocket will connect to the new downstream target
+										// - the gateway will continue reading messages from the client ws
+										//   (starting with the message that caused the hibernation to end)
+										let close = handler
+											.handle_websocket_hibernation(ws_handle.clone())
+											.await?;
+
+										if close {
+											tracing::debug!(
+												"hibernating websocket handler complete"
+											);
+											break;
+										}
+									} else if attempts > max_attempts
+										|| !is_retryable_ws_error(&err)
 									{
 										tracing::debug!(
 											?attempts,
@@ -1929,79 +1944,79 @@ impl ProxyService {
 
 										break;
 									} else {
-										if !ws_retry {
-											let backoff = ProxyService::calculate_backoff(
-												attempts,
-												initial_interval,
-											);
+										let backoff = ProxyService::calculate_backoff(
+											attempts,
+											initial_interval,
+										);
 
-											tracing::debug!(
-												?backoff,
-												"WebSocket attempt {attempts} failed (service unavailable)"
-											);
+										tracing::debug!(
+											?backoff,
+											"WebSocket attempt {attempts} failed (service unavailable)"
+										);
 
-											tokio::time::sleep(backoff).await;
-										}
+										// Apply backoff for retryable error
+										tokio::time::sleep(backoff).await;
+									}
 
-										match state
-											.resolve_route(
-												&req_host,
-												&req_path,
-												&req_method,
-												state.port_type.clone(),
-												&req_headers,
-												true,
-											)
-											.await
-										{
-											Ok(ResolveRouteOutput::CustomServe(new_handlers)) => {
-												handlers = new_handlers;
-												continue;
-											}
-											Ok(ResolveRouteOutput::Response(response)) => {
-												ws_handle
-													.send(to_hyper_close(Some(str_to_close_frame(
-														response.message.as_ref(),
-													))))
-													.await?;
-
-												// Flush to ensure close frame is sent
-												ws_handle.flush().await?;
-
-												// Keep TCP connection open briefly to allow client to process close
-												tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
-											}
-											Ok(ResolveRouteOutput::Target(_)) => {
-												ws_handle
-													.send(to_hyper_close(Some(err_to_close_frame(
-														errors::WebSocketTargetChanged.build(),
-														ray_id,
-													))))
-													.await?;
+									// Retry route resolution
+									match state
+										.resolve_route(
+											&req_host,
+											&req_path,
+											&req_method,
+											state.port_type.clone(),
+											&req_headers,
+											true,
+										)
+										.await
+									{
+										Ok(ResolveRouteOutput::CustomServe(new_handler)) => {
+											handler = new_handler;
+											continue;
+										}
+										Ok(ResolveRouteOutput::Response(response)) => {
+											ws_handle
+												.send(to_hyper_close(Some(str_to_close_frame(
+													response.message.as_ref(),
+												))))
+												.await?;
+
+											// Flush to ensure close frame is sent
+											ws_handle.flush().await?;
+
+											// Keep TCP connection open briefly to allow client to process close
+											tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
+										}
+										Ok(ResolveRouteOutput::Target(_)) => {
+											ws_handle
+												.send(to_hyper_close(Some(err_to_close_frame(
+													errors::WebSocketTargetChanged.build(),
+													ray_id,
+												))))
+												.await?;
 
-												// Flush to ensure close frame is sent
-												ws_handle.flush().await?;
+											// Flush to ensure close frame is sent
+											ws_handle.flush().await?;
 
-												// Keep TCP connection open briefly to allow client to process close
-												tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
+											// Keep TCP connection open briefly to allow client to process close
+											tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
 
-												break;
-											}
-											Err(err) => {
-												ws_handle
-													.send(to_hyper_close(Some(err_to_close_frame(
-														err, ray_id,
-													))))
-													.await?;
+											break;
+										}
+										Err(err) => {
+											ws_handle
+												.send(to_hyper_close(Some(err_to_close_frame(
+													err, ray_id,
+												))))
+												.await?;
 
-												// Flush to ensure close frame is sent
-												ws_handle.flush().await?;
+											// Flush to ensure close frame is sent
+											ws_handle.flush().await?;
 
-												// Keep TCP connection open briefly to allow client to process close
-												tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
+											// Keep TCP connection open briefly to allow client to process close
+											tokio::time::sleep(WEBSOCKET_CLOSE_LINGER).await;
 
-												break;
-											}
+											break;
 										}
 									}
 								}
@@ -2509,9 +2524,9 @@ fn is_retryable_ws_error(err: &anyhow::Error) -> bool {
 	}
 }
 
-fn is_ws_retry(err: &anyhow::Error) -> bool {
+fn is_ws_hibernate(err: &anyhow::Error) -> bool {
 	if let Some(rivet_err) = err.chain().find_map(|x| x.downcast_ref::<RivetError>()) {
-		rivet_err.group() == "guard" && rivet_err.code() == "websocket_service_retry"
+		rivet_err.group() == "guard" && rivet_err.code() == "websocket_service_hibernate"
 	} else {
 		false
 	}

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

@@ -1,5 +1,5 @@
 use anyhow::*;
-use futures_util::{SinkExt, StreamExt};
+use futures_util::{SinkExt, StreamExt, stream::Peekable};
 use hyper::upgrade::Upgraded;
 use hyper_tungstenite::HyperWebsocket;
 use hyper_tungstenite::tungstenite::Message as WsMessage;
@@ -8,7 +8,8 @@ use std::sync::Arc;
 use tokio::sync::Mutex;
 use tokio_tungstenite::WebSocketStream;
 
-pub type WebSocketReceiver = futures_util::stream::SplitStream<WebSocketStream<TokioIo<Upgraded>>>;
+pub type WebSocketReceiver =
+	Peekable<futures_util::stream::SplitStream<WebSocketStream<TokioIo<Upgraded>>>>;
 
 pub type WebSocketSender =
 	futures_util::stream::SplitSink<WebSocketStream<TokioIo<Upgraded>>, WsMessage>;
@@ -26,7 +27,7 @@ impl WebSocketHandle {
 
 		Ok(Self {
 			ws_tx: Arc::new(Mutex::new(ws_tx)),
-			ws_rx: Arc::new(Mutex::new(ws_rx)),
+			ws_rx: Arc::new(Mutex::new(ws_rx.peekable())),
 		})
 	}
 

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

@@ -102,6 +102,10 @@ impl CustomServeTrait for TestCustomServe {
 
 		Ok(None)
 	}
+
+	async fn handle_websocket_hibernation(&self, _websocket: WebSocketHandle) -> Result<bool> {
+		todo!();
+	}
 }
 
 // Create routing function that returns CustomServe

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

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use anyhow::*;
+use anyhow::{Context, Result};
 use async_trait::async_trait;
 use bytes::Bytes;
 use gas::prelude::*;
@@ -30,14 +30,14 @@ impl CustomServeTrait for ApiPublicService {
 		let response = service
 			.call(req)
 			.await
-			.map_err(|e| anyhow::anyhow!("Failed to call api-public service: {}", e))?;
+			.context("failed to call api-public service")?;
 
 		// Collect the body and convert to ResponseBody
 		let (parts, body) = response.into_parts();
 		let collected = body
 			.collect()
 			.await
-			.map_err(|e| anyhow::anyhow!("Failed to collect response body: {}", e))?;
+			.context("failed to collect response body")?;
 		let bytes = collected.to_bytes();
 		let response_body = ResponseBody::Full(Full::new(bytes));
 		let response = Response::from_parts(parts, response_body);
@@ -55,6 +55,10 @@ impl CustomServeTrait for ApiPublicService {
 	) -> Result<Option<CloseFrame>> {
 		bail!("api-public does not support WebSocket connections")
 	}
+
+	async fn handle_websocket_hibernation(&self, _client_ws: WebSocketHandle) -> Result<bool> {
+		bail!("api-public does not support WebSocket hibernation")
+	}
 }
 
 /// Route requests to the api-public service

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

@@ -235,6 +235,7 @@ async fn route_request_inner(
 
 	// Return pegboard-gateway instance with path
 	let gateway = pegboard_gateway::PegboardGateway::new(
+		ctx.clone(),
 		shared_state.pegboard_gateway.clone(),
 		runner_id,
 		actor_id,