fix(engine): ensure first keepalive tick is not skipped to prevent timeout on ws hibernation

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

@@ -1,5 +1,6 @@
 use anyhow::Result;
 use gas::prelude::*;
+use pegboard::tunnel::id as tunnel_id;
 use pegboard::tunnel::id::{GatewayId, RequestId};
 use rand::Rng;
 use std::time::Duration;
@@ -12,7 +13,6 @@ use crate::shared_state::SharedState;
 /// next actor start.
 ///
 /// Only ran for hibernating requests.
-
 pub async fn task(
 	shared_state: SharedState,
 	ctx: StandaloneCtx,
@@ -30,10 +30,6 @@ pub async fn task(
 	));
 	ping_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 
-	// Discard the first tick since it fires immediately and we've already called this
-	// above
-	ping_interval.tick().await;
-
 	loop {
 		tokio::select! {
 			_ = ping_interval.tick() => {}
@@ -46,6 +42,13 @@ pub async fn task(
 		let jitter = { rand::thread_rng().gen_range(0..128) };
 		tokio::time::sleep(Duration::from_millis(jitter)).await;
 
+		tracing::debug!(
+			%actor_id,
+			gateway_id=%tunnel_id::gateway_id_to_string(&gateway_id),
+			request_id=%tunnel_id::request_id_to_string(&request_id),
+			"updating hws keepalive"
+		);
+
 		tokio::try_join!(
 			ctx.op(pegboard::ops::actor::hibernating_request::upsert::Input {
 				actor_id,

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

@@ -8,20 +8,20 @@ use hyper::{Request, Response, StatusCode};
 use pegboard::tunnel::id::{self as tunnel_id, RequestId};
 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};
+use tokio::sync::{watch, Mutex};
 use tokio_tungstenite::tungstenite::{
+	protocol::frame::{coding::CloseCode, CloseFrame},
 	Message,
-	protocol::frame::{CloseFrame, coding::CloseCode},
 };
 
 use crate::shared_state::{InFlightRequestHandle, SharedState};
@@ -578,16 +578,6 @@ impl CustomServeTrait for PegboardGateway {
 		client_ws: WebSocketHandle,
 		request_id: RequestId,
 	) -> Result<HibernationResult> {
-		// Insert hibernating request entry before checking for pending messages
-		// This ensures the entry exists even if we immediately rewake the actor
-		self.ctx
-			.op(pegboard::ops::actor::hibernating_request::upsert::Input {
-				actor_id: self.actor_id,
-				gateway_id: self.shared_state.gateway_id(),
-				request_id,
-			})
-			.await?;
-
 		// Immediately rewake if we have pending messages
 		if self
 			.shared_state

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

@@ -2,8 +2,8 @@ use anyhow::Result;
 use gas::prelude::*;
 use pegboard::tunnel::id::{self as tunnel_id, GatewayId, RequestId};
 use rivet_guard_core::errors::WebSocketServiceTimeout;
-use rivet_runner_protocol::{self as protocol, PROTOCOL_VERSION, versioned};
-use scc::{HashMap, hash_map::Entry};
+use rivet_runner_protocol::{self as protocol, versioned, PROTOCOL_VERSION};
+use scc::{hash_map::Entry, HashMap};
 use std::{
 	ops::Deref,
 	sync::Arc,
@@ -13,7 +13,7 @@ use tokio::sync::{mpsc, watch};
 use universalpubsub::{NextOutput, PubSub, PublishOpts, Subscriber};
 use vbare::OwnedVersionedData;
 
-use crate::{WebsocketPendingLimitReached, metrics};
+use crate::{metrics, WebsocketPendingLimitReached};
 
 const GC_INTERVAL: Duration = Duration::from_secs(15);
 const TUNNEL_PING_TIMEOUT: i64 = util::duration::seconds(30);
@@ -512,7 +512,13 @@ impl SharedState {
 							}
 						}
 
-						if hs.last_ping.elapsed() > hibernation_timeout {
+						let hs_elapsed = hs.last_ping.elapsed();
+						tracing::debug!(
+							hs_elapsed=%hs_elapsed.as_secs_f64(),
+							timeout=%hibernation_timeout.as_secs_f64(),
+							"checking hibernating state elapsed time"
+						);
+						if hs_elapsed> hibernation_timeout {
 							break 'reason Some(MsgGcReason::HibernationTimeout);
 						}
 					} else if req.msg_tx.is_closed() {

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

@@ -1,9 +1,8 @@
 import type * as protocol from "@rivetkit/engine-runner-protocol";
+import { logger } from "./log";
 import type { PendingRequest } from "./tunnel";
+import { arraysEqual, idToStr, promiseWithResolvers } from "./utils";
 import type { WebSocketTunnelAdapter } from "./websocket-tunnel-adapter";
-import { arraysEqual, promiseWithResolvers } from "./utils";
-import { logger } from "./log";
-import * as tunnelId from "./tunnel-id";
 
 export interface ActorConfig {
 	name: string;
@@ -72,8 +71,8 @@ export class RunnerActor {
 		if (exists) {
 			logger()?.warn({
 				msg: "attempting to set pending request twice, replacing existing",
-				gatewayId: tunnelId.gatewayIdToString(gatewayId),
-				requestId: tunnelId.requestIdToString(requestId),
+				gatewayId: idToStr(gatewayId),
+				requestId: idToStr(requestId),
 			});
 			// Delete existing pending request before adding the new one
 			this.deletePendingRequest(gatewayId, requestId);
@@ -92,8 +91,8 @@ export class RunnerActor {
 		});
 		logger()?.debug({
 			msg: "added pending request",
-			gatewayId: tunnelId.gatewayIdToString(gatewayId),
-			requestId: tunnelId.requestIdToString(requestId),
+			gatewayId: idToStr(gatewayId),
+			requestId: idToStr(requestId),
 			length: this.pendingRequests.length,
 		});
 	}
@@ -109,8 +108,8 @@ export class RunnerActor {
 		if (exists) {
 			logger()?.warn({
 				msg: "attempting to set pending request twice, replacing existing",
-				gatewayId: tunnelId.gatewayIdToString(gatewayId),
-				requestId: tunnelId.requestIdToString(requestId),
+				gatewayId: idToStr(gatewayId),
+				requestId: idToStr(requestId),
 			});
 			// Delete existing pending request before adding the new one
 			this.deletePendingRequest(gatewayId, requestId);
@@ -130,8 +129,8 @@ export class RunnerActor {
 		});
 		logger()?.debug({
 			msg: "added pending request with stream controller",
-			gatewayId: tunnelId.gatewayIdToString(gatewayId),
-			requestId: tunnelId.requestIdToString(requestId),
+			gatewayId: idToStr(gatewayId),
+			requestId: idToStr(requestId),
 			length: this.pendingRequests.length,
 		});
 	}
@@ -149,8 +148,8 @@ export class RunnerActor {
 			this.pendingRequests.splice(index, 1);
 			logger()?.debug({
 				msg: "removed pending request",
-				gatewayId: tunnelId.gatewayIdToString(gatewayId),
-				requestId: tunnelId.requestIdToString(requestId),
+				gatewayId: idToStr(gatewayId),
+				requestId: idToStr(requestId),
 				length: this.pendingRequests.length,
 			});
 		}

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

@@ -13,8 +13,8 @@
 import { importWebSocket } from "./websocket.js";
 
 export type { HibernatingWebSocketMetadata };
-export * as tunnelId from "./tunnel-id";
 export { RunnerActor, type ActorConfig };
+export { idToStr } from "./utils";
 
 const KV_EXPIRE: number = 30_000;
 const PROTOCOL_VERSION: number = 3;
@@ -199,7 +199,7 @@
 	runnerId?: string;
 	#lastCommandIdx: number = -1;
 	#pingLoop?: NodeJS.Timeout;
 	#nextEventIdx: bigint = 0n;
 	#started: boolean = false;
 	#shutdown: boolean = false;
 	#shuttingDown: boolean = false;

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

@@ -1,4 +1,5 @@
 import type * as protocol from "@rivetkit/engine-runner-protocol";
+import { idToStr } from "./utils";
 
 /**
  * Helper function to stringify ArrayBuffer for logging
@@ -24,6 +25,13 @@ function stringifyMap(map: ReadonlyMap<string, string>): string {
 	return `Map(${map.size}){${entries}}`;
 }
 
+/**
+ * Helper function to stringify MessageId for logging
+ */
+function stringifyMessageId(messageId: protocol.MessageId): string {
+	return `MessageId{gatewayId: ${idToStr(messageId.gatewayId)}, requestId: ${idToStr(messageId.requestId)}, messageIndex: ${messageId.messageIndex}}`;
+}
+
 /**
  * Stringify ToServerTunnelMessageKind for logging
  * Handles ArrayBuffers, BigInts, and Maps that can't be JSON.stringified
@@ -111,15 +119,17 @@ export function stringifyToClientTunnelMessageKind(
 export function stringifyCommand(command: protocol.Command): string {
 	switch (command.tag) {
 		case "CommandStartActor": {
-			const { actorId, generation, config, hibernatingRequests } = command.val;
+			const { actorId, generation, config, hibernatingRequests } =
+				command.val;
 			const keyStr = config.key === null ? "null" : `"${config.key}"`;
 			const inputStr =
 				config.input === null
 					? "null"
 					: stringifyArrayBuffer(config.input);
-			const hibernatingRequestsStr = hibernatingRequests.length > 0
-				? `[${hibernatingRequests.map((hr) => `{gatewayId: ${stringifyArrayBuffer(hr.gatewayId)}, requestId: ${stringifyArrayBuffer(hr.requestId)}}`).join(", ")}]`
-				: "[]";
+			const hibernatingRequestsStr =
+				hibernatingRequests.length > 0
+					? `[${hibernatingRequests.map((hr) => `{gatewayId: ${idToStr(hr.gatewayId)}, requestId: ${idToStr(hr.requestId)}}`).join(", ")}]`
+					: "[]";
 			return `CommandStartActor{actorId: "${actorId}", generation: ${generation}, config: {name: "${config.name}", key: ${keyStr}, createTs: ${stringifyBigInt(config.createTs)}, input: ${inputStr}}, hibernatingRequests: ${hibernatingRequestsStr}}`;
 		}
 		case "CommandStopActor": {
@@ -193,8 +203,14 @@ export function stringifyEventWrapper(wrapper: protocol.EventWrapper): string {
 export function stringifyToServer(message: protocol.ToServer): string {
 	switch (message.tag) {
 		case "ToServerInit": {
-			const { name, version, totalSlots, lastCommandIdx, prepopulateActorNames, metadata } =
-				message.val;
+			const {
+				name,
+				version,
+				totalSlots,
+				lastCommandIdx,
+				prepopulateActorNames,
+				metadata,
+			} = message.val;
 			const lastCommandIdxStr =
 				lastCommandIdx === null
 					? "null"
@@ -227,8 +243,7 @@ export function stringifyToServer(message: protocol.ToServer): string {
 		}
 		case "ToServerTunnelMessage": {
 			const { messageId, messageKind } = message.val;
-			const messageIdStr = stringifyArrayBuffer(messageId);
-			return `ToServerTunnelMessage{messageId: ${messageIdStr}, messageKind: ${stringifyToServerTunnelMessageKind(messageKind)}}`;
+			return `ToServerTunnelMessage{messageId: ${stringifyMessageId(messageId)}, messageKind: ${stringifyToServerTunnelMessageKind(messageKind)}}`;
 		}
 	}
 }
@@ -261,8 +276,7 @@ export function stringifyToClient(message: protocol.ToClient): string {
 		}
 		case "ToClientTunnelMessage": {
 			const { messageId, messageKind } = message.val;
-			const messageIdStr = stringifyArrayBuffer(messageId);
-			return `ToClientTunnelMessage{messageId: ${messageIdStr}, messageKind: ${stringifyToClientTunnelMessageKind(messageKind)}}`;
+			return `ToClientTunnelMessage{messageId: ${stringifyMessageId(messageId)}, messageKind: ${stringifyToClientTunnelMessageKind(messageKind)}}`;
 		}
 	}
 }