diff --git a/codex-rs/thread-store/src/local/rollout_migration.rs b/codex-rs/thread-store/src/local/rollout_migration.rs index c259a6f034..3f54ec1fdb 100644 --- a/codex-rs/thread-store/src/local/rollout_migration.rs +++ b/codex-rs/thread-store/src/local/rollout_migration.rs @@ -7,11 +7,13 @@ //! The important invariant is that we always leave behind either the original legacy rollout or a //! recoverable paginated rollout. Once the rollout path is replaced, the durable `.pending` //! journal must be enough for a later migration run to finish SQLite recovery safely. +//! Compression and decompression workers retain these locks if their async caller is cancelled. use std::collections::HashMap; use std::io; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; use std::time::Duration; use chrono::DateTime; @@ -24,6 +26,7 @@ use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use codex_rollout::RolloutItem; use codex_rollout::RolloutLine; +use codex_rollout::RolloutMaintenanceGuard; use serde::Serialize; use tokio::fs::File; use tokio::io::AsyncBufReadExt; @@ -31,9 +34,11 @@ use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::io::BufWriter; +use tokio::sync::OwnedMutexGuard; use tokio::time::Instant; use super::LocalThreadStore; +use super::WriterLockGuard; use super::helpers::distinct_thread_metadata_title; use super::helpers::has_guardian_default_title; use super::thread_history; @@ -74,6 +79,14 @@ use telemetry::RolloutMigrationTrigger; const PROJECTION_BATCH_BYTES: u64 = 256 * 1024; const MAX_ROLLOUT_LINE_BYTES: usize = 16 * 1024 * 1024; +/// Shared with blocking workers so cancellation cannot release their file ownership early. +struct MigrationLocks { + _maintenance: Option>, + _writer: WriterLockGuard, + // Let local writers proceed only after releasing the cross-process locks. + _live_writer: OwnedMutexGuard<()>, +} + enum CanonicalizationAttempt { Complete { expected_length: u64, @@ -316,7 +329,7 @@ impl LocalThreadStore { paths: RolloutMigrationPaths, ) -> ThreadStoreResult { let mut limiter = RolloutMigrationRateLimiter::new(options.max_mib_per_second)?; - let _maintenance_guard = match options.mode { + let maintenance_guard = match options.mode { RolloutMigrationMode::DryRun => None, RolloutMigrationMode::Apply => Some( codex_rollout::try_acquire_rollout_maintenance_lock(&self.config.codex_home) @@ -326,7 +339,8 @@ impl LocalThreadStore { .to_string(), })?, ), - }; + } + .map(Arc::new); let mut paths = match paths { RolloutMigrationPaths::Discover => { find_all_rollout_paths(&self.config.codex_home).await? @@ -357,7 +371,13 @@ impl LocalThreadStore { for (index, path) in paths.into_iter().enumerate() { let outcome = self - .migrate_rollout_path(path, &options, &legacy_names, &mut limiter) + .migrate_rollout_path( + path, + &options, + &legacy_names, + &mut limiter, + maintenance_guard.as_ref(), + ) .await?; let outcome_status = outcome.as_ref().map(|outcome| outcome.status); if let Some(outcome) = outcome { @@ -379,6 +399,7 @@ impl LocalThreadStore { options: &RolloutMigrationOptions, legacy_names: &HashMap, limiter: &mut RolloutMigrationRateLimiter, + maintenance_guard: Option<&Arc>, ) -> ThreadStoreResult> { let mut retried_moved_path = false; let metadata = loop { @@ -482,7 +503,6 @@ impl LocalThreadStore { if metadata.meta.history_mode == ThreadHistoryMode::Paginated { let bytes_before = limiter.bytes_processed; let result = if pending_published_migration { - let _live_writer_guard = self.live_writer_locks.lock(thread_id).await; match self .recover_published_migration( thread_id, @@ -490,6 +510,7 @@ impl LocalThreadStore { &journal_path, legacy_names, limiter, + maintenance_guard, ) .await { @@ -535,8 +556,8 @@ impl LocalThreadStore { ))); } - let _live_writer_guard = self.live_writer_locks.lock(thread_id).await; - let _writer_guard = match self.acquire_writer_lock(thread_id) { + let live_writer_guard = self.live_writer_locks.lock(thread_id).await; + let writer_guard = match self.acquire_writer_lock(thread_id) { Ok(guard) => guard, Err(ThreadStoreError::Conflict { message }) => { return Ok(Some(skipped_busy_outcome( @@ -555,6 +576,11 @@ impl LocalThreadStore { ))); } }; + let locks = Arc::new(MigrationLocks { + _maintenance: maintenance_guard.cloned(), + _live_writer: live_writer_guard, + _writer: writer_guard, + }); // SessionMeta gives us the writer-lock id, so archiving can win between that read and // lock acquisition. Once the lock is ours, follow the same rollout to its current path. if !tokio::fs::try_exists(&path) @@ -567,7 +593,7 @@ impl LocalThreadStore { } let bytes_before = limiter.bytes_processed; let result = match self - .migrate_one_rollout(thread_id, &path, &journal_path, kind, legacy_names, limiter) + .migrate_one_rollout(thread_id, &path, kind, legacy_names, limiter, &locks) .await { Ok(()) => Ok(RolloutMigrationStatus::Migrated), @@ -602,10 +628,10 @@ impl LocalThreadStore { &self, thread_id: ThreadId, rollout_path: &Path, - journal_path: &Path, kind: RolloutMigrationKind, legacy_names: &HashMap, limiter: &mut RolloutMigrationRateLimiter, + locks: &Arc, ) -> ClassifiedMigrationResult<()> { if let Some(state_db) = &self.state_db && with_failure_reason( @@ -638,8 +664,9 @@ impl LocalThreadStore { thread_history::delete_thread(self, thread_id).await, RolloutMigrationFailureReason::SqliteMaterializationFailed, )?; + let journal_path = migration_journal_path(&self.config.codex_home, thread_id); with_failure_reason( - write_migration_journal(journal_path).await, + write_migration_journal(&journal_path).await, RolloutMigrationFailureReason::RolloutPublishFailed, )?; @@ -653,7 +680,7 @@ impl LocalThreadStore { let source_permissions = source_metadata.permissions(); let source_path = if let Some(decompressed_path) = decompressed_path.as_ref() { with_failure_reason( - decompress_rollout_to_path(rollout_path, decompressed_path).await, + decompress_rollout_to_path(rollout_path, decompressed_path, locks).await, RolloutMigrationFailureReason::RolloutReadFailed, )?; let decompressed_bytes = with_failure_reason( @@ -783,7 +810,7 @@ impl LocalThreadStore { let publish_result = async { let compressed_staged_path = if compressed { let path = compressed_staged_rollout_path(rollout_path)?; - compress_rollout_to_path(&staged_path, &path, source_permissions, source_modified) + compress_rollout_to_path(&staged_path, &path, source_permissions, source_modified, locks) .await?; Some(path) } else { @@ -817,7 +844,7 @@ impl LocalThreadStore { .map_err(migration_error)?; } sync_parent_directory(rollout_path).await?; - self.finish_published_migration(thread_id, journal_path, legacy_names) + self.finish_published_migration(thread_id, &journal_path, legacy_names) .await } .await; @@ -996,13 +1023,19 @@ impl LocalThreadStore { journal_path: &Path, legacy_names: &HashMap, limiter: &mut RolloutMigrationRateLimiter, + maintenance_guard: Option<&Arc>, ) -> ThreadStoreResult { - let _writer_guard = self.acquire_writer_lock(thread_id)?; + let live_writer_guard = self.live_writer_locks.lock(thread_id).await; + let locks = Arc::new(MigrationLocks { + _maintenance: maintenance_guard.cloned(), + _live_writer: live_writer_guard, + _writer: self.acquire_writer_lock(thread_id)?, + }); let decompressed_path = rollout_path_is_compressed(rollout_path) .then(|| decompressed_staged_rollout_path(rollout_path)) .transpose()?; let projection_path = if let Some(decompressed_path) = decompressed_path.as_ref() { - decompress_rollout_to_path(rollout_path, decompressed_path).await?; + decompress_rollout_to_path(rollout_path, decompressed_path, &locks).await?; let decompressed_bytes = tokio::fs::metadata(decompressed_path) .await .map_err(migration_error)? diff --git a/codex-rs/thread-store/src/local/rollout_migration/publish.rs b/codex-rs/thread-store/src/local/rollout_migration/publish.rs index 4c0f9656fb..731423714c 100644 --- a/codex-rs/thread-store/src/local/rollout_migration/publish.rs +++ b/codex-rs/thread-store/src/local/rollout_migration/publish.rs @@ -14,6 +14,7 @@ use std::io; use std::io::Write; use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; use std::time::SystemTime; #[cfg(unix)] @@ -29,6 +30,7 @@ use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::io::BufWriter; +use super::MigrationLocks; use super::migration_error; use crate::ThreadStoreResult; @@ -90,10 +92,13 @@ fn staged_path(rollout_path: &Path, suffix: &str) -> ThreadStoreResult pub(super) async fn decompress_rollout_to_path( compressed_path: &Path, plain_path: &Path, + locks: &Arc, ) -> ThreadStoreResult<()> { let compressed_path = compressed_path.to_path_buf(); let plain_path = plain_path.to_path_buf(); + let locks = Arc::clone(locks); tokio::task::spawn_blocking(move || -> io::Result<()> { + let _locks = locks; let input = std::fs::File::open(compressed_path)?; let mut decoder = zstd::stream::read::Decoder::new(input)?; let mut options = std::fs::OpenOptions::new(); @@ -116,10 +121,13 @@ pub(super) async fn compress_rollout_to_path( compressed_path: &Path, permissions: Permissions, modified_at: Option, + locks: &Arc, ) -> ThreadStoreResult<()> { let plain_path = plain_path.to_path_buf(); let compressed_path = compressed_path.to_path_buf(); + let locks = Arc::clone(locks); tokio::task::spawn_blocking(move || -> io::Result<()> { + let _locks = locks; let mut input = std::fs::File::open(plain_path)?; let output = std::fs::OpenOptions::new() .create(true) diff --git a/codex-rs/thread-store/src/local/rollout_migration_tests.rs b/codex-rs/thread-store/src/local/rollout_migration_tests.rs index e06cac6b02..4ddcbd596f 100644 --- a/codex-rs/thread-store/src/local/rollout_migration_tests.rs +++ b/codex-rs/thread-store/src/local/rollout_migration_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; #[cfg(unix)] use std::os::unix::fs::PermissionsExt; @@ -53,13 +54,14 @@ use serde_json::json; use tempfile::TempDir; use super::LocalThreadStore; +use super::MigrationLocks; use super::RolloutMigrationFailureReason; use super::RolloutMigrationMode; use super::RolloutMigrationOptions; use super::RolloutMigrationPaths; use super::RolloutMigrationProgress; use super::RolloutMigrationStatus; -#[cfg(unix)] +use super::compress_rollout_to_path; use super::decompress_rollout_to_path; use super::migration_journal_path; use super::telemetry::RolloutMigrationTrigger; @@ -2672,31 +2674,116 @@ async fn migration_repairs_a_missing_paginated_name_when_rerun() { ); } -#[cfg(unix)] -#[tokio::test] -async fn decompression_temporaries_are_owner_only() { - let home = TempDir::new().expect("create Codex home"); - let thread_id = ThreadId::new(); - let compressed_path = compress_rollout(&write_rollout( - home.path(), - thread_id, - SessionSource::Cli, - vec![user_message("compressed question")], - )); - let plain_path = home.path().join("decompressed.tmp"); +#[test] +fn cancelled_compression_workers_hold_migration_locks() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .expect("build runtime"); + runtime.block_on(async { + for compress in [true, false] { + let home = TempDir::new().expect("create Codex home"); + let thread_id = ThreadId::new(); + let plain_path = home.path().join("rollout.jsonl"); + let compressed_path = home.path().join("rollout.jsonl.zst"); + let contents = b"rollout contents\n"; + fs::write(&plain_path, contents).expect("write plain rollout"); + if !compress { + compress_rollout(&plain_path); + } + let source_path = if compress { + &plain_path + } else { + &compressed_path + }; + let permissions = fs::metadata(source_path) + .expect("source metadata") + .permissions(); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let locks = Arc::new(MigrationLocks { + _maintenance: Some(Arc::new( + codex_rollout::try_acquire_rollout_maintenance_lock(home.path()) + .expect("acquire maintenance lock") + .expect("maintenance lock is available"), + )), + _writer: store + .acquire_writer_lock(thread_id) + .expect("acquire writer lock"), + _live_writer: store.live_writer_locks.lock(thread_id).await, + }); + let coordination = store.live_writer_locks.coordination(thread_id).await; - decompress_rollout_to_path(&compressed_path, &plain_path) - .await - .expect("decompress rollout"); + // Hold the only blocking thread so the real file worker stays queued across cancellation. + let (reached_tx, reached_rx) = tokio::sync::oneshot::channel(); + let (resume_tx, resume_rx) = std::sync::mpsc::channel::<()>(); + let blocker = tokio::task::spawn_blocking(move || { + let _ = reached_tx.send(()); + // Dropping the sender after a failed assertion also lets runtime teardown finish. + let _ = resume_rx.recv(); + }); + reached_rx.await.expect("blocking thread occupied"); + let mut work = Box::pin(async { + if compress { + compress_rollout_to_path( + &plain_path, + &compressed_path, + permissions, + /*modified_at*/ None, + &locks, + ) + .await + } else { + decompress_rollout_to_path(&compressed_path, &plain_path, &locks).await + } + }); + assert!(futures::poll!(&mut work).is_pending()); + drop(work); + drop(locks); - assert_eq!( - fs::metadata(&plain_path) - .expect("read temporary metadata") - .permissions() - .mode() - & 0o777, - 0o600 - ); + assert!( + codex_rollout::try_acquire_rollout_maintenance_lock(home.path()) + .expect("probe maintenance lock") + .is_none() + ); + assert!(matches!( + store.acquire_writer_lock(thread_id), + Err(ThreadStoreError::Conflict { .. }) + )); + assert!(coordination.writer.try_lock().is_err()); + + drop(resume_tx); + blocker.await.expect("release blocking thread"); + // This guard is released last, after the worker's file I/O and other locks. + let _live_writer_guard = coordination.writer.lock().await; + assert!( + codex_rollout::try_acquire_rollout_maintenance_lock(home.path()) + .expect("probe released maintenance lock") + .is_some() + ); + store + .acquire_writer_lock(thread_id) + .expect("writer lock released"); + let decoded = if compress { + zstd::stream::decode_all( + fs::File::open(&compressed_path).expect("open compressed rollout"), + ) + .expect("decode complete rollout") + } else { + #[cfg(unix)] + assert_eq!( + fs::metadata(&plain_path) + .expect("temporary metadata") + .permissions() + .mode() + & 0o777, + 0o600 + ); + fs::read(&plain_path).expect("read decompressed rollout") + }; + assert_eq!(decoded, contents); + } + }); } #[tokio::test]