From a9bc7bebfa439e608543f45e9fc3415895c39153 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 6 Oct 2026 20:15:56 +0000 Subject: [PATCH] Load rollout history on a single blocking worker (#51499) ## What changed Move full rollout history reading and parsing into one blocking worker for both plain `.jsonl` and compressed `.jsonl.zst` files. Use the shared cancellation-aware line iterator so canceled loads stop before processing another record, while preserving rollout items, thread IDs, and parse-error counts. ## Testing Add regression tests for cancellation between records and equivalent history-loading results across plain and compressed files, including Unicode, CRLF line endings, malformed records, empty files, and invalid UTF-8. GitOrigin-RevId: 65e1bd7ebcdbbd0ab5f560e73b2c7f074f627f8f --- codex-rs/rollout/src/compression.rs | 1 + .../src/compression/blocking_reader.rs | 25 ++++ .../src/compression/blocking_reader_tests.rs | 30 ++++- codex-rs/rollout/src/compression_tests.rs | 43 +++++++ codex-rs/rollout/src/recorder.rs | 107 +++++++++--------- 5 files changed, 153 insertions(+), 53 deletions(-) diff --git a/codex-rs/rollout/src/compression.rs b/codex-rs/rollout/src/compression.rs index 3065b05373..dc4ffe1df5 100644 --- a/codex-rs/rollout/src/compression.rs +++ b/codex-rs/rollout/src/compression.rs @@ -21,6 +21,7 @@ mod error_metrics; mod path_metadata; mod read_metrics; +pub(crate) use blocking_reader::read_rollout_lines; use error_metrics::FailureMetric; pub(crate) use path_metadata::existing_rollout_with_metadata_sync; use read_metrics::ReadFailureSource; diff --git a/codex-rs/rollout/src/compression/blocking_reader.rs b/codex-rs/rollout/src/compression/blocking_reader.rs index 7b65a78d59..909e1e6518 100644 --- a/codex-rs/rollout/src/compression/blocking_reader.rs +++ b/codex-rs/rollout/src/compression/blocking_reader.rs @@ -2,6 +2,8 @@ //! Consumers may stop early; canceled reads do not contribute an EOF, failure, or duration. use std::io; +use std::io::BufRead; +use std::path::Path; use std::sync::Arc; use std::sync::Mutex; use std::sync::PoisonError; @@ -11,6 +13,29 @@ use std::time::Instant; use super::BlockingLineReader; use super::ReadFailureSource; use super::ReadMetrics; +use super::RolloutLineReaderInner; +use super::open_rollout_line_reader; + +/// Opens a fresh reader and processes its lines on the shared blocking worker. +pub(crate) async fn read_rollout_lines(path: &Path, read: F) -> io::Result +where + T: Send + 'static, + F: FnOnce(&mut dyn Iterator>) -> io::Result + Send + 'static, +{ + let reader = open_rollout_line_reader(path).await?; + let lines: BlockingLineReader = match reader.inner { + RolloutLineReaderInner::Plain(lines) => { + // The fresh reader has no buffered data to preserve when converting the file. + let file = lines.into_inner().into_inner().into_std().await; + io::BufReader::new(Box::new(file) as Box).lines() + } + RolloutLineReaderInner::Blocking(Some(lines)) => lines, + RolloutLineReaderInner::Blocking(None) => { + return Err(io::Error::other("compressed rollout reader is busy")); + } + }; + scan_lines(lines, reader.metrics, read).await +} /// Runs a consumer over measured lines, retaining metrics if the caller drops the future. pub(super) async fn scan_lines( diff --git a/codex-rs/rollout/src/compression/blocking_reader_tests.rs b/codex-rs/rollout/src/compression/blocking_reader_tests.rs index d1e9ce5356..1af510ba2d 100644 --- a/codex-rs/rollout/src/compression/blocking_reader_tests.rs +++ b/codex-rs/rollout/src/compression/blocking_reader_tests.rs @@ -1,4 +1,4 @@ -//! Exercises cancellation while the worker is blocked inside a read. +//! Exercises cancellation during reads and between records. use std::io; use std::io::BufRead; @@ -8,8 +8,10 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use pretty_assertions::assert_eq; +use tempfile::TempDir; use super::ReadMetrics; +use super::read_rollout_lines; use super::scan_lines; struct PausedRead { @@ -69,3 +71,29 @@ async fn cancellation_during_read_does_not_deliver_the_record() -> anyhow::Resul assert_eq!(calls.load(Ordering::SeqCst), 0); Ok(()) } + +#[tokio::test] +async fn cancelling_full_scan_stops_before_the_next_record() -> anyhow::Result<()> { + let home = TempDir::new()?; + let path = home.path().join("rollout.jsonl"); + std::fs::write(&path, "first\nsecond\n")?; + let (started_tx, started_rx) = tokio::sync::oneshot::channel(); + let (continue_tx, continue_rx) = std::sync::mpsc::channel(); + let (finished_tx, finished_rx) = tokio::sync::oneshot::channel(); + let task = tokio::spawn(async move { + read_rollout_lines(&path, move |lines| { + assert_eq!(lines.next().transpose()?, Some("first".to_string())); + let _ = started_tx.send(()); + continue_rx.recv().map_err(io::Error::other)?; + let _ = finished_tx.send(lines.next().transpose()); + Ok(()) + }) + .await + }); + started_rx.await?; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + continue_tx.send(())?; + assert_eq!(finished_rx.await??, None); + Ok(()) +} diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index b795c9c451..a4171801dd 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -50,6 +50,49 @@ async fn load_rollout_items_reads_compressed_rollout() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn full_history_load_preserves_records_and_errors_across_representations() +-> anyhow::Result<()> { + let home = TempDir::new()?; + let uuid = Uuid::from_u128(19); + let thread_id = ThreadId::from_string(&uuid.to_string())?; + let path = rollout_path(home.path(), "2025-01-03T12-00-00", uuid); + write_rollout(&path, thread_id, "message with unicode: π")?; + let original = fs::read_to_string(&path)?; + let expected = RolloutRecorder::load_rollout_items(&path).await?; + let contents = format!( + " \r\ninvalid JSON\n{{}}\n{}", + original.trim_end().replace('\n', "\r\n") + ); + fs::write(&path, &contents)?; + let expected = (expected.0, expected.1, 2); + assert_eq!( + serde_json::to_value(RolloutRecorder::load_rollout_items(&path).await?)?, + serde_json::to_value(&expected)? + ); + compress_now(&path)?; + assert_eq!( + serde_json::to_value(RolloutRecorder::load_rollout_items(&path).await?)?, + serde_json::to_value(&expected)? + ); + + for bytes in [b" \r\n\n".as_slice(), b"valid utf8\n\xff".as_slice()] { + fs::write(&path, bytes)?; + let plain = RolloutRecorder::load_rollout_items(&path) + .await + .unwrap_err(); + compress_now(&path)?; + let compressed = RolloutRecorder::load_rollout_items(&path) + .await + .unwrap_err(); + assert_eq!( + (compressed.kind(), compressed.to_string()), + (plain.kind(), plain.to_string()) + ); + } + Ok(()) +} + #[tokio::test] async fn read_session_meta_line_stops_before_invalid_utf8_tail() -> anyhow::Result<()> { let home = TempDir::new()?; diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 716d8f11c4..f99f2c5a89 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -1106,65 +1106,68 @@ impl RolloutRecorder { path: &Path, ) -> std::io::Result<(Vec, Option, usize)> { trace!("Resuming rollout from {path:?}"); - let mut items: Vec = Vec::new(); - let mut thread_id: Option = None; - let mut parse_errors = 0usize; - let mut reader = compression::open_rollout_line_reader(path).await?; - let mut saw_non_empty_line = false; - while let Some(line) = reader.next_line().await? { - if line.trim().is_empty() { - continue; - } - saw_non_empty_line = true; - let mut value: Value = match serde_json::from_str(&line) { - Ok(value) => value, - Err(e) => { - warn!("failed to parse line as JSON: {line:?}, error: {e}"); - parse_errors = parse_errors.saturating_add(1); + compression::read_rollout_lines(path, |reader| { + let mut items: Vec = Vec::new(); + let mut thread_id: Option = None; + let mut parse_errors = 0usize; + let mut saw_non_empty_line = false; + for line in reader { + let line = line?; + if line.trim().is_empty() { continue; } - }; - if strip_legacy_ghost_snapshot_rollout_line(&mut value) { - trace!("skipping legacy ghost_snapshot rollout line"); - continue; - } - if thread_id.is_none() { - // The first SessionMeta belongs to this rollout. Later SessionMeta lines - // can be copied from fork history, so only validate unknown history modes - // before we have parsed the rollout's own SessionMeta. - reject_unknown_thread_history_mode(&value)?; - } - - let rollout_line = match crate::decode_rollout_line(value) { - Ok(rollout_line) => rollout_line, - Err(e) => { - trace!("failed to parse rollout line: {e}"); - parse_errors = parse_errors.saturating_add(1); + saw_non_empty_line = true; + let mut value: Value = match serde_json::from_str(&line) { + Ok(value) => value, + Err(e) => { + warn!("failed to parse line as JSON: {line:?}, error: {e}"); + parse_errors = parse_errors.saturating_add(1); + continue; + } + }; + if strip_legacy_ghost_snapshot_rollout_line(&mut value) { + trace!("skipping legacy ghost_snapshot rollout line"); continue; } - }; + if thread_id.is_none() { + // The first SessionMeta belongs to this rollout. Later SessionMeta lines + // can be copied from fork history, so only validate unknown history modes + // before we have parsed the rollout's own SessionMeta. + reject_unknown_thread_history_mode(&value)?; + } - let item = rollout_line.item; - // Use the FIRST SessionMeta encountered in the file as the canonical - // thread id and main session information. Keep all items intact. - if thread_id.is_none() - && let RolloutItem::SessionMeta(session_meta_line) = &item - { - thread_id = Some(session_meta_line.meta.id); + let rollout_line = match crate::decode_rollout_line(value) { + Ok(rollout_line) => rollout_line, + Err(e) => { + trace!("failed to parse rollout line: {e}"); + parse_errors = parse_errors.saturating_add(1); + continue; + } + }; + + let item = rollout_line.item; + // Use the FIRST SessionMeta encountered in the file as the canonical + // thread id and main session information. Keep all items intact. + if thread_id.is_none() + && let RolloutItem::SessionMeta(session_meta_line) = &item + { + thread_id = Some(session_meta_line.meta.id); + } + items.push(item); + } + if !saw_non_empty_line { + return Err(IoError::other("empty session file")); } - items.push(item); - } - if !saw_non_empty_line { - return Err(IoError::other("empty session file")); - } - tracing::debug!( - "Resumed rollout with {} items, thread ID: {:?}, parse errors: {}", - items.len(), - thread_id, - parse_errors, - ); - Ok((items, thread_id, parse_errors)) + tracing::debug!( + "Resumed rollout with {} items, thread ID: {:?}, parse errors: {}", + items.len(), + thread_id, + parse_errors, + ); + Ok((items, thread_id, parse_errors)) + }) + .await } pub async fn get_rollout_history(path: &Path) -> std::io::Result {