mirror of
https://github.com/openai/codex.git
synced 2026-10-20 00:18:49 +00:00
Stream rollout records during thread history materialization (#54639)
## Why Thread history materialization allocated a buffer for the entire unread rollout suffix before processing records, increasing memory use for large rollouts. ## What changed Read records incrementally with a buffered reader bounded to the captured file end. Reuse the line buffer, preserve trailing partial records for the next pass, and report unexpected EOF if the rollout shrinks during reading. ## Testing Extend the partial-line boundary test with 64 KiB of padding per record to verify byte offsets and deferred processing until the final newline is durable. GitOrigin-RevId: c0426e3e364748b117285b248f6988a4fe064b94
This commit is contained in:
@@ -1,3 +1,5 @@
|
||||
use std::io::BufRead;
|
||||
use std::io::BufReader;
|
||||
use std::io::Read;
|
||||
use std::io::Seek;
|
||||
use std::io::SeekFrom;
|
||||
@@ -134,26 +136,31 @@ fn read_projection_steps_sync(
|
||||
.ok_or_else(|| ThreadStoreError::Internal {
|
||||
message: "durable rollout shrank before projection".to_string(),
|
||||
})?;
|
||||
let byte_count = usize::try_from(byte_count).map_err(|_| ThreadStoreError::Internal {
|
||||
message: "durable rollout append exceeds addressable memory".to_string(),
|
||||
})?;
|
||||
let mut bytes = vec![0; byte_count];
|
||||
file.seek(SeekFrom::Start(start_offset))
|
||||
.map_err(thread_store_io_error)?;
|
||||
file.read_exact(bytes.as_mut_slice())
|
||||
.map_err(thread_store_io_error)?;
|
||||
// Freeze this pass at the captured end, without retaining the entire durable suffix.
|
||||
let mut reader = BufReader::new(file.take(byte_count));
|
||||
let mut line_bytes = Vec::new();
|
||||
// Only project the newline-terminated prefix; leave a trailing partial record for the next
|
||||
// pass.
|
||||
let mut projections = Vec::new();
|
||||
let mut next_ordinal = expected_ordinal;
|
||||
let mut next_offset = start_offset;
|
||||
let mut line_start_offset = start_offset;
|
||||
for line_bytes in bytes.split_inclusive(|byte| *byte == b'\n') {
|
||||
loop {
|
||||
if is_cancelled() {
|
||||
return Err(thread_store_io_error(std::io::Error::other(
|
||||
"rollout read cancelled",
|
||||
)));
|
||||
}
|
||||
line_bytes.clear();
|
||||
if reader
|
||||
.read_until(b'\n', &mut line_bytes)
|
||||
.map_err(thread_store_io_error)?
|
||||
== 0
|
||||
{
|
||||
break;
|
||||
}
|
||||
if !line_bytes.ends_with(b"\n") {
|
||||
break;
|
||||
}
|
||||
@@ -171,7 +178,7 @@ fn read_projection_steps_sync(
|
||||
line_start_offset = line_end_offset;
|
||||
continue;
|
||||
}
|
||||
let value = match serde_json::from_slice::<serde_json::Value>(line_bytes) {
|
||||
let value = match serde_json::from_slice::<serde_json::Value>(&line_bytes) {
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -332,6 +339,12 @@ fn read_projection_steps_sync(
|
||||
next_offset = line_end_offset;
|
||||
line_start_offset = line_end_offset;
|
||||
}
|
||||
if reader.get_ref().limit() != 0 {
|
||||
return Err(thread_store_io_error(std::io::Error::new(
|
||||
std::io::ErrorKind::UnexpectedEof,
|
||||
"durable rollout shrank during projection",
|
||||
)));
|
||||
}
|
||||
Ok((projections, next_offset))
|
||||
}
|
||||
|
||||
|
||||
@@ -2204,6 +2204,7 @@ async fn synchronized_catch_up_does_not_replay_old_rows() {
|
||||
assert_eq!(status, "sentinel");
|
||||
}
|
||||
|
||||
/// Large records keep their byte positions and remain pending until their newline is durable.
|
||||
#[tokio::test]
|
||||
async fn catch_up_preserves_trailing_partial_line_boundaries() {
|
||||
let home = TempDir::new().expect("temp dir");
|
||||
@@ -2221,8 +2222,12 @@ async fn catch_up_preserves_trailing_partial_line_boundaries() {
|
||||
.await
|
||||
.expect("open thread history db");
|
||||
let before = projection_state(&pool, thread_id).await;
|
||||
let complete_line = rollout_line(Some(1), turn_started("turn-1"));
|
||||
let partial_line = rollout_line(Some(2), turn_completed("turn-1"));
|
||||
let padding = " ".repeat(64 * 1024);
|
||||
let complete_line = format!("{}{padding}", rollout_line(Some(1), turn_started("turn-1")));
|
||||
let partial_line = format!(
|
||||
"{}{padding}",
|
||||
rollout_line(Some(2), turn_completed("turn-1"))
|
||||
);
|
||||
let complete_suffix = format!("{complete_line}\n");
|
||||
let rollout_path = store
|
||||
.live_rollout_path(thread_id)
|
||||
@@ -2471,7 +2476,7 @@ async fn blank_and_rejected_rollout_lines_advance_projection() {
|
||||
.expect("queue valid retry");
|
||||
recorder.flush().await.expect("flush valid retry");
|
||||
|
||||
// Replay this small file, cancelling between records after its single read completes.
|
||||
// Replay this small file, cancelling between records.
|
||||
let checks = Cell::new(0);
|
||||
assert!(
|
||||
super::read_projection_steps_sync(
|
||||
|
||||
Reference in New Issue
Block a user