diff --git a/crates/worker/src/runtime/worker_allocation.rs b/crates/worker/src/runtime/worker_allocation.rs index 726c294ad..cb8b474b1 100644 --- a/crates/worker/src/runtime/worker_allocation.rs +++ b/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, diff --git a/crates/worker/src/runtime/worker_allocation/lifecycle.rs b/crates/worker/src/runtime/worker_allocation/lifecycle.rs index 2f758480c..6dd9c6537 100644 --- a/crates/worker/src/runtime/worker_allocation/lifecycle.rs +++ b/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, + segment_id: SegmentId, +) -> Result { + 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, #[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(); diff --git a/crates/worker/src/runtime/worker_allocation/table.rs b/crates/worker/src/runtime/worker_allocation/table.rs index b916cd178..d06dad5cd 100644 --- a/crates/worker/src/runtime/worker_allocation/table.rs +++ b/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> = 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"); diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 19d44b99d..fc4a21dcb 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/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>, service_records: std::sync::Mutex>, quarantined_service_cleanup: std::sync::Mutex>, + #[cfg(test)] + fail_service_stops: std::sync::Mutex>, internal_names: std::sync::Mutex>, 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 { + #[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() diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index bc8fc3afb..0c3d62aa3 100644 --- a/crates/worker/src/worker.rs +++ b/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(®istry)); + *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();