temp

Author: MasterPtato

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

@@ -2537,7 +2537,7 @@ fn is_retryable_ws_error(err: &anyhow::Error) -> bool {
 	}
 }
 
-fn is_ws_hibernate(err: &anyhow::Error) -> bool {
+pub 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_hibernate"
 	} else {

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

@@ -14,7 +14,7 @@ use rivet_guard_core::{
 		ServiceUnavailable, WebSocketServiceHibernate, WebSocketServiceTimeout,
 		WebSocketServiceUnavailable,
 	},
-	proxy_service::ResponseBody,
+	proxy_service::{ResponseBody, is_ws_hibernate},
 	request_context::RequestContext,
 	websocket_handle::WebSocketReceiver,
 };
@@ -559,28 +559,33 @@ impl CustomServeTrait for PegboardGateway {
 			(res, _) => res,
 		};
 
-		// Send WebSocket close message to runner
-		let (close_code, close_reason) = match &mut lifecycle_res {
-			// Taking here because it won't be used again
-			Ok(LifecycleResult::ClientClose(Some(close))) => {
-				(close.code, Some(std::mem::take(&mut close.reason)))
-			}
-			Ok(_) => (CloseCode::Normal.into(), None),
-			Err(_) => (CloseCode::Error.into(), Some("ws.downstream_closed".into())),
-		};
-		let close_message = protocol::ToClientTunnelMessageKind::ToClientWebSocketClose(
-			protocol::ToClientWebSocketClose {
-				code: Some(close_code.into()),
-				reason: close_reason.map(|x| x.as_str().to_string()),
-			},
-		);
-
-		if let Err(err) = self
-			.shared_state
-			.send_message(request_id, close_message)
-			.await
+		// Send close frame to runner if not hibernating
+		if lifecycle_res
+			.as_ref()
+			.map_or_else(is_ws_hibernate, |_| false)
 		{
-			tracing::error!(?err, "error sending close message");
+			let (close_code, close_reason) = match &mut lifecycle_res {
+				// Taking here because it won't be used again
+				Ok(LifecycleResult::ClientClose(Some(close))) => {
+					(close.code, Some(std::mem::take(&mut close.reason)))
+				}
+				Ok(_) => (CloseCode::Normal.into(), None),
+				Err(_) => (CloseCode::Error.into(), Some("ws.downstream_closed".into())),
+			};
+			let close_message = protocol::ToClientTunnelMessageKind::ToClientWebSocketClose(
+				protocol::ToClientWebSocketClose {
+					code: Some(close_code.into()),
+					reason: close_reason.map(|x| x.as_str().to_string()),
+				},
+			);
+
+			if let Err(err) = self
+				.shared_state
+				.send_message(request_id, close_message)
+				.await
+			{
+				tracing::error!(?err, "error sending close message");
+			}
 		}
 
 		// Send WebSocket close message to client