mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
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:
Generated
+1
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user