Search compressed rollouts concurrently

codex-rs/rollout/Cargo.toml

@@ -27,6 +27,7 @@ codex-utils-absolute-path = { workspace = true }
 codex-utils-output-truncation = { workspace = true }
 codex-utils-path = { workspace = true }
 codex-utils-string = { workspace = true }
+futures = { workspace = true, features = ["alloc"] }
 memchr = { workspace = true }
 regex = { workspace = true }
 serde = { workspace = true, features = ["derive"] }

codex-rs/rollout/src/search.rs

@@ -9,6 +9,7 @@ use codex_protocol::models::ContentItem;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::strip_user_message_prefix;
+use futures::future::join_all;
 use regex::Regex;
 use regex::RegexBuilder;
 use tokio::process::Command;
@@ -223,8 +224,8 @@ async fn scan_compressed_rollout_matches(
     search_term: &str,
 ) -> io::Result<RolloutSearchMatches> {
     let mut matches = HashMap::new();
+    let mut paths = Vec::new();
     let mut dirs = vec![root.to_path_buf()];
-    let mut patterns = None;
 
     while let Some(dir) = dirs.pop() {
         let mut entries = match tokio::fs::read_dir(dir).await {
@@ -248,26 +249,39 @@ async fn scan_compressed_rollout_matches(
             if !rollout_file.is_compressed() {
                 continue;
             }
+            paths.push(rollout_file.into_path());
+        }
+    }
+
+    const COMPRESSED_SEARCH_WORKERS: usize = 4;
+    let scan = async |paths: &[PathBuf]| {
+        let mut matches = HashMap::new();
+        let mut patterns = None;
+        for path in paths {
             if let Some(snippet) =
-                first_content_match_snippet(rollout_file.path(), search_term, &mut patterns)
+                first_content_match_snippet(path.as_path(), search_term, &mut patterns)
                     .await
                     .unwrap_or_else(|err| {
                         tracing::warn!(
-                            path = %rollout_file.path().display(),
+                            path = %path.as_path().display(),
                             %err,
                             "Failed to search compressed rollout"
                         );
                         None
                     })
             {
                 matches.insert(
-                    compression::plain_rollout_path(rollout_file.path()),
+                    compression::plain_rollout_path(path.as_path()),
                     Some(snippet),
                 );
             }
         }
+        matches
+    };
+    let chunk_size = paths.len().div_ceil(COMPRESSED_SEARCH_WORKERS).max(1);
+    for batch in join_all(paths.chunks(chunk_size).map(scan)).await {
+        matches.extend(batch);
     }
-
     Ok(matches)
 }
 

codex-rs/rollout/src/search_tests.rs

@@ -7,8 +7,19 @@ use tempfile::TempDir;
 #[tokio::test]
 async fn compressed_search_preserves_error_handling() -> io::Result<()> {
     let root = TempDir::new()?;
-    let path = root.path().join("rollout-0001.jsonl.zst");
-    std::fs::write(&path, zstd::stream::encode_all(&b""[..], 3)?)?;
+    let line = br#"{"timestamp":"2025-01-01T00:00:00Z","type":"event_msg","payload":{"type":"agent_message","message":"needle"}}"#;
+    let compressed = zstd::stream::encode_all(&line[..], 3)?;
+    let mut expected = RolloutSearchMatches::new();
+    for index in 0..7 {
+        let path = root.path().join(format!("rollout-{index:04}.jsonl"));
+        std::fs::write(path.with_extension("jsonl.zst"), &compressed)?;
+        expected.insert(path, Some("needle".to_string()));
+    }
+    std::fs::write(root.path().join("rollout-broken.jsonl.zst"), b"not zstd")?;
+    assert_eq!(
+        scan_compressed_rollout_matches(root.path(), "needle").await?,
+        expected
+    );
     let query = "a".repeat(500_000);
 
     assert!(case_insensitive_literal_regex(&query).is_err());