diff --git a/codex-rs/thread-store/src/local/benchmark.rs b/codex-rs/thread-store/src/local/benchmark.rs new file mode 100644 index 0000000000..9b446da8cc --- /dev/null +++ b/codex-rs/thread-store/src/local/benchmark.rs @@ -0,0 +1,184 @@ +use std::io::Write; +use std::time::Duration; +use std::time::Instant; + +use codex_protocol::ThreadId; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadMemoryMode; +use tempfile::TempDir; + +use super::LoadThreadHistoryParams; +use super::LocalThreadStore; +use super::ReadThreadParams; +use super::ResumeThreadParams; +use super::read_thread; +use super::test_support::test_config; +use super::test_support::write_session_file; +use crate::ThreadPersistenceMetadata; +use crate::ThreadStore; + +const BENCHMARK_HISTORY_ITEMS: usize = 10_000; +const BENCHMARK_RUNS: usize = 20; + +fn write_thread_history_benchmark_fixture( + root: &std::path::Path, + uuid: uuid::Uuid, +) -> std::path::PathBuf { + let timestamp = "2025-01-04T10-00-00"; + let rollout_path = write_session_file(root, timestamp, uuid).expect("write benchmark session"); + let message = "x".repeat(1_024); + let event = serde_json::json!({ + "timestamp": timestamp, + "type": "event_msg", + "payload": { + "type": "user_message", + "message": message, + "kind": "plain", + }, + }); + let mut file = std::fs::OpenOptions::new() + .append(true) + .open(&rollout_path) + .expect("open benchmark session"); + for _ in 0..BENCHMARK_HISTORY_ITEMS { + writeln!(file, "{event}").expect("append benchmark history item"); + } + rollout_path +} + +fn print_thread_history_benchmark( + label: &str, + mut durations: Vec, + rollout_bytes: u64, + read_work: &read_thread::read_work::ReadWork, +) { + durations.sort_unstable(); + let p50 = durations[durations.len() / 2]; + let p95 = durations[durations.len() * 95 / 100]; + let max = *durations.last().expect("benchmark duration"); + eprintln!( + "{BENCHMARK_HISTORY_ITEMS}-item thread history {label} benchmark: rollout_bytes={rollout_bytes} p50={p50:?} p95={p95:?} max={max:?} work={read_work:?}" + ); +} + +fn thread_metadata() -> ThreadPersistenceMetadata { + ThreadPersistenceMetadata { + cwd: Some(std::env::current_dir().expect("cwd")), + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + } +} + +#[tokio::test] +#[ignore = "release benchmark"] +async fn thread_history_benchmark_10_000_items() { + let home = TempDir::new().expect("temp dir"); + let uuid = uuid::Uuid::from_u128(408); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let rollout_path = write_thread_history_benchmark_fixture(home.path(), uuid); + let rollout_bytes = std::fs::metadata(&rollout_path) + .expect("benchmark rollout metadata") + .len(); + + let config = test_config(home.path()); + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let mut builder = codex_state::ThreadMetadataBuilder::new( + thread_id, + rollout_path.clone(), + chrono::Utc::now(), + SessionSource::Cli, + ); + builder.model_provider = Some(config.default_model_provider_id.clone()); + builder.cwd = home.path().to_path_buf(); + runtime + .upsert_thread(&builder.build(config.default_model_provider_id.as_str())) + .await + .expect("state db upsert should succeed"); + let sqlite_store = LocalThreadStore::new(config.clone(), Some(runtime)); + sqlite_store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: true, + }) + .await + .expect("warm SQLite history read"); + + let mut sqlite_durations = Vec::with_capacity(BENCHMARK_RUNS); + let mut sqlite_work = None; + for _ in 0..BENCHMARK_RUNS { + let started = Instant::now(); + let (thread, read_work) = + read_thread::read_work::measure(sqlite_store.read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: true, + })) + .await; + assert_eq!( + thread + .expect("read SQLite-backed thread") + .history + .expect("SQLite-backed history") + .items + .len(), + BENCHMARK_HISTORY_ITEMS + 2 + ); + sqlite_durations.push(started.elapsed()); + sqlite_work = Some(read_work); + } + print_thread_history_benchmark( + "SQLite", + sqlite_durations, + rollout_bytes, + &sqlite_work.expect("SQLite read work"), + ); + + let live_store = LocalThreadStore::new(config, /*state_db*/ None); + live_store + .resume_thread(ResumeThreadParams { + thread_id, + rollout_path: Some(rollout_path), + history: None, + include_archived: true, + metadata: thread_metadata(), + }) + .await + .expect("resume benchmark thread"); + live_store + .load_history(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect("warm active-writer history read"); + + let mut live_durations = Vec::with_capacity(BENCHMARK_RUNS); + let mut live_work = None; + for _ in 0..BENCHMARK_RUNS { + let started = Instant::now(); + let (history, read_work) = + read_thread::read_work::measure(live_store.load_history(LoadThreadHistoryParams { + thread_id, + include_archived: false, + })) + .await; + assert_eq!( + history.expect("load active-writer history").items.len(), + BENCHMARK_HISTORY_ITEMS + 2 + ); + live_durations.push(started.elapsed()); + live_work = Some(read_work); + } + print_thread_history_benchmark( + "active writer", + live_durations, + rollout_bytes, + &live_work.expect("active-writer read work"), + ); +} diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 19a0d2955b..bc4c1c3c9d 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -1,4 +1,6 @@ mod archive_thread; +#[cfg(test)] +mod benchmark; mod create_thread; mod delete_thread; mod helpers; @@ -182,17 +184,11 @@ impl LocalThreadStore { message: format!("thread {} is archived", params.thread_id), }); } - return read_thread::read_thread_by_rollout_path( - self, - rollout_path, - /*include_archived*/ true, - /*include_history*/ true, + return read_thread::load_history_from_rollout_path( + rollout_path.as_path(), + params.thread_id, ) - .await? - .history - .ok_or_else(|| ThreadStoreError::Internal { - message: format!("failed to load history for thread {}", params.thread_id), - }); + .await; } read_thread::read_thread( @@ -954,13 +950,13 @@ mod tests { .await .expect("flush live thread"); - let history = store - .load_history(LoadThreadHistoryParams { + let (history, read_work) = + read_thread::read_work::measure(store.load_history(LoadThreadHistoryParams { thread_id, include_archived: false, - }) - .await - .expect("load external live history"); + })) + .await; + let history = history.expect("load external live history"); assert!(history.items.iter().any(|item| { matches!( @@ -968,6 +964,13 @@ mod tests { RolloutItem::EventMsg(EventMsg::UserMessage(event)) if event.message == "external history item" ) })); + assert_eq!( + read_work, + read_thread::read_work::ReadWork { + summary_reads: 0, + history_reads: 1, + } + ); } #[tokio::test] diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index a1e1c59190..1d63344022 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -2,6 +2,7 @@ use chrono::DateTime; use chrono::Utc; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; use codex_rollout::RolloutRecorder; @@ -26,6 +27,54 @@ use crate::StoredThreadHistory; use crate::ThreadStoreError; use crate::ThreadStoreResult; +#[cfg(test)] +pub(crate) mod read_work { + use std::future::Future; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + #[derive(Clone, Default)] + struct ReadWorkRecorder { + summary_reads: Arc, + history_reads: Arc, + } + + #[derive(Debug, PartialEq, Eq)] + pub(crate) struct ReadWork { + pub(crate) summary_reads: usize, + pub(crate) history_reads: usize, + } + + tokio::task_local! { + static READ_WORK: ReadWorkRecorder; + } + + pub(super) fn record_summary_read() { + let _ = READ_WORK.try_with(|work| { + work.summary_reads.fetch_add(1, Ordering::Relaxed); + }); + } + + pub(super) fn record_history_read() { + let _ = READ_WORK.try_with(|work| { + work.history_reads.fetch_add(1, Ordering::Relaxed); + }); + } + + pub(crate) async fn measure(future: impl Future) -> (T, ReadWork) { + let recorder = ReadWorkRecorder::default(); + let result = READ_WORK.scope(recorder.clone(), future).await; + ( + result, + ReadWork { + summary_reads: recorder.summary_reads.load(Ordering::Relaxed), + history_reads: recorder.history_reads.load(Ordering::Relaxed), + }, + ) + } +} + pub(super) async fn read_thread( store: &LocalThreadStore, params: ReadThreadParams, @@ -38,36 +87,48 @@ pub(super) async fn read_thread( store.config.codex_home.as_path(), metadata.rollout_path.as_path(), ))) - && (!params.include_history - || sqlite_rollout_path_can_load_history_for_thread( - store, - &metadata.rollout_path, - thread_id, - ) - .await) { - let metadata_sandbox_policy = metadata.sandbox_policy.clone(); - let mut thread = stored_thread_from_sqlite_metadata(store, metadata).await; - if !params.include_history - && let Some(rollout_path) = thread.rollout_path.clone() - && let Ok(mut rollout_thread) = read_thread_from_rollout_path(store, rollout_path).await - && rollout_thread.thread_id == thread_id - && (params.include_archived || rollout_thread.archived_at.is_none()) - && !rollout_thread.preview.is_empty() - { - rollout_thread.recency_at = thread.recency_at; - if thread.name.is_some() { - rollout_thread.name = thread.name; + let history = if params.include_history { + // The full parse also validates the first SessionMeta thread ID, so do not precede it + // with a separate summary read on the common SQLite path. + load_history_from_rollout_path(metadata.rollout_path.as_path(), thread_id) + .await + .ok() + } else { + None + }; + if !params.include_history || history.is_some() { + let metadata_sandbox_policy = metadata.sandbox_policy.clone(); + let session_meta = history.as_ref().and_then(|history| { + history.items.iter().find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(meta_line.clone()), + _ => None, + }) + }); + let mut thread = + stored_thread_from_sqlite_metadata(store, metadata, session_meta).await; + if !params.include_history + && let Some(rollout_path) = thread.rollout_path.clone() + && let Ok(mut rollout_thread) = + read_thread_from_rollout_path(store, rollout_path).await + && rollout_thread.thread_id == thread_id + && (params.include_archived || rollout_thread.archived_at.is_none()) + && !rollout_thread.preview.is_empty() + { + rollout_thread.recency_at = thread.recency_at; + if thread.name.is_some() { + rollout_thread.name = thread.name; + } + rollout_thread.git_info = thread.git_info; + rollout_thread.permission_profile = permission_profile_from_metadata_value( + &metadata_sandbox_policy, + rollout_thread.cwd.as_path(), + ); + thread = rollout_thread; } - rollout_thread.git_info = thread.git_info; - rollout_thread.permission_profile = permission_profile_from_metadata_value( - &metadata_sandbox_policy, - rollout_thread.cwd.as_path(), - ); - thread = rollout_thread; + thread.history = history; + return Ok(thread); } - attach_history_if_requested(&mut thread, params.include_history).await?; - return Ok(thread); } let path = resolve_rollout_path(store, thread_id, params.include_archived) @@ -86,22 +147,6 @@ pub(super) async fn read_thread( Ok(thread) } -async fn sqlite_rollout_path_can_load_history_for_thread( - store: &LocalThreadStore, - path: &std::path::Path, - thread_id: codex_protocol::ThreadId, -) -> bool { - if codex_rollout::existing_rollout_path(path).await.is_none() { - return false; - } - // SQLite metadata can outlive a moved/recreated rollout path. When history is - // requested, verify the path still resolves to the requested thread before - // trusting it as the source replay. - read_thread_from_rollout_path(store, path.to_path_buf()) - .await - .is_ok_and(|thread| thread.thread_id == thread_id) -} - pub(super) async fn read_thread_by_rollout_path( store: &LocalThreadStore, rollout_path: std::path::PathBuf, @@ -190,8 +235,7 @@ async fn attach_history_if_requested( message: format!("failed to load thread history for thread {thread_id}"), }); }; - let items = load_history_items(&path).await?; - thread.history = Some(StoredThreadHistory { thread_id, items }); + thread.history = Some(load_history_from_rollout_path(&path, thread_id).await?); Ok(()) } @@ -248,6 +292,8 @@ async fn read_thread_from_rollout_path( store: &LocalThreadStore, path: std::path::PathBuf, ) -> ThreadStoreResult { + #[cfg(test)] + read_work::record_summary_read(); let Some(item) = read_thread_item_from_rollout(path.clone()).await else { return stored_thread_from_session_meta(store, path).await; }; @@ -280,15 +326,31 @@ async fn read_thread_from_rollout_path( Ok(thread) } -async fn load_history_items( +/// Loads one rollout as history and verifies that its first SessionMeta belongs to the caller's +/// thread before the caller trusts a state-database or live-writer path. +pub(super) async fn load_history_from_rollout_path( path: &std::path::Path, -) -> ThreadStoreResult> { - let (items, _, _) = RolloutRecorder::load_rollout_items(path) + expected_thread_id: codex_protocol::ThreadId, +) -> ThreadStoreResult { + #[cfg(test)] + read_work::record_history_read(); + let (items, thread_id, _) = RolloutRecorder::load_rollout_items(path) .await .map_err(|err| ThreadStoreError::Internal { message: format!("failed to load thread history {}: {err}", path.display()), })?; - Ok(items) + if thread_id != Some(expected_thread_id) { + return Err(ThreadStoreError::InvalidRequest { + message: format!( + "rollout {} does not contain history for thread {expected_thread_id}", + path.display() + ), + }); + } + Ok(StoredThreadHistory { + thread_id: expected_thread_id, + items, + }) } async fn read_sqlite_metadata( @@ -302,6 +364,7 @@ async fn read_sqlite_metadata( async fn stored_thread_from_sqlite_metadata( store: &LocalThreadStore, metadata: ThreadMetadata, + session_meta: Option, ) -> StoredThread { let name = match distinct_thread_metadata_title(&metadata) { Some(title) => Some(title), @@ -311,10 +374,13 @@ async fn stored_thread_from_sqlite_metadata( .flatten() .filter(|title| !title.trim().is_empty()), }; - let session_meta = read_session_meta_line(metadata.rollout_path.as_path()) - .await - .ok() - .map(|meta_line| meta_line.meta); + let session_meta = match session_meta { + Some(meta_line) => Some(meta_line.meta), + None => read_session_meta_line(metadata.rollout_path.as_path()) + .await + .ok() + .map(|meta_line| meta_line.meta), + }; let rollout_path = codex_rollout::plain_rollout_path(metadata.rollout_path.as_path()); let forked_from_id = session_meta.as_ref().and_then(|meta| meta.forked_from_id); let parent_thread_id = session_meta.as_ref().and_then(|meta| meta.parent_thread_id); @@ -960,14 +1026,13 @@ mod tests { .await .expect("state db upsert should succeed"); - let thread = store - .read_thread(ReadThreadParams { - thread_id, - include_archived: false, - include_history: true, - }) - .await - .expect("read thread"); + let (thread, read_work) = read_work::measure(store.read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: true, + })) + .await; + let thread = thread.expect("read thread"); assert_eq!(thread.thread_id, thread_id); assert_eq!(thread.rollout_path, Some(rollout_path)); @@ -979,6 +1044,13 @@ mod tests { let history = thread.history.expect("history should load"); assert_eq!(history.thread_id, thread_id); assert_eq!(history.items.len(), 1); + assert_eq!( + read_work, + read_work::ReadWork { + summary_reads: 0, + history_reads: 1, + } + ); } #[tokio::test]