fix(guard): handle actor stopped during ws/req

Author: MasterPtato

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

@@ -152,6 +152,11 @@ impl CustomServeTrait for PegboardGateway {
 			.context("failed to read body")?
 			.to_bytes();
 
+		let mut stopped_sub = self
+			.ctx
+			.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", self.actor_id))
+			.await?;
+
 		// Build subject to publish to
 		let tunnel_subject =
 			pegboard::pubsub_subjects::RunnerReceiverSubject::new(self.runner_id).to_string();
@@ -212,6 +217,10 @@ impl CustomServeTrait for PegboardGateway {
 							break;
 						}
 					}
+					_ = stopped_sub.next() => {
+						tracing::debug!("actor stopped while waiting for request response");
+						return Err(ServiceUnavailable.build());
+					}
 					_ = drop_rx.changed() => {
 						tracing::warn!("tunnel message timeout");
 						return Err(ServiceUnavailable.build());
@@ -278,6 +287,11 @@ impl CustomServeTrait for PegboardGateway {
 			}
 		}
 
+		let mut stopped_sub = self
+			.ctx
+			.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", self.actor_id))
+			.await?;
+
 		// Build subject to publish to
 		let tunnel_subject =
 			pegboard::pubsub_subjects::RunnerReceiverSubject::new(self.runner_id).to_string();
@@ -339,6 +353,10 @@ impl CustomServeTrait for PegboardGateway {
 								break;
 							}
 						}
+						_ = stopped_sub.next() => {
+							tracing::debug!("actor stopped while waiting for websocket open");
+							return Err(WebSocketServiceUnavailable.build());
+						}
 						_ = drop_rx.changed() => {
 							tracing::warn!("websocket open timeout");
 							return Err(WebSocketServiceUnavailable.build());
@@ -364,7 +382,7 @@ impl CustomServeTrait for PegboardGateway {
 			open_msg.can_hibernate
 		};
 
-		// Send reclaimed messages
+		// Send pending messages
 		self.shared_state
 			.resend_pending_websocket_messages(request_id)
 			.await?;
@@ -415,6 +433,15 @@ impl CustomServeTrait for PegboardGateway {
 								return Err(WebSocketServiceHibernate.build());
 							}
 						}
+						_ = stopped_sub.next() => {
+							tracing::debug!("actor stopped during websocket handler loop");
+
+							if can_hibernate {
+								return Err(WebSocketServiceHibernate.build());
+							} else {
+								return Err(WebSocketServiceUnavailable.build());
+							}
+						}
 						_ = drop_rx.changed() => {
 							tracing::warn!("websocket message timeout");
 							return Err(WebSocketServiceTimeout.build());
@@ -579,6 +606,15 @@ impl CustomServeTrait for PegboardGateway {
 		client_ws: WebSocketHandle,
 		unique_request_id: Uuid,
 	) -> Result<HibernationResult> {
+		// Immediately rewake if we have pending messages
+		if self
+			.shared_state
+			.has_pending_websocket_messages(unique_request_id.into_bytes())
+			.await?
+		{
+			return Ok(HibernationResult::Continue);
+		}
+
 		// Start keepalive task
 		let ctx = self.ctx.clone();
 		let actor_id = self.actor_id;

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

@@ -344,6 +344,18 @@ impl SharedState {
 		Ok(())
 	}
 
+	pub async fn has_pending_websocket_messages(&self, request_id: RequestId) -> Result<bool> {
+		let Some(req) = self.in_flight_requests.get_async(&request_id).await else {
+			bail!("request not in flight");
+		};
+
+		if let Some(hs) = &req.hibernation_state {
+			Ok(!hs.pending_ws_msgs.is_empty())
+		} else {
+			Ok(false)
+		}
+	}
+
 	pub async fn ack_pending_websocket_messages(
 		&self,
 		request_id: RequestId,

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

@@ -12,7 +12,7 @@ use rivet_guard_core::{
 use rivet_runner_protocol as protocol;
 use std::time::Duration;
 use tokio::sync::watch;
-use tokio_tungstenite::tungstenite::protocol::frame::{CloseFrame, coding::CloseCode};
+use tokio_tungstenite::tungstenite::protocol::frame::CloseFrame;
 use universalpubsub::PublishOpts;
 use vbare::OwnedVersionedData;
 
@@ -243,30 +243,19 @@ impl CustomServeTrait for PegboardRunnerWsCustomServe {
 			);
 		}
 
-		// Send WebSocket close messages to all remaining active requests
+		// Send close messages to all remaining active requests
 		let active_requests = conn.tunnel_active_requests.lock().await;
 		for (request_id, req) in &*active_requests {
-			let close_msg_kind = if req.is_ws {
-				let (close_code, close_reason) = if lifecycle_res.is_ok() {
-					(CloseCode::Normal.into(), None)
-				} else {
-					(CloseCode::Error.into(), Some("ws.upstream_closed".into()))
-				};
-
-				protocol::ToServerTunnelMessageKind::ToServerWebSocketClose(
-					protocol::ToServerWebSocketClose {
-						code: Some(close_code),
-						reason: close_reason,
-						hibernate: true,
-					},
-				)
-			} else {
-				protocol::ToServerTunnelMessageKind::ToServerResponseAbort
-			};
+			// Websockets are not ephemeral like requests. If the runner ws closes they are not informed;
+			// instead they wait for the actor itself to stop.
+			if req.is_ws {
+				continue;
+			}
+
 			let close_message = protocol::ToServerTunnelMessage {
 				request_id: request_id.clone(),
 				message_id: Uuid::new_v4().into_bytes(),
-				message_kind: close_msg_kind,
+				message_kind: protocol::ToServerTunnelMessageKind::ToServerResponseAbort,
 			};
 
 			let msg_serialized = protocol::versioned::ToGateway::wrap_latest(protocol::ToGateway {