Clean up implementation

This commit is contained in:
Javier Soto
2026-03-19 11:41:53 -07:00
parent a8b8f8e0e2
commit bf5a4924c7
2 changed files with 85 additions and 135 deletions

View File

@@ -5,12 +5,6 @@ use std::fs::{self};
use std::io::Error as IoError;
use std::path::Path;
use std::path::PathBuf;
#[cfg(test)]
use std::sync::Arc;
#[cfg(test)]
use std::sync::atomic::AtomicUsize;
#[cfg(test)]
use std::sync::atomic::Ordering;
use chrono::SecondsFormat;
use chrono::Utc;
@@ -761,7 +755,16 @@ async fn rollout_writer(
pending_items.extend(items);
if writer.is_none() {
continue;
if meta.is_some() {
continue;
}
match reopen_rollout_writer(&rollout_path) {
Ok(reopened_writer) => writer = Some(reopened_writer),
Err(err) => {
warn!("rollout reopen failed; keeping pending items queued: {err}");
continue;
}
}
}
write_and_reconcile_items(
@@ -774,18 +777,20 @@ async fn rollout_writer(
)
.await
.map(|()| pending_items.clear())
.unwrap_or_else(|err| warn!("rollout write failed; keeping writer alive: {err}"));
.unwrap_or_else(|err| {
writer = None;
warn!("rollout write failed; queued items will retry after reopen: {err}");
});
}
RolloutCmd::Persist { ack } => {
if writer.is_none() || meta.is_some() || !pending_items.is_empty() {
let result = async {
if writer.is_none() {
let Some(log_file_info) = deferred_log_file_info.as_ref() else {
return Err(IoError::other(
"deferred rollout recorder missing log file metadata",
));
};
let file = open_log_file(log_file_info.path.as_path())?;
let writer_path = deferred_log_file_info
.as_ref()
.map(|log_file_info| log_file_info.path.as_path())
.unwrap_or(rollout_path.as_path());
let file = open_log_file(writer_path)?;
writer = Some(JsonlWriter::new(tokio::fs::File::from_std(file)));
}
@@ -823,6 +828,7 @@ async fn rollout_writer(
.await;
if let Err(err) = result {
writer = None;
warn!("rollout persist failed; keeping writer alive: {err}");
let _ = ack.send(Err(err));
continue;
@@ -832,9 +838,10 @@ async fn rollout_writer(
}
RolloutCmd::Flush { ack } => {
// Deferred fresh threads may not have an initialized file yet.
if let Some(writer) = writer.as_mut()
&& let Err(e) = writer.file.flush().await
if let Some(current_writer) = writer.as_mut()
&& let Err(e) = current_writer.file.flush().await
{
writer = None;
warn!("rollout flush failed; keeping writer alive: {e}");
let _ = ack.send(Err(e));
continue;
@@ -960,10 +967,6 @@ async fn sync_thread_state_after_write(
struct JsonlWriter {
file: tokio::fs::File,
#[cfg(test)]
fail_next_writes: Option<Arc<AtomicUsize>>,
#[cfg(test)]
fail_next_flushes: Option<Arc<AtomicUsize>>,
}
#[derive(serde::Serialize)]
@@ -975,31 +978,7 @@ struct RolloutLineRef<'a> {
impl JsonlWriter {
fn new(file: tokio::fs::File) -> Self {
Self {
file,
#[cfg(test)]
fail_next_writes: None,
#[cfg(test)]
fail_next_flushes: None,
}
}
#[cfg(test)]
fn with_fail_next_writes(file: tokio::fs::File, fail_next_writes: Arc<AtomicUsize>) -> Self {
Self {
file,
fail_next_writes: Some(fail_next_writes),
fail_next_flushes: None,
}
}
#[cfg(test)]
fn with_fail_next_flushes(file: tokio::fs::File, fail_next_flushes: Arc<AtomicUsize>) -> Self {
Self {
file,
fail_next_writes: None,
fail_next_flushes: Some(fail_next_flushes),
}
Self { file }
}
async fn write_rollout_item(&mut self, rollout_item: &RolloutItem) -> std::io::Result<()> {
@@ -1008,17 +987,6 @@ impl JsonlWriter {
}
async fn write_rollout_items(&mut self, rollout_items: &[RolloutItem]) -> std::io::Result<()> {
#[cfg(test)]
if let Some(fail_next_writes) = self.fail_next_writes.as_ref()
&& fail_next_writes
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| {
(remaining > 0).then(|| remaining - 1)
})
.is_ok()
{
return Err(IoError::other("injected rollout write failure"));
}
let file_len_before_write = self.file.metadata().await?.len();
let mut json = String::new();
for rollout_item in rollout_items {
@@ -1028,16 +996,6 @@ impl JsonlWriter {
let result = async {
self.file.write_all(json.as_bytes()).await?;
#[cfg(test)]
if let Some(fail_next_flushes) = self.fail_next_flushes.as_ref()
&& fail_next_flushes
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |remaining| {
(remaining > 0).then(|| remaining - 1)
})
.is_ok()
{
return Err(IoError::other("injected rollout flush failure"));
}
self.file.flush().await
}
.await;
@@ -1070,6 +1028,12 @@ impl JsonlWriter {
}
}
fn reopen_rollout_writer(rollout_path: &Path) -> std::io::Result<JsonlWriter> {
open_log_file(rollout_path)
.map(tokio::fs::File::from_std)
.map(JsonlWriter::new)
}
impl From<codex_state::ThreadsPage> for ThreadsPage {
fn from(db_page: codex_state::ThreadsPage) -> Self {
let items = db_page

View File

@@ -16,9 +16,6 @@ use std::fs::{self};
use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::time::Duration;
use tempfile::TempDir;
use uuid::Uuid;
@@ -83,14 +80,10 @@ async fn flush_test_rollout_writer(tx: &mpsc::Sender<RolloutCmd>) -> std::io::Re
.map_err(|e| IoError::other(format!("flush ack should be sent: {e}")))?
}
async fn wait_for_test_fault_to_fire(fault: &AtomicUsize) {
tokio::time::timeout(Duration::from_secs(5), async {
while fault.load(Ordering::Acquire) != 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("injected rollout fault should be consumed");
async fn open_read_only_rollout_writer(path: &Path) -> std::io::Result<JsonlWriter> {
std::fs::write(path, "")?;
let file = tokio::fs::OpenOptions::new().read(true).open(path).await?;
Ok(JsonlWriter::new(file))
}
fn assert_rollout_message_count(text: &str, message: &str, expected: usize) {
@@ -183,17 +176,11 @@ async fn recorder_materializes_only_after_explicit_persist() -> std::io::Result<
}
#[tokio::test]
async fn rollout_writer_recovers_after_transient_write_failure() -> std::io::Result<()> {
async fn rollout_writer_reopens_after_initial_write_error_and_retries_pending_items()
-> std::io::Result<()> {
let home = TempDir::new().expect("temp dir");
let rollout_path = home.path().join("rollout.jsonl");
let fail_next_writes = Arc::new(AtomicUsize::new(1));
let file = tokio::fs::OpenOptions::new()
.create(true)
.append(true)
.open(&rollout_path)
.await?;
let writer = JsonlWriter::with_fail_next_writes(file, fail_next_writes.clone());
let writer = open_read_only_rollout_writer(&rollout_path).await?;
let (tx, rx) = mpsc::channel::<RolloutCmd>(256);
let writer_task = tokio::spawn(rollout_writer(
Some(writer),
@@ -211,11 +198,10 @@ async fn rollout_writer_recovers_after_transient_write_failure() -> std::io::Res
queue_test_rollout_agent_message(&tx, "first-write-fails")
.await
.expect("first write should queue");
wait_for_test_fault_to_fire(&fail_next_writes).await;
let _ = queue_test_rollout_agent_message(&tx, "second-write-succeeds").await;
let _ = flush_test_rollout_writer(&tx).await;
queue_test_rollout_agent_message(&tx, "second-write-succeeds")
.await
.expect("second write should queue after reopen");
flush_test_rollout_writer(&tx).await?;
drop(tx);
@@ -238,59 +224,59 @@ async fn rollout_writer_recovers_after_transient_write_failure() -> std::io::Res
}
#[tokio::test]
async fn rollout_writer_rolls_back_partial_batch_and_retries_without_duplicates()
-> std::io::Result<()> {
async fn recorder_retries_persist_after_materialization_failure() -> std::io::Result<()> {
let home = TempDir::new().expect("temp dir");
let rollout_path = home.path().join("rollout.jsonl");
let fail_next_flushes = Arc::new(AtomicUsize::new(1));
let file = tokio::fs::OpenOptions::new()
.create(true)
.append(true)
.open(&rollout_path)
let config = ConfigBuilder::default()
.codex_home(home.path().to_path_buf())
.build()
.await?;
let writer = JsonlWriter::with_fail_next_flushes(file, fail_next_flushes.clone());
let (tx, rx) = mpsc::channel::<RolloutCmd>(256);
let writer_task = tokio::spawn(rollout_writer(
Some(writer),
None,
rx,
None,
home.path().to_path_buf(),
rollout_path.clone(),
let recorder = RolloutRecorder::new(
&config,
RolloutRecorderParams::new(
ThreadId::new(),
None,
SessionSource::Exec,
BaseInstructions::default(),
Vec::new(),
EventPersistenceMode::Limited,
),
None,
None,
"test-provider".to_string(),
false,
));
)
.await?;
let rollout_path = recorder.rollout_path().to_path_buf();
std::fs::create_dir_all(
rollout_path
.parent()
.expect("rollout path should have a parent directory"),
)?;
std::fs::create_dir(&rollout_path)?;
tx.send(RolloutCmd::AddItems(vec![
test_rollout_agent_message("batch-item-1"),
test_rollout_agent_message("batch-item-2"),
]))
.await
.expect("failed batch should queue");
wait_for_test_fault_to_fire(&fail_next_flushes).await;
queue_test_rollout_agent_message(&tx, "batch-item-3")
recorder
.record_items(&[test_rollout_agent_message("buffered-before-failure")])
.await?;
let persist_error = recorder
.persist()
.await
.expect("retry trigger should queue");
flush_test_rollout_writer(&tx).await?;
drop(tx);
.expect_err("materialization should fail while the rollout path is a directory");
assert!(
persist_error.to_string().contains("Is a directory"),
"expected a real open failure, got: {persist_error}"
);
std::fs::remove_dir(&rollout_path)?;
recorder
.record_items(&[test_rollout_agent_message("buffered-after-failure")])
.await?;
recorder.persist().await?;
recorder.flush().await?;
let text = std::fs::read_to_string(&rollout_path)?;
assert_rollout_message_count(&text, "batch-item-1", 1);
assert_rollout_message_count(&text, "batch-item-2", 1);
assert_rollout_message_count(&text, "batch-item-3", 1);
for line in text.lines() {
serde_json::from_str::<serde_json::Value>(line).expect("rollout line should be valid JSON");
}
writer_task
.await
.expect("writer task should join cleanly")
.expect("writer task should exit cleanly after channel closes");
assert_rollout_message_count(&text, "buffered-before-failure", 1);
assert_rollout_message_count(&text, "buffered-after-failure", 1);
recorder.shutdown().await?;
Ok(())
}