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

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

@@ -4,38 +4,39 @@ use bytes::Bytes;
 use futures_util::TryStreamExt;
 use gas::prelude::*;
 use http_body_util::{BodyExt, Full};
-use hyper::{Request, Response, StatusCode};
+use hyper::{Request, Response, StatusCode, body::Body};
 use rivet_error::*;
 use rivet_guard_core::{
-	WebSocketHandle,
+	ResponseBody, WebSocketHandle,
 	custom_serve::{CustomServeTrait, HibernationResult},
 	errors::{ServiceUnavailable, WebSocketServiceUnavailable},
-	proxy_service::{ResponseBody, is_ws_hibernate},
-	request_context::RequestContext,
+	request_context::{CorsConfig, RequestContext},
+	utils::is_ws_hibernate,
 	websocket_handle::WebSocketReceiver,
 };
-use rivet_runner_protocol as protocol;
+use rivet_runner_protocol::{self as protocol, PROTOCOL_MK1_VERSION};
 use rivet_util::serde::HashableMap;
-use std::{sync::Arc, time::Duration};
+use std::{
+	sync::{Arc, atomic::AtomicU64},
+	time::Duration,
+};
 use tokio::sync::{Mutex, watch};
 use tokio_tungstenite::tungstenite::{
 	Message,
 	protocol::frame::{CloseFrame, coding::CloseCode},
 };
+use universaldb::utils::IsolationLevel::*;
 
 use crate::shared_state::{InFlightRequestHandle, SharedState};
 
 mod keepalive_task;
 mod metrics;
+mod metrics_task;
 mod ping_task;
 pub mod shared_state;
 mod tunnel_to_ws_task;
 mod ws_to_tunnel_task;
 
-const WEBSOCKET_OPEN_TIMEOUT: Duration = Duration::from_secs(15);
-const RESPONSE_START_TIMEOUT: Duration = Duration::from_secs(15);
-const UPDATE_PING_INTERVAL: Duration = Duration::from_secs(3);
-
 #[derive(RivetError, Serialize, Deserialize)]
 #[error(
 	"guard",
@@ -44,9 +45,11 @@ const UPDATE_PING_INTERVAL: Duration = Duration::from_secs(3);
 )]
 pub struct WebsocketPendingLimitReached;
 
+const UPDATE_METRICS_INTERVAL: Duration = Duration::from_secs(15);
+
 #[derive(Debug)]
 enum LifecycleResult {
-	ServerClose(protocol::ToServerWebSocketClose),
+	ServerClose(protocol::mk2::ToServerWebSocketClose),
 	ClientClose(Option<CloseFrame>),
 	Aborted,
 }
@@ -78,17 +81,16 @@ impl PegboardGateway {
 	}
 }
 
-#[async_trait]
-impl CustomServeTrait for PegboardGateway {
-	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id, runner_id=?self.runner_id))]
-	async fn handle_request(
+impl PegboardGateway {
+	async fn handle_request_inner(
 		&self,
+		ctx: &StandaloneCtx,
 		req: Request<Full<Bytes>>,
-		_request_context: &mut RequestContext,
-		request_id: protocol::RequestId,
+		req_ctx: &mut RequestContext,
 	) -> Result<Response<ResponseBody>> {
 		// Use the actor ID from the gateway instance
 		let actor_id = self.actor_id.to_string();
+		let request_id = req_ctx.in_flight_request_id()?;
 
 		// Extract origin for CORS (before consuming request)
 		// When credentials: true, we must echo back the actual origin, not "*"
@@ -100,19 +102,20 @@ impl CustomServeTrait for PegboardGateway {
 			.to_string();
 
 		// Extract request parts
-		let mut headers = HashableMap::new();
-		for (name, value) in req.headers() {
-			if let Result::Ok(value_str) = value.to_str() {
-				headers.insert(name.to_string(), value_str.to_string());
-			}
-		}
-
-		// Extract method and path before consuming the request
-		let method = req.method().to_string();
+		let headers = req
+			.headers()
+			.iter()
+			.filter_map(|(name, value)| {
+				value
+					.to_str()
+					.ok()
+					.map(|value_str| (name.to_string(), value_str.to_string()))
+			})
+			.collect::<HashableMap<_, _>>();
 
 		// Handle CORS preflight OPTIONS requests at gateway level
 		//
-		// We need to do this in Guard because there is no way of sending an OPTIONS request to the
+		// We need to do this in the gateway because there is no way of sending an OPTIONS request to the
 		// actor since we don't have the `x-rivet-token` header. This implementation allows
 		// requests from anywhere and lets the actor handle CORS manually in `onBeforeConnect`.
 		// This had the added benefit of also applying to WebSockets.
@@ -126,26 +129,32 @@ impl CustomServeTrait for PegboardGateway {
 				.and_then(|v| v.to_str().ok())
 				.unwrap_or("*");
 
-			let mut response = Response::builder()
-				.status(StatusCode::NO_CONTENT)
-				.header("access-control-allow-origin", &origin)
-				.header("access-control-allow-credentials", "true")
-				.header(
-					"access-control-allow-methods",
-					"GET, POST, PUT, DELETE, OPTIONS, PATCH",
-				)
-				.header("access-control-allow-headers", requested_headers)
-				.header("access-control-expose-headers", "*")
-				.header("access-control-max-age", "86400");
-
-			// Add Vary header to prevent cache poisoning when echoing origin
-			if origin != "*" {
-				response = response.header("vary", "Origin");
-			}
+			req_ctx.set_cors(CorsConfig {
+				allow_origin: origin.clone(),
+				allow_credentials: true,
+				expose_headers: "*".to_string(),
+				allow_methods: Some("GET, POST, PUT, DELETE, OPTIONS, PATCH".to_string()),
+				allow_headers: Some(requested_headers.to_string()),
+				max_age: Some(86400),
+			});
 
-			return Ok(response.body(ResponseBody::Full(Full::new(Bytes::new())))?);
+			return Ok(Response::builder()
+				.status(StatusCode::NO_CONTENT)
+				.body(ResponseBody::Full(Full::new(Bytes::new())))?);
 		}
 
+		// Set CORS headers through guard
+		req_ctx.set_cors(CorsConfig {
+			allow_origin: origin.clone(),
+			allow_credentials: true,
+			expose_headers: "*".to_string(),
+			// Not an options req, not required
+			allow_methods: None,
+			allow_headers: None,
+			max_age: None,
+		});
+
+		// NOTE: Size constraints have already been applied by guard
 		let body_bytes = req
 			.into_body()
 			.collect()
@@ -153,10 +162,10 @@ impl CustomServeTrait for PegboardGateway {
 			.context("failed to read body")?
 			.to_bytes();
 
-		let mut stopped_sub = self
-			.ctx
-			.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", self.actor_id))
-			.await?;
+		let (mut stopped_sub, runner_protocol_version) = tokio::try_join!(
+			ctx.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", self.actor_id)),
+			get_runner_protocol_version(&ctx, self.runner_id),
+		)?;
 
 		// Build subject to publish to
 		let tunnel_subject =
@@ -169,14 +178,14 @@ impl CustomServeTrait for PegboardGateway {
 			..
 		} = self
 			.shared_state
-			.start_in_flight_request(tunnel_subject, request_id)
+			.start_in_flight_request(tunnel_subject, runner_protocol_version, request_id)
 			.await;
 
 		// Start request
-		let message = protocol::ToClientTunnelMessageKind::ToClientRequestStart(
-			protocol::ToClientRequestStart {
+		let message = protocol::mk2::ToClientTunnelMessageKind::ToClientRequestStart(
+			protocol::mk2::ToClientRequestStart {
 				actor_id: actor_id.clone(),
-				method,
+				method: req_ctx.method().to_string(),
 				path: self.path.clone(),
 				headers,
 				body: if body_bytes.is_empty() {
@@ -197,12 +206,12 @@ impl CustomServeTrait for PegboardGateway {
 					res = msg_rx.recv() => {
 						if let Some(msg) = res {
 							match msg {
-								protocol::ToServerTunnelMessageKind::ToServerResponseStart(
+								protocol::mk2::ToServerTunnelMessageKind::ToServerResponseStart(
 									response_start,
 								) => {
 									return anyhow::Ok(response_start);
 								}
-								protocol::ToServerTunnelMessageKind::ToServerResponseAbort => {
+								protocol::mk2::ToServerTunnelMessageKind::ToServerResponseAbort => {
 									tracing::warn!("request aborted");
 									return Err(ServiceUnavailable.build());
 								}
@@ -223,7 +232,7 @@ impl CustomServeTrait for PegboardGateway {
 						return Err(ServiceUnavailable.build());
 					}
 					_ = drop_rx.changed() => {
-						tracing::warn!("tunnel message timeout");
+						tracing::warn!(reason=?drop_rx.borrow(), "tunnel message timeout");
 						return Err(ServiceUnavailable.build());
 					}
 				}
@@ -231,7 +240,13 @@ impl CustomServeTrait for PegboardGateway {
 
 			Err(ServiceUnavailable.build())
 		};
-		let response_start = tokio::time::timeout(RESPONSE_START_TIMEOUT, fut)
+		let response_start_timeout = Duration::from_millis(
+			self.ctx
+				.config()
+				.pegboard()
+				.gateway_response_start_timeout_ms(),
+		);
+		let response_start = tokio::time::timeout(response_start_timeout, fut)
 			.await
 			.map_err(|_| {
 				tracing::warn!("timed out waiting for response start from runner");
@@ -249,17 +264,6 @@ impl CustomServeTrait for PegboardGateway {
 			response_builder = response_builder.header(key, value);
 		}
 
-		// Add CORS headers to actual request
-		response_builder = response_builder
-			.header("access-control-allow-origin", &origin)
-			.header("access-control-allow-credentials", "true")
-			.header("access-control-expose-headers", "*");
-
-		// Add Vary header to prevent cache poisoning when echoing origin
-		if origin != "*" {
-			response_builder = response_builder.header("vary", "Origin");
-		}
-
 		// Add body
 		let body = response_start.body.unwrap_or_default();
 		let response = response_builder.body(ResponseBody::Full(Full::new(Bytes::from(body))))?;
@@ -267,31 +271,27 @@ impl CustomServeTrait for PegboardGateway {
 		Ok(response)
 	}
 
-	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id, runner_id=?self.runner_id, request_id=%protocol::util::id_to_string(&request_id)))]
-	async fn handle_websocket(
+	async fn handle_websocket_inner(
 		&self,
+		ctx: &StandaloneCtx,
+		req_ctx: &mut RequestContext,
 		client_ws: WebSocketHandle,
-		headers: &hyper::HeaderMap,
-		_path: &str,
-		_request_context: &mut RequestContext,
-		request_id: protocol::RequestId,
 		after_hibernation: bool,
 	) -> Result<Option<CloseFrame>> {
-		// Use the actor ID from the gateway instance
-		let actor_id = self.actor_id.to_string();
+		let request_id = req_ctx.in_flight_request_id()?;
 
 		// Extract headers
 		let mut request_headers = HashableMap::new();
-		for (name, value) in headers {
+		for (name, value) in req_ctx.headers() {
 			if let Result::Ok(value_str) = value.to_str() {
 				request_headers.insert(name.to_string(), value_str.to_string());
 			}
 		}
 
-		let mut stopped_sub = self
-			.ctx
-			.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", self.actor_id))
-			.await?;
+		let (mut stopped_sub, runner_protocol_version) = tokio::try_join!(
+			ctx.subscribe::<pegboard::workflows::actor::Stopped>(("actor_id", self.actor_id)),
+			get_runner_protocol_version(&ctx, self.runner_id),
+		)?;
 
 		// Build subject to publish to
 		let tunnel_subject =
@@ -304,7 +304,7 @@ impl CustomServeTrait for PegboardGateway {
 			new,
 		} = self
 			.shared_state
-			.start_in_flight_request(tunnel_subject.clone(), request_id)
+			.start_in_flight_request(tunnel_subject.clone(), runner_protocol_version, request_id)
 			.await;
 
 		ensure!(
@@ -317,9 +317,9 @@ impl CustomServeTrait for PegboardGateway {
 			true
 		} else {
 			// Send WebSocket open message
-			let open_message = protocol::ToClientTunnelMessageKind::ToClientWebSocketOpen(
-				protocol::ToClientWebSocketOpen {
-					actor_id: actor_id.clone(),
+			let open_message = protocol::mk2::ToClientTunnelMessageKind::ToClientWebSocketOpen(
+				protocol::mk2::ToClientWebSocketOpen {
+					actor_id: self.actor_id.to_string(),
 					path: self.path.clone(),
 					headers: request_headers,
 				},
@@ -338,10 +338,10 @@ impl CustomServeTrait for PegboardGateway {
 						res = msg_rx.recv() => {
 							if let Some(msg) = res {
 								match msg {
-									protocol::ToServerTunnelMessageKind::ToServerWebSocketOpen(msg) => {
+									protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketOpen(msg) => {
 										return anyhow::Ok(msg);
 									}
-									protocol::ToServerTunnelMessageKind::ToServerWebSocketClose(close) => {
+									protocol::mk2::ToServerTunnelMessageKind::ToServerWebSocketClose(close) => {
 										tracing::warn!(?close, "websocket closed before opening");
 										return Err(WebSocketServiceUnavailable.build());
 									}
@@ -364,7 +364,7 @@ impl CustomServeTrait for PegboardGateway {
 							return Err(WebSocketServiceUnavailable.build());
 						}
 						_ = drop_rx.changed() => {
-							tracing::warn!("websocket open timeout");
+							tracing::warn!(reason=?drop_rx.borrow(), "websocket open timeout");
 							return Err(WebSocketServiceUnavailable.build());
 						}
 					}
@@ -373,7 +373,13 @@ impl CustomServeTrait for PegboardGateway {
 				Err(WebSocketServiceUnavailable.build())
 			};
 
-			let open_msg = tokio::time::timeout(WEBSOCKET_OPEN_TIMEOUT, fut)
+			let websocket_open_timeout = Duration::from_millis(
+				self.ctx
+					.config()
+					.pegboard()
+					.gateway_websocket_open_timeout_ms(),
+			);
+			let open_msg = tokio::time::timeout(websocket_open_timeout, fut)
 				.await
 				.map_err(|_| {
 					tracing::warn!("timed out waiting for websocket open from runner");
@@ -388,6 +394,9 @@ impl CustomServeTrait for PegboardGateway {
 			open_msg.can_hibernate
 		};
 
+		let ingress_bytes = Arc::new(AtomicU64::new(0));
+		let egress_bytes = Arc::new(AtomicU64::new(0));
+
 		// Send pending messages
 		self.shared_state
 			.resend_pending_websocket_messages(request_id)
@@ -399,6 +408,7 @@ impl CustomServeTrait for PegboardGateway {
 		let (ws_to_tunnel_abort_tx, ws_to_tunnel_abort_rx) = watch::channel(());
 		let (ping_abort_tx, ping_abort_rx) = watch::channel(());
 		let (keepalive_abort_tx, keepalive_abort_rx) = watch::channel(());
+		let (metrics_abort_tx, metrics_abort_rx) = watch::channel(());
 
 		let tunnel_to_ws = tokio::spawn(tunnel_to_ws_task::task(
 			self.shared_state.clone(),
@@ -408,23 +418,40 @@ impl CustomServeTrait for PegboardGateway {
 			msg_rx,
 			drop_rx,
 			can_hibernate,
+			egress_bytes.clone(),
 			tunnel_to_ws_abort_rx,
 		));
 		let ws_to_tunnel = tokio::spawn(ws_to_tunnel_task::task(
 			self.shared_state.clone(),
 			request_id,
 			ws_rx,
+			ingress_bytes.clone(),
 			ws_to_tunnel_abort_rx,
 		));
+		let update_ping_interval = Duration::from_millis(
+			self.ctx
+				.config()
+				.pegboard()
+				.gateway_update_ping_interval_ms(),
+		);
 		let ping = tokio::spawn(ping_task::task(
 			self.shared_state.clone(),
 			request_id,
 			ping_abort_rx,
+			update_ping_interval,
+		));
+		let metrics = tokio::spawn(metrics_task::task(
+			ctx.clone(),
+			self.actor_id,
+			self.runner_id,
+			ingress_bytes,
+			egress_bytes,
+			metrics_abort_rx,
 		));
 		let keepalive = if can_hibernate {
 			Some(tokio::spawn(keepalive_task::task(
 				self.shared_state.clone(),
-				self.ctx.clone(),
+				ctx.clone(),
 				self.actor_id,
 				self.shared_state.gateway_id(),
 				request_id,
@@ -434,12 +461,8 @@ impl CustomServeTrait for PegboardGateway {
 			None
 		};
 
-		let tunnel_to_ws_abort_tx2 = tunnel_to_ws_abort_tx.clone();
-		let ws_to_tunnel_abort_tx2 = ws_to_tunnel_abort_tx.clone();
-		let ping_abort_tx2 = ping_abort_tx.clone();
-
 		// Wait for all tasks to complete
-		let (tunnel_to_ws_res, ws_to_tunnel_res, ping_res, keepalive_res) = tokio::join!(
+		let (tunnel_to_ws_res, ws_to_tunnel_res, ping_res, keepalive_res, metrics_res) = tokio::join!(
 			async {
 				let res = tunnel_to_ws.await?;
 
@@ -450,6 +473,7 @@ impl CustomServeTrait for PegboardGateway {
 					let _ = ping_abort_tx.send(());
 					let _ = ws_to_tunnel_abort_tx.send(());
 					let _ = keepalive_abort_tx.send(());
+					let _ = metrics_abort_tx.send(());
 				} else {
 					tracing::debug!(?res, "tunnel to ws task completed");
 				}
@@ -463,9 +487,10 @@ impl CustomServeTrait for PegboardGateway {
 				if !matches!(res, Ok(LifecycleResult::Aborted)) {
 					tracing::debug!(?res, "ws to tunnel task completed, aborting counterpart");
 
-					let _ = ping_abort_tx2.send(());
+					let _ = ping_abort_tx.send(());
 					let _ = tunnel_to_ws_abort_tx.send(());
 					let _ = keepalive_abort_tx.send(());
+					let _ = metrics_abort_tx.send(());
 				} else {
 					tracing::debug!(?res, "ws to tunnel task completed");
 				}
@@ -479,9 +504,10 @@ impl CustomServeTrait for PegboardGateway {
 				if !matches!(res, Ok(LifecycleResult::Aborted)) {
 					tracing::debug!(?res, "ping task completed, aborting others");
 
-					let _ = ws_to_tunnel_abort_tx2.send(());
-					let _ = tunnel_to_ws_abort_tx2.send(());
+					let _ = ws_to_tunnel_abort_tx.send(());
+					let _ = tunnel_to_ws_abort_tx.send(());
 					let _ = keepalive_abort_tx.send(());
+					let _ = metrics_abort_tx.send(());
 				} else {
 					tracing::debug!(?res, "ping task completed");
 				}
@@ -499,30 +525,54 @@ impl CustomServeTrait for PegboardGateway {
 				if !matches!(res, Ok(LifecycleResult::Aborted)) {
 					tracing::debug!(?res, "keepalive task completed, aborting others");
 
-					let _ = ws_to_tunnel_abort_tx2.send(());
-					let _ = tunnel_to_ws_abort_tx2.send(());
-					let _ = ping_abort_tx2.send(());
+					let _ = ws_to_tunnel_abort_tx.send(());
+					let _ = tunnel_to_ws_abort_tx.send(());
+					let _ = ping_abort_tx.send(());
+					let _ = metrics_abort_tx.send(());
 				} else {
 					tracing::debug!(?res, "keepalive task completed");
 				}
 
+				res
+			},
+			async {
+				let res = metrics.await?;
+
+				// Abort others if not aborted
+				if !matches!(res, Ok(LifecycleResult::Aborted)) {
+					tracing::debug!(?res, "metrics task completed, aborting others");
+
+					let _ = ws_to_tunnel_abort_tx.send(());
+					let _ = tunnel_to_ws_abort_tx.send(());
+					let _ = ping_abort_tx.send(());
+					let _ = keepalive_abort_tx.send(());
+				} else {
+					tracing::debug!(?res, "metrics task completed");
+				}
+
 				res
 			},
 		);
 
 		// Determine single result from all tasks
-		let mut lifecycle_res = match (tunnel_to_ws_res, ws_to_tunnel_res, ping_res, keepalive_res)
-		{
+		let mut lifecycle_res = match (
+			tunnel_to_ws_res,
+			ws_to_tunnel_res,
+			ping_res,
+			keepalive_res,
+			metrics_res,
+		) {
 			// Prefer error
-			(Err(err), _, _, _) => Err(err),
-			(_, Err(err), _, _) => Err(err),
-			(_, _, Err(err), _) => Err(err),
-			(_, _, _, Err(err)) => Err(err),
+			(Err(err), _, _, _, _) => Err(err),
+			(_, Err(err), _, _, _) => Err(err),
+			(_, _, Err(err), _, _) => Err(err),
+			(_, _, _, Err(err), _) => Err(err),
+			(_, _, _, _, Err(err)) => Err(err),
 			// Prefer non aborted result if all succeed
-			(Ok(res), Ok(LifecycleResult::Aborted), _, _) => Ok(res),
-			(Ok(LifecycleResult::Aborted), Ok(res), _, _) => Ok(res),
+			(Ok(res), Ok(LifecycleResult::Aborted), _, _, _) => Ok(res),
+			(Ok(LifecycleResult::Aborted), Ok(res), _, _, _) => Ok(res),
 			// Unlikely case
-			(res, _, _, _) => res,
+			(res, _, _, _, _) => res,
 		};
 
 		// Send close frame to runner if not hibernating
@@ -538,8 +588,8 @@ impl CustomServeTrait for PegboardGateway {
 				Ok(_) => (CloseCode::Normal.into(), None),
 				Err(_) => (CloseCode::Error.into(), Some("ws.downstream_closed".into())),
 			};
-			let close_message = protocol::ToClientTunnelMessageKind::ToClientWebSocketClose(
-				protocol::ToClientWebSocketClose {
+			let close_message = protocol::mk2::ToClientTunnelMessageKind::ToClientWebSocketClose(
+				protocol::mk2::ToClientWebSocketClose {
 					code: Some(close_code.into()),
 					reason: close_reason.map(|x| x.as_str().to_string()),
 				},
@@ -570,13 +620,100 @@ impl CustomServeTrait for PegboardGateway {
 			Err(err) => Err(err),
 		}
 	}
+}
 
-	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id, request_id=%protocol::util::id_to_string(&request_id)))]
+#[async_trait]
+impl CustomServeTrait for PegboardGateway {
+	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id, runner_id=?self.runner_id))]
+	async fn handle_request(
+		&self,
+		req: Request<Full<Bytes>>,
+		req_ctx: &mut RequestContext,
+	) -> Result<Response<ResponseBody>> {
+		let ctx = self.ctx.with_ray(req_ctx.ray_id(), req_ctx.req_id())?;
+		let req_body_size_hint = req.body().size_hint();
+
+		let (res, metrics_res) = tokio::join!(
+			self.handle_request_inner(&ctx, req, req_ctx),
+			record_req_metrics(
+				&ctx,
+				self.runner_id,
+				self.actor_id,
+				Metric::HttpIngress(
+					req_body_size_hint
+						.upper()
+						.unwrap_or(req_body_size_hint.lower()) as usize
+				),
+			),
+		);
+
+		let response_size = match &res {
+			Ok(res) => res.size_hint().upper().unwrap_or(res.size_hint().lower()),
+			Err(_) => 0,
+		};
+
+		if let Err(err) = metrics_res {
+			tracing::error!(?err, "http req ingress metrics failed");
+		} else {
+			if let Err(err) = record_req_metrics(
+				&ctx,
+				self.runner_id,
+				self.actor_id,
+				Metric::HttpEgress(response_size as usize),
+			)
+			.await
+			{
+				tracing::error!(
+					?err,
+					runner_id=?self.runner_id,
+					"http req egress metrics failed, likely corrupt now",
+				);
+			}
+		}
+
+		res
+	}
+
+	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id, runner_id=?self.runner_id))]
+	async fn handle_websocket(
+		&self,
+		req_ctx: &mut RequestContext,
+		client_ws: WebSocketHandle,
+		after_hibernation: bool,
+	) -> Result<Option<CloseFrame>> {
+		let ctx = self.ctx.with_ray(req_ctx.ray_id(), req_ctx.req_id())?;
+		let (res, metrics_res) = tokio::join!(
+			self.handle_websocket_inner(&ctx, req_ctx, client_ws, after_hibernation),
+			record_req_metrics(&ctx, self.runner_id, self.actor_id, Metric::WebsocketOpen),
+		);
+
+		if let Err(err) = metrics_res {
+			tracing::error!(?err, "ws open metrics failed");
+		} else {
+			if let Err(err) =
+				record_req_metrics(&ctx, self.runner_id, self.actor_id, Metric::WebsocketClose)
+					.await
+			{
+				tracing::error!(
+					?err,
+					runner_id=?self.runner_id,
+					"ws close metrics failed, likely corrupt now",
+				);
+			}
+		}
+
+		res
+	}
+
+	#[tracing::instrument(skip_all, fields(actor_id=?self.actor_id))]
 	async fn handle_websocket_hibernation(
 		&self,
+		req_ctx: &mut RequestContext,
 		client_ws: WebSocketHandle,
-		request_id: protocol::RequestId,
 	) -> Result<HibernationResult> {
+		let ctx = self.ctx.with_ray(req_ctx.ray_id(), req_ctx.req_id())?;
+		let request_id = req_ctx.in_flight_request_id()?;
+
 		// Immediately rewake if we have pending messages
 		if self
 			.shared_state
@@ -592,31 +729,66 @@ impl CustomServeTrait for PegboardGateway {
 		let (keepalive_abort_tx, keepalive_abort_rx) = watch::channel(());
 		let keepalive_handle = tokio::spawn(keepalive_task::task(
 			self.shared_state.clone(),
-			self.ctx.clone(),
+			ctx.clone(),
 			self.actor_id,
 			self.shared_state.gateway_id(),
 			request_id,
 			keepalive_abort_rx,
 		));
 
-		let res = self.handle_websocket_hibernation_inner(client_ws).await;
+		let (res, metrics_res) = tokio::join!(
+			self.handle_websocket_hibernation_inner(client_ws),
+			record_req_metrics(
+				&ctx,
+				self.runner_id,
+				self.actor_id,
+				Metric::WebsocketHibernate
+			)
+		);
 
 		let _ = keepalive_abort_tx.send(());
 		let _ = keepalive_handle.await;
 
-		match &res {
-			Ok(HibernationResult::Continue) => {}
-			Ok(HibernationResult::Close) | Err(_) => {
-				// No longer an active hibernating request, delete entry
-				self.ctx
-					.op(pegboard::ops::actor::hibernating_request::delete::Input {
-						actor_id: self.actor_id,
-						gateway_id: self.shared_state.gateway_id(),
-						request_id,
-					})
-					.await?;
-			}
-		}
+		let (delete_res, _) = tokio::join!(
+			async {
+				match &res {
+					Ok(HibernationResult::Continue) => {}
+					Ok(HibernationResult::Close) | Err(_) => {
+						// No longer an active hibernating request, delete entry
+						ctx.op(pegboard::ops::actor::hibernating_request::delete::Input {
+							actor_id: self.actor_id,
+							gateway_id: self.shared_state.gateway_id(),
+							request_id,
+						})
+						.await?;
+					}
+				}
+
+				anyhow::Ok(())
+			},
+			async {
+				if let Err(err) = metrics_res {
+					tracing::error!(?err, "ws hibernate metrics failed");
+				} else {
+					if let Err(err) = record_req_metrics(
+						&ctx,
+						self.runner_id,
+						self.actor_id,
+						Metric::WebsocketStopHibernate,
+					)
+					.await
+					{
+						tracing::error!(
+							?err,
+							runner_id=?self.runner_id,
+							"ws stop hibernate metrics failed, likely corrupt now"
+						);
+					}
+				}
+			},
+		);
+
+		delete_res?;
 
 		res
 	}
@@ -696,3 +868,209 @@ async fn hibernate_ws(ws_rx: Arc<Mutex<WebSocketReceiver>>) -> Result<Hibernatio
 		}
 	}
 }
+
+async fn get_runner_protocol_version(ctx: &StandaloneCtx, runner_id: Id) -> Result<u16> {
+	ctx.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(pegboard::keys::subspace());
+
+			let protocol_version_entry = tx
+				.read_opt(
+					&pegboard::keys::runner::ProtocolVersionKey::new(runner_id),
+					Serializable,
+				)
+				.await?;
+
+			Ok(protocol_version_entry.unwrap_or(PROTOCOL_MK1_VERSION))
+		})
+		.await
+}
+
+enum Metric {
+	HttpIngress(usize),
+	HttpEgress(usize),
+	WebsocketOpen,
+	// Ingress, Egress
+	WebsocketTransfer(usize, usize),
+	WebsocketClose,
+	WebsocketHibernate,
+	WebsocketStopHibernate,
+}
+
+async fn record_req_metrics(
+	ctx: &StandaloneCtx,
+	runner_id: Id,
+	actor_id: Id,
+	metric: Metric,
+) -> Result<u16> {
+	let metric = &metric;
+	// Read runner protocol version
+	let (protocol_version, has_name, actor_workflow_id) = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(pegboard::keys::subspace());
+
+			let protocol_version_key = pegboard::keys::runner::ProtocolVersionKey::new(runner_id);
+			let namespace_id_key = pegboard::keys::runner::NamespaceIdKey::new(runner_id);
+			let actor_name_key = pegboard::keys::actor::NameKey::new(actor_id);
+			let actor_workflow_id_key = pegboard::keys::actor::WorkflowIdKey::new(actor_id);
+			let (protocol_version_entry, namespace_id, actor_name_entry, actor_workflow_id) = tokio::try_join!(
+				tx.read_opt(&protocol_version_key, Serializable),
+				tx.read(&namespace_id_key, Serializable),
+				tx.read_opt(&actor_name_key, Serializable),
+				tx.read(&actor_workflow_id_key, Serializable),
+			)?;
+			let has_name = actor_name_entry.is_some();
+
+			if let Some(name) = &actor_name_entry {
+				metric_inc(&tx, namespace_id, name, metric);
+			}
+
+			Ok((
+				protocol_version_entry.unwrap_or(PROTOCOL_MK1_VERSION),
+				has_name,
+				actor_workflow_id,
+			))
+		})
+		.await?;
+
+	// NOTE: The name key was added via backfill. If the actor has not backfilled the key yet (key is none),
+	// we need to fetch it from the actor state
+	if !has_name {
+		let wf = ctx
+			.workflow::<pegboard::workflows::actor::Input>(actor_workflow_id)
+			.get()
+			.await?
+			.context("actor not found")?;
+		let actor_state = &wf
+			.parse_state::<Option<pegboard::workflows::actor::State>>()?
+			.context("actor did not initialize state yet")?;
+
+		// Record metrics
+		ctx.udb()?
+			.run(|tx| async move {
+				let tx = tx.with_subspace(pegboard::keys::subspace());
+				metric_inc(&tx, actor_state.namespace_id, &actor_state.name, metric);
+
+				Ok(())
+			})
+			.await?;
+	}
+
+	Ok(protocol_version)
+}
+
+fn metric_inc(tx: &universaldb::Transaction, namespace_id: Id, name: &str, metric: &Metric) {
+	match metric {
+		Metric::HttpIngress(size) => {
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::GatewayIngress(
+					name.to_string(),
+					"http".to_string(),
+				),
+				(*size).try_into().unwrap_or_default(),
+			);
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::Requests(name.to_string(), "http".to_string()),
+				1,
+			);
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::ActiveRequests(
+					name.to_string(),
+					"http".to_string(),
+				),
+				1,
+			);
+		}
+		Metric::HttpEgress(size) => {
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::GatewayEgress(
+					name.to_string(),
+					"http".to_string(),
+				),
+				(*size).try_into().unwrap_or_default(),
+			);
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::ActiveRequests(
+					name.to_string(),
+					"http".to_string(),
+				),
+				-1,
+			);
+		}
+		Metric::WebsocketOpen => {
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::Requests(name.to_string(), "ws".to_string()),
+				1,
+			);
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::ActiveRequests(name.to_string(), "ws".to_string()),
+				1,
+			);
+		}
+		Metric::WebsocketTransfer(ingress_size, egress_size) => {
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::GatewayIngress(name.to_string(), "ws".to_string()),
+				(*ingress_size).try_into().unwrap_or_default(),
+			);
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::GatewayEgress(name.to_string(), "ws".to_string()),
+				(*egress_size).try_into().unwrap_or_default(),
+			);
+		}
+		Metric::WebsocketClose => {
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::ActiveRequests(name.to_string(), "ws".to_string()),
+				-1,
+			);
+		}
+		Metric::WebsocketHibernate => {
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::Requests(name.to_string(), "hws".to_string()),
+				1,
+			);
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::ActiveRequests(
+					name.to_string(),
+					"hws".to_string(),
+				),
+				1,
+			);
+		}
+		Metric::WebsocketStopHibernate => {
+			namespace::keys::metric::inc(
+				&tx.with_subspace(namespace::keys::subspace()),
+				namespace_id,
+				namespace::keys::metric::Metric::ActiveRequests(
+					name.to_string(),
+					"hws".to_string(),
+				),
+				-1,
+			);
+		}
+	}
+}