From b27877c99b10a360a640280e4489fca4f291a98d Mon Sep 17 00:00:00 2001 From: Zyy Date: Tue, 6 Oct 2026 15:47:25 +0800 Subject: [PATCH] fix(memory): release request copies and preallocate XZ buffers Release preliminary Responses decodes before middleware execution and move WebSocket frame ownership without cloning the request body. Reserve XZ block buffers from validated lengths while retaining overflow detection. Add allocation regression coverage for 24 MiB and 32 MiB XZ blocks. Refs #338 Co-authored-by: Codex --- backend/Cargo.lock | 1 + .../gateway-api/src/openai/responses/http.rs | 2 ++ .../src/openai/responses/websocket/mod.rs | 9 ++++--- backend/crates/providers/openai/Cargo.toml | 1 + .../transport/profile/platform_release/xz.rs | 2 ++ .../transport/profile/platform_release/xz.rs | 26 ++++++++++++++++++- 6 files changed, 36 insertions(+), 5 deletions(-) diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 729b48f15..b0a057344 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -3238,6 +3238,7 @@ dependencies = [ name = "provider-openai" version = "0.0.0" dependencies = [ + "allocation-counter", "async-stream", "async-trait", "aws-lc-rs", diff --git a/backend/crates/gateway-api/src/openai/responses/http.rs b/backend/crates/gateway-api/src/openai/responses/http.rs index 72470bc37..afed0b4f6 100644 --- a/backend/crates/gateway-api/src/openai/responses/http.rs +++ b/backend/crates/gateway-api/src/openai/responses/http.rs @@ -134,6 +134,8 @@ pub(crate) async fn execute_prepared_responses( ClientTransport::HttpJson }; let model_hint = Some(decoded.metadata().requested_model().to_owned()); + // 终端会重新校验中间件改写后的正文，初步解码不必保留到响应交付 + drop(decoded); let execution = service.execution(); let request_id = prepared.request_id().clone(); let cancellation = prepared.cancellation(); diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs index 20696965f..342479b25 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs @@ -214,8 +214,9 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke } request_count = request_count.saturating_add(1); let correlation_id = Arc::::from(service.next_request_id()); - let decoded = match decode_response_create_with_context(&payload, &request_headers) { - Ok(decoded) => decoded, + // 初步解码只保留路由事实，避免整份正文跨越准入等待和响应交付 + let model_hint = match decode_response_create_with_context(&payload, &request_headers) { + Ok(decoded) => Some(decoded.metadata().requested_model().to_owned()), Err(error) => { trace_rejected_request( &correlation_id, @@ -298,9 +299,9 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke protocol: "openai".to_owned(), operation: Some(OperationKind::Generate), transport: ClientTransport::WebSocket, - model_hint: Some(decoded.metadata().requested_model().to_owned()), + model_hint, headers: encode_headers(&raw_headers), - body: Bytes::from(payload.clone()), + body: Bytes::from(payload), }; let service_for_terminal = service.clone(); let replay_for_terminal = replay.clone(); diff --git a/backend/crates/providers/openai/Cargo.toml b/backend/crates/providers/openai/Cargo.toml index 86efa7666..5ae2d6ac6 100644 --- a/backend/crates/providers/openai/Cargo.toml +++ b/backend/crates/providers/openai/Cargo.toml @@ -72,6 +72,7 @@ openssl = "=0.10.75" openssl-sys = { version = "=0.9.111", features = ["vendored"] } [dev-dependencies] +allocation-counter = "0.8.1" tempfile = "3.27.0" tokio = { version = "1.52.3", default-features = false, features = ["rt-multi-thread", "test-util"] } tracing-subscriber = { version = "0.3.23", features = ["json"] } diff --git a/backend/crates/providers/openai/src/transport/profile/platform_release/xz.rs b/backend/crates/providers/openai/src/transport/profile/platform_release/xz.rs index 0edb2ea8b..4651fc6d9 100644 --- a/backend/crates/providers/openai/src/transport/profile/platform_release/xz.rs +++ b/backend/crates/providers/openai/src/transport/profile/platform_release/xz.rs @@ -140,6 +140,8 @@ impl Read for IndexedXz { .map_err(|_| invalid())?; let decoder = xz2::read::XzDecoder::new_stream(bytes.as_slice(), stream); self.buffer.clear(); + // 索引长度已受 MAX_BLOCK 限制，多留一字节检查越界，避免 read_to_end 倍增容量 + self.buffer.reserve_exact(block.unpacked as usize + 1); decoder .take(block.unpacked + 1) .read_to_end(&mut self.buffer)?; diff --git a/backend/crates/providers/openai/tests/transport/profile/platform_release/xz.rs b/backend/crates/providers/openai/tests/transport/profile/platform_release/xz.rs index 7280ece44..55e48e871 100644 --- a/backend/crates/providers/openai/tests/transport/profile/platform_release/xz.rs +++ b/backend/crates/providers/openai/tests/transport/profile/platform_release/xz.rs @@ -1,4 +1,4 @@ -//! 验证 XZ 索引读取的定位、跨块访问与无关块跳过 +//! 验证 XZ 索引读取的定位、跨块访问、无关块跳过与大块缓冲保留量 use provider_openai::transport::profile::platform_release::xz::IndexedXz; use std::io::{Cursor, Read, Seek, SeekFrom, Write}; @@ -8,6 +8,30 @@ pub(super) fn compress(bytes: &[u8]) -> Vec { encoder.write_all(bytes).unwrap(); encoder.finish().unwrap() } + +#[test] +fn indexed_xz_keeps_large_block_buffer_close_to_unpacked_size() { + // 同时覆盖实际制品的 24 MiB 块和会因 EOF 探测再次扩容的二次幂边界 + for size in [24 * 1024 * 1024, 32 * 1024 * 1024] { + let compressed = compress(&vec![0x5a; size]); + let mut reader = + IndexedXz::new(Cursor::new(&compressed), 0, compressed.len() as u64).unwrap(); + let mut byte = [0]; + let allocations = allocation_counter::measure(|| reader.read_exact(&mut byte).unwrap()); + + assert_eq!(byte, [0x5a]); + // 测量读取器仍存活时的 Rust 分配；为块元数据留余量，不等同于进程 RSS + assert!( + allocations.bytes_current <= size as i64 + 64 * 1024, + "size={size}, {allocations:?}" + ); + reader.seek(SeekFrom::End(-1)).unwrap(); + reader.read_exact(&mut byte).unwrap(); + assert_eq!(byte, [0x5a]); + assert_eq!(reader.read(&mut byte).unwrap(), 0); + } +} + #[test] fn indexed_xz_supports_backward_and_end_relative_seeks() { let bytes: Vec = (0..100_000).map(|i| (i % 251) as u8).collect();