openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Batch thread history materialization into SQLite (#55244)

Browse Source 
## Why

History catch-up collected the entire durable rollout prefix before applying it in one SQLite transaction. Large rollouts therefore required buffering all projection steps and committing them together.

## What changed

Stream projection batches through a bounded queue and commit each batch with its checkpoint. Flush between complete records after reaching 1 MiB of source data or 256 projection steps, keeping ordinal gaps with their following records and allowing oversized records to remain intact.

Wait for batches before opening SQLite transactions so an idle reader does not hold the writer. Preserve committed progress across later failures or cancellation, allowing retries to resume from the last checkpoint.

## Testing

Add coverage for reader and writer failures under backpressure, atomic batch rollback, retries without replaying committed records, and unrelated writes while projection is paused. Verify that producer disconnection and consumer cancellation retain committed progress.

GitOrigin-RevId: 06e79fa4cf273de9abd95faf7eed370f18cedc33
This commit is contained in:
jif
2026-10-19 10:48:13 +00:00
committed by copyberry
parent eb354a928b
commit 47db419304
3 changed files with 376 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/thread-store/src/local/thread_history_materialization.rs
+30 -52
View File
@@ -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<u64>,
) -> ThreadStoreResult<(Vec<RolloutProjectionStep>, 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<u64>,
is_cancelled: impl Fn() -> bool,
) -> ThreadStoreResult<(Vec<RolloutProjectionStep>, u64)> {
mut emit: impl FnMut(Vec<RolloutProjectionStep>, 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)]
codex-rs/thread-store/src/local/thread_history_materialization_pipeline.rs
+107
View File
@@ -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<RolloutProjectionStep>,
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<u64>,
) -> 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<ProjectionMessage>,
) -> 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(),
})
}
codex-rs/thread-store/src/local/thread_history_materialization_tests.rs
+239 -1
View File
@@ -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()
);