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

Author: NathanFlurry

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() {