fix(pegboard-gateway): write hws immediately on socket hibernate

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

@@ -6,7 +6,7 @@
 use http_body_util::{BodyExt, Full};
 use hyper::{Request, Response, StatusCode};
 use pegboard::tunnel::id::{self as tunnel_id, RequestId};
 use rand::Rng;
 use rivet_error::*;
 use rivet_guard_core::{
 	custom_serve::{CustomServeTrait, HibernationResult},
@@ -18,13 +18,13 @@
 };
 use rivet_runner_protocol as protocol;
 use rivet_util::serde::HashableMap;
 use std::{sync::Arc, time::Duration};
 use tokio::{
 	sync::{watch, Mutex},
 	task::JoinHandle,
 };
 use tokio_tungstenite::tungstenite::{
 	protocol::frame::{coding::CloseCode, CloseFrame},
 	Message,
 };
 
@@ -539,6 +539,16 @@
 	) -> Result<HibernationResult> {
 		let request_id = unique_request_id;
 
+		// 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
@@ -549,6 +559,7 @@
 				?unique_request_id,
 				"detected pending requests on websocket hibernation, rewaking actor"
 			);
+
 			return Ok(HibernationResult::Continue);
 		}
 
@@ -567,6 +578,10 @@
 			));
 			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 {
 				ping_interval.tick().await;