From 4891934519fb8dd07f04c6ec2de8f488146db08c Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Mon, 6 Apr 2026 23:34:28 -0700 Subject: [PATCH] Durably flush shutdown state --- codex-rs/core/src/codex.rs | 5 ++ codex-rs/rollout/src/recorder.rs | 82 ++++++++++++++++++++------- codex-rs/rollout/src/session_index.rs | 1 + codex-rs/state/src/runtime.rs | 11 ++++ 4 files changed, 78 insertions(+), 21 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 6e5529800d..e99c1379e6 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -5646,6 +5646,11 @@ mod handlers { }; sess.send_event_raw(event).await; } + if let Some(state_db) = sess.services.state_db.as_deref() + && let Err(e) = state_db.checkpoint_wal().await + { + warn!("failed to checkpoint state db WAL during shutdown: {e}"); + } let event = Event { id: sub_id, diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 3554cb3797..2b94eb5275 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -98,10 +98,10 @@ enum RolloutCmd { }, /// Ensure all prior writes are processed; respond when flushed. Flush { - ack: oneshot::Sender<()>, + ack: oneshot::Sender>, }, Shutdown { - ack: oneshot::Sender<()>, + ack: oneshot::Sender>, }, } @@ -525,7 +525,7 @@ impl RolloutRecorder { .await .map_err(|e| IoError::other(format!("failed to queue rollout flush: {e}")))?; rx.await - .map_err(|e| IoError::other(format!("failed waiting for rollout flush: {e}"))) + .map_err(|e| IoError::other(format!("failed waiting for rollout flush: {e}")))? } pub async fn load_rollout_items( @@ -613,9 +613,9 @@ impl RolloutRecorder { pub async fn shutdown(&self) -> std::io::Result<()> { let (tx_done, rx_done) = oneshot::channel(); match self.tx.send(RolloutCmd::Shutdown { ack: tx_done }).await { - Ok(_) => rx_done - .await - .map_err(|e| IoError::other(format!("failed waiting for rollout shutdown: {e}")))?, + Ok(_) => rx_done.await.map_err(|e| { + IoError::other(format!("failed waiting for rollout shutdown: {e}")) + })??, Err(e) => { warn!("failed to send rollout shutdown command: {e}"); return Err(IoError::other(format!( @@ -805,38 +805,78 @@ async fn rollout_writer( buffered_items.clear(); } + flush_writer(writer.as_mut()).await?; Ok(()) } .await; - if let Err(err) = result { - let kind = err.kind(); - let message = err.to_string(); - let _ = ack.send(Err(IoError::new(kind, message.clone()))); - return Err(IoError::new(kind, message)); + match result { + Ok(()) => { + let _ = ack.send(Ok(())); + } + Err(err) => { + let return_err = clone_io_error(&err); + let _ = ack.send(Err(err)); + return Err(return_err); + } } + } else if let Err(err) = flush_writer(writer.as_mut()).await { + let return_err = clone_io_error(&err); + let _ = ack.send(Err(err)); + return Err(return_err); + } else { + let _ = ack.send(Ok(())); } - let _ = ack.send(Ok(())); } 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 - { - let _ = ack.send(()); - return Err(e); + match flush_writer(writer.as_mut()).await { + Ok(()) => { + let _ = ack.send(Ok(())); + } + Err(err) => { + let return_err = clone_io_error(&err); + let _ = ack.send(Err(err)); + return Err(return_err); + } } - let _ = ack.send(()); - } - RolloutCmd::Shutdown { ack } => { - let _ = ack.send(()); } + RolloutCmd::Shutdown { ack } => match sync_writer(writer.as_mut()).await { + Ok(()) => { + let _ = ack.send(Ok(())); + break; + } + Err(err) => { + let return_err = clone_io_error(&err); + let _ = ack.send(Err(err)); + return Err(return_err); + } + }, } } Ok(()) } +fn clone_io_error(err: &IoError) -> IoError { + IoError::new(err.kind(), err.to_string()) +} + +async fn sync_writer(writer: Option<&mut JsonlWriter>) -> std::io::Result<()> { + if let Some(writer) = writer { + writer.file.flush().await?; + writer.file.sync_all().await?; + } + Ok(()) +} + +async fn flush_writer(writer: Option<&mut JsonlWriter>) -> std::io::Result<()> { + if let Some(writer) = writer { + writer.file.flush().await?; + } + Ok(()) +} + #[allow(clippy::too_many_arguments)] async fn write_session_meta( mut writer: Option<&mut JsonlWriter>, diff --git a/codex-rs/rollout/src/session_index.rs b/codex-rs/rollout/src/session_index.rs index f1acdd8124..a3b12fa62f 100644 --- a/codex-rs/rollout/src/session_index.rs +++ b/codex-rs/rollout/src/session_index.rs @@ -60,6 +60,7 @@ pub async fn append_session_index_entry( line.push('\n'); file.write_all(line.as_bytes()).await?; file.flush().await?; + file.sync_all().await?; Ok(()) } diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index f71b6adf05..e45506e372 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -139,6 +139,17 @@ impl StateRuntime { pub fn codex_home(&self) -> &Path { self.codex_home.as_path() } + + /// Checkpoint both runtime SQLite databases so WAL contents are persisted to the main files. + pub async fn checkpoint_wal(&self) -> anyhow::Result<()> { + sqlx::query("PRAGMA wal_checkpoint(TRUNCATE)") + .execute(self.pool.as_ref()) + .await?; + sqlx::query("PRAGMA wal_checkpoint(TRUNCATE)") + .execute(self.logs_pool.as_ref()) + .await?; + Ok(()) + } } fn base_sqlite_options(path: &Path) -> SqliteConnectOptions {