Retain rollout migration locks across worker cancellation (#53794)

## Why

Cancelling an async migration caller can leave its blocking compression or decompression worker running. Releasing migration locks before that worker finishes allows other writers or migrations to access files still being processed.

## What changed

Share maintenance, cross-process writer, and local writer locks with blocking workers during migration and recovery. Keep those locks held until file work completes, releasing cross-process locks before allowing local writers to proceed.

## Testing

Add a regression test for cancelled compression and decompression callers that verifies locks remain held until the worker finishes and the output is complete. Preserve the Unix check that decompression temporaries have owner-only permissions.

GitOrigin-RevId: 74855250d89ecab81eaadc2956e328d37b65e72e

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<Arc<RolloutMaintenanceGuard>>,
+    _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<RolloutMigrationReport> {
         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<ThreadId, String>,
         limiter: &mut RolloutMigrationRateLimiter,
+        maintenance_guard: Option<&Arc<RolloutMaintenanceGuard>>,
     ) -> ThreadStoreResult<Option<RolloutMigrationOutcome>> {
         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<ThreadId, String>,
         limiter: &mut RolloutMigrationRateLimiter,
+        locks: &Arc<MigrationLocks>,
     ) -> 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<ThreadId, String>,
         limiter: &mut RolloutMigrationRateLimiter,
+        maintenance_guard: Option<&Arc<RolloutMaintenanceGuard>>,
     ) -> ThreadStoreResult<PathBuf> {
-        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)?

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<PathBuf>
 pub(super) async fn decompress_rollout_to_path(
     compressed_path: &Path,
     plain_path: &Path,
+    locks: &Arc<MigrationLocks>,
 ) -> 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<SystemTime>,
+    locks: &Arc<MigrationLocks>,
 ) -> 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)

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]