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 <noreply@openai.com>
This commit is contained in:
Zyy
2026-10-06 15:47:25 +08:00
co-authored by Codex
parent 7505459ece
commit b27877c99b
6 changed files with 36 additions and 5 deletions
+1
View File
@@ -3238,6 +3238,7 @@ dependencies = [
name = "provider-openai"
version = "0.0.0"
dependencies = [
"allocation-counter",
"async-stream",
"async-trait",
"aws-lc-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();
@@ -214,8 +214,9 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke
}
request_count = request_count.saturating_add(1);
let correlation_id = Arc::<str>::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();
@@ -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"] }
@@ -140,6 +140,8 @@ impl<R: Read + Seek> Read for IndexedXz<R> {
.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)?;
@@ -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<u8> {
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<u8> = (0..100_000).map(|i| (i % 251) as u8).collect();