From 7b2a18a6c7cef501932380ccb3a4e886d8b0c79b Mon Sep 17 00:00:00 2001 From: anais Date: Mon, 15 Jun 2026 21:49:54 -0700 Subject: [PATCH] speed up rollout history loading --- codex-rs/rollout/src/list.rs | 124 +++++++++++---- codex-rs/rollout/src/recorder.rs | 265 +++++++++++++++++++++++++------ 2 files changed, 305 insertions(+), 84 deletions(-) diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 845145aaef..93e359b0a2 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1,7 +1,12 @@ #![allow(warnings, clippy::all)] +use chrono::DateTime; +use chrono::Datelike; +use chrono::Timelike; +use chrono::Utc; use codex_utils_path as path_utils; use std::cmp::Reverse; +use std::collections::BinaryHeap; use std::ffi::OsStr; use std::io; use std::num::NonZero; @@ -23,7 +28,6 @@ use crate::state_db; use codex_file_search as file_search; use codex_protocol::ThreadId; use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::USER_MESSAGE_BEGIN; @@ -928,20 +932,62 @@ pub(crate) fn parse_timestamp_uuid_from_filename(name: &str) -> Option<(OffsetDa // Expected: rollout-YYYY-MM-DDThh-mm-ss-.jsonl[.zst] let name = compression::parse_rollout_file_name(name)?; let core = name.strip_prefix("rollout-")?.strip_suffix(".jsonl")?; - - // Scan from the right for a '-' such that the suffix parses as a UUID. - let (sep_idx, uuid) = core - .match_indices('-') - .rev() - .find_map(|(i, _)| Uuid::parse_str(&core[i + 1..]).ok().map(|u| (i, u)))?; - - let ts_str = &core[..sep_idx]; - let format: &[FormatItem] = - format_description!("[year]-[month]-[day]T[hour]-[minute]-[second]"); - let ts = PrimitiveDateTime::parse(ts_str, format).ok()?.assume_utc(); + let uuid_start = core.len().checked_sub(36)?; + let (ts_with_sep, uuid_str) = core.split_at(uuid_start); + let ts_str = ts_with_sep.strip_suffix('-')?; + let uuid = Uuid::parse_str(uuid_str).ok()?; + let ts = parse_rollout_filename_timestamp(ts_str)?; Some((ts, uuid)) } +fn parse_rollout_filename_timestamp(ts_str: &str) -> Option { + let bytes = ts_str.as_bytes(); + if bytes.len() != 19 + || bytes[4] != b'-' + || bytes[7] != b'-' + || bytes[10] != b'T' + || bytes[13] != b'-' + || bytes[16] != b'-' + { + return None; + } + + let year = parse_decimal_i32(&bytes[0..4])?; + let month = parse_decimal_u8(&bytes[5..7])?; + let day = parse_decimal_u8(&bytes[8..10])?; + let hour = parse_decimal_u8(&bytes[11..13])?; + let minute = parse_decimal_u8(&bytes[14..16])?; + let second = parse_decimal_u8(&bytes[17..19])?; + let date = + time::Date::from_calendar_date(year, time::Month::try_from(month).ok()?, day).ok()?; + let time = time::Time::from_hms(hour, minute, second).ok()?; + Some(time::PrimitiveDateTime::new(date, time).assume_utc()) +} + +fn parse_decimal_i32(bytes: &[u8]) -> Option { + let mut value = 0i32; + for &byte in bytes { + let digit = byte.checked_sub(b'0')?; + if digit > 9 { + return None; + } + value = value.checked_mul(10)?.checked_add(i32::from(digit))?; + } + Some(value) +} + +fn parse_decimal_u8(bytes: &[u8]) -> Option { + let mut value = 0u8; + for &byte in bytes { + let digit = byte.checked_sub(b'0')?; + if digit > 9 { + return None; + } + value = value.checked_mul(10)?.checked_add(digit)?; + } + Some(value) +} + struct ThreadCandidate { path: PathBuf, id: Uuid, @@ -1087,10 +1133,10 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result = serde_json::from_str(trimmed); - let Ok(rollout_line) = parsed else { continue }; + let parsed: Result = serde_json::from_str(trimmed); + let Ok(rollout_item) = parsed else { continue }; - match rollout_line.item { + match rollout_item { RolloutItem::SessionMeta(session_meta_line) => { if !summary.saw_session_meta { summary.source = Some(session_meta_line.meta.source.clone()); @@ -1118,9 +1164,8 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { - summary - .created_at - .get_or_insert_with(|| rollout_line.timestamp.clone()); + // The listing path only returns files with session metadata, which carries + // the canonical creation timestamp. } RolloutItem::TurnContext(_) => { // Not included in `head`; skip. @@ -1167,8 +1212,8 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result(trimmed) { - match rollout_line.item { + if let Ok(rollout_item) = serde_json::from_str::(trimmed) { + match rollout_item { RolloutItem::SessionMeta(session_meta_line) => { if let Ok(value) = serde_json::to_value(session_meta_line) { head.push(value); @@ -1229,19 +1274,32 @@ fn event_msg_preview(event: &EventMsg) -> Option { /// Read the SessionMetaLine from the head of a rollout file for reuse by /// callers that need the session metadata (e.g. to derive a cwd for config). pub async fn read_session_meta_line(path: &Path) -> io::Result { - let head = read_head_for_summary(path).await?; - let Some(first) = head.first() else { - return Err(io::Error::other(format!( - "rollout at {} is empty", - path.display() - ))); - }; - serde_json::from_value::(first.clone()).map_err(|_| { - io::Error::other(format!( - "rollout at {} does not start with session metadata", - path.display() - )) - }) + let mut lines = compression::open_rollout_line_reader(path).await?; + while let Some(line) = lines.next_line().await? { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + let Ok(rollout_item) = serde_json::from_str::(trimmed) else { + continue; + }; + return match rollout_item { + RolloutItem::SessionMeta(session_meta_line) => Ok(session_meta_line), + RolloutItem::ResponseItem(_) | RolloutItem::InterAgentCommunication(_) => { + Err(io::Error::other(format!( + "rollout at {} does not start with session metadata", + path.display() + ))) + } + RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => { + continue; + } + }; + } + Err(io::Error::other(format!( + "rollout at {} is empty", + path.display() + ))) } async fn file_modified_time(path: &Path) -> io::Result> { diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 5dc207049f..783b6525bd 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -1,9 +1,13 @@ //! Persist Codex session rollouts (.jsonl) so sessions can be replayed or inspected later. use std::collections::HashSet; +use std::ffi::OsStr; use std::fs; use std::fs::File; +use std::io::BufRead; +use std::io::BufReader; use std::io::Error as IoError; +use std::io::ErrorKind; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; @@ -13,6 +17,7 @@ use chrono::SecondsFormat; use codex_protocol::ThreadId; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ResponseItem; use serde_json::Value; use time::OffsetDateTime; use time::format_description::FormatItem; @@ -55,7 +60,6 @@ use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::ResumedHistory; use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; @@ -63,6 +67,9 @@ use codex_protocol::protocol::ThreadSource; use codex_state::StateRuntime; use codex_utils_path as path_utils; +const ROLLOUT_HISTORY_READ_BUFFER_SIZE: usize = 64 * 1024; +const ROLLOUT_HISTORY_BYTES_PER_ITEM_HINT: u64 = 128; + /// Writes canonical session rollout items to JSONL. /// /// Rollouts are recorded as JSONL and can be inspected with tools such as: @@ -847,59 +854,22 @@ 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 v: Value = match serde_json::from_str(&line) { - Ok(v) => v, - 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 v) { - trace!("skipping legacy ghost_snapshot rollout line"); - continue; - } - - // Parse the rollout line structure - match serde_json::from_value::(v.clone()) { - Ok(rollout_line) => { - 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); - } - Err(e) => { - trace!("failed to parse rollout line: {e}"); - parse_errors = parse_errors.saturating_add(1); + let requested_path = path.to_path_buf(); + for _ in 0..3 { + let path = requested_path.clone(); + let result = tokio::task::spawn_blocking(move || load_rollout_items_blocking(path)) + .await + .map_err(IoError::other)?; + match result { + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + tokio::time::sleep(std::time::Duration::from_millis(50)).await; } + other => return other, } } - 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)) + tokio::task::spawn_blocking(move || load_rollout_items_blocking(requested_path)) + .await + .map_err(IoError::other)? } pub async fn get_rollout_history(path: &Path) -> std::io::Result { @@ -947,6 +917,199 @@ impl RolloutRecorder { } } +fn load_rollout_items_blocking( + path: PathBuf, +) -> std::io::Result<(Vec, Option, usize)> { + let (file, path) = open_rollout_file_for_history(path)?; + let capacity_hint = file + .metadata() + .ok() + .map(|metadata| rollout_items_capacity_hint(metadata.len())) + .unwrap_or_default(); + if is_compressed_rollout_path(path.as_path()) { + let decoder = zstd::stream::read::Decoder::new(file)?; + load_rollout_items_from_reader( + BufReader::with_capacity(ROLLOUT_HISTORY_READ_BUFFER_SIZE, decoder), + capacity_hint, + ) + } else { + load_rollout_items_from_reader( + BufReader::with_capacity(ROLLOUT_HISTORY_READ_BUFFER_SIZE, file), + capacity_hint, + ) + } +} + +fn open_rollout_file_for_history(path: PathBuf) -> std::io::Result<(File, PathBuf)> { + let plain_path = compression::plain_rollout_path(path.as_path()); + if let Some(file) = open_regular_file(plain_path.as_path())? { + return Ok((file, plain_path)); + } + + let compressed_path = compressed_rollout_path(plain_path.as_path()); + if let Some(file) = open_regular_file(compressed_path.as_path())? { + return Ok((file, compressed_path)); + } + + Err(IoError::new( + ErrorKind::NotFound, + format!("rollout file not found: {}", path.display()), + )) +} + +fn open_regular_file(path: &Path) -> std::io::Result> { + let file = match File::open(path) { + Ok(file) => file, + Err(err) if err.kind() == ErrorKind::NotFound => return Ok(None), + Err(err) => return Err(err), + }; + if file.metadata().is_ok_and(|metadata| metadata.is_file()) { + return Ok(Some(file)); + } + Ok(None) +} + +fn compressed_rollout_path(path: &Path) -> PathBuf { + if is_compressed_rollout_path(path) { + return path.to_path_buf(); + } + let mut file_name = path + .file_name() + .map(OsStr::to_os_string) + .unwrap_or_else(|| OsStr::new("rollout.jsonl").to_os_string()); + file_name.push(".zst"); + path.with_file_name(file_name) +} + +fn is_compressed_rollout_path(path: &Path) -> bool { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| name.ends_with(".jsonl.zst")) +} + +fn load_rollout_items_from_reader( + mut reader: R, + capacity_hint: usize, +) -> std::io::Result<(Vec, Option, usize)> { + let mut items: Vec = Vec::with_capacity(capacity_hint); + let mut thread_id: Option = None; + let mut parse_errors = 0usize; + let mut saw_non_empty_line = false; + let mut line = Vec::new(); + loop { + line.clear(); + if reader.read_until(b'\n', &mut line)? == 0 { + break; + } + let trimmed_line = trim_ascii_whitespace(line.as_slice()); + if trimmed_line.is_empty() { + continue; + } + saw_non_empty_line = true; + // Parse the rollout line structure + match parse_rollout_item_preserving_legacy_ghost_snapshot_filter(trimmed_line) { + Ok(None) => continue, + Ok(Some(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); + } + Err(e) => { + if matches!( + e.classify(), + serde_json::error::Category::Eof | serde_json::error::Category::Syntax + ) { + warn!("failed to parse line as JSON: {line:?}, error: {e}"); + } else { + trace!("failed to parse rollout line: {e}"); + } + parse_errors = parse_errors.saturating_add(1); + } + } + } + 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)) +} + +fn rollout_items_capacity_hint(file_size_bytes: u64) -> usize { + usize::try_from(file_size_bytes / ROLLOUT_HISTORY_BYTES_PER_ITEM_HINT) + .unwrap_or(usize::MAX) + .saturating_add(1) + .min(16_384) +} + +fn trim_ascii_whitespace(bytes: &[u8]) -> &[u8] { + let mut start = 0usize; + let mut end = bytes.len(); + while start < end && bytes[start].is_ascii_whitespace() { + start += 1; + } + while end > start && bytes[end - 1].is_ascii_whitespace() { + end -= 1; + } + &bytes[start..end] +} + +fn parse_rollout_item_preserving_legacy_ghost_snapshot_filter( + line: &[u8], +) -> Result, serde_json::Error> { + match serde_json::from_slice::(line) { + Ok(item) => { + if rollout_item_needs_legacy_ghost_snapshot_filter(&item) + && contains_bytes(line, b"\"ghost_snapshot\"") + { + return parse_rollout_item_after_legacy_ghost_snapshot_strip(line); + } + Ok(Some(item)) + } + Err(err) => { + if contains_bytes(line, b"\"ghost_snapshot\"") { + parse_rollout_item_after_legacy_ghost_snapshot_strip(line) + } else { + Err(err) + } + } + } +} + +fn rollout_item_needs_legacy_ghost_snapshot_filter(item: &RolloutItem) -> bool { + matches!( + item, + RolloutItem::ResponseItem(ResponseItem::Other) | RolloutItem::Compacted(_) + ) +} + +fn parse_rollout_item_after_legacy_ghost_snapshot_strip( + line: &[u8], +) -> Result, serde_json::Error> { + let mut v: Value = serde_json::from_slice(line)?; + if strip_legacy_ghost_snapshot_rollout_line(&mut v) { + trace!("skipping legacy ghost_snapshot rollout line"); + return Ok(None); + } + serde_json::from_value::(v).map(Some) +} + +fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool { + haystack + .windows(needle.len()) + .any(|window| window == needle) +} + fn strip_legacy_ghost_snapshot_rollout_line(value: &mut Value) -> bool { match value.get("type").and_then(Value::as_str) { Some("response_item") => value