diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 2df77f7239..326cec8139 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -153,6 +153,8 @@ pub struct LocalThreadStore { state_db: Option, thread_history_db: Arc>, thread_data_cleanup: Option>, + #[cfg(test)] + startup_migration_retry_gate: Option>, } type ThreadDataCleanup = dyn Fn(Vec) -> ThreadStoreFuture<'static, ()> + Send + Sync; @@ -227,6 +229,14 @@ impl LiveWriterLocks { } } +// Per-store synchronization for exercising startup contention through the public entry point. +#[cfg(test)] +#[derive(Default)] +struct StartupMigrationRetryGate { + retry_started: tokio::sync::Notify, + retry_allowed: tokio::sync::Notify, +} + /// Process-scoped configuration for local thread storage. /// /// This describes where local storage lives. New-thread rollout metadata such @@ -271,6 +281,8 @@ impl LocalThreadStore { state_db, thread_history_db: Arc::new(OnceCell::new()), thread_data_cleanup: None, + #[cfg(test)] + startup_migration_retry_gate: None, } } diff --git a/codex-rs/thread-store/src/local/rollout_migration/startup.rs b/codex-rs/thread-store/src/local/rollout_migration/startup.rs index dbc48ed8c0..21cba720e0 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/startup.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/startup.rs @@ -208,6 +208,12 @@ async fn run_startup_migration( codex_rollout::try_acquire_rollout_maintenance_lock(&store.config.codex_home) .map_err(migration_error)? else { + #[cfg(test)] + if let Some(gate) = &store.startup_migration_retry_gate { + gate.retry_started.notify_one(); + gate.retry_allowed.notified().await; + continue; + } tokio::time::sleep(MAINTENANCE_RETRY_DELAY).await; continue; }; diff --git a/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs b/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs index d2e144ca33..a2b7d34f81 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/startup_tests.rs @@ -2,6 +2,7 @@ use std::fs; use std::io::Write; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; use std::time::Duration; use codex_protocol::ThreadId; @@ -22,6 +23,7 @@ use super::super::publish::migration_journal_path; use super::super::publish::write_migration_journal; use super::super::thread_history; use super::LocalThreadStore; +use crate::local::StartupMigrationRetryGate; use crate::local::test_support::test_config; const TIMESTAMP: &str = "2025-01-03T12:00:00Z"; @@ -297,7 +299,9 @@ async fn waits_for_rollout_maintenance_before_migrating() { let home = TempDir::new().expect("create Codex home"); let thread_id = ThreadId::new(); let path = write_rollout(home.path(), thread_id, ThreadHistoryMode::Legacy); - let store = indexed_store(home.path()).await; + let mut store = indexed_store(home.path()).await; + let gate = Arc::new(StartupMigrationRetryGate::default()); + store.startup_migration_retry_gate = Some(Arc::clone(&gate)); let maintenance_guard = codex_rollout::try_acquire_rollout_maintenance_lock(home.path()) .expect("acquire rollout maintenance lock") .expect("claim rollout maintenance lock"); @@ -309,13 +313,27 @@ async fn waits_for_rollout_maintenance_before_migrating() { .expect("migrate startup rollouts"); }); - assert!( - tokio::time::timeout(Duration::from_millis(50), &mut migration) + // Observe real lock contention instead of assuming the task reached it after 50 ms. + tokio::time::timeout(Duration::from_secs(2), async { + tokio::select! { + result = &mut migration => panic!("migration should wait for rollout maintenance: {result:?}"), + _ = gate.retry_started.notified() => {} + } + }) + .await + .expect("migration should encounter rollout maintenance"); + assert_eq!( + codex_rollout::read_session_meta_line(&path) .await - .is_err(), - "migration should wait for rollout maintenance" + .expect("read metadata while maintenance is held") + .meta + .history_mode, + ThreadHistoryMode::Legacy ); + drop(maintenance_guard); + // Drive the retry explicitly so the completion budget does not include the production backoff. + gate.retry_allowed.notify_one(); tokio::time::timeout(Duration::from_secs(2), migration) .await .expect("migration should retry rollout maintenance")