From 0e41a5c4a8f51304ea076dcdedf1d876f1fd939e Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 6 Mar 2026 18:58:14 +0000 Subject: [PATCH] chore: improve DB flushing (#13620) This branch: * Avoid flushing DB when not necessary * Filter events for which we perfom an `upsert` into the DB * Add a dedicated update function of the `thread:updated_at` that is lighter This should significantly reduce the DB lock contention. If it is not sufficient, we can de-sync the flush of the DB for `updated_at` --- codex-rs/core/src/codex.rs | 2 +- codex-rs/core/src/rollout/recorder.rs | 233 +++++++++++++++++++++++--- codex-rs/core/src/state_db.rs | 31 +++- codex-rs/state/src/extract.rs | 11 ++ codex-rs/state/src/lib.rs | 1 + codex-rs/state/src/runtime/threads.rs | 116 ++++++++++++- 6 files changed, 366 insertions(+), 28 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index c993eb3c6..ca351e883 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -1694,7 +1694,7 @@ impl Session { self.services.state_db.clone() } - /// Ensure all rollout writes are durably flushed. + /// Ensure rollout file writes are durably flushed. pub(crate) async fn flush_rollout(&self) { let recorder = { let guard = self.services.rollout.lock().await; diff --git a/codex-rs/core/src/rollout/recorder.rs b/codex-rs/core/src/rollout/recorder.rs index 466c97497..df150bf2b 100644 --- a/codex-rs/core/src/rollout/recorder.rs +++ b/codex-rs/core/src/rollout/recorder.rs @@ -7,6 +7,7 @@ use std::path::Path; use std::path::PathBuf; use chrono::SecondsFormat; +use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::BaseInstructions; @@ -448,7 +449,6 @@ impl RolloutRecorder { // future will yield, which is fine – we only need to ensure we do not // perform *blocking* I/O on the caller's thread. let (tx, rx) = mpsc::channel::(256); - // Spawn a Tokio task that owns the file handle and performs async // writes. Using `tokio::fs::File` keeps everything on the async I/O // driver instead of blocking the runtime. @@ -614,14 +614,15 @@ impl RolloutRecorder { 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}"))), + .map_err(|e| IoError::other(format!("failed waiting for rollout shutdown: {e}")))?, Err(e) => { warn!("failed to send rollout shutdown command: {e}"); - Err(IoError::other(format!( + return Err(IoError::other(format!( "failed to send rollout shutdown command: {e}" - ))) + ))); } - } + }; + Ok(()) } } @@ -744,25 +745,21 @@ async fn rollout_writer( while let Some(cmd) = rx.recv().await { match cmd { RolloutCmd::AddItems(items) => { - let mut persisted_items = Vec::new(); - for item in items { - persisted_items.push(item); - } - if persisted_items.is_empty() { + if items.is_empty() { continue; } if writer.is_none() { - buffered_items.extend(persisted_items); + buffered_items.extend(items); continue; } write_and_reconcile_items( writer.as_mut(), - persisted_items.as_slice(), + items.as_slice(), &rollout_path, state_db_ctx.as_deref(), - &mut state_builder, + state_builder.as_ref(), default_provider.as_str(), ) .await?; @@ -800,7 +797,7 @@ async fn rollout_writer( buffered_items.as_slice(), &rollout_path, state_db_ctx.as_deref(), - &mut state_builder, + state_builder.as_ref(), default_provider.as_str(), ) .await?; @@ -861,13 +858,12 @@ async fn write_session_meta( if let Some(writer) = writer.as_mut() { writer.write_rollout_item(&rollout_item).await?; } - state_db::reconcile_rollout( + sync_thread_state_after_write( state_db_ctx, rollout_path, - default_provider, state_builder.as_ref(), std::slice::from_ref(&rollout_item), - None, + default_provider, (!generate_memories).then_some("disabled"), ) .await; @@ -879,7 +875,7 @@ async fn write_and_reconcile_items( items: &[RolloutItem], rollout_path: &Path, state_db_ctx: Option<&StateRuntime>, - state_builder: &mut Option, + state_builder: Option<&ThreadMetadataBuilder>, default_provider: &str, ) -> std::io::Result<()> { if let Some(writer) = writer.as_mut() { @@ -887,20 +883,65 @@ async fn write_and_reconcile_items( writer.write_rollout_item(item).await?; } } - if let Some(builder) = state_builder.as_mut() { - builder.rollout_path = rollout_path.to_path_buf(); + sync_thread_state_after_write( + state_db_ctx, + rollout_path, + state_builder, + items, + default_provider, + None, + ) + .await; + Ok(()) +} + +async fn sync_thread_state_after_write( + state_db_ctx: Option<&StateRuntime>, + rollout_path: &Path, + state_builder: Option<&ThreadMetadataBuilder>, + items: &[RolloutItem], + default_provider: &str, + new_thread_memory_mode: Option<&str>, +) { + let updated_at = Utc::now(); + if new_thread_memory_mode.is_some() + || items + .iter() + .any(codex_state::rollout_item_affects_thread_metadata) + { + state_db::apply_rollout_items( + state_db_ctx, + rollout_path, + default_provider, + state_builder, + items, + "rollout_writer", + new_thread_memory_mode, + Some(updated_at), + ) + .await; + return; + } + + let thread_id = state_builder + .map(|builder| builder.id) + .or_else(|| metadata::builder_from_items(items, rollout_path).map(|builder| builder.id)); + if state_db::touch_thread_updated_at(state_db_ctx, thread_id, updated_at, "rollout_writer") + .await + { + return; } state_db::apply_rollout_items( state_db_ctx, rollout_path, default_provider, - state_builder.as_ref(), + state_builder, items, "rollout_writer", - None, + new_thread_memory_mode, + Some(updated_at), ) .await; - Ok(()) } struct JsonlWriter { @@ -1079,6 +1120,7 @@ mod tests { use std::io::Write; use std::path::Path; use std::path::PathBuf; + use std::time::Duration; use tempfile::TempDir; use uuid::Uuid; @@ -1200,6 +1242,151 @@ mod tests { Ok(()) } + #[tokio::test] + async fn metadata_irrelevant_events_touch_state_db_updated_at() -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let mut config = ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .build() + .await?; + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow sqlite"); + + let state_db = codex_state::StateRuntime::init( + home.path().to_path_buf(), + config.model_provider_id.clone(), + None, + ) + .await + .expect("state db should initialize"); + state_db + .mark_backfill_complete(None) + .await + .expect("backfill should be complete"); + + let thread_id = ThreadId::new(); + let recorder = RolloutRecorder::new( + &config, + RolloutRecorderParams::new( + thread_id, + None, + SessionSource::Cli, + BaseInstructions::default(), + Vec::new(), + EventPersistenceMode::Limited, + ), + Some(state_db.clone()), + None, + ) + .await?; + + recorder + .record_items(&[RolloutItem::EventMsg(EventMsg::UserMessage( + UserMessageEvent { + message: "first-user-message".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + ))]) + .await?; + recorder.persist().await?; + recorder.flush().await?; + let initial_thread = state_db + .get_thread(thread_id) + .await + .expect("thread should load") + .expect("thread should exist"); + let initial_updated_at = initial_thread.updated_at; + let initial_title = initial_thread.title.clone(); + let initial_first_user_message = initial_thread.first_user_message.clone(); + + tokio::time::sleep(Duration::from_secs(1)).await; + + recorder + .record_items(&[RolloutItem::EventMsg(EventMsg::AgentMessage( + AgentMessageEvent { + message: "assistant text".to_string(), + phase: None, + }, + ))]) + .await?; + recorder.flush().await?; + + let updated_thread = state_db + .get_thread(thread_id) + .await + .expect("thread should load after agent message") + .expect("thread should still exist"); + + assert!(updated_thread.updated_at > initial_updated_at); + assert_eq!(updated_thread.title, initial_title); + assert_eq!( + updated_thread.first_user_message, + initial_first_user_message + ); + + recorder.shutdown().await?; + Ok(()) + } + + #[tokio::test] + async fn metadata_irrelevant_events_fall_back_to_upsert_when_thread_missing() + -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let mut config = ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .build() + .await?; + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow sqlite"); + + let state_db = codex_state::StateRuntime::init( + home.path().to_path_buf(), + config.model_provider_id.clone(), + None, + ) + .await + .expect("state db should initialize"); + let thread_id = ThreadId::new(); + let rollout_path = home.path().join("rollout.jsonl"); + let builder = ThreadMetadataBuilder::new( + thread_id, + rollout_path.clone(), + Utc::now(), + SessionSource::Cli, + ); + let items = vec![RolloutItem::EventMsg(EventMsg::AgentMessage( + AgentMessageEvent { + message: "assistant text".to_string(), + phase: None, + }, + ))]; + + sync_thread_state_after_write( + Some(state_db.as_ref()), + rollout_path.as_path(), + Some(&builder), + items.as_slice(), + config.model_provider_id.as_str(), + None, + ) + .await; + + let thread = state_db + .get_thread(thread_id) + .await + .expect("thread should load after fallback") + .expect("thread should be inserted after fallback"); + assert_eq!(thread.id, thread_id); + + Ok(()) + } + #[tokio::test] async fn list_threads_db_disabled_does_not_skip_paginated_items() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/core/src/state_db.rs b/codex-rs/core/src/state_db.rs index 30140eb1c..2a9fa5e45 100644 --- a/codex-rs/core/src/state_db.rs +++ b/codex-rs/core/src/state_db.rs @@ -361,6 +361,7 @@ pub async fn reconcile_rollout( items, "reconcile_rollout", new_thread_memory_mode, + None, ) .await; return; @@ -491,6 +492,7 @@ pub async fn read_repair_rollout_path( } /// Apply rollout items incrementally to SQLite. +#[allow(clippy::too_many_arguments)] pub async fn apply_rollout_items( context: Option<&codex_state::StateRuntime>, rollout_path: &Path, @@ -499,6 +501,7 @@ pub async fn apply_rollout_items( items: &[RolloutItem], stage: &str, new_thread_memory_mode: Option<&str>, + updated_at_override: Option>, ) { let Some(ctx) = context else { return; @@ -520,7 +523,13 @@ pub async fn apply_rollout_items( builder.rollout_path = rollout_path.to_path_buf(); builder.cwd = normalize_cwd_for_state_db(&builder.cwd); if let Err(err) = ctx - .apply_rollout_items(&builder, items, None, new_thread_memory_mode) + .apply_rollout_items( + &builder, + items, + None, + new_thread_memory_mode, + updated_at_override, + ) .await { warn!( @@ -530,6 +539,26 @@ pub async fn apply_rollout_items( } } +pub async fn touch_thread_updated_at( + context: Option<&codex_state::StateRuntime>, + thread_id: Option, + updated_at: DateTime, + stage: &str, +) -> bool { + let Some(ctx) = context else { + return false; + }; + let Some(thread_id) = thread_id else { + return false; + }; + ctx.touch_thread_updated_at(thread_id, updated_at) + .await + .unwrap_or_else(|err| { + warn!("state db touch_thread_updated_at failed during {stage} for {thread_id}: {err}"); + false + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 1a7fb3724..ba425adbe 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -29,6 +29,17 @@ pub fn apply_rollout_item( } } +/// Return whether this rollout item can mutate thread metadata stored in SQLite. +pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool { + match item { + RolloutItem::SessionMeta(_) | RolloutItem::TurnContext(_) => true, + RolloutItem::EventMsg(EventMsg::TokenCount(_) | EventMsg::UserMessage(_)) => true, + RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) => { + false + } + } +} + fn apply_session_meta_from_item(metadata: &mut ThreadMetadata, meta_line: &SessionMetaLine) { if metadata.id != meta_line.meta.id { // Ignore session_meta lines that don't match the canonical thread ID, diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index fa0d74c59..84b4462ce 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -23,6 +23,7 @@ pub use runtime::StateRuntime; /// /// Most consumers should prefer [`StateRuntime`]. pub use extract::apply_rollout_item; +pub use extract::rollout_item_affects_thread_metadata; pub use model::AgentJob; pub use model::AgentJobCreateParams; pub use model::AgentJobItem; diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 2a1f44e1e..e671dfaf6 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -278,6 +278,19 @@ ON CONFLICT(id) DO NOTHING Ok(result.rows_affected() > 0) } + pub async fn touch_thread_updated_at( + &self, + thread_id: ThreadId, + updated_at: DateTime, + ) -> anyhow::Result { + let result = sqlx::query("UPDATE threads SET updated_at = ? WHERE id = ?") + .bind(datetime_to_epoch_seconds(updated_at)) + .bind(thread_id.to_string()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + pub async fn update_thread_git_info( &self, thread_id: ThreadId, @@ -436,6 +449,7 @@ ON CONFLICT(thread_id, position) DO NOTHING items: &[RolloutItem], otel: Option<&OtelManager>, new_thread_memory_mode: Option<&str>, + updated_at_override: Option>, ) -> anyhow::Result<()> { if items.is_empty() { return Ok(()); @@ -451,7 +465,11 @@ ON CONFLICT(thread_id, position) DO NOTHING if let Some(existing_metadata) = existing_metadata.as_ref() { metadata.prefer_existing_git_info(existing_metadata); } - if let Some(updated_at) = file_modified_time_utc(builder.rollout_path.as_path()).await { + let updated_at = match updated_at_override { + Some(updated_at) => Some(updated_at), + None => file_modified_time_utc(builder.rollout_path.as_path()).await, + }; + if let Some(updated_at) = updated_at { metadata.updated_at = updated_at; } // Keep the thread upsert before dynamic tools to satisfy the foreign key constraint: @@ -649,6 +667,7 @@ mod tests { use super::*; use crate::runtime::test_support::test_thread_metadata; use crate::runtime::test_support::unique_temp_dir; + use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; @@ -735,7 +754,7 @@ mod tests { })]; runtime - .apply_rollout_items(&builder, &items, None, None) + .apply_rollout_items(&builder, &items, None, None, None) .await .expect("apply_rollout_items should succeed"); @@ -793,7 +812,7 @@ mod tests { })]; runtime - .apply_rollout_items(&builder, &items, None, None) + .apply_rollout_items(&builder, &items, None, None, None) .await .expect("apply_rollout_items should succeed"); @@ -947,4 +966,95 @@ mod tests { assert_eq!(persisted.git_branch, None); assert_eq!(persisted.git_origin_url, None); } + + #[tokio::test] + async fn touch_thread_updated_at_updates_only_updated_at() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None) + .await + .expect("state db should initialize"); + let thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000791").expect("valid thread id"); + let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); + metadata.title = "original title".to_string(); + metadata.first_user_message = Some("first-user-message".to_string()); + + runtime + .upsert_thread(&metadata) + .await + .expect("initial upsert should succeed"); + + let touched_at = DateTime::::from_timestamp(1_700_001_111, 0).expect("timestamp"); + let touched = runtime + .touch_thread_updated_at(thread_id, touched_at) + .await + .expect("touch should succeed"); + assert!(touched); + + let persisted = runtime + .get_thread(thread_id) + .await + .expect("thread should load") + .expect("thread should exist"); + assert_eq!(persisted.updated_at, touched_at); + assert_eq!(persisted.title, "original title"); + assert_eq!( + persisted.first_user_message.as_deref(), + Some("first-user-message") + ); + } + + #[tokio::test] + async fn apply_rollout_items_uses_override_updated_at_when_provided() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None) + .await + .expect("state db should initialize"); + let thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000792").expect("valid thread id"); + let metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); + + runtime + .upsert_thread(&metadata) + .await + .expect("initial upsert should succeed"); + + let builder = ThreadMetadataBuilder::new( + thread_id, + metadata.rollout_path.clone(), + metadata.created_at, + SessionSource::Cli, + ); + let items = vec![RolloutItem::EventMsg(EventMsg::TokenCount( + codex_protocol::protocol::TokenCountEvent { + info: Some(codex_protocol::protocol::TokenUsageInfo { + total_token_usage: codex_protocol::protocol::TokenUsage { + input_tokens: 0, + cached_input_tokens: 0, + output_tokens: 0, + reasoning_output_tokens: 0, + total_tokens: 321, + }, + last_token_usage: codex_protocol::protocol::TokenUsage::default(), + model_context_window: None, + }), + rate_limits: None, + }, + ))]; + let override_updated_at = + DateTime::::from_timestamp(1_700_001_234, 0).expect("timestamp"); + + runtime + .apply_rollout_items(&builder, &items, None, None, Some(override_updated_at)) + .await + .expect("apply_rollout_items should succeed"); + + let persisted = runtime + .get_thread(thread_id) + .await + .expect("thread should load") + .expect("thread should exist"); + assert_eq!(persisted.tokens_used, 321); + assert_eq!(persisted.updated_at, override_updated_at); + } }