i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 17:37:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: harden concurrency limits and high-RPM runtime paths

Browse Source 
Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
This commit is contained in:
elky
2026-09-10 08:14:58 +08:00
parent 361952ada9
commit ecc16673eb
149 changed files with 27963 additions and 1926 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/execution_runtime/attempt_cancellation.rs
+2
View File
@@ -234,7 +234,9 @@ impl Drop for AttemptCancellationGuard {
);
return;
};
let usage_producer = state.usage_runtime.track_producer();
handle.spawn(async move {
let _usage_producer = usage_producer;
settle_cancelled_attempt(state, armed, error_type, error_message).await;
});
}
apps/aether-gateway/src/execution_runtime/attempt_lifecycle.rs
+2
View File
@@ -674,8 +674,10 @@ impl ExecutionAttemptLifecycle {
let billing_void = settlement.billing.is_void();
let usage_runtime = Arc::clone(&state.usage_runtime);
let usage_data = Arc::clone(state.usage_lifecycle_data_state());
let usage_producer = usage_runtime.track_producer();
self.stage_guard
.await_detachable_stage(self.trace_id.as_str(), "usage_terminal", async move {
let _usage_producer = usage_producer;
usage_runtime
.record_stream_terminal(
usage_data.as_ref(),
apps/aether-gateway/src/execution_runtime/mod.rs
+1
View File
@@ -20,6 +20,7 @@ mod response_header_rules;
mod server;
pub(crate) mod stream;
mod stream_pump;
mod stream_read_timeout;
pub(crate) mod submission;
pub(crate) mod sync;
pub(crate) mod transport;
apps/aether-gateway/src/execution_runtime/stream/capture_budget.rs
+251
View File
@@ -0,0 +1,251 @@
use std::ops::Deref;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, LazyLock};
const DEFAULT_STREAM_CAPTURE_MEMORY_BUDGET_BYTES: usize = 128 * 1024 * 1024;
const STREAM_CAPTURE_MEMORY_BUDGET_ENV: &str = "AETHER_GATEWAY_STREAM_CAPTURE_MEMORY_BUDGET_BYTES";
static STREAM_CAPTURE_BUDGET: LazyLock<Arc<StreamCaptureBudget>> = LazyLock::new(|| {
StreamCaptureBudget::new(
std::env::var(STREAM_CAPTURE_MEMORY_BUDGET_ENV)
.ok()
.and_then(|value| value.trim().parse().ok())
.unwrap_or(DEFAULT_STREAM_CAPTURE_MEMORY_BUDGET_BYTES),
)
});
#[derive(Debug)]
pub(super) struct StreamCaptureBudget {
available: AtomicUsize,
}
impl StreamCaptureBudget {
pub(super) fn new(bytes: usize) -> Arc<Self> {
Arc::new(Self {
available: AtomicUsize::new(bytes),
})
}
fn reserve_up_to(&self, wanted: usize, minimum: usize) -> usize {
let mut available = self.available.load(Ordering::Relaxed);
loop {
let reserved = wanted.min(available);
if reserved < minimum {
return 0;
}
match self.available.compare_exchange_weak(
available,
available - reserved,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return reserved,
Err(current) => available = current,
}
}
}
fn release(&self, bytes: usize) {
self.available.fetch_add(bytes, Ordering::Relaxed);
}
}
/// Only retained diagnostic bytes belong here. Protocol and billing observers
/// must consume the original chunks independently of capture admission.
#[derive(Debug)]
pub(super) struct StreamBodyCapture {
bytes: Vec<u8>,
budget: Arc<StreamCaptureBudget>,
}
impl Default for StreamBodyCapture {
fn default() -> Self {
Self::with_budget(Arc::clone(&STREAM_CAPTURE_BUDGET))
}
}
impl StreamBodyCapture {
pub(super) fn with_budget(budget: Arc<StreamCaptureBudget>) -> Self {
Self {
bytes: Vec::new(),
budget,
}
}
pub(super) fn append(&mut self, chunk: &[u8], limit: usize, truncated: &mut bool) {
if chunk.is_empty() || *truncated {
return;
}
let wanted_len = self.bytes.len().saturating_add(chunk.len()).min(limit);
if wanted_len > self.bytes.capacity() {
// Keep the old allocation charged until its replacement has been
// allocated and copied, including their overlap during growth.
let wanted_capacity = wanted_len
.max(self.bytes.capacity().saturating_mul(2))
.min(limit);
let reserved = self
.budget
.reserve_up_to(wanted_capacity, self.bytes.capacity().saturating_add(1));
if reserved > 0 {
let mut replacement = Vec::new();
if replacement.try_reserve_exact(reserved).is_ok() {
let extra = replacement.capacity().saturating_sub(reserved);
if extra == 0 || self.budget.reserve_up_to(extra, extra) == extra {
replacement.extend_from_slice(&self.bytes);
let old = std::mem::replace(&mut self.bytes, replacement);
let old_capacity = old.capacity();
drop(old);
self.budget.release(old_capacity);
} else {
drop(replacement);
self.budget.release(reserved);
}
} else {
self.budget.release(reserved);
}
}
}
let keep = wanted_len
.min(self.bytes.capacity())
.saturating_sub(self.bytes.len());
self.bytes.extend_from_slice(&chunk[..keep]);
// Once bytes are omitted, never append a later suffix to this prefix.
*truncated = keep < chunk.len();
}
}
impl Deref for StreamBodyCapture {
type Target = [u8];
fn deref(&self) -> &Self::Target {
&self.bytes
}
}
impl Drop for StreamBodyCapture {
fn drop(&mut self) {
let bytes = std::mem::take(&mut self.bytes);
let capacity = bytes.capacity();
drop(bytes);
self.budget.release(capacity);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn stream_capture_budget_is_shared_and_released_on_drop() {
let budget = StreamCaptureBudget::new(12);
let mut provider = StreamBodyCapture::with_budget(Arc::clone(&budget));
let mut client = StreamBodyCapture::with_budget(Arc::clone(&budget));
let mut provider_truncated = false;
let mut client_truncated = false;
provider.append(b"12345678", 64, &mut provider_truncated);
client.append(b"abcdefgh", 64, &mut client_truncated);
assert_eq!(&*provider, b"12345678");
assert_eq!(&*client, b"abcd");
assert!(!provider_truncated);
assert!(client_truncated);
assert_eq!(budget.available.load(Ordering::Relaxed), 0);
drop(provider);
client.append(b"later", 64, &mut client_truncated);
assert_eq!(&*client, b"abcd");
drop(client);
assert_eq!(budget.available.load(Ordering::Relaxed), 12);
}
#[test]
fn stream_capture_budget_charges_capacity_and_reallocation_overlap() {
let budget = StreamCaptureBudget::new(16);
let mut capture = StreamBodyCapture::with_budget(Arc::clone(&budget));
let mut truncated = false;
capture.append(b"1234", 64, &mut truncated);
capture.append(b"5", 64, &mut truncated);
assert_eq!(capture.bytes.capacity(), 8);
assert_eq!(budget.available.load(Ordering::Relaxed), 8);
capture.append(b"6789", 64, &mut truncated);
assert_eq!(&*capture, b"12345678");
assert!(truncated);
assert_eq!(budget.available.load(Ordering::Relaxed), 8);
drop(capture);
assert_eq!(budget.available.load(Ordering::Relaxed), 16);
}
#[test]
fn stream_capture_budget_zero_disables_capture_without_allocating() {
let budget = StreamCaptureBudget::new(0);
let mut capture = StreamBodyCapture::with_budget(budget);
let mut truncated = false;
capture.append(b"data", 64, &mut truncated);
assert!(capture.is_empty());
assert_eq!(capture.bytes.capacity(), 0);
assert!(truncated);
}
#[test]
fn stream_capture_budget_exhaustion_uses_existing_spare_capacity() {
let budget = StreamCaptureBudget::new(14);
let mut capture = StreamBodyCapture::with_budget(budget);
let mut truncated = false;
capture.append(b"1234", 64, &mut truncated);
capture.append(b"5", 64, &mut truncated);
assert_eq!(capture.bytes.capacity(), 8);
capture.append(b"6789", 64, &mut truncated);
assert_eq!(&*capture, b"12345678");
assert_eq!(capture.bytes.capacity(), 8);
assert!(truncated);
}
#[test]
fn stream_capture_budget_local_limit_keeps_a_contiguous_prefix() {
let budget = StreamCaptureBudget::new(128);
let mut capture = StreamBodyCapture::with_budget(Arc::clone(&budget));
let mut truncated = false;
capture.append(b"abcdef", 3, &mut truncated);
assert_eq!(&*capture, b"abc");
assert!(truncated);
assert_eq!(budget.available.load(Ordering::Relaxed), 125);
}
#[test]
fn stream_capture_budget_concurrent_growth_and_drop_never_exceeds_capacity() {
const LIMIT: usize = 256;
const THREADS: usize = 8;
let budget = StreamCaptureBudget::new(LIMIT);
let barrier = std::sync::Barrier::new(THREADS);
let held = AtomicUsize::new(0);
std::thread::scope(|scope| {
for index in 0..THREADS {
let budget = &budget;
let barrier = &barrier;
let held = &held;
scope.spawn(move || {
for _ in 0..32 {
let mut capture = StreamBodyCapture::with_budget(Arc::clone(budget));
let mut truncated = false;
barrier.wait();
capture.append(&[1; 16], LIMIT, &mut truncated);
capture.append(&[2; 48], LIMIT, &mut truncated);
held.fetch_add(capture.bytes.capacity(), Ordering::SeqCst);
barrier.wait();
if index == 0 {
let retained = held.load(Ordering::SeqCst);
assert!(retained <= LIMIT);
assert_eq!(retained + budget.available.load(Ordering::Relaxed), LIMIT,);
}
barrier.wait();
drop(capture);
barrier.wait();
if index == 0 {
assert_eq!(budget.available.load(Ordering::Relaxed), LIMIT);
held.store(0, Ordering::SeqCst);
}
barrier.wait();
}
});
}
});
}
}
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+604 -88
View File
File diff suppressed because it is too large Load Diff
apps/aether-gateway/src/execution_runtime/stream/mod.rs
+3
View File
@@ -1,7 +1,10 @@
mod capture_budget;
mod commit_policy;
mod error;
mod execution;
mod usage_fallback;
pub(crate) use execution::{
execute_execution_runtime_stream, execute_execution_runtime_stream_with_retry_scope,
ClientVisibleStreamCompletionTracker,
};
apps/aether-gateway/src/execution_runtime/stream/usage_fallback.rs
+1163
View File
File diff suppressed because it is too large Load Diff
apps/aether-gateway/src/execution_runtime/stream_pump.rs
+244 -538
View File
@@ -12,7 +12,6 @@ use async_stream::stream;
use axum::body::Bytes;
use base64::Engine as _;
use futures_util::{Stream, StreamExt};
use http_body_util::BodyExt;
use serde_json::Value;
use tracing::warn;
@@ -21,9 +20,14 @@ use crate::ai_serving::api::{
normalize_provider_private_report_context, StreamingStandardTerminalObserver,
};
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::stream::ClientVisibleStreamCompletionTracker;
use crate::execution_runtime::stream_read_timeout::{
await_stream_idle_read, stream_idle_timeout_message,
};
use crate::execution_runtime::transport::{
append_upstream_response_body_chunk, decode_response_body_bytes,
stream_first_byte_timeout_message, DirectUpstreamResponse,
direct_upstream_response_byte_stream, stream_first_byte_timeout_message,
DirectUpstreamResponse,
};
use crate::execution_runtime::DirectUpstreamStreamExecution;
use crate::GatewayError;
@@ -31,6 +35,15 @@ use crate::GatewayError;
const STREAM_USAGE_OBSERVER_MAX_LINE_BYTES: usize = 1024 * 1024;
const UPSTREAM_STREAM_READ_ERROR_MESSAGE: &str = "Upstream response stream failed";
fn upstream_stream_error_category(response: &DirectUpstreamResponse) -> &'static str {
match response {
DirectUpstreamResponse::Reqwest(_) => "reqwest_body_read_failed",
DirectUpstreamResponse::HyperH2c(_) => "hyper_body_read_failed",
DirectUpstreamResponse::BrowserWreq(_) => "browser_body_read_failed",
DirectUpstreamResponse::LocalTunnel(_) => "tunnel_body_read_failed",
}
}
pub(crate) fn build_direct_execution_frame_stream(
execution: DirectUpstreamStreamExecution,
) -> impl Stream<Item = Result<Bytes, IoError>> + Send + 'static {
@@ -49,9 +62,11 @@ pub(crate) fn build_direct_execution_frame_stream(
started_at,
response_observation,
stream_first_byte_timeout,
stream_idle_timeout,
upstream_target_permit,
} = execution;
let _upstream_target_permit = upstream_target_permit;
let upstream_error_category = upstream_stream_error_category(&response);
let mut observer_context = stream_summary_report_context;
if observer_context
@@ -74,6 +89,7 @@ pub(crate) fn build_direct_execution_frame_stream(
let mut private_stream_normalizer =
maybe_build_provider_private_stream_normalizer(Some(&observer_context));
let mut stream_terminal_observer = StreamingStandardTerminalObserver::default();
let mut stream_completion = ClientVisibleStreamCompletionTracker::default();
let mut observer_buffered = Vec::new();
if should_buffer_non_stream_response(
@@ -87,6 +103,7 @@ pub(crate) fn build_direct_execution_frame_stream(
response,
started_at,
stream_first_byte_timeout,
stream_idle_timeout,
)
.await
{
@@ -166,6 +183,7 @@ pub(crate) fn build_direct_execution_frame_stream(
ttfb_ms,
upstream_bytes,
first_byte_timeout,
idle_timeout,
}) => {
match encode_headers_frame(
status_code,
@@ -180,6 +198,8 @@ pub(crate) fn build_direct_execution_frame_stream(
}
let error_frame = if let Some(timeout) = first_byte_timeout {
encode_first_byte_timeout_frame(timeout)
} else if let Some(timeout) = idle_timeout {
encode_idle_timeout_frame(timeout)
} else {
encode_error_frame(message)
};
@@ -228,7 +248,9 @@ pub(crate) fn build_direct_execution_frame_stream(
let mut prefetched_body_failed = false;
for item in prefetched_body {
match item {
Ok(chunk) if chunk.is_empty() => continue,
Ok(chunk) => {
stream_completion.observe_chunk(&chunk);
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
@@ -280,328 +302,97 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
if !prefetched_body_failed {
match response {
DirectUpstreamResponse::Reqwest(response) => {
let mut bytes_stream = response.bytes_stream();
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
match encode_first_byte_timeout_frame(timeout) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
break;
}
let mut bytes_stream = direct_upstream_response_byte_stream(VecDeque::new(), response);
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(), started_at, stream_first_byte_timeout,
).await {
Ok(item) => item,
Err(timeout) => {
match encode_first_byte_timeout_frame(timeout) {
Ok(frame) => yield Ok(frame),
Err(err) => yield Err(err),
}
} else {
bytes_stream.next().await
};
let Some(item) = item else {
break;
};
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
}
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error_category = "reqwest_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
}
break;
}
}
}
}
DirectUpstreamResponse::HyperH2c(response) => {
let mut bytes_stream = response.into_body().into_data_stream();
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(),
started_at,
stream_first_byte_timeout,
)
.await
} else {
match await_stream_idle_read(bytes_stream.next(), stream_idle_timeout).await {
Ok(item) => item,
Err(timeout) => {
drop(bytes_stream);
if stream_completion.successful_completion()
|| (!stream_completion.observed_terminal()
&& stream_terminal_observer.latest_summary().is_some_and(|summary| {
summary.observed_finish && summary.parser_error.is_none()
&& summary.finish_reason.as_deref() != Some("error")
}))
{
Ok(item) => item,
Err(timeout) => {
match encode_first_byte_timeout_frame(timeout) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
break;
}
}
} else {
bytes_stream.next().await
};
let Some(item) = item else {
break;
};
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
}
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error_category = "hyper_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
}
break;
}
if stream_terminal_observer.latest_summary().is_some_and(|summary| {
summary.observed_finish && summary.parser_error.is_some()
}) {
// The terminal summary carries the original provider failure.
break;
}
match encode_idle_timeout_frame(timeout) {
Ok(frame) => yield Ok(frame),
Err(err) => yield Err(err),
}
break;
}
}
}
DirectUpstreamResponse::BrowserWreq(response) => {
let mut bytes_stream = response.bytes_stream();
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
match encode_first_byte_timeout_frame(timeout) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
break;
}
}
} else {
bytes_stream.next().await
};
let Some(item) = item else {
break;
};
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
}
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error_category = "browser_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
}
break;
}
};
let Some(item) = item else { break };
match item {
Ok(chunk) => {
stream_completion.observe_chunk(&chunk);
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
}
}
DirectUpstreamResponse::LocalTunnel(mut response) => loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
response.next_chunk(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
match encode_first_byte_timeout_frame(timeout) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
break;
}
}
} else {
response.next_chunk().await
};
match item {
Ok(Some(chunk)) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
if !first_chunk_telemetry_emitted {
match encode_telemetry_frame(ttfb_ms, ttfb_ms, upstream_bytes) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
Ok(None) => break,
Err(_message) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error_category = "tunnel_body_read_failed",
"upstream body stream read error"
);
match encode_error_frame(UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string()) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
return;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
break;
}
}
Err(_) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
status_code,
upstream_bytes,
error_category = upstream_error_category,
"upstream body stream read error"
);
match encode_error_frame(UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string()) {
Ok(frame) => yield Ok(frame),
Err(err) => yield Err(err),
}
break;
}
}
}
}
@@ -704,6 +495,22 @@ fn encode_first_byte_timeout_frame(timeout: Duration) -> Result<Bytes, IoError>
})
}
fn encode_idle_timeout_frame(timeout: Duration) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::ReadTimeout,
phase: ExecutionPhase::StreamRead,
message: stream_idle_timeout_message(timeout),
upstream_status: None,
retryable: true,
failover_recommended: true,
},
},
})
}
async fn await_stream_first_byte<T, F>(
future: F,
started_at: Instant,
@@ -737,6 +544,7 @@ struct BufferedUpstreamBodyError {
ttfb_ms: Option<u64>,
upstream_bytes: u64,
first_byte_timeout: Option<Duration>,
idle_timeout: Option<Duration>,
}
fn append_buffered_upstream_body_chunk(
@@ -752,6 +560,7 @@ fn append_buffered_upstream_body_chunk(
ttfb_ms,
upstream_bytes: *upstream_bytes,
first_byte_timeout: None,
idle_timeout: None,
}
})
}
@@ -816,16 +625,52 @@ fn should_buffer_non_stream_response(
}
async fn buffer_non_sse_upstream_body(
mut prefetched_body: VecDeque<Result<Bytes, String>>,
prefetched_body: VecDeque<Result<Bytes, String>>,
response: DirectUpstreamResponse,
started_at: Instant,
stream_first_byte_timeout: Option<Duration>,
stream_idle_timeout: Option<Duration>,
) -> Result<BufferedUpstreamBody, BufferedUpstreamBodyError> {
let mut body_bytes = Vec::new();
let mut upstream_bytes = 0u64;
let mut ttfb_ms = None;
while let Some(item) = prefetched_body.pop_front() {
let upstream_error_category = upstream_stream_error_category(&response);
let mut bytes_stream = direct_upstream_response_byte_stream(prefetched_body, response);
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
return Err(BufferedUpstreamBodyError {
message: stream_first_byte_timeout_message(timeout),
ttfb_ms,
upstream_bytes,
first_byte_timeout: Some(timeout),
idle_timeout: None,
})
}
}
} else {
match await_stream_idle_read(bytes_stream.next(), stream_idle_timeout).await {
Ok(item) => item,
Err(timeout) => {
return Err(BufferedUpstreamBodyError {
message: stream_idle_timeout_message(timeout),
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
idle_timeout: Some(timeout),
})
}
}
};
let Some(item) = item else { break };
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
@@ -838,246 +683,24 @@ async fn buffer_non_sse_upstream_body(
&mut upstream_bytes,
)?;
}
Err(_message) => {
Err(_) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error_category = upstream_error_category,
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
message: UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string(),
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
idle_timeout: None,
});
}
}
}
match response {
DirectUpstreamResponse::Reqwest(response) => {
let mut bytes_stream = response.bytes_stream();
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
return Err(BufferedUpstreamBodyError {
message: stream_first_byte_timeout_message(timeout),
ttfb_ms,
upstream_bytes,
first_byte_timeout: Some(timeout),
});
}
}
} else {
bytes_stream.next().await
};
let Some(item) = item else {
break;
};
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_buffered_upstream_body_chunk(
&mut body_bytes,
&chunk,
ttfb_ms,
&mut upstream_bytes,
)?;
}
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error_category = "reqwest_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
message,
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
});
}
}
}
}
DirectUpstreamResponse::HyperH2c(response) => {
let mut bytes_stream = response.into_body().into_data_stream();
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
return Err(BufferedUpstreamBodyError {
message: stream_first_byte_timeout_message(timeout),
ttfb_ms,
upstream_bytes,
first_byte_timeout: Some(timeout),
});
}
}
} else {
bytes_stream.next().await
};
let Some(item) = item else {
break;
};
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_buffered_upstream_body_chunk(
&mut body_bytes,
&chunk,
ttfb_ms,
&mut upstream_bytes,
)?;
}
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error_category = "hyper_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
message,
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
});
}
}
}
}
DirectUpstreamResponse::BrowserWreq(response) => {
let mut bytes_stream = response.bytes_stream();
loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
bytes_stream.next(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
return Err(BufferedUpstreamBodyError {
message: stream_first_byte_timeout_message(timeout),
ttfb_ms,
upstream_bytes,
first_byte_timeout: Some(timeout),
});
}
}
} else {
bytes_stream.next().await
};
let Some(item) = item else {
break;
};
match item {
Ok(chunk) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_buffered_upstream_body_chunk(
&mut body_bytes,
&chunk,
ttfb_ms,
&mut upstream_bytes,
)?;
}
Err(_err) => {
let message = UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string();
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error_category = "browser_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
message,
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
});
}
}
}
}
DirectUpstreamResponse::LocalTunnel(mut response) => loop {
let item = if ttfb_ms.is_none() {
match await_stream_first_byte(
response.next_chunk(),
started_at,
stream_first_byte_timeout,
)
.await
{
Ok(item) => item,
Err(timeout) => {
return Err(BufferedUpstreamBodyError {
message: stream_first_byte_timeout_message(timeout),
ttfb_ms,
upstream_bytes,
first_byte_timeout: Some(timeout),
});
}
}
} else {
response.next_chunk().await
};
match item {
Ok(Some(chunk)) => {
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
append_buffered_upstream_body_chunk(
&mut body_bytes,
&chunk,
ttfb_ms,
&mut upstream_bytes,
)?;
}
Ok(None) => break,
Err(_message) => {
warn!(
event_name = "stream_pump_body_read_error",
log_type = "ops",
upstream_bytes,
error_category = "tunnel_body_read_failed",
"upstream body stream read error"
);
return Err(BufferedUpstreamBodyError {
message: UPSTREAM_STREAM_READ_ERROR_MESSAGE.to_string(),
ttfb_ms,
upstream_bytes,
first_byte_timeout: None,
});
}
}
},
}
Ok(BufferedUpstreamBody {
body_bytes,
ttfb_ms,
@@ -1605,6 +1228,89 @@ mod tests {
assert_eq!(error.get("failover_recommended"), Some(&Value::Bool(true)));
}
#[tokio::test]
async fn direct_execution_frame_stream_enforces_idle_timeout_after_first_byte() {
for (content_type, first_chunk, expect_timeout, provider_format) in [
("text/event-stream", "data: hello\n\n", true, "openai:chat"),
("application/json", "{\"message\":", true, "openai:chat"),
("text/event-stream", "data: [DONE]\n\n", false, "openai:chat"),
("text/event-stream", "event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\"}}\n\n", false, "openai:responses"),
("text/event-stream", "event: response.incomplete\ndata: {\"type\":\"response.incomplete\",\"response\":{}}\n\n", true, "openai:responses"),
] {
let listener = crate::test_support::bind_loopback_listener().await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.unwrap();
let mut request = [0_u8; 4096];
socket.read(&mut request).await.unwrap();
let response = if content_type == "application/json" {
format!("HTTP/1.1 200 OK\r\ncontent-type: {content_type}\r\ncontent-length: 1024\r\n\r\n{first_chunk}")
} else {
format!(
"HTTP/1.1 200 OK\r\ncontent-type: {content_type}\r\ntransfer-encoding: chunked\r\n\r\n{:x}\r\n{first_chunk}\r\n",
first_chunk.len(),
)
};
socket.write_all(response.as_bytes()).await.unwrap();
socket.flush().await.unwrap();
tokio::time::sleep(Duration::from_secs(5)).await;
});
let execution = DirectSyncExecutionRuntime::new()
.execute_stream(&ExecutionPlan {
request_id: "req-stream-idle-timeout".into(),
candidate_id: Some("cand-stream-idle-timeout".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: format!("http://{addr}/chat"),
headers: BTreeMap::from([("content-type".into(), "application/json".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(serde_json::json!({"stream": true})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: provider_format.into(),
model_name: Some("gpt-5".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
first_byte_ms: Some(1_000),
read_ms: Some(10),
..ExecutionTimeouts::default()
}),
})
.await
.expect("stream response headers");
let frames = tokio::time::timeout(
Duration::from_secs(1),
build_direct_execution_frame_stream(execution).collect::<Vec<_>>(),
)
.await;
server.abort();
let frames = frames
.expect("idle timeout must terminate both SSE and buffered JSON")
.into_iter()
.map(|line| serde_json::from_slice::<Value>(&line.unwrap()).unwrap())
.collect::<Vec<_>>();
let errors = frames
.iter()
.filter(|frame| frame["type"] == "error")
.collect::<Vec<_>>();
assert_eq!(errors.len(), usize::from(expect_timeout));
if expect_timeout {
assert_eq!(errors[0]["payload"]["error"]["kind"], "read_timeout");
assert_eq!(errors[0]["payload"]["error"]["phase"], "stream_read");
}
assert!(frames.iter().any(|frame| frame["type"] == "eof"));
assert_eq!(
frames.iter().any(|frame| frame["type"] == "data"),
content_type == "text/event-stream"
);
}
}
#[tokio::test]
async fn direct_execution_frame_stream_emits_telemetry_before_first_data_frame() {
let listener = crate::test_support::bind_loopback_listener()
apps/aether-gateway/src/execution_runtime/stream_read_timeout.rs
+187
View File
@@ -0,0 +1,187 @@
use std::future::Future;
use std::time::Duration;
use aether_contracts::ExecutionPlan;
use axum::body::Bytes;
use futures_util::{Stream, StreamExt};
const STREAM_IDLE_TIMEOUT_MS_ENV: &str = "AETHER_GATEWAY_UPSTREAM_STREAM_IDLE_TIMEOUT_MS";
const DEFAULT_STREAM_IDLE_TIMEOUT_MS: u64 = 300_000;
pub(crate) fn resolve_stream_idle_timeout(plan: &ExecutionPlan) -> Option<Duration> {
if !plan.stream {
return None;
}
let configured = std::env::var(STREAM_IDLE_TIMEOUT_MS_ENV).ok();
stream_idle_timeout_from_config(
plan.timeouts.as_ref().and_then(|timeouts| timeouts.read_ms),
configured.as_deref(),
)
}
fn stream_idle_timeout_from_config(
read_ms: Option<u64>,
configured: Option<&str>,
) -> Option<Duration> {
let timeout_ms = read_ms
.or_else(|| configured.and_then(|value| value.trim().parse::<u64>().ok()))
.unwrap_or(DEFAULT_STREAM_IDLE_TIMEOUT_MS);
// Zero explicitly disables the idle limit for providers with long silent reasoning phases.
(timeout_ms > 0).then(|| Duration::from_millis(timeout_ms))
}
pub(crate) fn stream_idle_timeout_message(timeout: Duration) -> String {
format!(
"provider stream idle read timeout after {} ms",
timeout.as_millis()
)
}
pub(crate) async fn await_stream_idle_read<T>(
future: impl Future<Output = T>,
timeout: Option<Duration>,
) -> Result<T, Duration> {
match timeout {
Some(timeout) => tokio::time::timeout(timeout, future)
.await
.map_err(|_| timeout),
None => Ok(future.await),
}
}
pub(crate) fn skip_empty_upstream_chunks<E: Send + 'static>(
upstream: impl Stream<Item = Result<Bytes, E>> + Send + 'static,
) -> impl Stream<Item = Result<Bytes, E>> + Send {
async_stream::stream! {
tokio::pin!(upstream);
while let Some(item) = upstream.next().await {
match item {
Ok(chunk) if chunk.is_empty() => {
// Empty frames are not progress; yield so an always-ready source cannot
// monopolize the executor or prevent its enclosing timeout from firing.
tokio::task::yield_now().await;
}
item => yield item,
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
#[test]
fn idle_timeout_configuration_preserves_provider_override_and_explicit_disable() {
assert_eq!(
stream_idle_timeout_from_config(None, None),
Some(Duration::from_secs(300))
);
assert_eq!(
stream_idle_timeout_from_config(None, Some(" invalid ")),
Some(Duration::from_secs(300))
);
assert_eq!(
stream_idle_timeout_from_config(None, Some(" 600000 ")),
Some(Duration::from_secs(600))
);
assert_eq!(
stream_idle_timeout_from_config(Some(120_000), Some("600000")),
Some(Duration::from_secs(120))
);
assert_eq!(
stream_idle_timeout_from_config(Some(0), Some("600000")),
None
);
assert_eq!(stream_idle_timeout_from_config(None, Some("0")), None);
}
#[tokio::test]
async fn idle_timeout_cancels_the_pending_upstream_read() {
struct DropMarker(Arc<AtomicBool>);
impl Drop for DropMarker {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
let dropped = Arc::new(AtomicBool::new(false));
let marker = DropMarker(Arc::clone(&dropped));
let outcome = await_stream_idle_read(
async move {
let _marker = marker;
std::future::pending::<()>().await;
},
Some(Duration::from_millis(5)),
)
.await;
assert_eq!(outcome, Err(Duration::from_millis(5)));
assert!(dropped.load(Ordering::SeqCst));
}
#[tokio::test]
async fn idle_timeout_allows_progressing_stream_to_outlive_one_timeout() {
for _ in 0..3 {
assert_eq!(
await_stream_idle_read(
async {
tokio::time::sleep(Duration::from_millis(10)).await;
1
},
Some(Duration::from_millis(25))
)
.await,
Ok(1)
);
}
assert_eq!(
await_stream_idle_read(
async {
tokio::time::sleep(Duration::from_millis(30)).await;
2
},
None
)
.await,
Ok(2)
);
}
#[tokio::test]
async fn empty_upstream_chunks_do_not_reset_idle_timeout() {
let upstream = futures_util::stream::repeat(Ok::<_, ()>(Bytes::new()));
let filtered = skip_empty_upstream_chunks(upstream);
tokio::pin!(filtered);
let outcome = tokio::time::timeout(
Duration::from_secs(1),
await_stream_idle_read(filtered.next(), Some(Duration::from_millis(5))),
)
.await
.expect("empty ready chunks must yield to the idle timer");
assert_eq!(outcome, Err(Duration::from_millis(5)));
}
#[tokio::test]
async fn downstream_keepalive_ticks_do_not_reset_pending_upstream_idle_timeout() {
let read = await_stream_idle_read(
std::future::pending::<()>(),
Some(Duration::from_millis(30)),
);
tokio::pin!(read);
let mut keepalive = tokio::time::interval(Duration::from_millis(2));
let mut ticks = 0;
loop {
tokio::select! {
result = &mut read => {
assert_eq!(result, Err(Duration::from_millis(30)));
assert!(ticks > 0);
break;
}
_ = keepalive.tick() => { ticks += 1; }
}
}
}
}
apps/aether-gateway/src/execution_runtime/sync/execution.rs
+2
View File
@@ -270,7 +270,9 @@ impl Drop for SyncAttemptTerminalGuard {
let candidate_started_unix_ms = self.candidate_started_unix_ms;
let candidate_started_at = self.candidate_started_at;
if let Ok(handle) = tokio::runtime::Handle::try_current() {
let usage_producer = state.usage_runtime.track_producer();
handle.spawn(async move {
let _usage_producer = usage_producer;
record_sync_attempt_forced_terminal_state(
state,
plan,
apps/aether-gateway/src/execution_runtime/transport.rs
+73 -5
View File
@@ -36,7 +36,7 @@ use hyper::body::Incoming as HyperIncomingBody;
use hyper::client::conn::http2::SendRequest as HyperH2cSendRequest;
use hyper_util::client::legacy::connect::HttpConnector;
use hyper_util::client::legacy::Client as HyperLegacyClient;
use hyper_util::rt::{TokioExecutor, TokioIo};
use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use reqwest::redirect::Policy;
use serde::Serialize;
@@ -50,6 +50,7 @@ use tokio::sync::OnceCell as TokioOnceCell;
use crate::ai_serving::api::extract_provider_private_stream_error_body;
#[cfg(test)]
use crate::execution_runtime::remote_compat::execute_sync_plan_via_remote_execution_runtime;
use crate::execution_runtime::stream_read_timeout::resolve_stream_idle_timeout;
use crate::execution_runtime::windsurf::maybe_execute_windsurf_sync;
use crate::frontdoor_loop_guard::{
configured_gateway_frontdoor_base_url, gateway_frontdoor_self_loop_guard_error,
@@ -109,7 +110,7 @@ const DIRECT_REQWEST_PREWARM_SYNC_CLIENTS_ENV: &str =
"AETHER_GATEWAY_DIRECT_REQWEST_PREWARM_SYNC_CLIENTS";
const DEFAULT_H2_TARGET_STREAMS_PER_CLIENT: usize = 8;
const DEFAULT_HTTP1_TARGET_STREAMS_PER_CLIENT: usize = 512;
const DEFAULT_DIRECT_H2C_POOL_MAX_IDLE_PER_HOST: usize = 512;
const DEFAULT_DIRECT_H2C_POOL_MAX_IDLE_PER_HOST: usize = 32;
const DEFAULT_DIRECT_H2C_TARGET_STREAMS_PER_CLIENT: usize = 128;
const DEFAULT_DIRECT_H2C_SENDER_SELECT_WINDOW: usize = 4;
const MAX_DIRECT_H2C_DRIVER_RUNTIME_THREADS: usize = 16;
@@ -382,6 +383,7 @@ static DIRECT_H2C_SENDER_CACHE: LazyLock<
static DIRECT_H2C_POOL_MAX_IDLE_PER_HOST: LazyLock<usize> = LazyLock::new(|| {
env_positive_usize(DIRECT_H2C_POOL_MAX_IDLE_PER_HOST_ENV)
.unwrap_or(DEFAULT_DIRECT_H2C_POOL_MAX_IDLE_PER_HOST)
.min(1024)
});
static DIRECT_H2C_SENDER_SELECT_WINDOW: LazyLock<usize> = LazyLock::new(|| {
@@ -1198,6 +1200,42 @@ pub(crate) enum DirectUpstreamResponse {
LocalTunnel(tunnel::DirectRelayResponse),
}
pub(crate) fn direct_upstream_response_byte_stream(
prefetched_body: VecDeque<Result<Bytes, String>>,
response: DirectUpstreamResponse,
) -> futures_util::stream::BoxStream<'static, Result<Bytes, String>> {
let response_stream = match response {
DirectUpstreamResponse::Reqwest(response) => response
.bytes_stream()
.map(|item| item.map_err(|err| format_upstream_request_error(&err)))
.boxed(),
DirectUpstreamResponse::HyperH2c(response) => response
.into_body()
.into_data_stream()
.map(|item| item.map_err(|err| format_hyper_error_chain(&err)))
.boxed(),
DirectUpstreamResponse::BrowserWreq(response) => response
.bytes_stream()
.map(|item| item.map_err(|err| format_wreq_upstream_request_error(&err)))
.boxed(),
DirectUpstreamResponse::LocalTunnel(mut response) => async_stream::stream! {
loop {
match response.next_chunk().await {
Ok(Some(chunk)) => yield Ok(chunk),
Ok(None) => break,
Err(err) => {
yield Err(err);
break;
}
}
}
}
.boxed(),
};
let upstream = futures_util::stream::iter(prefetched_body).chain(response_stream);
crate::execution_runtime::stream_read_timeout::skip_empty_upstream_chunks(upstream).boxed()
}
pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) request_id: String,
pub(crate) candidate_id: Option<String>,
@@ -1214,6 +1252,7 @@ pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) started_at: Instant,
pub(crate) response_observation: ExecutionResponseObservation,
pub(crate) stream_first_byte_timeout: Option<Duration>,
pub(crate) stream_idle_timeout: Option<Duration>,
pub(crate) upstream_target_permit: Option<UpstreamTargetAdmissionPermit>,
}
@@ -1347,6 +1386,7 @@ impl DirectSyncExecutionRuntime {
request_order_id,
},
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
stream_idle_timeout: resolve_stream_idle_timeout(plan),
upstream_target_permit: None,
})
}
@@ -1494,6 +1534,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
request_order_id,
},
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
stream_idle_timeout: resolve_stream_idle_timeout(plan),
upstream_target_permit: None,
}))
}
@@ -2593,6 +2634,8 @@ fn build_direct_h2c_client_from_cache_key(
builder.http2_only(true);
builder.http2_adaptive_window(true);
builder.pool_max_idle_per_host(cache_key.pool_max_idle_per_host);
builder.pool_timer(TokioTimer::new());
builder.pool_idle_timeout(Duration::from_millis(upstream_pool_idle_timeout_ms()));
builder.build(connector)
}
@@ -2854,8 +2897,18 @@ async fn send_via_browser_wreq_transport(
let profile = plan.transport_profile.as_ref().ok_or_else(|| {
ExecutionRuntimeTransportError::UnsupportedTransportProfile(String::new())
})?;
let mut client_timeouts = plan.timeouts.clone();
if plan.stream {
if let Some(timeouts) = client_timeouts.as_mut() {
// Streamed responses use the shared idle reader; sync collectors retain
// their existing client read timeout. Zero explicitly disables either.
if apply_request_total_timeout || timeouts.read_ms == Some(0) {
timeouts.read_ms = None;
}
}
}
let client = build_browser_wreq_client(
plan.timeouts.as_ref(),
client_timeouts.as_ref(),
plan.proxy.as_ref(),
profile,
transport_controls,
@@ -4214,6 +4267,7 @@ fn build_direct_reqwest_client_from_cache_key(
&HttpClientConfig {
connect_timeout_ms: cache_key.connect_timeout_ms,
pool_max_idle_per_host: Some(direct_reqwest_pool_max_idle_per_host()),
pool_idle_timeout_ms: Some(upstream_pool_idle_timeout_ms()),
..HttpClientConfig::default()
},
);
@@ -4233,12 +4287,22 @@ fn build_direct_reqwest_client_from_cache_key(
}
fn direct_reqwest_pool_max_idle_per_host() -> usize {
const DEFAULT_MAX_IDLE_PER_HOST: usize = 1024;
const DEFAULT_MAX_IDLE_PER_HOST: usize = 32;
std::env::var("AETHER_GATEWAY_UPSTREAM_POOL_MAX_IDLE_PER_HOST")
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_MAX_IDLE_PER_HOST)
.min(1024)
}
fn upstream_pool_idle_timeout_ms() -> u64 {
std::env::var("AETHER_GATEWAY_UPSTREAM_POOL_IDLE_TIMEOUT_MS")
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.filter(|value| *value > 0)
.unwrap_or(15_000)
.min(300_000)
}
pub(crate) fn direct_reqwest_client_cache_metric_samples() -> Vec<MetricSample> {
@@ -4585,7 +4649,11 @@ pub(crate) fn build_browser_wreq_client(
) -> Result<wreq::Client, ExecutionRuntimeTransportError> {
let emulation = browser_wreq_emulation_from_profile(transport_profile)?;
let proxy_url = resolve_proxy_url(proxy)?;
let mut builder = wreq::Client::builder().no_proxy().emulation(emulation);
let mut builder = wreq::Client::builder()
.no_proxy()
.emulation(emulation)
.pool_max_idle_per_host(direct_reqwest_pool_max_idle_per_host())
.pool_idle_timeout(Duration::from_millis(upstream_pool_idle_timeout_ms()));
if proxy_url.is_none() {
builder = builder.dns_resolver(ExecutionSafeDnsResolver);
}