chore(rivetkit): fix hibernation implementation

Author: NathanFlurry

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

@@ -9,23 +9,23 @@ 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, 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};
@@ -486,7 +486,7 @@ impl CustomServeTrait for PegboardGateway {
 		};
 
 		// Send close frame to runner if not hibernating
-		if lifecycle_res
+		if !&lifecycle_res
 			.as_ref()
 			.map_or_else(is_ws_hibernate, |_| false)
 		{

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

@@ -124,7 +124,6 @@ impl SharedState {
 				entry.msg_tx = msg_tx;
 				entry.drop_tx = drop_tx;
 				entry.opened = false;
-				entry.message_index = 0;
 
 				if entry.stopping {
 					entry.hibernation_state = None;
@@ -194,6 +193,11 @@ impl SharedState {
 			};
 
 			hs.pending_ws_msgs.push(pending_ws_msg);
+			tracing::debug!(
+				index=current_message_index,
+				new_count=hs.pending_ws_msgs.len(),
+				"pushed pending websocket message"
+			);
 		}
 
 		self.ups
@@ -391,7 +395,6 @@ impl SharedState {
 
 		let len_after = hs.pending_ws_msgs.len();
 		tracing::debug!(
-			request_id=?tunnel_id::request_id_to_string(&request_id),
 			ack_index,
 			removed_count = len_before - len_after,
 			remaining_count = len_after,

engine/sdks/typescript/runner/src/actor.ts

@@ -1,7 +1,9 @@
 import type * as protocol from "@rivetkit/engine-runner-protocol";
 import type { PendingRequest } from "./tunnel";
 import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
-import { arraysEqual } from "./utils";
+import { arraysEqual, promiseWithResolvers } from "./utils";
+import { logger } from "./log";
+import * as tunnelId from "./tunnel-id";
 
 export interface ActorConfig {
 	name: string;
@@ -24,11 +26,28 @@ export class RunnerActor {
 		requestId: protocol.RequestId;
 		ws: WebSocketTunnelAdapter;
 	}> = [];
+	actorStartPromise: ReturnType<typeof promiseWithResolvers<void>>;
 
-	constructor(actorId: string, generation: number, config: ActorConfig) {
+	/**
+	 * If restoreHibernatingRequests has been called. This is used to assert
+	 * that the caller is implemented correctly.
+	 **/
+	hibernationRestored: boolean = false;
+
+	constructor(
+		actorId: string,
+		generation: number,
+		config: ActorConfig,
+		/**
+		 * List of hibernating requests provided by the gateway on actor start.
+		 * This represents the WebSocket connections that the gateway knows about.
+		 **/
+		public hibernatingRequests: readonly protocol.HibernatingRequest[],
+	) {
 		this.actorId = actorId;
 		this.generation = generation;
 		this.config = config;
+		this.actorStartPromise = promiseWithResolvers();
 	}
 
 	// Pending request methods
@@ -43,13 +62,78 @@ export class RunnerActor {
 		)?.request;
 	}
 
-	setPendingRequest(
+	createPendingRequest(
+		gatewayId: protocol.GatewayId,
+		requestId: protocol.RequestId,
+		clientMessageIndex: number,
+	) {
+		const exists =
+			this.getPendingRequest(gatewayId, requestId) !== undefined;
+		if (exists) {
+			logger()?.warn({
+				msg: "attempting to set pending request twice, replacing existing",
+				gatewayId: tunnelId.gatewayIdToString(gatewayId),
+				requestId: tunnelId.requestIdToString(requestId),
+			});
+			// Delete existing pending request before adding the new one
+			this.deletePendingRequest(gatewayId, requestId);
+		}
+		this.pendingRequests.push({
+			gatewayId,
+			requestId,
+			request: {
+				resolve: () => {},
+				reject: () => {},
+				actorId: this.actorId,
+				gatewayId: gatewayId,
+				requestId: requestId,
+				clientMessageIndex,
+			},
+		});
+		logger()?.debug({
+			msg: "added pending request",
+			gatewayId: tunnelId.gatewayIdToString(gatewayId),
+			requestId: tunnelId.requestIdToString(requestId),
+			length: this.pendingRequests.length,
+		});
+	}
+
+	createPendingRequestWithStreamController(
 		gatewayId: protocol.GatewayId,
 		requestId: protocol.RequestId,
-		request: PendingRequest,
+		clientMessageIndex: number,
+		streamController: ReadableStreamDefaultController<Uint8Array>,
 	) {
-		this.deletePendingRequest(gatewayId, requestId);
-		this.pendingRequests.push({ gatewayId, requestId, request });
+		const exists =
+			this.getPendingRequest(gatewayId, requestId) !== undefined;
+		if (exists) {
+			logger()?.warn({
+				msg: "attempting to set pending request twice, replacing existing",
+				gatewayId: tunnelId.gatewayIdToString(gatewayId),
+				requestId: tunnelId.requestIdToString(requestId),
+			});
+			// Delete existing pending request before adding the new one
+			this.deletePendingRequest(gatewayId, requestId);
+		}
+		this.pendingRequests.push({
+			gatewayId,
+			requestId,
+			request: {
+				resolve: () => {},
+				reject: () => {},
+				actorId: this.actorId,
+				gatewayId: gatewayId,
+				requestId: requestId,
+				clientMessageIndex,
+				streamController,
+			},
+		});
+		logger()?.debug({
+			msg: "added pending request with stream controller",
+			gatewayId: tunnelId.gatewayIdToString(gatewayId),
+			requestId: tunnelId.requestIdToString(requestId),
+			length: this.pendingRequests.length,
+		});
 	}
 
 	deletePendingRequest(
@@ -63,6 +147,12 @@ export class RunnerActor {
 		);
 		if (index !== -1) {
 			this.pendingRequests.splice(index, 1);
+			logger()?.debug({
+				msg: "removed pending request",
+				gatewayId: tunnelId.gatewayIdToString(gatewayId),
+				requestId: tunnelId.requestIdToString(requestId),
+				length: this.pendingRequests.length,
+			});
 		}
 	}
 
@@ -83,7 +173,11 @@ export class RunnerActor {
 		requestId: protocol.RequestId,
 		ws: WebSocketTunnelAdapter,
 	) {
-		this.deleteWebSocket(gatewayId, requestId);
+		const exists = this.getWebSocket(gatewayId, requestId) !== undefined;
+		if (exists) {
+			logger()?.warn({ msg: "attempting to set websocket twice" });
+			return;
+		}
 		this.webSockets.push({ gatewayId, requestId, ws });
 	}