diff --git a/codex-rs/rollout/src/compression.rs b/codex-rs/rollout/src/compression.rs index ba7eb24c5b..47c86520f4 100644 --- a/codex-rs/rollout/src/compression.rs +++ b/codex-rs/rollout/src/compression.rs @@ -16,6 +16,7 @@ use std::os::unix::fs::OpenOptionsExt; #[cfg(unix)] use std::os::unix::fs::PermissionsExt; +mod blocking_reader; mod error_metrics; mod read_metrics; @@ -298,6 +299,30 @@ impl RolloutLineReader { } result } + + /// Keeps a compressed scan on one worker while retaining this reader's format and I/O metrics. + pub(crate) async fn find_map( + mut self, + mut find: impl FnMut(&str) -> Option + Send + 'static, + ) -> io::Result> { + let RolloutLineReaderInner::Blocking(Some(reader)) = self.inner else { + while let Some(line) = self.next_line().await? { + if let Some(found) = find(&line) { + return Ok(Some(found)); + } + } + return Ok(None); + }; + blocking_reader::scan_lines(reader, self.metrics, move |lines| { + for line in lines { + if let Some(found) = find(&line?) { + return Ok(Some(found)); + } + } + Ok(None) + }) + .await + } } type BlockingLineReader = std::io::Lines>>; diff --git a/codex-rs/rollout/src/compression/blocking_reader.rs b/codex-rs/rollout/src/compression/blocking_reader.rs new file mode 100644 index 0000000000..7b65a78d59 --- /dev/null +++ b/codex-rs/rollout/src/compression/blocking_reader.rs @@ -0,0 +1,93 @@ +//! Runs a rollout scan on one blocking worker, with cancellation and one reader observation. +//! Consumers may stop early; canceled reads do not contribute an EOF, failure, or duration. + +use std::io; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::time::Duration; +use std::time::Instant; + +use super::BlockingLineReader; +use super::ReadFailureSource; +use super::ReadMetrics; + +/// Runs a consumer over measured lines, retaining metrics if the caller drops the future. +pub(super) async fn scan_lines( + lines: BlockingLineReader, + metrics: ReadMetrics, + scan: F, +) -> io::Result +where + T: Send + 'static, + F: FnOnce(&mut dyn Iterator>) -> io::Result + Send + 'static, +{ + // Include the initial queue wait, as the per-line reader does, but exclude consumer work. + let queued_at = Instant::now(); + let metrics = Arc::new(Mutex::new(metrics)); + let worker_metrics = Arc::clone(&metrics); + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); + let span = tracing::Span::current(); + let task = tokio::task::spawn_blocking(move || { + let queue_duration = queued_at.elapsed(); + span.in_scope(|| { + let mut metrics = worker_metrics + .lock() + .unwrap_or_else(PoisonError::into_inner); + scan(&mut MeasuredLines { + lines, + metrics: &mut metrics, + stop, + queue_duration, + }) + }) + }); + let result = task.await.map_err(|err| { + let err = io::Error::other(err); + metrics + .lock() + .unwrap_or_else(PoisonError::into_inner) + .failed("read", ReadFailureSource::TaskJoin, &err); + err + }); + drop(keep_running); + result? +} + +struct MeasuredLines<'a> { + lines: BlockingLineReader, + metrics: &'a mut ReadMetrics, + stop: tokio::sync::oneshot::Sender<()>, + queue_duration: Duration, +} + +impl Iterator for MeasuredLines<'_> { + type Item = io::Result; + + fn next(&mut self) -> Option { + if self.stop.is_closed() { + return None; + } + let started_at = Instant::now(); + let line = self.lines.next(); + if self.stop.is_closed() { + return None; + } + let metrics = &mut self.metrics; + metrics.duration = metrics + .duration + .saturating_add(std::mem::take(&mut self.queue_duration)) + .saturating_add(started_at.elapsed()); + metrics.reached_eof = line.is_none(); + match &line { + Some(Ok(_)) => metrics.read_any_line = true, + Some(Err(err)) => metrics.failed("read", ReadFailureSource::Stream, err), + None => {} + } + line + } +} + +#[cfg(test)] +#[path = "blocking_reader_tests.rs"] +mod tests; diff --git a/codex-rs/rollout/src/compression/blocking_reader_tests.rs b/codex-rs/rollout/src/compression/blocking_reader_tests.rs new file mode 100644 index 0000000000..d1e9ce5356 --- /dev/null +++ b/codex-rs/rollout/src/compression/blocking_reader_tests.rs @@ -0,0 +1,71 @@ +//! Exercises cancellation while the worker is blocked inside a read. + +use std::io; +use std::io::BufRead; +use std::io::Read; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use pretty_assertions::assert_eq; + +use super::ReadMetrics; +use super::scan_lines; + +struct PausedRead { + contents: io::Cursor<&'static [u8]>, + started: Option>, + resume: std::sync::mpsc::Receiver<()>, + finished: Option>, +} + +impl Read for PausedRead { + fn read(&mut self, buffer: &mut [u8]) -> io::Result { + if let Some(started) = self.started.take() { + let _ = started.send(()); + self.resume.recv().map_err(io::Error::other)?; + } + self.contents.read(buffer) + } +} + +impl Drop for PausedRead { + fn drop(&mut self) { + if let Some(finished) = self.finished.take() { + let _ = finished.send(()); + } + } +} + +#[tokio::test] +async fn cancellation_during_read_does_not_deliver_the_record() -> anyhow::Result<()> { + let (started, waiting) = tokio::sync::oneshot::channel(); + let (resume, paused) = std::sync::mpsc::channel(); + let (finished, done) = tokio::sync::oneshot::channel(); + let reader = Box::new(PausedRead { + contents: io::Cursor::new(b"first\n".as_slice()), + started: Some(started), + resume: paused, + finished: Some(finished), + }) as Box; + let calls = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&calls); + let task = tokio::spawn(scan_lines( + io::BufReader::new(reader).lines(), + ReadMetrics::default(), + move |lines| { + for line in lines { + line?; + seen.fetch_add(1, Ordering::SeqCst); + } + Ok(()) + }, + )); + waiting.await?; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + resume.send(())?; + done.await?; + assert_eq!(calls.load(Ordering::SeqCst), 0); + Ok(()) +} diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index 71843ce373..b795c9c451 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -26,6 +26,7 @@ use crate::RolloutLine; use crate::RolloutRecorder; use crate::RolloutRecorderParams; use crate::append_rollout_item_to_path; +use crate::first_rollout_content_match_snippet; use crate::read_session_meta_line; use crate::search_rollout_matches; @@ -263,6 +264,117 @@ async fn search_rollout_matches_uses_logical_path_for_compressed_rollout() -> an Ok(()) } +#[tokio::test] +async fn compressed_search_preserves_first_visible_snippet_and_no_match() -> anyhow::Result<()> { + let home = TempDir::new()?; + let uuid = Uuid::from_u128(22); + 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, "an unrelated first message")?; + let mut contents = fs::read(&path)?; + contents.extend_from_slice(b"{malformed needle raw-only}\n"); + fs::write(&path, contents)?; + for message in [ + "the first Needle [x] is visible", + "the second needle is visible", + ] { + append_rollout_item_to_path( + &path, + &RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: message.to_string(), + ..Default::default() + })), + ) + .await?; + } + + for compressed in [false, true] { + if compressed { + compress_now(&path)?; + } + for (query, expected) in [ + ("NEEDLE", Some("the first Needle [x] is visible")), + ("[x]", Some("the first Needle [x] is visible")), + ("raw-only", None), + ("absent", None), + ] { + assert_eq!( + first_rollout_content_match_snippet(&path, query).await?, + expected.map(str::to_string), + "compressed={compressed}, query={query}" + ); + } + } + Ok(()) +} + +#[tokio::test] +async fn compressed_search_only_reports_stream_errors_before_a_match() -> anyhow::Result<()> { + let home = TempDir::new()?; + let uuid = Uuid::from_u128(23); + 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, "needle before unreadable tail")?; + let mut contents = fs::read(&path)?; + contents.extend_from_slice(b"\xff\n"); + fs::write(&path, contents)?; + compress_now(&path)?; + assert_eq!( + first_rollout_content_match_snippet(&path, "needle").await?, + Some("needle before unreadable tail".to_owned()) + ); + assert_eq!( + first_rollout_content_match_snippet(&path, "absent") + .await + .unwrap_err() + .kind(), + io::ErrorKind::InvalidData + ); + Ok(()) +} + +#[tokio::test] +async fn compressed_search_stops_between_lines_when_cancelled() -> anyhow::Result<()> { + struct NotifyOnDrop(Option>); + impl Drop for NotifyOnDrop { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } + } + let home = TempDir::new()?; + let uuid = Uuid::from_u128(24); + 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, "multiple lines are available")?; + compress_now(&path)?; + let reader = open_rollout_line_reader(&path).await?; + let calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let seen = calls.clone(); + let (entered, waiting) = tokio::sync::oneshot::channel(); + let (release, blocked) = std::sync::mpsc::channel(); + let (finished, done) = tokio::sync::oneshot::channel(); + let guard = NotifyOnDrop(Some(finished)); + let mut entered = Some(entered); + let handle = tokio::spawn(reader.find_map(move |_| { + let _guard = &guard; + seen.fetch_add(1, Ordering::SeqCst); + if let Some(entered) = entered.take() { + let _ = entered.send(()); + let _ = blocked.recv(); + } + None::<()> + })); + waiting.await?; + handle.abort(); + assert!(handle.await.unwrap_err().is_cancelled()); + release.send(())?; + tokio::time::timeout(Duration::from_secs(5), done).await??; + assert_eq!(calls.load(Ordering::SeqCst), 1); + Ok(()) +} + #[tokio::test] async fn worker_compresses_old_active_and_archived_rollouts() -> anyhow::Result<()> { let home = TempDir::new()?; diff --git a/codex-rs/rollout/src/search.rs b/codex-rs/rollout/src/search.rs index b9c2bc7343..b0f1a68f87 100644 --- a/codex-rs/rollout/src/search.rs +++ b/codex-rs/rollout/src/search.rs @@ -187,17 +187,18 @@ pub async fn first_rollout_content_match_snippet( path: &Path, search_term: &str, ) -> io::Result> { - let mut lines = compression::open_rollout_line_reader(path).await?; + let lines = compression::open_rollout_line_reader(path).await?; let json_search_term = case_insensitive_literal_regex(json_escaped_search_term(search_term)?)?; let search_term = case_insensitive_literal_regex(search_term)?; - while let Some(line) = lines.next_line().await? { - if json_search_term.is_match(line.as_str()) - && let Some(snippet) = content_match_snippet(line.as_str(), &search_term) - { - return Ok(Some(snippet)); - } - } - Ok(None) + lines + .find_map(move |line| { + if json_search_term.is_match(line) { + content_match_snippet(line, &search_term) + } else { + None + } + }) + .await } async fn scan_compressed_rollout_matches(