From fc2fa3958f1166b347c9361688959e27ce52d62c Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Thu, 30 Jul 2026 15:56:17 -0700 Subject: [PATCH] fix(http): make request handoff ordered and non-resumable --- .../pegboard-envoy/src/tunnel_to_ws_task.rs | 62 ++++++++++++++----- .../tests/support/tunnel_to_ws_delivery.rs | 50 +++++++++++++++ .../sdks/rust/envoy-client/src/actor/http.rs | 24 ++++++- engine/sdks/rust/envoy-client/src/tunnel.rs | 30 ++++++--- .../tests/support/actor_http_stream.rs | 36 +++++++++++ 5 files changed, 178 insertions(+), 24 deletions(-) create mode 100644 engine/packages/pegboard-envoy/tests/support/tunnel_to_ws_delivery.rs diff --git a/engine/packages/pegboard-envoy/src/tunnel_to_ws_task.rs b/engine/packages/pegboard-envoy/src/tunnel_to_ws_task.rs index 05a87709de..09ca4ac912 100644 --- a/engine/packages/pegboard-envoy/src/tunnel_to_ws_task.rs +++ b/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,7 +260,7 @@ 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(); @@ -254,6 +268,22 @@ async fn handle_message( Ok(false) } +async fn ordered_handoff( + handoff: impl Future>, + reply: impl Future>, +) -> 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; diff --git a/engine/packages/pegboard-envoy/tests/support/tunnel_to_ws_delivery.rs b/engine/packages/pegboard-envoy/tests/support/tunnel_to_ws_delivery.rs new file mode 100644 index 0000000000..f7a24f8d10 --- /dev/null +++ b/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)); +} diff --git a/engine/sdks/rust/envoy-client/src/actor/http.rs b/engine/sdks/rust/envoy-client/src/actor/http.rs index 6228a3c184..c5d0d1d5fc 100644 --- a/engine/sdks/rust/envoy-client/src/actor/http.rs +++ b/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; } } diff --git a/engine/sdks/rust/envoy-client/src/tunnel.rs b/engine/sdks/rust/envoy-client/src/tunnel.rs index ab6505cc48..844d60657a 100644 --- a/engine/sdks/rust/envoy-client/src/tunnel.rs +++ b/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( diff --git a/engine/sdks/rust/envoy-client/tests/support/actor_http_stream.rs b/engine/sdks/rust/envoy-client/tests/support/actor_http_stream.rs index d21b319df8..a7d30ceaa4 100644 --- a/engine/sdks/rust/envoy-client/tests/support/actor_http_stream.rs +++ b/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();