From 8ad185c741b901092888df3fcf7c2955d986ea90 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 29 May 2026 11:51:30 +0100 Subject: [PATCH] make it stronger --- codex-rs/rollout/src/compression.rs | 60 ++++++++++++++++------- codex-rs/rollout/src/compression_tests.rs | 22 +++++++++ codex-rs/rollout/src/list.rs | 11 ++++- 3 files changed, 75 insertions(+), 18 deletions(-) diff --git a/codex-rs/rollout/src/compression.rs b/codex-rs/rollout/src/compression.rs index 9932175649..17da4b05c7 100644 --- a/codex-rs/rollout/src/compression.rs +++ b/codex-rs/rollout/src/compression.rs @@ -3,6 +3,7 @@ use std::fs::File; use std::fs::FileTimes; use std::fs::Permissions; use std::io; +use std::io::BufRead; use std::io::Read; use std::io::Write; use std::path::Path; @@ -30,6 +31,7 @@ const TEMP_SUFFIX: &str = ".tmp"; const COMPRESSION_LEVEL: i32 = 3; const MIN_ROLLOUT_AGE: Duration = Duration::from_secs(7 * 24 * 60 * 60); const GLOBAL_LOCK_STALE_AFTER: Duration = Duration::from_secs(6 * 60 * 60); +const TEMP_FILE_STALE_AFTER: Duration = GLOBAL_LOCK_STALE_AFTER; const WORKER_MAX_RUNTIME: Duration = Duration::from_secs(5 * 60 * 60); const LOCK_FILE_NAME: &str = "rollout-compression.lock"; static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); @@ -103,7 +105,13 @@ pub async fn open_rollout_line_reader(path: &Path) -> io::Result Ok(reader), Err(err) if err.kind() == io::ErrorKind::NotFound => { - open_rollout_line_reader_alternate(path).await + match open_rollout_line_reader_once(path).await { + Ok(reader) => Ok(reader), + Err(err) if err.kind() == io::ErrorKind::NotFound => { + open_rollout_line_reader_alternate(path).await + } + Err(err) => Err(err), + } } Err(err) => Err(err), } @@ -209,7 +217,7 @@ pub struct RolloutLineReader { enum RolloutLineReaderInner { Plain(tokio::io::Lines>), - Memory(std::io::Lines>), + Blocking(Option), } impl RolloutLineReader { @@ -217,11 +225,23 @@ impl RolloutLineReader { pub async fn next_line(&mut self) -> io::Result> { match &mut self.inner { RolloutLineReaderInner::Plain(lines) => lines.next_line().await, - RolloutLineReaderInner::Memory(lines) => lines.next().transpose(), + RolloutLineReaderInner::Blocking(slot) => { + let Some(mut reader) = slot.take() else { + return Err(io::Error::other("compressed rollout reader is busy")); + }; + let (line, reader) = + tokio::task::spawn_blocking(move || (reader.next().transpose(), reader)) + .await + .map_err(io::Error::other)?; + *slot = Some(reader); + line + } } } } +type BlockingLineReader = std::io::Lines>>; + #[derive(Default)] struct CompressionStats { scanned: usize, @@ -385,7 +405,7 @@ fn compress_rollout_if_cold_blocking(path: &Path) -> io::Result { std::fs::remove_file(path)?; Ok(true) })(); - if result.is_err() { + if !matches!(result, Ok(true)) { let _ = std::fs::remove_file(temp_path.as_path()); } result @@ -524,6 +544,16 @@ async fn cleanup_stale_temps_in_root(root: &Path) -> io::Result<()> { .and_then(OsStr::to_str) .is_some_and(|name| name.ends_with(TEMP_SUFFIX)) { + let stale = entry + .metadata() + .await + .ok() + .and_then(|metadata| metadata.modified().ok()) + .and_then(|modified| SystemTime::now().duration_since(modified).ok()) + .is_some_and(|age| age >= TEMP_FILE_STALE_AFTER); + if !stale { + continue; + } match tokio::fs::remove_file(path.as_path()).await { Ok(()) => {} Err(err) if err.kind() == io::ErrorKind::NotFound => {} @@ -584,20 +614,16 @@ async fn open_rollout_line_reader_alternate(path: &Path) -> io::Result io::Result { - let text = tokio::task::spawn_blocking(move || decode_zstd_to_string(path.as_path())) - .await - .map_err(io::Error::other)??; - Ok(RolloutLineReader { - inner: RolloutLineReaderInner::Memory(std::io::BufRead::lines(io::Cursor::new(text))), + let reader = tokio::task::spawn_blocking(move || { + let input = File::open(path.as_path())?; + let decoder = zstd::stream::read::Decoder::new(input)?; + Ok::<_, io::Error>(io::BufReader::new(Box::new(decoder) as Box).lines()) + }) + .await + .map_err(io::Error::other)??; + Ok(RolloutLineReader { + inner: RolloutLineReaderInner::Blocking(Some(reader)), }) -} - -fn decode_zstd_to_string(path: &Path) -> io::Result { - let input = File::open(path)?; - let mut decoder = zstd::stream::read::Decoder::new(input)?; - let mut text = String::new(); - decoder.read_to_string(&mut text)?; - Ok(text) } fn temp_path_for(path: &Path, operation: &str) -> PathBuf { diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index 10028556fd..c05c5cd7e6 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -94,6 +94,10 @@ async fn worker_compresses_old_active_and_archived_rollouts() -> anyhow::Result< let stale_temp = active_path.with_file_name("rollout-stale.jsonl.zst.tmp"); fs::write(&stale_temp, "stale temp")?; + set_old_mtime(&stale_temp)?; + + let fresh_temp = active_path.with_file_name("rollout-fresh.jsonl.zst.tmp"); + fs::write(&fresh_temp, "fresh temp")?; run_rollout_compression_worker(home.path().to_path_buf()).await?; @@ -104,6 +108,7 @@ async fn worker_compresses_old_active_and_archived_rollouts() -> anyhow::Result< assert!(fresh_path.exists()); assert!(!compressed_rollout_path(&fresh_path).exists()); assert!(!stale_temp.exists()); + assert!(fresh_temp.exists()); Ok(()) } @@ -188,6 +193,23 @@ async fn find_thread_path_by_id_handles_compressed_rollout_filenames() -> anyhow Ok(()) } +#[tokio::test] +async fn find_thread_path_by_id_ignores_compression_temp_matches() -> anyhow::Result<()> { + let home = TempDir::new()?; + let uuid = Uuid::from_u128(9); + let thread_id = ThreadId::from_string(&uuid.to_string())?; + let temp_path = rollout_path(home.path(), "2025-01-03T12-00-00", uuid).with_file_name(format!( + "rollout-2025-01-03T12-00-00-{uuid}.jsonl.zst.compress.1.0.tmp" + )); + write_rollout(&temp_path, thread_id, "temporary file should not resolve")?; + + assert_eq!( + crate::find_thread_path_by_id_str(home.path(), &uuid.to_string(), None).await?, + None + ); + Ok(()) +} + fn rollout_path(home: &std::path::Path, ts: &str, uuid: Uuid) -> std::path::PathBuf { home.join("sessions/2025/01/03") .join(format!("rollout-{ts}-{uuid}.jsonl")) diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index cc13ecd884..0091bbce48 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1372,7 +1372,16 @@ async fn find_thread_path_by_id_str_in_subdir( ) .map_err(|e| io::Error::other(format!("file search failed: {e}")))?; - let found = match results.matches.into_iter().next().map(|m| m.full_path()) { + let found = match results + .matches + .into_iter() + .map(|m| m.full_path()) + .find(|path| { + path.file_name() + .and_then(OsStr::to_str) + .is_some_and(|name| !name.ends_with(".tmp")) + && !compression::should_skip_compressed_sibling(path) + }) { Some(path) => Some(path), None => find_rollout_path_by_id_from_filenames(root.as_path(), id_str).await?, };