fix(http): harden streaming lifecycle and observability

engine/packages/guard-core/src/response_body.rs

@@ -45,7 +45,8 @@ pub enum ResponseBody {
 }
 
 impl ResponseBody {
-	pub(crate) fn with_completion(self, callback: impl FnOnce() + Send + 'static) -> Self {
+	#[doc(hidden)]
+	pub fn with_completion(self, callback: impl FnOnce() + Send + 'static) -> Self {
 		Self::WithCompletion {
 			body: Box::new(self),
 			completion: CompletionGuard(Some(Box::new(callback))),

engine/packages/pegboard-gateway2/src/http_stream/handler.rs

@@ -1,4 +1,11 @@
-use std::{collections::HashMap, time::Duration};
+use std::{
+	collections::HashMap,
+	sync::{
+		Arc,
+		atomic::{AtomicU64, Ordering},
+	},
+	time::Duration,
+};
 
 use anyhow::{Result, anyhow};
 use bytes::Bytes;
@@ -11,7 +18,7 @@ use rivet_guard_core::{
 	errors::{ActorStoppedWhileWaiting, InvalidRequestBody},
 	request_context::RequestContext,
 };
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, watch};
 use tracing::Instrument;
 
 use super::{
@@ -24,7 +31,7 @@ use super::{
 	send_http_request_abort,
 };
 use crate::{
-	PegboardGateway2,
+	PegboardGateway2, metrics_task,
 	request_metrics::{RequestKind, RequestMetrics},
 	shared_state::{InFlightRequestCtx, RequestProtocol, RequestStopResult},
 };
@@ -42,37 +49,72 @@ impl PegboardGateway2 {
 		B::Error: std::error::Error + Send + Sync + 'static,
 	{
 		let ctx = self.ctx.with_ray(req_ctx.ray_id(), req_ctx.req_id())?;
-		let req_body_size_hint = req.body().size_hint();
-		let ingress_size_hint = req_body_size_hint
-			.upper()
-			.unwrap_or(req_body_size_hint.lower()) as usize;
+		let ingress_bytes = Arc::new(AtomicU64::new(0));
+		let egress_bytes = Arc::new(AtomicU64::new(0));
 		let request_metrics = RequestMetrics::new(
 			ctx.clone(),
 			self.actor_id,
 			self.namespace_id,
 			self.envoy_key.clone(),
 			RequestKind::Http,
 		);
+		let (metrics_abort_tx, metrics_abort_rx) = watch::channel(());
+		let transfer_metrics = request_metrics.clone();
+		let transfer_ingress_bytes = ingress_bytes.clone();
+		let transfer_egress_bytes = egress_bytes.clone();
+		tokio::spawn(
+			async move {
+				if let Err(error) = metrics_task::task(
+					transfer_metrics,
+					transfer_ingress_bytes,
+					transfer_egress_bytes,
+					metrics_abort_rx,
+				)
+				.await
+				{
+					tracing::error!(?error, "HTTP transfer metrics task failed");
+				}
+			}
+			.in_current_span(),
+		);
 
 		let (res, active_metrics) = tokio::join!(
-			self.handle_request_inner(&ctx, req, req_ctx),
-			request_metrics.begin(ingress_size_hint),
+			self.handle_request_inner(
+				&ctx,
+				req,
+				req_ctx,
+				ingress_bytes.clone(),
+				egress_bytes.clone(),
+			),
+			request_metrics.begin(0),
 		);
 
-		let egress_size_hint = res
-			.as_ref()
-			.map(|res| res.size_hint().upper().unwrap_or(res.size_hint().lower()) as usize)
-			.unwrap_or_default();
-		active_metrics.finish_in_background(egress_size_hint);
-
-		res
+		match res {
+			Ok(response) => {
+				let (parts, body) = response.into_parts();
+				Ok(Response::from_parts(
+					parts,
+					body.with_completion(move || {
+						let _ = metrics_abort_tx.send(());
+						active_metrics.finish_in_background(0);
+					}),
+				))
+			}
+			Err(error) => {
+				let _ = metrics_abort_tx.send(());
+				active_metrics.finish_in_background(0);
+				Err(error)
+			}
+		}
 	}
 
 	async fn handle_request_inner<B>(
 		&self,
 		ctx: &StandaloneCtx,
 		req: Request<B>,
 		req_ctx: &mut RequestContext,
+		ingress_bytes: Arc<AtomicU64>,
+		egress_bytes: Arc<AtomicU64>,
 	) -> Result<Response<ResponseBody>>
 	where
 		B: Body<Data = Bytes> + Unpin,
@@ -119,6 +161,7 @@ impl PegboardGateway2 {
 				None,
 			)
 		};
+		ingress_bytes.fetch_add(body_bytes.len() as u64, Ordering::AcqRel);
 
 		let mut stopped_sub = ctx
 			.subscribe::<pegboard::workflows::actor2::Stopped>(("actor_id", self.actor_id))
@@ -222,6 +265,7 @@ impl PegboardGateway2 {
 					request_id,
 					body,
 					max_request_body_size,
+					ingress_bytes,
 					response_start_deadline,
 					response_start_timeout,
 				)
@@ -273,13 +317,15 @@ impl PegboardGateway2 {
 						expected_message_index,
 						self.actor_id,
 						idle_timeout,
+						egress_bytes,
 					)
 					.in_current_span(),
 				);
 
 				response_builder.body(ResponseBody::Channel(body_rx))?
 			} else {
 				let body = response_start.body.unwrap_or_default();
+				egress_bytes.fetch_add(body.len() as u64, Ordering::AcqRel);
 				let response =
 					response_builder.body(ResponseBody::Full(Full::new(Bytes::from(body))))?;
 

engine/packages/pegboard-gateway2/src/http_stream/request.rs

@@ -8,7 +8,13 @@ use rivet_guard_core::errors::{
 	ActorStoppedWhileWaiting, GatewayResponseStartTimeout, InvalidRequestBody,
 	TunnelMessageTimeout, TunnelRequestAborted, TunnelResponseClosed,
 };
-use std::time::Duration;
+use std::{
+	sync::{
+		Arc,
+		atomic::{AtomicU64, Ordering},
+	},
+	time::Duration,
+};
 use tokio::sync::{mpsc, watch};
 
 use crate::shared_state::{InFlightRequestHandle, InFlightTunnelMessage, MsgGcReason};
@@ -64,6 +70,7 @@ async fn send_streaming_http_request_body_chunks<B>(
 	in_flight_req: &InFlightRequestHandle,
 	mut body: B,
 	max_body_size: usize,
+	ingress_bytes: Arc<AtomicU64>,
 ) -> Result<()>
 where
 	B: Body<Data = Bytes> + Unpin,
@@ -85,7 +92,12 @@ where
 				frame = body.frame() => frame,
 				_ = tokio::time::sleep_until(deadline) => {
 					if let Some(chunk) = chunker.flush() {
-						send_http_request_body_chunk(in_flight_req, chunk, false).await?;
+						send_http_request_body_chunk(
+							in_flight_req,
+							chunk,
+							false,
+						)
+						.await?;
 					}
 					flush_deadline = None;
 					continue;
@@ -112,6 +124,7 @@ where
 		let Ok(data) = frame.into_data() else {
 			continue;
 		};
+		ingress_bytes.fetch_add(data.len() as u64, Ordering::AcqRel);
 		let Some(next_body_size) = next_request_body_size(body_size, data.len(), max_body_size)
 		else {
 			super::send_http_request_abort(
@@ -240,6 +253,7 @@ pub(super) async fn stream_http_request_and_wait_for_response<B>(
 	request_id: protocol::RequestId,
 	body: B,
 	max_body_size: usize,
+	ingress_bytes: Arc<AtomicU64>,
 	response_start_deadline: tokio::time::Instant,
 	response_start_timeout: Duration,
 ) -> Result<(protocol::MessageId, protocol::ToRivetResponseStart)>
@@ -250,7 +264,8 @@ where
 	// Upload ingress, cumulative byte accounting, and bounded-latency
 	// coalescing run in one future. Response start/abort is observed
 	// concurrently so a slow or rejected upload cannot retain the request.
-	let upload = send_streaming_http_request_body_chunks(in_flight_req, body, max_body_size);
+	let upload =
+		send_streaming_http_request_body_chunks(in_flight_req, body, max_body_size, ingress_bytes);
 	tokio::pin!(upload);
 	tokio::select! {
 		upload_result = &mut upload => {

engine/packages/pegboard-gateway2/src/http_stream/response.rs

@@ -1,6 +1,10 @@
 use bytes::Bytes;
 use gas::prelude::*;
 use rivet_envoy_protocol as protocol;
+use std::sync::{
+	Arc,
+	atomic::{AtomicU64, Ordering},
+};
 use std::time::Duration;
 use tokio::sync::{mpsc, watch};
 
@@ -31,10 +35,12 @@ async fn send_http_response_body_bytes(
 	actor_id: Id,
 	body: Vec<u8>,
 	detail: &'static str,
+	egress_bytes: &AtomicU64,
 ) -> bool {
 	let body = Bytes::from(body);
 	for offset in (0..body.len()).step_by(HTTP_BODY_CHUNK_SIZE) {
 		let chunk = body.slice(offset..(offset + HTTP_BODY_CHUNK_SIZE).min(body.len()));
+		let chunk_len = chunk.len();
 		let delivery = tokio::select! {
 			biased;
 			_ = body_tx.closed() => Err((RequestStopResult::ClientDisconnect, detail.to_owned())),
@@ -84,6 +90,7 @@ async fn send_http_response_body_bytes(
 			in_flight_req.stop(stop_result).await;
 			return false;
 		}
+		egress_bytes.fetch_add(chunk_len as u64, Ordering::AcqRel);
 	}
 
 	true
@@ -99,6 +106,7 @@ pub(super) async fn drain_http_response_stream(
 	mut expected_message_index: protocol::MessageIndex,
 	actor_id: Id,
 	idle_timeout: Option<Duration>,
+	egress_bytes: Arc<AtomicU64>,
 ) {
 	if let Some(body) = initial_body.filter(|body| !body.is_empty()) {
 		if !send_http_response_body_bytes(
@@ -109,6 +117,7 @@ pub(super) async fn drain_http_response_stream(
 			actor_id,
 			body,
 			"client dropped response before initial body was sent",
+			&egress_bytes,
 		)
 		.await
 		{
@@ -159,6 +168,7 @@ pub(super) async fn drain_http_response_stream(
 							actor_id,
 							chunk.body,
 							"client dropped streaming response body",
+							&egress_bytes,
 						).await {
 							return;
 						}

engine/packages/pegboard-gateway2/src/metrics_task.rs

@@ -21,7 +21,7 @@ pub async fn task(
 			_ = tokio::time::sleep(UPDATE_METRICS_INTERVAL) => {}
 			_ = metrics_abort_rx.changed() => {
 				// Record final values before abort
-				record_ws_transfer(
+				record_transfer(
 					&metrics,
 					&ingress_bytes,
 					&egress_bytes,
@@ -33,7 +33,7 @@ pub async fn task(
 			}
 		}
 
-		record_ws_transfer(
+		record_transfer(
 			&metrics,
 			&ingress_bytes,
 			&egress_bytes,
@@ -44,7 +44,7 @@ pub async fn task(
 	}
 }
 
-async fn record_ws_transfer(
+async fn record_transfer(
 	metrics: &RequestMetrics,
 	ingress_bytes: &AtomicU64,
 	egress_bytes: &AtomicU64,

engine/sdks/rust/envoy-client/src/actor/http.rs

@@ -151,7 +151,7 @@
 		.http_requests
 		.pending
 		.get_mut(&[&message_id.gateway_id, &message_id.request_id])
 	{
 		pending.task_abort_handle = Some(task_abort_handle);
 	}
 
@@ -215,16 +215,7 @@
 								),
 							};
 							tracing::warn!("streamed request body channel overloaded");
-							if let Some(body_abort_tx) = pending.body_abort_tx.take() {
-								body_abort_tx.send_replace(Some(HttpRequestBodyError {
-									reason: reason.clone(),
-								}));
-							}
-							if let Some(task_abort_handle) = pending.task_abort_handle.take() {
-								task_abort_handle.abort();
-							}
-							pending.body_tx = None;
-							pending.body_rejected = true;
+							reject_pending_request(pending, &reason);
 							response_aborted = true;
 
 							let shared = ctx.shared.clone();
@@ -258,6 +249,22 @@
 	}
 }
 
+fn reject_pending_request(
+	pending: &mut PendingHttpRequest,
+	reason: &protocol::HttpStreamAbortReason,
+) {
+	if let Some(body_abort_tx) = pending.body_abort_tx.take() {
+		body_abort_tx.send_replace(Some(HttpRequestBodyError {
+			reason: reason.clone(),
+		}));
+	}
+	if let Some(task_abort_handle) = pending.task_abort_handle.take() {
+		task_abort_handle.abort();
+	}
+	pending.body_tx = None;
+	pending.body_rejected = true;
+}
+
 pub(super) fn handle_req_complete(ctx: &mut ActorContext, message_id: protocol::MessageId) {
 	complete_phase(ctx, message_id, RequestPhase::Response);
 }