worker: test activation fault recovery

crates/worker/src/runtime/worker_allocation.rs

@@ -22,6 +22,10 @@ pub use lifecycle::{
     ScopeAllocationGuard, SegmentActivationCommitError, SegmentActivationGuard, SegmentLockInfo,
     adopt_allocation, install_top_level, install_top_level_with_deny, lookup_segment,
 };
+#[cfg(test)]
+pub(crate) use lifecycle::install_top_level_at_for_test;
+#[cfg(test)]
+pub(crate) use table::fail_next_save_for_test;
 pub use mutate::{
     delegate_scope, reclaim_delegated_scope, reclaim_stale, reclaim_stale_with, register_worker,
     register_worker_with_deny, release_worker,

crates/worker/src/runtime/worker_allocation/lifecycle.rs

@@ -201,6 +201,31 @@ pub fn install_top_level_with_deny(
     })
 }
 
+#[cfg(test)]
+pub(crate) fn install_top_level_at_for_test(
+    lock_path: PathBuf,
+    worker_name: String,
+    pid: u32,
+    socket: PathBuf,
+    scope_allow: Vec<ScopeRule>,
+    segment_id: SegmentId,
+) -> Result<ScopeAllocationGuard, ScopeLockError> {
+    let mut guard = LockFileGuard::open(&lock_path)?;
+    super::mutate::register_worker_with_deny(
+        &mut guard,
+        worker_name.clone(),
+        pid,
+        socket,
+        scope_allow,
+        Vec::new(),
+        segment_id,
+    )?;
+    Ok(ScopeAllocationGuard {
+        worker_name,
+        lock_path,
+    })
+}
+
 /// Take ownership of an existing allocation that was pre-registered by
 /// a spawning Worker.
 ///
@@ -260,7 +285,7 @@ pub fn lookup_segment(segment_id: SegmentId) -> Result<Option<SegmentLockInfo>,
 
 #[cfg(test)]
 mod tests {
-    use super::super::table::Allocation;
+    use super::super::table::{Allocation, LockFile};
     use super::super::test_util::*;
     use super::*;
     use tempfile::TempDir;
@@ -464,7 +489,7 @@ mod tests {
             } if expected == old_segment
         ));
         // The authority is returned rather than dropped on failure. Production
-        // stores it on SegmentState to keep restore admission blocked.
+        // stores it on Worker to keep restore admission blocked.
         let (sent, received) = std::sync::mpsc::channel();
         let lookup = std::thread::spawn(move || {
             sent.send(lookup_segment(replacement).unwrap()).unwrap();
@@ -480,6 +505,39 @@ mod tests {
         lookup.join().unwrap();
     }
 
+    #[test]
+    fn allocation_save_failure_preserves_source_table_and_returns_authority() {
+        let dir = TempDir::new().unwrap();
+        let _sandbox = RuntimeDirSandbox::new(dir.path());
+        let lock_path = dir.path().join("workers.json");
+        let old_segment = sid();
+        let replacement = sid();
+        let guard = install_top_level(
+            "activation".into(),
+            std::process::id(),
+            sock("activation"),
+            vec![write_rule("/work", true)],
+            old_segment,
+        )
+        .unwrap();
+        let activation = guard
+            .begin_segment_activation(old_segment, replacement)
+            .unwrap();
+        crate::runtime::worker_allocation::table::fail_next_save_for_test(&lock_path);
+
+        let error = activation.commit().unwrap_err();
+        let (source, authority) = error.into_parts();
+        assert!(matches!(source, ScopeLockError::Io(_)));
+        let persisted: LockFile = serde_json::from_slice(&std::fs::read(&lock_path).unwrap()).unwrap();
+        assert_eq!(
+            persisted
+                .find("activation")
+                .and_then(|allocation| allocation.segment_id),
+            Some(old_segment)
+        );
+        drop(authority);
+    }
+
     #[test]
     fn dropping_segment_activation_preserves_source_allocation() {
         let dir = TempDir::new().unwrap();

crates/worker/src/runtime/worker_allocation/table.rs

@@ -7,6 +7,14 @@ use std::path::{Path, PathBuf};
 use std::thread;
 use std::time::{Duration, Instant};
 
+#[cfg(test)]
+static FAIL_NEXT_SAVE_PATH: std::sync::Mutex<Option<PathBuf>> = std::sync::Mutex::new(None);
+
+#[cfg(test)]
+pub(crate) fn fail_next_save_for_test(path: &Path) {
+    *FAIL_NEXT_SAVE_PATH.lock().unwrap() = Some(path.to_path_buf());
+}
+
 use fs4::fs_std::FileExt;
 use manifest::{ScopeRule, paths};
 use serde::{Deserialize, Serialize};
@@ -209,6 +217,14 @@ impl LockFileGuard {
     /// Persist with atomic replacement while the separate allocation lock stays
     /// held. A crash exposes either the old complete table or the new one.
     pub fn save(&mut self) -> io::Result<()> {
+        #[cfg(test)]
+        {
+            let mut fail_path = FAIL_NEXT_SAVE_PATH.lock().unwrap();
+            if fail_path.as_ref() == Some(&self.data_path) {
+                *fail_path = None;
+                return Err(io::Error::other("injected allocation save failure"));
+            }
+        }
         let json = serde_json::to_vec_pretty(&self.data).map_err(io::Error::other)?;
         let mut temp_os = self.data_path.as_os_str().to_owned();
         temp_os.push(".tmp");

crates/worker/src/spawn/registry.rs

@@ -9,6 +9,8 @@
 //! rules to the parent scope.
 
 use std::collections::{BTreeMap, HashSet};
+#[cfg(test)]
+use std::collections::HashMap;
 use std::io;
 use std::sync::{
     Arc, Mutex,
@@ -312,6 +314,8 @@ pub struct SpawnedWorkerRegistry {
     internal_records: std::sync::Mutex<Vec<InternalSpawnedWorkerRecord>>,
     service_records: std::sync::Mutex<Vec<InternalServiceWorkerRecord>>,
     quarantined_service_cleanup: std::sync::Mutex<HashSet<String>>,
+    #[cfg(test)]
+    fail_service_stops: std::sync::Mutex<HashMap<String, usize>>,
     internal_names: std::sync::Mutex<HashSet<String>>,
     internal_shutting_down: AtomicBool,
     pending_internal_spawns: AtomicUsize,
@@ -333,6 +337,8 @@ impl SpawnedWorkerRegistry {
             internal_records: std::sync::Mutex::new(Vec::new()),
             service_records: std::sync::Mutex::new(Vec::new()),
             quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
+            #[cfg(test)]
+            fail_service_stops: std::sync::Mutex::new(HashMap::new()),
             internal_names: std::sync::Mutex::new(HashSet::new()),
             internal_shutting_down: AtomicBool::new(false),
             pending_internal_spawns: AtomicUsize::new(0),
@@ -349,6 +355,8 @@ impl SpawnedWorkerRegistry {
             internal_records: std::sync::Mutex::new(Vec::new()),
             service_records: std::sync::Mutex::new(Vec::new()),
             quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
+            #[cfg(test)]
+            fail_service_stops: std::sync::Mutex::new(HashMap::new()),
             internal_names: std::sync::Mutex::new(HashSet::new()),
             internal_shutting_down: AtomicBool::new(false),
             pending_internal_spawns: AtomicUsize::new(0),
@@ -364,6 +372,8 @@ impl SpawnedWorkerRegistry {
             internal_records: std::sync::Mutex::new(Vec::new()),
             service_records: std::sync::Mutex::new(Vec::new()),
             quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
+            #[cfg(test)]
+            fail_service_stops: std::sync::Mutex::new(HashMap::new()),
             internal_names: std::sync::Mutex::new(HashSet::new()),
             internal_shutting_down: AtomicBool::new(false),
             pending_internal_spawns: AtomicUsize::new(0),
@@ -448,6 +458,8 @@ impl SpawnedWorkerRegistry {
                 internal_records: std::sync::Mutex::new(Vec::new()),
                 service_records: std::sync::Mutex::new(Vec::new()),
             quarantined_service_cleanup: std::sync::Mutex::new(HashSet::new()),
+            #[cfg(test)]
+            fail_service_stops: std::sync::Mutex::new(HashMap::new()),
                 internal_names: std::sync::Mutex::new(HashSet::new()),
                 internal_shutting_down: AtomicBool::new(false),
                 pending_internal_spawns: AtomicUsize::new(0),
@@ -540,9 +552,35 @@ impl SpawnedWorkerRegistry {
         Ok(worker_ref)
     }
 
+    #[cfg(test)]
+    pub(crate) fn fail_service_stops_for_test(&self, session_id: &str, count: usize) {
+        self.fail_service_stops
+            .lock()
+            .unwrap()
+            .insert(session_id.to_owned(), count);
+    }
+
     /// Stop and remove one parent-owned service Worker. This is host-only and is
     /// intentionally separate from the SubWorker control surface.
+    #[cfg(test)]
+    pub(crate) fn is_service_cleanup_quarantined_for_test(&self, session_id: &str) -> bool {
+        self.quarantined_service_cleanup
+            .lock()
+            .unwrap()
+            .contains(session_id)
+    }
+
     pub(crate) async fn stop_service(&self, session_id: &str) -> io::Result<bool> {
+        #[cfg(test)]
+        {
+            let mut failures = self.fail_service_stops.lock().unwrap();
+            if let Some(remaining) = failures.get_mut(session_id) {
+                if *remaining > 0 {
+                    *remaining -= 1;
+                    return Err(io::Error::other("injected service stop failure"));
+                }
+            }
+        }
         let record = self
             .service_records
             .lock()

crates/worker/src/worker.rs

@@ -8584,6 +8584,100 @@ mod build_summary_prompt_tests {
         assert_eq!(restored.active_run_turn_count, None);
     }
 
+    #[tokio::test]
+    async fn failed_activation_worker_drop_releases_nested_authority_before_allocation() {
+        let dir = tempfile::tempdir().unwrap();
+        let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
+        let mut worker = Worker::new(
+            minimal_manifest(),
+            Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
+            store,
+            WorkerWorkspaceContext::no_workspace(),
+            WorkerFilesystemAuthority::None,
+            Scope::empty(),
+        )
+        .await
+        .unwrap();
+        worker.ensure_segment_head().await.unwrap();
+        let old_segment = worker.segment_id();
+        let replacement = session_store::new_segment_id();
+        let lock_path = dir.path().join("workers.json");
+        let allocation = worker_allocation::install_top_level_at_for_test(
+            lock_path.clone(),
+            "test-worker".into(),
+            std::process::id(),
+            dir.path().join("worker.sock"),
+            Vec::new(),
+            old_segment,
+        )
+        .unwrap();
+        let activation = allocation
+            .begin_segment_activation(old_segment, replacement)
+            .unwrap();
+        worker.scope_allocation = Some(allocation);
+        worker_allocation::fail_next_save_for_test(&lock_path);
+        let (source, authority) = activation.commit().unwrap_err().into_parts();
+        assert!(matches!(source, ScopeLockError::Io(_)));
+        worker.fail_closed_segment_activation(authority);
+        assert!(worker.segment_state.ensure_append_allowed().is_err());
+
+        let (sent, received) = std::sync::mpsc::channel();
+        let teardown = std::thread::spawn(move || {
+            drop(worker);
+            sent.send(()).unwrap();
+        });
+        received
+            .recv_timeout(Duration::from_secs(1))
+            .expect("Worker teardown must not deadlock on nested allocation locks");
+        teardown.join().unwrap();
+    }
+
+    #[tokio::test]
+    async fn shutdown_cleanup_failures_transfer_service_authority_to_quarantine() {
+        let dir = tempfile::tempdir().unwrap();
+        let store = session_store::FsStore::new(dir.path().join("sessions")).unwrap();
+        let mut worker = Worker::new(
+            minimal_manifest(),
+            Engine::<_, Mutable, SessionHistoryMetadata>::new_annotated(NoopClient),
+            store,
+            WorkerWorkspaceContext::no_workspace(),
+            WorkerFilesystemAuthority::None,
+            Scope::empty(),
+        )
+        .await
+        .unwrap();
+        let registry =
+            crate::spawn::registry::SpawnedWorkerRegistry::new_for_internal_services();
+        registry.fail_service_stops_for_test("cleanup-session", 3);
+        worker.internal_worker_registry = Some(Arc::clone(&registry));
+        *worker
+            .pending_compaction_cleanup
+            .lock()
+            .expect("pending cleanup lock poisoned") = Some(PendingCompactionCleanup {
+            session_id: "cleanup-session".into(),
+            compaction_id: "compaction".into(),
+        });
+
+        tokio::time::timeout(
+            Duration::from_secs(1),
+            worker.finish_pending_compaction_cleanup_for_shutdown(),
+        )
+        .await
+        .expect("shutdown cleanup must be bounded")
+        .unwrap();
+        assert!(!worker.has_pending_compaction_cleanup());
+        tokio::time::timeout(Duration::from_secs(1), async {
+            loop {
+                if !registry.is_service_cleanup_quarantined_for_test("cleanup-session") {
+                    break;
+                }
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .unwrap();
+    }
+
     #[tokio::test]
     async fn auto_fork_checkpoints_interrupted_run_budget_for_restore() {
         let dir = tempfile::tempdir().unwrap();