Keep compressed rollout snippet searches on one blocking worker (#49692)

## Why

Compressed snippet searches schedule a blocking task for every line read. Keeping the scan on one worker avoids repeated task scheduling.

## What changed

Use `RolloutLineReader::find_map` to scan compressed rollouts on a single blocking worker, stopping at the first matching snippet. Preserve read metrics and stop delivering records when the search is canceled, including cancellation during a read.

## Testing

Add tests for first-match snippets, case-insensitive literal matching, malformed records, missing matches, stream errors before a match, and cancellation between lines or during a read.

GitOrigin-RevId: 295206b1a00e861bc44606bd843f6ed8c3e6aae1

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<T: Send + 'static>(
+        mut self,
+        mut find: impl FnMut(&str) -> Option<T> + Send + 'static,
+    ) -> io::Result<Option<T>> {
+        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<std::io::BufReader<Box<dyn Read + Send>>>;

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<T, F>(
+    lines: BlockingLineReader,
+    metrics: ReadMetrics,
+    scan: F,
+) -> io::Result<T>
+where
+    T: Send + 'static,
+    F: FnOnce(&mut dyn Iterator<Item = io::Result<String>>) -> io::Result<T> + 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<String>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        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;

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<tokio::sync::oneshot::Sender<()>>,
+    resume: std::sync::mpsc::Receiver<()>,
+    finished: Option<tokio::sync::oneshot::Sender<()>>,
+}
+
+impl Read for PausedRead {
+    fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
+        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<dyn Read + Send>;
+    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(())
+}

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<tokio::sync::oneshot::Sender<()>>);
+    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()?;

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<Option<String>> {
-    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(