From 4baff97c24df3ca13590241dcc955b968ca90993 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Wed, 29 Jul 2026 16:10:18 -0700 Subject: [PATCH] fix(http): address streaming review findings --- CLAUDE.md | 2 + engine/artifacts/config-schema.json | 31 ++++- engine/packages/config/src/config/pegboard.rs | 23 ++++ .../packages/guard-core/src/response_body.rs | 1 + .../src/http_stream/handler.rs | 12 +- .../pegboard-gateway2/src/http_stream/mod.rs | 1 - .../src/http_stream/request.rs | 114 +++++++++--------- .../src/http_stream/response.rs | 44 ++++--- .../src/http_stream/response_queue.rs | 17 ++- .../pegboard-gateway2/src/shared_state.rs | 10 +- .../tests/support/http_response_queue.rs | 25 ++-- .../tests/support/http_stream_request.rs | 15 ++- 12 files changed, 198 insertions(+), 97 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 1f21c4c751..23f9537969 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -341,6 +341,8 @@ When the user asks to track something in a note, store it in `~/.agents/notes/` - Follow existing patterns in neighboring files. - Always check existing imports and dependencies before adding new ones. - **Always add imports at the top of the file instead of inline within a function.** +- Prefer explicit `if` or `match` statements over calling `.then(...)` on booleans. +- Warning and error logs should briefly identify the likely operational cause or direction when known. ### Comments diff --git a/engine/artifacts/config-schema.json b/engine/artifacts/config-schema.json index 33b6c2e1e5..54f1dee13d 100644 --- a/engine/artifacts/config-schema.json +++ b/engine/artifacts/config-schema.json @@ -1004,6 +1004,24 @@ "format": "uint", "minimum": 0.0 }, + "gateway_http_response_body_channel_capacity": { + "description": "Number of body chunks buffered between a streaming response handler and its HTTP client.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0.0 + }, + "gateway_http_response_queue_max_messages": { + "description": "Maximum number of envoy response messages buffered per HTTP request.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0.0 + }, "gateway_hws_max_pending_size": { "description": "Max pending message buffer size for hibernating WebSockets in bytes.", "type": [ @@ -1023,7 +1041,7 @@ "minimum": 0.0 }, "gateway_response_chunk_idle_timeout_ms": { - "description": "Timeout between streaming HTTP response chunks in milliseconds.", + "description": "Timeout between streaming HTTP response chunks in milliseconds.\n\nDisabled when unset so long-lived streams such as SSE may remain idle.", "type": [ "integer", "null" @@ -1040,6 +1058,15 @@ "format": "uint64", "minimum": 0.0 }, + "gateway_streaming_http_response_queue_max_bytes": { + "description": "Maximum streaming response bytes buffered per HTTP request.", + "type": [ + "integer", + "null" + ], + "format": "uint", + "minimum": 0.0 + }, "gateway_tunnel_ping_timeout_ms": { "description": "Tunnel ping timeout in milliseconds.", "type": [ @@ -1535,4 +1562,4 @@ "additionalProperties": false } } -} +} \ No newline at end of file diff --git a/engine/packages/config/src/config/pegboard.rs b/engine/packages/config/src/config/pegboard.rs index cc85c3bdea..9a1f0f17af 100644 --- a/engine/packages/config/src/config/pegboard.rs +++ b/engine/packages/config/src/config/pegboard.rs @@ -121,6 +121,12 @@ pub struct Pegboard { pub gateway_hws_max_pending_size: Option, /// Max HTTP request body size in bytes for requests to actors. pub gateway_http_max_request_body_size: Option, + /// Number of body chunks buffered between a streaming response handler and its HTTP client. + pub gateway_http_response_body_channel_capacity: Option, + /// Maximum number of envoy response messages buffered per HTTP request. + pub gateway_http_response_queue_max_messages: Option, + /// Maximum streaming response bytes buffered per HTTP request. + pub gateway_streaming_http_response_queue_max_bytes: Option, // === Envoy Settings === /// How long to wait before considering an envoy lost and evicting all of its actors. @@ -309,6 +315,23 @@ impl Pegboard { .unwrap_or(128 * 1024 * 1024) // 128 MiB } + pub fn gateway_http_response_body_channel_capacity(&self) -> usize { + self.gateway_http_response_body_channel_capacity + .unwrap_or(16) + .max(1) + } + + pub fn gateway_http_response_queue_max_messages(&self) -> usize { + // A 20 MiB response occupies 320 64 KiB chunks. Leave room for + // control frames and smaller chunks while still bounding amplification. + self.gateway_http_response_queue_max_messages.unwrap_or(384) + } + + pub fn gateway_streaming_http_response_queue_max_bytes(&self) -> usize { + self.gateway_streaming_http_response_queue_max_bytes + .unwrap_or(20 * 1024 * 1024) // 20 MiB + } + pub fn runner_max_response_payload_body_size(&self) -> usize { self.runner_max_response_payload_body_size .unwrap_or(20 * 1024 * 1024) // 20 MiB diff --git a/engine/packages/guard-core/src/response_body.rs b/engine/packages/guard-core/src/response_body.rs index 9a3e28dbc7..bdf97fa89b 100644 --- a/engine/packages/guard-core/src/response_body.rs +++ b/engine/packages/guard-core/src/response_body.rs @@ -46,6 +46,7 @@ pub enum ResponseBody { impl ResponseBody { #[doc(hidden)] + /// Runs `callback` exactly once when the body reaches EOF, errors, or is dropped. pub fn with_completion(self, callback: impl FnOnce() + Send + 'static) -> Self { Self::WithCompletion { body: Box::new(self), diff --git a/engine/packages/pegboard-gateway2/src/http_stream/handler.rs b/engine/packages/pegboard-gateway2/src/http_stream/handler.rs index fccfc69ea6..4191abfdb4 100644 --- a/engine/packages/pegboard-gateway2/src/http_stream/handler.rs +++ b/engine/packages/pegboard-gateway2/src/http_stream/handler.rs @@ -22,7 +22,7 @@ use tokio::sync::{mpsc, watch}; use tracing::Instrument; use super::{ - HTTP_RESPONSE_BODY_CHANNEL_CAPACITY, ResponseBodyError, + ResponseBodyError, request::{ should_stream_http_request_body_hint, stream_http_request_and_wait_for_response, wait_for_http_response_start, @@ -293,9 +293,13 @@ impl PegboardGateway2 { } let response = if response_start.stream { - let (body_tx, body_rx) = mpsc::channel::>( - HTTP_RESPONSE_BODY_CHANNEL_CAPACITY, - ); + let body_channel_capacity = self + .ctx + .config() + .pegboard() + .gateway_http_response_body_channel_capacity(); + let (body_tx, body_rx) = + mpsc::channel::>(body_channel_capacity); let idle_timeout = self .ctx .config() diff --git a/engine/packages/pegboard-gateway2/src/http_stream/mod.rs b/engine/packages/pegboard-gateway2/src/http_stream/mod.rs index bc1470f718..12894ab883 100644 --- a/engine/packages/pegboard-gateway2/src/http_stream/mod.rs +++ b/engine/packages/pegboard-gateway2/src/http_stream/mod.rs @@ -11,7 +11,6 @@ pub(crate) use response_queue::{ HttpResponseQueueBudget, HttpResponseQueueOverloaded, HttpResponseQueuePermit, }; -pub(super) const HTTP_RESPONSE_BODY_CHANNEL_CAPACITY: usize = 16; pub(super) type ResponseBodyError = Box; pub(super) async fn send_http_request_abort( diff --git a/engine/packages/pegboard-gateway2/src/http_stream/request.rs b/engine/packages/pegboard-gateway2/src/http_stream/request.rs index d1faf2c80b..f2d2dcf3b0 100644 --- a/engine/packages/pegboard-gateway2/src/http_stream/request.rs +++ b/engine/packages/pegboard-gateway2/src/http_stream/request.rs @@ -24,13 +24,23 @@ const HTTP_BODY_CHUNK_SIZE: usize = 64 * 1024; const HTTP_BODY_CHUNK_FLUSH_INTERVAL: Duration = Duration::from_millis(10); pub(super) fn should_stream_http_request_body_hint(size_hint: &SizeHint) -> bool { - size_hint - .upper() - .map_or(true, |body_len| body_len as usize > HTTP_BODY_CHUNK_SIZE) + match size_hint.upper() { + Some(body_len) => body_len as usize > HTTP_BODY_CHUNK_SIZE, + None => true, + } +} + +#[derive(Debug, PartialEq)] +enum RequestBodySize { + WithinLimit(usize), + ExceedsLimit, } -fn next_request_body_size(current: usize, chunk: usize, limit: usize) -> Option { - current.checked_add(chunk).filter(|size| *size <= limit) +fn next_request_body_size(current: usize, chunk: usize, limit: usize) -> RequestBodySize { + match current.checked_add(chunk) { + Some(size) if size <= limit => RequestBodySize::WithinLimit(size), + Some(_) | None => RequestBodySize::ExceedsLimit, + } } #[derive(Default)] @@ -106,12 +116,17 @@ where } else { body.frame().await }; - let Some(frame) = frame else { - break; - }; - let frame = match frame { - Ok(frame) => frame, - Err(error) => { + let data = match frame { + // The client request body reached normal EOF. + None => break, + Some(Ok(frame)) => { + let Ok(data) = frame.into_data() else { + continue; + }; + data + } + Some(Err(error)) => { + tracing::warn!(%error, "failed to read streaming request body from client"); super::send_http_request_abort( in_flight_req, protocol::HttpStreamAbortReasonKind::ClientDisconnect, @@ -121,26 +136,20 @@ where return Err(anyhow!("failed to read streaming request body: {error}")); } }; - 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( - in_flight_req, - protocol::HttpStreamAbortReasonKind::BodyTooLarge, - Some(format!( - "request body exceeded the {max_body_size}-byte limit" - )), - ) - .await; - return Err(InvalidRequestBody { - reason: format!("request body exceeded the {max_body_size}-byte limit"), + body_size = match next_request_body_size(body_size, data.len(), max_body_size) { + RequestBodySize::WithinLimit(next_body_size) => next_body_size, + RequestBodySize::ExceedsLimit => { + let reason = format!("request body exceeded the {max_body_size}-byte limit"); + super::send_http_request_abort( + in_flight_req, + protocol::HttpStreamAbortReasonKind::BodyTooLarge, + Some(reason.clone()), + ) + .await; + return Err(InvalidRequestBody { reason }.build()); } - .build()); }; - body_size = next_body_size; let was_empty = chunker.is_empty(); for chunk in chunker.push(&data) { @@ -153,12 +162,10 @@ where } } - // Normal EOF flushes any partial protocol chunk, then sends exactly one final - // marker so actor-side upload state and request routing can be released. - 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, Vec::new(), true).await + // Normal EOF sends exactly one final protocol chunk so actor-side upload + // state and request routing can be released. + let final_chunk = chunker.flush().unwrap_or_default(); + send_http_request_body_chunk(in_flight_req, final_chunk, true).await } async fn send_http_request_body_chunk( @@ -267,35 +274,28 @@ where let upload = send_streaming_http_request_body_chunks(in_flight_req, body, max_body_size, ingress_bytes); tokio::pin!(upload); + let response_start = wait_for_http_response_start( + msg_rx, + drop_rx, + stopped_sub, + actor_id, + request_id, + response_start_deadline, + response_start_timeout, + ); + tokio::pin!(response_start); tokio::select! { upload_result = &mut upload => { // Normal completion already sent the final upload marker. Continue - // waiting under the original response-start deadline. + // polling the same response future under the original deadline. upload_result?; - wait_for_http_response_start( - msg_rx, - drop_rx, - stopped_sub, - actor_id, - request_id, - response_start_deadline, - response_start_timeout, - ) - .await + response_start.await } - response_start = wait_for_http_response_start( - msg_rx, - drop_rx, - stopped_sub, - actor_id, - request_id, - response_start_deadline, - response_start_timeout, - ) => { + response_start = &mut response_start => { let response_start = response_start?; - // An early successful response makes the unread client body - // irrelevant. Drop the upload future and send a clean final marker - // so actor request routing can be released. + // HTTP handlers may intentionally respond before consuming the full + // upload. Stop reading the client and send a clean final marker so + // actor request routing can be released. send_http_request_body_chunk(in_flight_req, Vec::new(), true).await?; Ok(response_start) } diff --git a/engine/packages/pegboard-gateway2/src/http_stream/response.rs b/engine/packages/pegboard-gateway2/src/http_stream/response.rs index 3f3af352e4..962a75b203 100644 --- a/engine/packages/pegboard-gateway2/src/http_stream/response.rs +++ b/engine/packages/pegboard-gateway2/src/http_stream/response.rs @@ -15,7 +15,10 @@ use crate::shared_state::{ use super::{ResponseBodyError, send_http_request_abort}; const HTTP_BODY_CHUNK_SIZE: usize = 64 * 1024; +const HTTP_RESPONSE_QUEUE_OVERLOADED_DETAIL: &str = + "actor response producer outpaced downstream delivery and filled the gateway buffer"; +/// Advances the expected tunnel sequence while rejecting missing or reordered response frames. fn advance_http_stream_message_index( expected: protocol::MessageIndex, actual: protocol::MessageIndex, @@ -54,7 +57,7 @@ async fn send_http_response_body_bytes( send_http_request_abort( in_flight_req, protocol::HttpStreamAbortReasonKind::Overloaded, - Some("gateway streaming response queue overloaded".to_owned()), + Some(HTTP_RESPONSE_QUEUE_OVERLOADED_DETAIL.to_owned()), ) .await; } @@ -160,17 +163,21 @@ pub(super) async fn drain_http_response_stream( match msg.message_kind { protocol::ToRivetTunnelMessageKind::ToRivetResponseChunk(chunk) => { - if !chunk.body.is_empty() && !send_http_response_body_bytes( - &in_flight_req, - &body_tx, - &mut drop_rx, - &mut stopped_sub, - actor_id, - chunk.body, - "client dropped streaming response body", - &egress_bytes, - ).await { - return; + if !chunk.body.is_empty() { + let delivered = send_http_response_body_bytes( + &in_flight_req, + &body_tx, + &mut drop_rx, + &mut stopped_sub, + actor_id, + chunk.body, + "client dropped streaming response body", + &egress_bytes, + ) + .await; + if !delivered { + return; + } } if chunk.finish { @@ -223,7 +230,7 @@ pub(super) async fn drain_http_response_stream( send_http_request_abort( &in_flight_req, protocol::HttpStreamAbortReasonKind::Overloaded, - Some("gateway streaming response queue overloaded".to_owned()), + Some(HTTP_RESPONSE_QUEUE_OVERLOADED_DETAIL.to_owned()), ) .await; } @@ -277,7 +284,16 @@ fn send_http_body_error( body_tx: &mpsc::Sender>, message: impl Into, ) { - let _ = body_tx.try_send(Err(Box::new(std::io::Error::other(message.into())))); + let error: Result = + Err(Box::new(std::io::Error::other(message.into()))); + match body_tx.try_send(error) { + Ok(()) | Err(mpsc::error::TrySendError::Closed(_)) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + tracing::warn!( + "could not deliver streaming response error because the downstream buffer is full" + ); + } + } } #[cfg(test)] diff --git a/engine/packages/pegboard-gateway2/src/http_stream/response_queue.rs b/engine/packages/pegboard-gateway2/src/http_stream/response_queue.rs index a6feb3eab7..d0927a14a3 100644 --- a/engine/packages/pegboard-gateway2/src/http_stream/response_queue.rs +++ b/engine/packages/pegboard-gateway2/src/http_stream/response_queue.rs @@ -5,9 +5,6 @@ use std::sync::{ use rivet_envoy_protocol::ToRivetTunnelMessageKind; -const HTTP_RESPONSE_QUEUE_MAX_MESSAGES: usize = 384; -const STREAMING_HTTP_RESPONSE_QUEUE_MAX_BYTES: usize = 20 * 1024 * 1024; - /// Bounds response data waiting between pubsub delivery and the downstream HTTP client. /// /// Permits are attached to queued messages and release their capacity on delivery or drop. The @@ -17,14 +14,22 @@ pub(crate) struct HttpResponseQueueBudget { messages: AtomicUsize, bytes: AtomicUsize, buffered_response_max_bytes: usize, + max_messages: usize, + streaming_response_max_bytes: usize, } impl HttpResponseQueueBudget { - pub(crate) fn new(buffered_response_max_bytes: usize) -> Self { + pub(crate) fn new( + buffered_response_max_bytes: usize, + max_messages: usize, + streaming_response_max_bytes: usize, + ) -> Self { Self { messages: AtomicUsize::new(0), bytes: AtomicUsize::new(0), buffered_response_max_bytes, + max_messages, + streaming_response_max_bytes, } } @@ -34,7 +39,7 @@ impl HttpResponseQueueBudget { streaming: bool, ) -> Option { let max_bytes = if streaming { - STREAMING_HTTP_RESPONSE_QUEUE_MAX_BYTES + self.streaming_response_max_bytes } else { self.buffered_response_max_bytes }; @@ -43,7 +48,7 @@ impl HttpResponseQueueBudget { } let previous_messages = self.messages.fetch_add(1, Ordering::AcqRel); - if previous_messages >= HTTP_RESPONSE_QUEUE_MAX_MESSAGES { + if previous_messages >= self.max_messages { self.messages.fetch_sub(1, Ordering::AcqRel); return None; } diff --git a/engine/packages/pegboard-gateway2/src/shared_state.rs b/engine/packages/pegboard-gateway2/src/shared_state.rs index e4cdf077a0..edafa2821b 100644 --- a/engine/packages/pegboard-gateway2/src/shared_state.rs +++ b/engine/packages/pegboard-gateway2/src/shared_state.rs @@ -160,6 +160,8 @@ pub struct SharedStateInner { hws_message_ack_timeout: Duration, hws_max_pending_size: u64, buffered_http_response_max_bytes: usize, + http_response_queue_max_messages: usize, + streaming_http_response_queue_max_bytes: usize, } #[derive(Clone)] @@ -188,6 +190,10 @@ impl SharedState { ), hws_max_pending_size: pegboard_config.gateway_hws_max_pending_size(), buffered_http_response_max_bytes: pegboard_config.envoy_max_response_payload_size(), + http_response_queue_max_messages: pegboard_config + .gateway_http_response_queue_max_messages(), + streaming_http_response_queue_max_bytes: pegboard_config + .gateway_streaming_http_response_queue_max_bytes(), })) } @@ -386,6 +392,8 @@ impl SharedState { let http_response_queue_budget = if matches!(protocol, RequestProtocol::Http) { Some(Arc::new(HttpResponseQueueBudget::new( self.buffered_http_response_max_bytes, + self.http_response_queue_max_messages, + self.streaming_http_response_queue_max_bytes, ))) } else { None @@ -1501,7 +1509,7 @@ fn forward_tunnel_message( request_id=%display_id(&message_id.request_id), message_index=message_id.message_index, http_response_bytes = bytes, - "HTTP response queue exceeded its in-flight limit" + "actor response producer outpaced downstream delivery and filled the HTTP response queue" ); drop_tx.send_replace(Some(MsgGcReason::HttpResponseQueueOverloaded)); return None; diff --git a/engine/packages/pegboard-gateway2/tests/support/http_response_queue.rs b/engine/packages/pegboard-gateway2/tests/support/http_response_queue.rs index 60e25f6f3c..988b89f7e2 100644 --- a/engine/packages/pegboard-gateway2/tests/support/http_response_queue.rs +++ b/engine/packages/pegboard-gateway2/tests/support/http_response_queue.rs @@ -1,9 +1,16 @@ use super::*; +const MAX_MESSAGES: usize = 384; +const MAX_BYTES: usize = 20 * 1024 * 1024; + #[test] fn http_response_queue_budget_bounds_messages_and_releases_capacity() { - let budget = Arc::new(HttpResponseQueueBudget::new(20 * 1024 * 1024)); - let permits = (0..HTTP_RESPONSE_QUEUE_MAX_MESSAGES) + let budget = Arc::new(HttpResponseQueueBudget::new( + MAX_BYTES, + MAX_MESSAGES, + MAX_BYTES, + )); + let permits = (0..MAX_MESSAGES) .map(|_| budget.try_reserve(0, true).expect("message should fit")) .collect::>(); @@ -14,16 +21,16 @@ fn http_response_queue_budget_bounds_messages_and_releases_capacity() { #[test] fn http_response_queue_budget_bounds_bytes_without_limiting_total_stream_size() { - let budget = Arc::new(HttpResponseQueueBudget::new(20 * 1024 * 1024)); + let budget = Arc::new(HttpResponseQueueBudget::new( + MAX_BYTES, + MAX_MESSAGES, + MAX_BYTES, + )); let permit = budget - .try_reserve(STREAMING_HTTP_RESPONSE_QUEUE_MAX_BYTES, true) + .try_reserve(MAX_BYTES, true) .expect("queue-sized chunk should fit"); assert!(budget.try_reserve(1, true).is_none()); drop(permit); - assert!( - budget - .try_reserve(STREAMING_HTTP_RESPONSE_QUEUE_MAX_BYTES, true) - .is_some() - ); + assert!(budget.try_reserve(MAX_BYTES, true).is_some()); } diff --git a/engine/packages/pegboard-gateway2/tests/support/http_stream_request.rs b/engine/packages/pegboard-gateway2/tests/support/http_stream_request.rs index 4407873319..290e1d192d 100644 --- a/engine/packages/pegboard-gateway2/tests/support/http_stream_request.rs +++ b/engine/packages/pegboard-gateway2/tests/support/http_stream_request.rs @@ -14,9 +14,18 @@ fn request_body_streams_only_after_one_chunk() { #[test] fn streaming_request_body_size_is_cumulative_and_overflow_safe() { - assert_eq!(next_request_body_size(6, 4, 10), Some(10)); - assert_eq!(next_request_body_size(7, 4, 10), None); - assert_eq!(next_request_body_size(usize::MAX, 1, usize::MAX), None); + assert_eq!( + next_request_body_size(6, 4, 10), + RequestBodySize::WithinLimit(10) + ); + assert_eq!( + next_request_body_size(7, 4, 10), + RequestBodySize::ExceedsLimit + ); + assert_eq!( + next_request_body_size(usize::MAX, 1, usize::MAX), + RequestBodySize::ExceedsLimit + ); } #[test]