@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/packages/pegboard-gateway/src/metrics.rs

@@ -1,14 +1,10 @@
-use rivet_metrics::{
-	BUCKETS,
-	otel::{global::*, metrics::*},
-};
+use rivet_metrics::{BUCKETS, REGISTRY, prometheus::*};
 
 lazy_static::lazy_static! {
-	static ref METER: Meter = meter("rivet-gateway");
-
-	/// Has no expected attributes
-	pub static ref TUNNEL_PING_DURATION: Histogram<f64> = METER.f64_histogram("rivet_gateway_tunnel_ping_duration")
-		.with_description("RTT of messages from gateway to pegboard.")
-		.with_boundaries(BUCKETS.to_vec())
-		.build();
+	pub static ref TUNNEL_PING_DURATION: Histogram = register_histogram_with_registry!(
+		"gateway_tunnel_ping_duration",
+		"RTT of messages from gateway to pegboard.",
+		BUCKETS.to_vec(),
+		*REGISTRY
+	).unwrap();
 }

package/packages/pegboard-gateway/src/metrics_task.rs

@@ -0,0 +1,80 @@
+use std::sync::{Arc, atomic::AtomicU64};
+
+use anyhow::Result;
+use gas::prelude::*;
+use tokio::sync::watch;
+
+use super::{LifecycleResult, Metric, UPDATE_METRICS_INTERVAL, record_req_metrics};
+
+pub async fn task(
+	ctx: StandaloneCtx,
+	actor_id: Id,
+	runner_id: Id,
+	ingress_bytes: Arc<AtomicU64>,
+	egress_bytes: Arc<AtomicU64>,
+	mut metrics_abort_rx: watch::Receiver<()>,
+) -> Result<LifecycleResult> {
+	let mut last_ingress_bytes = 0;
+	let mut last_egress_bytes = 0;
+
+	loop {
+		tokio::select! {
+			_ = tokio::time::sleep(UPDATE_METRICS_INTERVAL) => {}
+			_ = metrics_abort_rx.changed() => {
+				// Record final values before abort
+				record_ws_transfer(
+					&ctx,
+					actor_id,
+					runner_id,
+					&ingress_bytes,
+					&egress_bytes,
+					&mut last_ingress_bytes,
+					&mut last_egress_bytes,
+				).await?;
+
+				return Ok(LifecycleResult::Aborted);
+			}
+		}
+
+		record_ws_transfer(
+			&ctx,
+			actor_id,
+			runner_id,
+			&ingress_bytes,
+			&egress_bytes,
+			&mut last_ingress_bytes,
+			&mut last_egress_bytes,
+		)
+		.await?;
+	}
+}
+
+pub async fn record_ws_transfer(
+	ctx: &StandaloneCtx,
+	actor_id: Id,
+	runner_id: Id,
+	ingress_bytes: &AtomicU64,
+	egress_bytes: &AtomicU64,
+	last_ingress_bytes: &mut u64,
+	last_egress_bytes: &mut u64,
+) -> Result<()> {
+	let new_ingress_bytes = ingress_bytes.load(std::sync::atomic::Ordering::Acquire);
+	let new_egress_bytes = egress_bytes.load(std::sync::atomic::Ordering::Acquire);
+	let ingress_diff = new_ingress_bytes - *last_ingress_bytes;
+	let egress_diff = new_egress_bytes - *last_egress_bytes;
+
+	if ingress_diff > 0 || egress_diff > 0 {
+		record_req_metrics(
+			ctx,
+			runner_id,
+			actor_id,
+			Metric::WebsocketTransfer(ingress_diff as usize, egress_diff as usize),
+		)
+		.await?;
+	}
+
+	*last_ingress_bytes = new_ingress_bytes;
+	*last_egress_bytes = new_egress_bytes;
+
+	Ok(())
+}

package/packages/pegboard-gateway/src/ping_task.rs

@@ -1,23 +1,30 @@
 use anyhow::Result;
+use rand::Rng;
 use rivet_runner_protocol as protocol;
+use std::time::Duration;
 use tokio::sync::watch;
 
-use super::{LifecycleResult, UPDATE_PING_INTERVAL};
+use super::LifecycleResult;
 use crate::shared_state::SharedState;
 
 pub async fn task(
 	shared_state: SharedState,
 	request_id: protocol::RequestId,
 	mut ping_abort_rx: watch::Receiver<()>,
+	update_ping_interval: Duration,
 ) -> Result<LifecycleResult> {
 	loop {
 		tokio::select! {
-			_ = tokio::time::sleep(UPDATE_PING_INTERVAL) => {}
+			_ = tokio::time::sleep(update_ping_interval) => {}
 			_ = ping_abort_rx.changed() => {
 				return Ok(LifecycleResult::Aborted);
 			}
 		}
 
+		// Jitter sleep to prevent stampeding herds
+		let jitter = { rand::thread_rng().gen_range(0..128) };
+		tokio::time::sleep(Duration::from_millis(jitter)).await;
+
 		shared_state.send_and_check_ping(request_id).await?;
 	}
 }

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

@@ -1,7 +1,9 @@
 use anyhow::Result;
 use gas::prelude::*;
 use rivet_guard_core::errors::WebSocketServiceTimeout;
-use rivet_runner_protocol::{self as protocol, PROTOCOL_VERSION, versioned};
+use rivet_runner_protocol::{
+	self as protocol, PROTOCOL_MK1_VERSION, PROTOCOL_MK2_VERSION, versioned,
+};
 use scc::{HashMap, hash_map::Entry};
 use std::{
 	ops::Deref,
@@ -14,32 +16,28 @@ use vbare::OwnedVersionedData;
 
 use crate::{WebsocketPendingLimitReached, metrics};
 
-const GC_INTERVAL: Duration = Duration::from_secs(15);
-const TUNNEL_PING_TIMEOUT: i64 = util::duration::seconds(30);
-const HWS_MESSAGE_ACK_TIMEOUT: Duration = Duration::from_secs(30);
-const HWS_MAX_PENDING_MSGS_SIZE_PER_REQ: u64 = util::size::mebibytes(1);
-
 pub struct InFlightRequestHandle {
-	pub msg_rx: mpsc::Receiver<protocol::ToServerTunnelMessageKind>,
+	pub msg_rx: mpsc::Receiver<protocol::mk2::ToServerTunnelMessageKind>,
 	/// Used to check if the request handler has been dropped.
 	///
 	/// This is separate from `msg_rx` there may still be messages that need to be sent to the
 	/// request after `msg_rx` has dropped.
-	pub drop_rx: watch::Receiver<()>,
+	pub drop_rx: watch::Receiver<Option<MsgGcReason>>,
 	pub new: bool,
 }
 
 struct InFlightRequest {
 	/// UPS subject to send messages to for this request.
 	receiver_subject: String,
+	protocol_version: u16,
 	/// Sender for incoming messages to this request.
-	msg_tx: mpsc::Sender<protocol::ToServerTunnelMessageKind>,
+	msg_tx: mpsc::Sender<protocol::mk2::ToServerTunnelMessageKind>,
 	/// Used to check if the request handler has been dropped.
-	drop_tx: watch::Sender<()>,
+	drop_tx: watch::Sender<Option<MsgGcReason>>,
 	/// True once first message for this request has been sent (so runner learned reply_to).
 	opened: bool,
 	/// Message index counter for this request.
-	message_index: protocol::MessageIndex,
+	message_index: protocol::mk2::MessageIndex,
 	hibernation_state: Option<HibernationState>,
 	stopping: bool,
 	last_pong: i64,
@@ -55,15 +53,36 @@ struct HibernationState {
 pub struct PendingWebsocketMessage {
 	payload: Vec<u8>,
 	send_instant: Instant,
-	message_index: protocol::MessageIndex,
+	message_index: protocol::mk2::MessageIndex,
+}
+
+#[derive(Debug)]
+pub enum MsgGcReason {
+	/// Gateway channel is closed and is not hibernating
+	GatewayClosed,
+	/// WebSocket pending messages (ToServerWebSocketMessageAck)
+	WebSocketMessageNotAcked {
+		#[allow(dead_code)]
+		first_msg_index: u16,
+		#[allow(dead_code)]
+		last_msg_index: u16,
+	},
+	/// The gateway has not kept alive the in flight request during hibernation for the given timeout
+	/// duration.
+	HibernationTimeout,
 }
 
 pub struct SharedStateInner {
 	ups: PubSub,
-	gateway_id: protocol::GatewayId,
+	gateway_id: protocol::mk2::GatewayId,
 	receiver_subject: String,
-	in_flight_requests: HashMap<protocol::RequestId, InFlightRequest>,
+	in_flight_requests: HashMap<protocol::mk2::RequestId, InFlightRequest>,
 	hibernation_timeout: i64,
+	// Config values
+	gc_interval: Duration,
+	tunnel_ping_timeout: i64,
+	hws_message_ack_timeout: Duration,
+	hws_max_pending_size: u64,
 }
 
 #[derive(Clone)]
@@ -76,16 +95,23 @@ impl SharedState {
 		let receiver_subject =
 			pegboard::pubsub_subjects::GatewayReceiverSubject::new(gateway_id).to_string();
 
+		let pegboard_config = config.pegboard();
 		Self(Arc::new(SharedStateInner {
 			ups,
 			gateway_id,
 			receiver_subject,
 			in_flight_requests: HashMap::new(),
-			hibernation_timeout: config.pegboard().hibernating_request_eligible_threshold(),
+			hibernation_timeout: pegboard_config.hibernating_request_eligible_threshold(),
+			gc_interval: Duration::from_millis(pegboard_config.gateway_gc_interval_ms()),
+			tunnel_ping_timeout: pegboard_config.gateway_tunnel_ping_timeout_ms(),
+			hws_message_ack_timeout: Duration::from_millis(
+				pegboard_config.gateway_hws_message_ack_timeout_ms(),
+			),
+			hws_max_pending_size: pegboard_config.gateway_hws_max_pending_size(),
 		}))
 	}
 
-	pub fn gateway_id(&self) -> protocol::GatewayId {
+	pub fn gateway_id(&self) -> protocol::mk2::GatewayId {
 		self.gateway_id
 	}
 
@@ -106,15 +132,17 @@ impl SharedState {
 	pub async fn start_in_flight_request(
 		&self,
 		receiver_subject: String,
-		request_id: protocol::RequestId,
+		protocol_version: u16,
+		request_id: protocol::mk2::RequestId,
 	) -> InFlightRequestHandle {
 		let (msg_tx, msg_rx) = mpsc::channel(128);
-		let (drop_tx, drop_rx) = watch::channel(());
+		let (drop_tx, drop_rx) = watch::channel(None);
 
 		let new = match self.in_flight_requests.entry_async(request_id).await {
 			Entry::Vacant(entry) => {
 				entry.insert_entry(InFlightRequest {
 					receiver_subject,
+					protocol_version,
 					msg_tx,
 					drop_tx,
 					opened: false,
@@ -126,11 +154,13 @@ impl SharedState {
 
 				true
 			}
+			// If the entry already exists it means we transition from hibernating to active
 			Entry::Occupied(mut entry) => {
 				entry.receiver_subject = receiver_subject;
 				entry.msg_tx = msg_tx;
 				entry.drop_tx = drop_tx;
 				entry.opened = false;
+				entry.last_pong = util::timestamp::now();
 
 				if entry.stopping {
 					entry.hibernation_state = None;
@@ -151,8 +181,8 @@ impl SharedState {
 	#[tracing::instrument(skip_all, fields(request_id=%protocol::util::id_to_string(&request_id)))]
 	pub async fn send_message(
 		&self,
-		request_id: protocol::RequestId,
-		message_kind: protocol::ToClientTunnelMessageKind,
+		request_id: protocol::mk2::RequestId,
+		message_kind: protocol::mk2::ToClientTunnelMessageKind,
 	) -> Result<()> {
 		let mut req = self
 			.in_flight_requests
@@ -161,7 +191,7 @@ impl SharedState {
 			.context("request not in flight")?;
 
 		// Generate message ID
-		let message_id = protocol::MessageId {
+		let message_id = protocol::mk2::MessageId {
 			gateway_id: self.gateway_id,
 			request_id,
 			message_index: req.message_index,
@@ -180,23 +210,30 @@ impl SharedState {
 		// Check if this is a WebSocket message for hibernation tracking
 		let is_ws_message = matches!(
 			message_kind,
-			protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage(_)
+			protocol::mk2::ToClientTunnelMessageKind::ToClientWebSocketMessage(_)
 		);
 
-		let payload = protocol::ToClientTunnelMessage {
+		let payload = protocol::mk2::ToClientTunnelMessage {
 			message_id,
 			message_kind,
 		};
 
-		// Send message
-		let message = protocol::ToRunner::ToClientTunnelMessage(payload);
-		let message_serialized = versioned::ToRunner::wrap_latest(message)
-			.serialize_with_embedded_version(PROTOCOL_VERSION)?;
+		let message_serialized = if protocol::is_mk2(req.protocol_version) {
+			let message = protocol::mk2::ToRunner::ToClientTunnelMessage(payload);
+			versioned::ToRunnerMk2::wrap_latest(message)
+				.serialize_with_embedded_version(PROTOCOL_MK2_VERSION)?
+		} else {
+			let message = protocol::ToRunner::ToClientTunnelMessage(
+				versioned::to_client_tunnel_message_mk2_to_mk1(payload),
+			);
+			versioned::ToRunner::wrap_latest(message)
+				.serialize_with_embedded_version(PROTOCOL_MK1_VERSION)?
+		};
 
 		if let (Some(hs), true) = (&mut req.hibernation_state, is_ws_message) {
 			hs.total_pending_ws_msgs_size += message_serialized.len() as u64;
 
-			if hs.total_pending_ws_msgs_size > HWS_MAX_PENDING_MSGS_SIZE_PER_REQ
+			if hs.total_pending_ws_msgs_size > self.hws_max_pending_size
 				|| hs.pending_ws_msgs.len() >= u16::MAX as usize
 			{
 				return Err(WebsocketPendingLimitReached {}.build());
@@ -228,7 +265,7 @@ impl SharedState {
 	}
 
 	#[tracing::instrument(skip_all, fields(request_id=%protocol::util::id_to_string(&request_id)))]
-	pub async fn send_and_check_ping(&self, request_id: protocol::RequestId) -> Result<()> {
+	pub async fn send_and_check_ping(&self, request_id: protocol::mk2::RequestId) -> Result<()> {
 		let req = self
 			.in_flight_requests
 			.get_async(&request_id)
@@ -238,19 +275,28 @@ impl SharedState {
 		let now = util::timestamp::now();
 
 		// Verify ping timeout
-		if now.saturating_sub(req.last_pong) > TUNNEL_PING_TIMEOUT {
-			tracing::warn!("tunnel timeout");
+		if now.saturating_sub(req.last_pong) > self.tunnel_ping_timeout {
+			tracing::warn!(runner_topic=%req.receiver_subject, "tunnel timeout");
 			return Err(WebSocketServiceTimeout.build());
 		}
 
-		// Send message
-		let message = protocol::ToRunner::ToRunnerPing(protocol::ToRunnerPing {
-			gateway_id: self.gateway_id,
-			request_id,
-			ts: now,
-		});
-		let message_serialized = versioned::ToRunner::wrap_latest(message)
-			.serialize_with_embedded_version(PROTOCOL_VERSION)?;
+		let message_serialized = if protocol::is_mk2(req.protocol_version) {
+			let message = protocol::mk2::ToRunner::ToRunnerPing(protocol::mk2::ToRunnerPing {
+				gateway_id: self.gateway_id,
+				request_id,
+				ts: now,
+			});
+			versioned::ToRunnerMk2::wrap_latest(message)
+				.serialize_with_embedded_version(PROTOCOL_MK2_VERSION)?
+		} else {
+			let message = protocol::ToRunner::ToRunnerPing(protocol::ToRunnerPing {
+				gateway_id: self.gateway_id,
+				request_id,
+				ts: now,
+			});
+			versioned::ToRunner::wrap_latest(message)
+				.serialize_with_embedded_version(PROTOCOL_MK1_VERSION)?
+		};
 
 		self.ups
 			.publish(
@@ -264,7 +310,7 @@ impl SharedState {
 	}
 
 	#[tracing::instrument(skip_all, fields(request_id=%protocol::util::id_to_string(&request_id)))]
-	pub async fn keepalive_hws(&self, request_id: protocol::RequestId) -> Result<()> {
+	pub async fn keepalive_hws(&self, request_id: protocol::mk2::RequestId) -> Result<()> {
 		let mut req = self
 			.in_flight_requests
 			.get_async(&request_id)
@@ -289,7 +335,7 @@ impl SharedState {
 			);
 
 			match versioned::ToGateway::deserialize_with_embedded_version(&msg.payload) {
-				Ok(protocol::ToGateway::ToGatewayPong(pong)) => {
+				Ok(protocol::mk2::ToGateway::ToGatewayPong(pong)) => {
 					let Some(mut in_flight) =
 						self.in_flight_requests.get_async(&pong.request_id).await
 					else {
@@ -304,9 +350,9 @@ impl SharedState {
 					in_flight.last_pong = now;
 
 					let rtt = now.saturating_sub(pong.ts);
-					metrics::TUNNEL_PING_DURATION.record(rtt as f64 * 0.001, &[]);
+					metrics::TUNNEL_PING_DURATION.observe(rtt as f64 * 0.001);
 				}
-				Ok(protocol::ToGateway::ToServerTunnelMessage(msg)) => {
+				Ok(protocol::mk2::ToGateway::ToServerTunnelMessage(msg)) => {
 					let message_id = msg.message_id;
 
 					let Some(in_flight) = self
@@ -324,10 +370,17 @@ impl SharedState {
 					};
 
 					// Send message to the request handler to emulate the real network action
+					let inner_size = match &msg.message_kind {
+						protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketMessage(
+							ws_msg,
+						) => ws_msg.data.len(),
+						_ => 0,
+					};
 					tracing::debug!(
 						gateway_id=%protocol::util::id_to_string(&message_id.gateway_id),
 						request_id=%protocol::util::id_to_string(&message_id.request_id),
 						message_index=message_id.message_index,
+						inner_size,
 						"forwarding message to request handler"
 					);
 					let _ = in_flight.msg_tx.send(msg.message_kind.clone()).await;
@@ -342,7 +395,7 @@ impl SharedState {
 	#[tracing::instrument(skip_all, fields(request_id=%protocol::util::id_to_string(&request_id), %enable))]
 	pub async fn toggle_hibernation(
 		&self,
-		request_id: protocol::RequestId,
+		request_id: protocol::mk2::RequestId,
 		enable: bool,
 	) -> Result<()> {
 		let mut req = self
@@ -370,7 +423,7 @@ impl SharedState {
 	#[tracing::instrument(skip_all, fields(request_id=%protocol::util::id_to_string(&request_id)))]
 	pub async fn resend_pending_websocket_messages(
 		&self,
-		request_id: protocol::RequestId,
+		request_id: protocol::mk2::RequestId,
 	) -> Result<()> {
 		let Some(mut req) = self.in_flight_requests.get_async(&request_id).await else {
 			bail!("request not in flight");
@@ -396,7 +449,7 @@ impl SharedState {
 	#[tracing::instrument(skip_all, fields(request_id=%protocol::util::id_to_string(&request_id)))]
 	pub async fn has_pending_websocket_messages(
 		&self,
-		request_id: protocol::RequestId,
+		request_id: protocol::mk2::RequestId,
 	) -> Result<bool> {
 		let Some(req) = self.in_flight_requests.get_async(&request_id).await else {
 			bail!("request not in flight");
@@ -412,7 +465,7 @@ impl SharedState {
 	#[tracing::instrument(skip_all, fields(request_id=%protocol::util::id_to_string(&request_id), %ack_index))]
 	pub async fn ack_pending_websocket_messages(
 		&self,
-		request_id: protocol::RequestId,
+		request_id: protocol::mk2::RequestId,
 		ack_index: u16,
 	) -> Result<()> {
 		let Some(mut req) = self.in_flight_requests.get_async(&request_id).await else {
@@ -441,7 +494,7 @@ impl SharedState {
 
 	#[tracing::instrument(skip_all)]
 	async fn gc(&self) {
-		let mut interval = tokio::time::interval(GC_INTERVAL);
+		let mut interval = tokio::time::interval(self.gc_interval);
 		interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 
 		loop {
@@ -473,22 +526,6 @@ impl SharedState {
 	/// 2a. Remove all requests where it was flagged as stopping and `drop_rx` has been dropped
 	#[tracing::instrument(skip_all)]
 	async fn gc_in_flight_requests(&self) {
-		#[derive(Debug)]
-		enum MsgGcReason {
-			/// Gateway channel is closed and is not hibernating
-			GatewayClosed,
-			/// WebSocket pending messages (ToServerWebSocketMessageAck)
-			WebSocketMessageNotAcked {
-				#[allow(dead_code)]
-				first_msg_index: u16,
-				#[allow(dead_code)]
-				last_msg_index: u16,
-			},
-			/// The gateway has not kept alive the in flight request during hibernation for the given timeout
-			/// duration.
-			HibernationTimeout,
-		}
-
 		let now = Instant::now();
 		let hibernation_timeout =
 			Duration::from_millis(self.hibernation_timeout.try_into().unwrap_or(90_000));
@@ -497,7 +534,8 @@ impl SharedState {
 		// message ack
 		self.in_flight_requests
 			.iter_mut_async(|mut entry| {
-				let (request_id, req) = &mut *entry;
+				let request_id = entry.key().clone();
+				let req = &mut *entry;
 
 				if req.stopping {
 					return true;
@@ -506,12 +544,10 @@ impl SharedState {
 				let reason = 'reason: {
 					if let Some(hs) = &req.hibernation_state {
 						if let Some(earliest_pending_ws_msg) = hs.pending_ws_msgs.first() {
-							if now.duration_since(earliest_pending_ws_msg.send_instant)
-								> HWS_MESSAGE_ACK_TIMEOUT
-							{
+							if now.duration_since(earliest_pending_ws_msg.send_instant) > self.hws_message_ack_timeout {
 								break 'reason Some(MsgGcReason::WebSocketMessageNotAcked {
 									first_msg_index: earliest_pending_ws_msg.message_index,
-									last_msg_index: req.message_index
+									last_msg_index: req.message_index,
 								});
 							}
 						}
@@ -522,7 +558,7 @@ impl SharedState {
 							timeout=%hibernation_timeout.as_secs_f64(),
 							"checking hibernating state elapsed time"
 						);
-						if hs_elapsed> hibernation_timeout {
+						if hs_elapsed > hibernation_timeout {
 							break 'reason Some(MsgGcReason::HibernationTimeout);
 						}
 					} else if req.msg_tx.is_closed() {
@@ -532,15 +568,15 @@ impl SharedState {
 					None
 				};
 
-				if let Some(reason) = &reason {
+				if let Some(reason) = reason {
 					tracing::debug!(
-						request_id=%protocol::util::id_to_string(request_id),
+						request_id=%protocol::util::id_to_string(&request_id),
 						?reason,
 						"gc stopping in flight request"
 					);
 
-					if req.drop_tx.send(()).is_err() {
-						tracing::debug!(request_id=%protocol::util::id_to_string(request_id), "failed to send timeout msg to tunnel");
+					if req.drop_tx.send(Some(reason)).is_err() {
+						tracing::debug!(request_id=%protocol::util::id_to_string(&request_id), "failed to send timeout msg to tunnel");
 					}
 
 					// Mark req as stopping to skip this loop next time the gc is run

package/packages/pegboard-gateway/src/tunnel_to_ws_task.rs

@@ -1,3 +1,8 @@
+use std::sync::{
+	Arc,
+	atomic::{AtomicU64, Ordering},
+};
+
 use anyhow::Result;
 use gas::prelude::*;
 use rivet_guard_core::{
@@ -9,16 +14,17 @@ use tokio::sync::{mpsc, watch};
 use tokio_tungstenite::tungstenite::Message;
 
 use super::LifecycleResult;
-use crate::shared_state::SharedState;
+use crate::shared_state::{MsgGcReason, SharedState};
 
 pub async fn task(
 	shared_state: SharedState,
 	client_ws: WebSocketHandle,
 	request_id: protocol::RequestId,
 	mut stopped_sub: message::SubscriptionHandle<pegboard::workflows::actor::Stopped>,
-	mut msg_rx: mpsc::Receiver<protocol::ToServerTunnelMessageKind>,
-	mut drop_rx: watch::Receiver<()>,
+	mut msg_rx: mpsc::Receiver<protocol::mk2::ToServerTunnelMessageKind>,
+	mut drop_rx: watch::Receiver<Option<MsgGcReason>>,
 	can_hibernate: bool,
+	egress_bytes: Arc<AtomicU64>,
 	mut tunnel_to_ws_abort_rx: watch::Receiver<()>,
 ) -> Result<LifecycleResult> {
 	loop {
@@ -26,7 +32,13 @@ pub async fn task(
 			res = msg_rx.recv() => {
 				if let Some(msg) = res {
 					match msg {
-						protocol::ToServerTunnelMessageKind::ToServerWebSocketMessage(ws_msg) => {
+						protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketMessage(ws_msg) => {
+							tracing::trace!(
+								request_id=%protocol::util::id_to_string(&request_id),
+								data_len=ws_msg.data.len(),
+								binary=ws_msg.binary,
+								"forwarding websocket message to client"
+							);
 							let msg = if ws_msg.binary {
 								Message::Binary(ws_msg.data.into())
 							} else {
@@ -34,9 +46,11 @@ pub async fn task(
 									String::from_utf8_lossy(&ws_msg.data).into_owned().into(),
 								)
 							};
+
+							egress_bytes.fetch_add(msg.len() as u64, Ordering::AcqRel);
 							client_ws.send(msg).await?;
 						}
-						protocol::ToServerTunnelMessageKind::ToServerWebSocketMessageAck(ack) => {
+						protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketMessageAck(ack) => {
 							tracing::debug!(
 								request_id=%protocol::util::id_to_string(&request_id),
 								ack_index=?ack.index,
@@ -46,7 +60,7 @@ pub async fn task(
 								.ack_pending_websocket_messages(request_id, ack.index)
 								.await?;
 						}
-						protocol::ToServerTunnelMessageKind::ToServerWebSocketClose(close) => {
+						protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketClose(close) => {
 							tracing::debug!(?close, "server closed websocket");
 
 							if can_hibernate && close.hibernate {
@@ -73,7 +87,7 @@ pub async fn task(
 				}
 			}
 			_ = drop_rx.changed() => {
-				tracing::warn!("garbage collected");
+				tracing::warn!(reason=?drop_rx.borrow().as_ref(), "garbage collected");
 				return Err(WebSocketServiceTimeout.build());
 			}
 			_ = tunnel_to_ws_abort_rx.changed() => {

package/packages/pegboard-gateway/src/ws_to_tunnel_task.rs

@@ -2,7 +2,10 @@ use anyhow::Result;
 use futures_util::TryStreamExt;
 use rivet_guard_core::websocket_handle::WebSocketReceiver;
 use rivet_runner_protocol as protocol;
-use std::sync::Arc;
+use std::sync::{
+	Arc,
+	atomic::{AtomicU64, Ordering},
+};
 use tokio::sync::{Mutex, watch};
 use tokio_tungstenite::tungstenite::Message;
 
@@ -11,8 +14,9 @@ use crate::shared_state::SharedState;
 
 pub async fn task(
 	shared_state: SharedState,
-	request_id: protocol::RequestId,
+	request_id: protocol::mk2::RequestId,
 	ws_rx: Arc<Mutex<WebSocketReceiver>>,
+	ingress_bytes: Arc<AtomicU64>,
 	mut ws_to_tunnel_abort_rx: watch::Receiver<()>,
 ) -> Result<LifecycleResult> {
 	let mut ws_rx = ws_rx.lock().await;
@@ -21,11 +25,13 @@ pub async fn task(
 		tokio::select! {
 			res = ws_rx.try_next() => {
 				if let Some(msg) = res? {
+					ingress_bytes.fetch_add(msg.len() as u64, Ordering::AcqRel);
+
 					match msg {
 						Message::Binary(data) => {
 							let ws_message =
-								protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage(
-									protocol::ToClientWebSocketMessage {
+								protocol::mk2::ToClientTunnelMessageKind::ToClientWebSocketMessage(
+									protocol::mk2::ToClientWebSocketMessage {
 										data: data.into(),
 										binary: true,
 									},
@@ -36,8 +42,8 @@ pub async fn task(
 						}
 						Message::Text(text) => {
 							let ws_message =
-								protocol::ToClientTunnelMessageKind::ToClientWebSocketMessage(
-									protocol::ToClientWebSocketMessage {
+								protocol::mk2::ToClientTunnelMessageKind::ToClientWebSocketMessage(
+									protocol::mk2::ToClientWebSocketMessage {
 										data: text.as_bytes().to_vec(),
 										binary: false,
 									},