fix(http): make request handoff ordered and non-resumable

engine/packages/pegboard-envoy/src/tunnel_to_ws_task.rs

@@ -3,7 +3,7 @@ use gas::prelude::*;
 use hyper_tungstenite::tungstenite::Message;
 use pegboard::pubsub_subjects::GatewayReceiverSubject;
 use rivet_envoy_protocol::{self as protocol, PROTOCOL_VERSION, versioned};
-use std::{sync::Arc, time::Instant};
+use std::{future::Future, sync::Arc, time::Instant};
 use tokio::sync::watch;
 use universalpubsub as ups;
 use universalpubsub::{NextOutput, PublishOpts, Subscriber};
@@ -101,7 +101,7 @@ async fn handle_message(
 	);
 
 	// Parse message
-	let start = Instant::now();
+	let ack_start = Instant::now();
 	let msg = match versioned::ToEnvoyConn::deserialize_with_embedded_version(&tunnel_msg.payload) {
 		Result::Ok(x) => x,
 		Err(err) => {
@@ -110,14 +110,17 @@ async fn handle_message(
 		}
 	};
 
-	// Need to reply to tunnel request so it can continue
-	tunnel_msg.reply(&[]).await?;
-
-	metrics::ACK_MSG_DURATION
-		.with_label_values(&[conn.namespace_id.to_string().as_str(), &conn.pool_name])
-		.observe(start.elapsed().as_secs_f64());
+	// Tunnel messages are acknowledged only after they have been written to the
+	// actor WebSocket. This keeps the gateway's sequential send loop behind the
+	// final transport handoff and turns a disconnect into a request failure
+	// instead of acknowledging data that the actor never received.
+	let reply_after_websocket_handoff =
+		matches!(&msg, protocol::ToEnvoyConn::ToEnvoyTunnelMessage(_));
+	if !reply_after_websocket_handoff {
+		reply_to_gateway(conn, &tunnel_msg, ack_start).await?;
+	}
 
-	let start = Instant::now();
+	let process_start = Instant::now();
 
 	// Convert to ToEnvoy types
 	let mut tunnel_message_meta = None;
@@ -226,10 +229,21 @@ async fn handle_message(
 		);
 	}
 	let _in_flight = ws_to_tunnel_task::WsResponseInFlightGuard::new();
-	conn.ws_handle
-		.send(ws_msg)
-		.await
-		.context("failed to send message to WebSocket")?;
+	let websocket_handoff = async {
+		conn.ws_handle
+			.send(ws_msg)
+			.await
+			.context("failed to send message to WebSocket")
+	};
+	if reply_after_websocket_handoff {
+		ordered_handoff(
+			websocket_handoff,
+			reply_to_gateway(conn, &tunnel_msg, ack_start),
+		)
+		.await?;
+	} else {
+		websocket_handoff.await?;
+	}
 	drop(_in_flight);
 	if let Some((gateway_id, request_id, message_index, message_kind, inner_data_len)) =
 		&tunnel_message_meta
@@ -246,14 +260,30 @@ async fn handle_message(
 
 	metrics::PROCESS_MSG_DURATION
 		.with_label_values(&[conn.namespace_id.to_string().as_str(), &conn.pool_name])
-		.observe(start.elapsed().as_secs_f64());
+		.observe(process_start.elapsed().as_secs_f64());
 	metrics::MSG_PROCESSED_TOTAL
 		.with_label_values(&[conn.namespace_id.to_string().as_str(), &conn.pool_name])
 		.inc();
 
 	Ok(false)
 }
 
+async fn ordered_handoff(
+	handoff: impl Future<Output = Result<()>>,
+	reply: impl Future<Output = Result<()>>,
+) -> Result<()> {
+	handoff.await?;
+	reply.await
+}
+
+async fn reply_to_gateway(conn: &Conn, tunnel_msg: &ups::Message, start: Instant) -> Result<()> {
+	tunnel_msg.reply(&[]).await?;
+	metrics::ACK_MSG_DURATION
+		.with_label_values(&[conn.namespace_id.to_string().as_str(), &conn.pool_name])
+		.observe(start.elapsed().as_secs_f64());
+	Ok(())
+}
+
 fn to_envoy_tunnel_message_kind_name(kind: &protocol::ToEnvoyTunnelMessageKind) -> &'static str {
 	match kind {
 		protocol::ToEnvoyTunnelMessageKind::ToEnvoyRequestStart(_) => "ToEnvoyRequestStart",
@@ -283,3 +313,7 @@ fn to_envoy_tunnel_message_inner_data_len(kind: &protocol::ToEnvoyTunnelMessageK
 #[cfg(test)]
 #[path = "../tests/support/tunnel_to_ws_payload_accounting.rs"]
 mod payload_accounting_tests;
+
+#[cfg(test)]
+#[path = "../tests/support/tunnel_to_ws_delivery.rs"]
+mod delivery_tests;

engine/packages/pegboard-envoy/tests/support/tunnel_to_ws_delivery.rs

@@ -0,0 +1,50 @@
+use std::sync::{
+	Arc,
+	atomic::{AtomicBool, Ordering},
+};
+
+use anyhow::{Result, anyhow};
+use tokio::sync::oneshot;
+
+use super::ordered_handoff;
+
+#[tokio::test]
+async fn gateway_reply_waits_for_websocket_handoff() {
+	let replied = Arc::new(AtomicBool::new(false));
+	let replied_for_task = replied.clone();
+	let (handoff_tx, handoff_rx) = oneshot::channel();
+
+	let task = tokio::spawn(ordered_handoff(
+		async move {
+			handoff_rx.await.expect("handoff sender dropped");
+			Ok(())
+		},
+		async move {
+			replied_for_task.store(true, Ordering::Release);
+			Ok(())
+		},
+	));
+
+	tokio::task::yield_now().await;
+	assert!(!replied.load(Ordering::Acquire));
+
+	handoff_tx.send(()).expect("handoff receiver dropped");
+	task.await.expect("handoff task panicked").unwrap();
+	assert!(replied.load(Ordering::Acquire));
+}
+
+#[tokio::test]
+async fn failed_websocket_handoff_is_not_acknowledged() {
+	let replied = Arc::new(AtomicBool::new(false));
+	let replied_for_task = replied.clone();
+
+	let result: Result<()> =
+		ordered_handoff(async { Err(anyhow!("websocket closed")) }, async move {
+			replied_for_task.store(true, Ordering::Release);
+			Ok(())
+		})
+		.await;
+
+	assert!(result.is_err());
+	assert!(!replied.load(Ordering::Acquire));
+}

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

@@ -154,7 +154,6 @@ pub(super) fn handle_req_start(
 	{
 		pending.task_abort_handle = Some(task_abort_handle);
 	}
-
 }
 
 pub(super) fn handle_task_result(result: Result<(), JoinError>) {
@@ -236,7 +235,28 @@ pub(super) fn handle_req_chunk(
 			}
 		}
 		None => {
-			tracing::warn!("received request chunk without an active request");
+			tracing::warn!(
+				gateway_id = ?message_id.gateway_id,
+				request_id = ?message_id.request_id,
+				message_index = message_id.message_index,
+				"received request chunk without an active request"
+			);
+			let shared = ctx.shared.clone();
+			let gateway_id = message_id.gateway_id;
+			let request_id = message_id.request_id;
+			spawn_detached(async move {
+				send_response_abort(
+					&shared,
+					gateway_id,
+					request_id,
+					0,
+					protocol::HttpStreamAbortReason {
+						kind: protocol::HttpStreamAbortReasonKind::HandlerError,
+						detail: Some("request start was not delivered".to_owned()),
+					},
+				)
+				.await;
+			});
 			return;
 		}
 	}

engine/sdks/rust/envoy-client/src/tunnel.rs

@@ -27,7 +27,7 @@ pub async fn handle_tunnel_message(ctx: &mut EnvoyContext, msg: protocol::ToEnvo
 			handle_request_start(ctx, message_id, req).await;
 		}
 		protocol::ToEnvoyTunnelMessageKind::ToEnvoyRequestChunk(chunk) => {
-			handle_request_chunk(ctx, message_id, chunk);
+			handle_request_chunk(ctx, message_id, chunk).await;
 		}
 		protocol::ToEnvoyTunnelMessageKind::ToEnvoyRequestAbort(abort) => {
 			handle_request_abort(ctx, message_id, abort.reason);
@@ -54,7 +54,14 @@ async fn handle_request_start(
 
 	if !has_actor {
 		tracing::warn!(actor_id = %actor_id, "received request for unknown actor");
-		send_error_response(ctx, message_id.gateway_id, message_id.request_id).await;
+		send_error_response(
+			ctx,
+			message_id.gateway_id,
+			message_id.request_id,
+			"envoy.actor_not_found",
+			"Actor not found",
+		)
+		.await;
 		return;
 	}
 
@@ -71,7 +78,7 @@ async fn handle_request_start(
 	});
 }
 
-fn handle_request_chunk(
+async fn handle_request_chunk(
 	ctx: &mut EnvoyContext,
 	message_id: protocol::MessageId,
 	chunk: protocol::ToEnvoyRequestChunk,
@@ -97,6 +104,14 @@ fn handle_request_chunk(
 			message_index = message_id.message_index,
 			"received request chunk without request start"
 		);
+		send_error_response(
+			ctx,
+			message_id.gateway_id,
+			message_id.request_id,
+			"envoy.request_not_found",
+			"Request start was not delivered",
+		)
+		.await;
 	}
 }
 
@@ -277,13 +292,12 @@ async fn send_error_response(
 	ctx: &EnvoyContext,
 	gateway_id: protocol::GatewayId,
 	request_id: protocol::RequestId,
+	error_code: &str,
+	message: &str,
 ) {
-	let body = b"Actor not found".to_vec();
+	let body = message.as_bytes().to_vec();
 	let mut headers = HashMap::new();
-	headers.insert(
-		"x-rivet-error".to_string(),
-		"envoy.actor_not_found".to_string(),
-	);
+	headers.insert("x-rivet-error".to_string(), error_code.to_owned());
 	headers.insert("content-length".to_string(), body.len().to_string());
 
 	ws_send(

engine/sdks/rust/envoy-client/tests/support/actor_http_stream.rs

@@ -21,6 +21,42 @@ use crate::{
 	http::{HTTP_BODY_MAX_CHUNK_SIZE, ResponseChunk},
 };
 
+#[tokio::test]
+async fn request_chunk_without_start_is_rejected_instead_of_buffered() {
+	let (shared, _envoy_rx) = build_shared_context(Arc::new(TestCallbacks::idle()));
+	let (ws_tx, mut ws_rx) = mpsc::unbounded_channel();
+	*shared.ws_tx.lock().await = Some(ws_tx);
+	let (actor_tx, _) = create_actor(
+		shared,
+		"actor-missing-request-start".to_string(),
+		1,
+		actor_config(),
+		Vec::new(),
+		None,
+	);
+
+	actor_tx
+		.send(ToActor::ReqChunk {
+			message_id: message_id(),
+			chunk: protocol::ToEnvoyRequestChunk {
+				body: vec![1, 2, 3],
+				finish: false,
+			},
+		})
+		.expect("failed to send request chunk");
+
+	let abort = recv_ws_tunnel_msg(&mut ws_rx).await;
+	assert!(matches!(
+		abort.message_kind,
+		protocol::ToRivetTunnelMessageKind::ToRivetResponseAbort(protocol::ToRivetResponseAbort {
+			reason: protocol::HttpStreamAbortReason {
+				kind: protocol::HttpStreamAbortReasonKind::HandlerError,
+				..
+			}
+		})
+	));
+}
+
 #[tokio::test]
 async fn streamed_request_remains_cancellable_after_upload_finishes() {
 	let (fetch_started_tx, fetch_started_rx) = oneshot::channel();