diff --git a/codex-rs/thread-store/src/local/thread_history_materialization.rs b/codex-rs/thread-store/src/local/thread_history_materialization.rs index 06ac3e8072..e28fff1152 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization.rs @@ -1,3 +1,5 @@ +//! Reads durable history into bounded batches ending at complete records, each with a checkpoint. + use std::io::BufRead; use std::io::BufReader; use std::io::Read; @@ -18,8 +20,13 @@ use super::thread_history::RolloutProjectionStep; use crate::ThreadStoreError; use crate::ThreadStoreResult; +#[path = "thread_history_materialization_pipeline.rs"] +mod pipeline; + const SQLITE_PROJECTION_METRIC: &str = "codex.thread_history.sqlite_projection"; const SQLITE_PROJECTION_ANOMALY_METRIC: &str = "codex.thread_history.sqlite_projection.anomaly"; +const BATCH_SOURCE_BYTES: u64 = 1024 * 1024; +const BATCH_STEPS: usize = 256; pub(super) async fn materialize_to_sqlite( store: &LocalThreadStore, @@ -61,71 +68,32 @@ async fn materialize_to_sqlite_with_state_db( let expected_ordinal = projection_state .as_ref() .map_or(initial_ordinal, |state| state.next_ordinal); - let (projections, next_offset) = read_projection_steps( + pipeline::read_and_apply( + store, + thread_id, rollout_path, start_offset, expected_ordinal, - thread_id, + initial_ordinal, subagent_history_start_ordinal, ) - .await?; - // Empty valid records can still consume bytes through blank complete lines. - if projections.is_empty() && start_offset == next_offset { - return Ok(()); - } - super::thread_history::apply_projection( - store, - thread_id, - start_offset, - next_offset, - initial_ordinal, - projections, - ) .await } -async fn read_projection_steps( - rollout_path: &Path, - start_offset: u64, - expected_ordinal: u64, - thread_id: ThreadId, - subagent_history_start_ordinal: Option, -) -> ThreadStoreResult<(Vec, u64)> { - let path = rollout_path.to_path_buf(); - let span = tracing::Span::current(); - let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); - let result = tokio::task::spawn_blocking(move || { - let _entered = span.enter(); - read_projection_steps_sync( - &path, - start_offset, - expected_ordinal, - thread_id, - subagent_history_start_ordinal, - || stop.is_closed(), - ) - }) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to join rollout projection read: {err}"), - })?; - drop(keep_running); - result -} - /// Read and project the complete durable prefix in one blocking task. -fn read_projection_steps_sync( +fn read_projection_batches( rollout_path: &Path, start_offset: u64, expected_ordinal: u64, thread_id: ThreadId, subagent_history_start_ordinal: Option, is_cancelled: impl Fn() -> bool, -) -> ThreadStoreResult<(Vec, u64)> { + mut emit: impl FnMut(Vec, u64) -> ThreadStoreResult<()>, +) -> ThreadStoreResult<()> { let mut file = match codex_rollout::open_rollout_seekable_reader(rollout_path, &is_cancelled) { Ok(file) => file, Err(err) if err.kind() == std::io::ErrorKind::NotFound && start_offset == 0 => { - return Ok((Vec::new(), 0)); + return Ok(()); } Err(err) => return Err(thread_store_io_error(err)), }; @@ -143,16 +111,23 @@ fn read_projection_steps_sync( let mut line_bytes = Vec::new(); // Only project the newline-terminated prefix; leave a trailing partial record for the next // pass. - let mut projections = Vec::new(); let mut next_ordinal = expected_ordinal; let mut next_offset = start_offset; let mut line_start_offset = start_offset; + let mut batch_start_offset = start_offset; + let mut batch = Vec::new(); loop { if is_cancelled() { return Err(thread_store_io_error(std::io::Error::other( "rollout read cancelled", ))); } + // Flush only between complete records: an ordinal gap and its following line must + // share a checkpoint. Byte-only progress from skipped records also bounds batches. + if next_offset - batch_start_offset >= BATCH_SOURCE_BYTES || batch.len() >= BATCH_STEPS { + emit(std::mem::take(&mut batch), next_offset)?; + batch_start_offset = next_offset; + } line_bytes.clear(); if reader .read_until(b'\n', &mut line_bytes) @@ -285,7 +260,7 @@ fn read_projection_steps_sync( .ok_or_else(|| ThreadStoreError::Internal { message: "rollout ordinal exceeds SQLite integer range".to_string(), })?; - projections.push(RolloutProjectionStep::SkippedOrdinalRange { + batch.push(RolloutProjectionStep::SkippedOrdinalRange { start_ordinal: next_ordinal, end_ordinal_exclusive, }); @@ -311,7 +286,7 @@ fn read_projection_steps_sync( "skipping missing rollout ordinal range during projection" ); record_projection_anomaly(ProjectionAnomaly::ForwardOrdinalGap); - projections.push(RolloutProjectionStep::SkippedOrdinalRange { + batch.push(RolloutProjectionStep::SkippedOrdinalRange { start_ordinal: next_ordinal, end_ordinal_exclusive: ordinal, }); @@ -322,7 +297,7 @@ fn read_projection_steps_sync( .ok_or_else(|| ThreadStoreError::Internal { message: "rollout ordinal exceeds SQLite integer range".to_string(), })?; - projections.push(RolloutProjectionStep::Line(Box::new( + batch.push(RolloutProjectionStep::Line(Box::new( ProjectedRolloutLine { ordinal, start_byte_offset: line_start_offset, @@ -345,7 +320,10 @@ fn read_projection_steps_sync( "durable rollout shrank during projection", ))); } - Ok((projections, next_offset)) + if next_offset != batch_start_offset { + emit(batch, next_offset)?; + } + Ok(()) } #[derive(Clone, Copy)] diff --git a/codex-rs/thread-store/src/local/thread_history_materialization_pipeline.rs b/codex-rs/thread-store/src/local/thread_history_materialization_pipeline.rs new file mode 100644 index 0000000000..c972ed1e1f --- /dev/null +++ b/codex-rs/thread-store/src/local/thread_history_materialization_pipeline.rs @@ -0,0 +1,107 @@ +//! Pipes bounded rollout batches into short SQLite transactions, committing each checkpoint. + +use std::path::Path; + +use codex_protocol::ThreadId; + +use super::super::LocalThreadStore; +use super::super::thread_history::RolloutProjectionStep; +use super::super::thread_history::apply_projection; +use super::read_projection_batches; +use crate::ThreadStoreError; +use crate::ThreadStoreResult; + +const QUEUED_BATCHES: usize = 2; + +pub(super) enum ProjectionMessage { + Batch { + projections: Vec, + next_offset: u64, + }, + Complete, + Failed(ThreadStoreError), +} + +pub(super) async fn read_and_apply( + store: &LocalThreadStore, + thread_id: ThreadId, + rollout_path: &Path, + start_offset: u64, + expected_ordinal: u64, + initial_ordinal: u64, + subagent_history_start_ordinal: Option, +) -> ThreadStoreResult<()> { + // Two queued batches, one being produced and one being consumed. Each is limited by + // source bytes and step count; an oversized individual record remains indivisible. + let (sender, receiver) = tokio::sync::mpsc::channel(QUEUED_BATCHES); + let (stop, keep_running) = tokio::sync::oneshot::channel::<()>(); + let path = rollout_path.to_path_buf(); + let span = tracing::Span::current(); + let worker = tokio::task::spawn_blocking(move || { + let _entered = span.enter(); + let read = read_projection_batches( + &path, + start_offset, + expected_ordinal, + thread_id, + subagent_history_start_ordinal, + || stop.is_closed(), + |projections, next_offset| { + sender + .blocking_send(ProjectionMessage::Batch { + projections, + next_offset, + }) + .map_err(|_| ThreadStoreError::Internal { + message: "projection consumer stopped".to_string(), + }) + }, + ); + let completion = match read { + Ok(()) => ProjectionMessage::Complete, + Err(error) => ProjectionMessage::Failed(error), + }; + let _ = sender.blocking_send(completion); + }); + let result = apply_batches(store, thread_id, start_offset, initial_ordinal, receiver).await; + // apply_batches drops the receiver on error, unblocking any pending producer send. + drop(keep_running); + worker.await.map_err(|err| ThreadStoreError::Internal { + message: format!("projection producer failed: {err}"), + })?; + result +} + +pub(super) async fn apply_batches( + store: &LocalThreadStore, + thread_id: ThreadId, + mut start_offset: u64, + initial_ordinal: u64, + mut receiver: tokio::sync::mpsc::Receiver, +) -> ThreadStoreResult<()> { + // Receive before opening the transaction so a slow reader cannot hold the SQLite writer. + while let Some(message) = receiver.recv().await { + match message { + ProjectionMessage::Batch { + projections, + next_offset, + } => { + apply_projection( + store, + thread_id, + start_offset, + next_offset, + initial_ordinal, + projections, + ) + .await?; + start_offset = next_offset; + } + ProjectionMessage::Complete => return Ok(()), + ProjectionMessage::Failed(error) => return Err(error), + } + } + Err(ThreadStoreError::Internal { + message: "projection producer stopped before successful completion".to_string(), + }) +} diff --git a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs index e3299fed59..213735805b 100644 --- a/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs +++ b/codex-rs/thread-store/src/local/thread_history_materialization_tests.rs @@ -2,6 +2,7 @@ use std::cell::Cell; use std::fs; use std::io::Write; use std::path::Path; +use std::sync::Arc; use std::time::Duration; use chrono::Utc; @@ -2362,6 +2363,242 @@ async fn jsonl_failure_does_not_create_projection_database() { ); } +/// A failed batch rolls back atomically; retries start after the last committed complete record. +#[tokio::test] +async fn catch_up_preserves_committed_batches_on_reader_or_writer_failure() { + enum Failure { + Reader, + Writer, + } + + for failure in [Failure::Reader, Failure::Writer] { + let home = TempDir::new().expect("temp dir"); + let store = projection_store(home.path()).await; + let thread_id = ThreadId::default(); + create_paginated_thread(&store, thread_id).await; + store + .persist_thread(thread_id, PersistContext::Standard) + .await + .expect("persist session metadata"); + let pool = codex_state::open_thread_history_db( + &codex_state::SqliteConfig::new_for_testing(home.path().abs()), + ) + .await + .expect("open thread history db"); + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("rollout path"); + // More than the producer, queue and consumer can hold: failures exercise backpressure. + append_suffix(rollout_path.as_path(), "\nnot-json\n"); + for index in 1..=2049 { + // The gap and following record cross the step limit, but must commit together. + let ordinal = if index >= 256 { index + 1 } else { index }; + append_suffix( + rollout_path.as_path(), + &format!( + "{}\n", + rollout_line(Some(ordinal), turn_started(&format!("turn-{index}"))) + ), + ); + } + append_suffix(rollout_path.as_path(), "\n"); + let valid_len = fs::metadata(rollout_path.as_path()) + .expect("rollout metadata") + .len(); + match failure { + Failure::Reader => append_suffix( + rollout_path.as_path(), + &format!( + "{}\n", + rollout_line(Some(u64::MAX), turn_started("overflow")) + ), + ), + Failure::Writer => { + sqlx::query( + "CREATE TRIGGER reject_projection BEFORE INSERT ON thread_turns + WHEN NEW.turn_id = 'turn-300' + BEGIN SELECT RAISE(FAIL, 'reject projection'); END", + ) + .execute(&pool) + .await + .expect("reject a later batch"); + } + } + super::materialize_to_sqlite(&store, thread_id, rollout_path.as_path()) + .await + .expect_err("failed batch must not commit"); + let turn_count = + sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM thread_turns WHERE thread_id = ?") + .bind(thread_id.to_string()) + .fetch_one(&pool) + .await + .expect("count committed turns"); + let committed_turns = match failure { + Failure::Reader => 2048, + Failure::Writer => 256, + }; + let committed_ordinal = committed_turns + 1; + let (_, committed_offset) = + rollout_line_byte_offsets(rollout_path.as_path(), committed_ordinal); + assert_eq!( + (projection_state(&pool, thread_id).await, turn_count), + ( + (committed_offset, committed_ordinal as i64 + 1), + committed_turns as i64 + ) + ); + + match failure { + Failure::Reader => fs::OpenOptions::new() + .write(true) + .open(rollout_path.as_path()) + .expect("open rollout") + .set_len(valid_len) + .expect("remove invalid suffix"), + Failure::Writer => { + sqlx::query("DROP TRIGGER reject_projection") + .execute(&pool) + .await + .expect("allow retry"); + } + } + sqlx::query( + "CREATE TRIGGER reject_replay BEFORE INSERT ON thread_turns + WHEN NEW.turn_id = 'turn-1' + BEGIN SELECT RAISE(FAIL, 'replayed committed prefix'); END", + ) + .execute(&pool) + .await + .expect("reject replay of committed rows"); + super::materialize_to_sqlite(&store, thread_id, rollout_path.as_path()) + .await + .expect("retry uncommitted suffix"); + let turn_count = + sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM thread_turns WHERE thread_id = ?") + .bind(thread_id.to_string()) + .fetch_one(&pool) + .await + .expect("count committed turns"); + assert_eq!( + (projection_state(&pool, thread_id).await, turn_count), + ( + (i64::try_from(valid_len).expect("rollout length"), 2051), + 2049 + ) + ); + } +} + +#[tokio::test] +async fn paused_projection_releases_writer_and_retains_progress_when_stopped() { + enum Stop { + Disconnect, + Cancel, + } + + for stop in [Stop::Disconnect, Stop::Cancel] { + let home = TempDir::new().expect("temp dir"); + let store = Arc::new(projection_store(home.path()).await); + let thread_id = ThreadId::default(); + let other_thread_id = ThreadId::default(); + for id in [thread_id, other_thread_id] { + create_paginated_thread(&store, id).await; + store + .persist_thread(id, PersistContext::Standard) + .await + .expect("persist session metadata"); + } + let pool = store.thread_history_db().await.expect("history db"); + let before = projection_state(pool, thread_id).await; + let rollout_path = store + .live_rollout_path(thread_id) + .await + .expect("rollout path"); + append_suffix( + rollout_path.as_path(), + &format!("{}\n", rollout_line(Some(1), turn_started("turn-1"))), + ); + let (_, next_offset) = + rollout_line_byte_offsets(rollout_path.as_path(), /*ordinal*/ 1); + let (sender, receiver) = tokio::sync::mpsc::channel(2); + super::read_projection_batches( + rollout_path.as_path(), + before.0 as u64, + before.1 as u64, + thread_id, + /*subagent_history_start_ordinal*/ None, + || false, + |projections, next_offset| { + assert!( + sender + .try_send(super::pipeline::ProjectionMessage::Batch { + projections, + next_offset, + }) + .is_ok(), + "queue the complete record" + ); + Ok(()) + }, + ) + .expect("read batch"); + let consumer_store = Arc::clone(&store); + let consumer = tokio::spawn(async move { + super::pipeline::apply_batches( + &consumer_store, + thread_id, + before.0 as u64, + /*initial_ordinal*/ 0, + receiver, + ) + .await + }); + // Keep the producer connected but idle until its first batch is visibly committed. + while projection_state(pool, thread_id).await == before { + tokio::task::yield_now().await; + } + assert_eq!(projection_state(pool, thread_id).await, (next_offset, 2)); + store + .append_items(AppendThreadItemsParams { + thread_id: other_thread_id, + items: vec![turn_started("unrelated-turn")], + }) + .await + .expect("append while the reader is idle"); + store + .flush_thread(other_thread_id) + .await + .expect("flush unrelated writer"); + assert_eq!(projection_state(pool, other_thread_id).await.1, 2); + + match stop { + Stop::Disconnect => { + drop(sender); + consumer + .await + .expect("join consumer") + .expect_err("missing completion"); + } + Stop::Cancel => { + consumer.abort(); + assert!( + consumer + .await + .expect_err("cancelled consumer") + .is_cancelled() + ); + sender.closed().await; + } + } + assert_eq!(projection_state(pool, thread_id).await, (next_offset, 2)); + super::materialize_to_sqlite(&store, thread_id, rollout_path.as_path()) + .await + .expect("resume after stopped producer"); + assert_eq!(projection_state(pool, thread_id).await, (next_offset, 2)); + } +} + #[tokio::test] async fn catch_up_rejects_missing_rollout_after_projection() { let home = TempDir::new().expect("temp dir"); @@ -2479,13 +2716,14 @@ async fn blank_and_rejected_rollout_lines_advance_projection() { // Replay this small file, cancelling between records. let checks = Cell::new(0); assert!( - super::read_projection_steps_sync( + super::read_projection_batches( rollout_path.as_path(), /*start_offset*/ 0, /*expected_ordinal*/ 0, thread_id, /*subagent_history_start_ordinal*/ None, || checks.replace(checks.get() + 1) >= 6, + |_batch, _next_offset| Ok(()), ) .is_err() );