From 9d8e0114f562d6199af82c0b5dd4fd177d59406d Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 17 Sep 2026 00:14:33 +0900 Subject: [PATCH] fix: serialize runtime worker lifecycle --- Cargo.lock | 1 + crates/worker-runtime/Cargo.toml | 3 +- crates/worker-runtime/src/error.rs | 3 + crates/worker-runtime/src/fs_store.rs | 152 ++++++- crates/worker-runtime/src/http_server.rs | 2 + crates/worker-runtime/src/runtime.rs | 450 +++++++++++++++----- crates/worker-runtime/src/worker_backend.rs | 344 +++++++++++---- crates/worker/src/bootstrap.rs | 11 +- crates/worker/src/controller.rs | 119 ++++-- crates/worker/src/ipc/server.rs | 49 ++- crates/worker/tests/controller_test.rs | 27 ++ crates/workspace-server/src/hosts.rs | 5 + 12 files changed, 942 insertions(+), 224 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c787079d..b429be87 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6070,6 +6070,7 @@ dependencies = [ "base64 0.22.1", "decodal", "flow", + "fs4", "futures", "manifest", "protocol", diff --git a/crates/worker-runtime/Cargo.toml b/crates/worker-runtime/Cargo.toml index a838b48a..dabdf140 100644 --- a/crates/worker-runtime/Cargo.toml +++ b/crates/worker-runtime/Cargo.toml @@ -18,7 +18,7 @@ required-features = ["ws-server", "fs-store"] [features] default = ["ws-server", "fs-store"] -fs-store = [] +fs-store = ["dep:fs4"] http-server = ["dep:axum", "dep:tower", "dep:reqwest"] ws-server = ["http-server", "axum/ws", "dep:futures", "tokio/sync"] @@ -29,6 +29,7 @@ axum = { workspace = true, optional = true } futures = { workspace = true, optional = true } decodal.workspace = true flow = { path = "../flow" } +fs4 = { workspace = true, optional = true } manifest.workspace = true protocol.workspace = true serde = { workspace = true, features = ["derive"] } diff --git a/crates/worker-runtime/src/error.rs b/crates/worker-runtime/src/error.rs index 1de44f83..1f51b5a1 100644 --- a/crates/worker-runtime/src/error.rs +++ b/crates/worker-runtime/src/error.rs @@ -79,6 +79,9 @@ pub enum RuntimeError { name: String, }, + #[error("Runtime store is already owned by another process")] + RuntimeStoreAlreadyOpen { path: PathBuf }, + #[error("runtime store {operation} failed at {}: {source}", path.display())] StoreIo { operation: &'static str, diff --git a/crates/worker-runtime/src/fs_store.rs b/crates/worker-runtime/src/fs_store.rs index f6f295af..393696b2 100644 --- a/crates/worker-runtime/src/fs_store.rs +++ b/crates/worker-runtime/src/fs_store.rs @@ -8,11 +8,13 @@ use crate::identity::{ LegacyWorkerIdentityMapping, WorkerId, WorkerRef, legacy_worker_identity_mapping_digest, }; use crate::management::{RuntimeBackendKind, RuntimeStatus}; +use fs4::fs_std::FileExt; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; use std::fs::{self, File, OpenOptions}; use std::io::{BufReader, Write}; use std::path::{Path, PathBuf}; +use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; const SCHEMA_VERSION: u32 = 6; @@ -50,25 +52,47 @@ impl FsRuntimeStoreOptions { } } +#[derive(Debug)] +struct RuntimeStoreOwnerLock { + file: File, +} + +impl Drop for RuntimeStoreOwnerLock { + fn drop(&mut self) { + let _ = FileExt::unlock(&self.file); + } +} + /// Filesystem persistence boundary for one Worker Runtime state. /// /// Authority is the Workspace-owned typed Worker identity. Legacy pod paths, socket /// paths, and session paths are deliberately not part of the layout or lookup API. -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug)] pub struct FsRuntimeStore { root: PathBuf, + _owner_lock: Option>, } +impl PartialEq for FsRuntimeStore { + fn eq(&self, other: &Self) -> bool { + self.root == other.root + } +} + +impl Eq for FsRuntimeStore {} + impl FsRuntimeStore { pub fn migration_plan( options: &FsRuntimeStoreOptions, ) -> Result { + let _owner_lock = acquire_runtime_store_owner_lock(&options.root)?; plan_runtime_store_migration(&options.root, &options.runtime_id).map(|(plan, _)| plan) } pub fn migrate( options: &FsRuntimeStoreOptions, ) -> Result { + let _owner_lock = acquire_runtime_store_owner_lock(&options.root)?; migrate_runtime_store(&options.root, &options.runtime_id) } @@ -93,6 +117,15 @@ impl FsRuntimeStore { }); } + if !existed { + fs::create_dir_all(&root).map_err(|source| RuntimeError::StoreIo { + operation: "create runtime store root", + path: root.clone(), + source, + })?; + } + let owner_lock = acquire_runtime_store_owner_lock(&root)?; + fs::create_dir_all(root.join(WORKERS_DIR)).map_err(|source| RuntimeError::StoreIo { operation: "create runtime store", path: root.join(WORKERS_DIR), @@ -114,7 +147,10 @@ impl FsRuntimeStore { if existed { migrate_runtime_store(&root, runtime_id)?; } - let store = Self { root }; + let store = Self { + root, + _owner_lock: Some(owner_lock), + }; let state = if existed { Some(store.load_runtime_state()?) } else { @@ -256,6 +292,61 @@ impl FsRuntimeStore { } } +fn acquire_runtime_store_owner_lock( + root: &Path, +) -> Result, RuntimeError> { + let canonical_root = fs::canonicalize(root).map_err(|source| RuntimeError::StoreIo { + operation: "resolve runtime store owner lock", + path: root.to_path_buf(), + source, + })?; + let parent = canonical_root + .parent() + .ok_or_else(|| RuntimeError::StoreCorrupt { + operation: "resolve runtime store owner lock", + path: root.to_path_buf(), + message: "runtime store root has no parent".to_string(), + })?; + let name = canonical_root + .file_name() + .ok_or_else(|| RuntimeError::StoreCorrupt { + operation: "resolve runtime store owner lock", + path: root.to_path_buf(), + message: "runtime store root has no file name".to_string(), + })?; + let mut lock_name = std::ffi::OsString::from("."); + lock_name.push(name); + lock_name.push(".runtime-owner.lock"); + let lock_path = parent.join(lock_name); + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(&lock_path) + .map_err(|source| RuntimeError::StoreIo { + operation: "open runtime store owner lock", + path: root.to_path_buf(), + source, + })?; + match file.try_lock_exclusive() { + Ok(true) => Ok(Arc::new(RuntimeStoreOwnerLock { file })), + Ok(false) => Err(RuntimeError::RuntimeStoreAlreadyOpen { + path: root.to_path_buf(), + }), + Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => { + Err(RuntimeError::RuntimeStoreAlreadyOpen { + path: root.to_path_buf(), + }) + } + Err(source) => Err(RuntimeError::StoreIo { + operation: "acquire runtime store owner lock", + path: root.to_path_buf(), + source, + }), + } +} + fn remove_legacy_observations(worker_dir: &Path) { let path = worker_dir.join(LEGACY_OBSERVATIONS_FILE); if path.is_file() { @@ -1193,6 +1284,7 @@ fn migrate_runtime_store( }; let staged_store = FsRuntimeStore { root: staging.clone(), + _owner_lock: None, }; if let Err(error) = staged_store.load_runtime_state() { let _ = fs::remove_dir_all(&staging); @@ -1628,6 +1720,62 @@ fn sync_directory(path: &Path, operation: &'static str) -> Result<(), RuntimeErr mod tests { use super::*; + #[test] + fn runtime_store_owner_lock_child_probe() { + let Some(root) = std::env::var_os("YOI_TEST_RUNTIME_STORE_LOCK_ROOT") else { + return; + }; + assert!(matches!( + FsRuntimeStore::open_or_create(PathBuf::from(root), "runtime-test").unwrap_err(), + RuntimeError::RuntimeStoreAlreadyOpen { .. } + )); + } + + #[test] + fn second_runtime_store_open_conflicts_before_store_mutation() { + let parent = tempfile::tempdir().unwrap(); + let root = parent.path().join("runtime-store"); + let first = FsRuntimeStore::open_or_create(root.clone(), "runtime-test").unwrap(); + let legacy_events = root.join("events.jsonl"); + fs::write(&legacy_events, b"must remain").unwrap(); + + let error = FsRuntimeStore::open_or_create(root.clone(), "runtime-test").unwrap_err(); + + assert!(matches!( + error, + RuntimeError::RuntimeStoreAlreadyOpen { path } if path == root + )); + assert_eq!(fs::read(&legacy_events).unwrap(), b"must remain"); + let child = std::process::Command::new(std::env::current_exe().unwrap()) + .arg("fs_store::tests::runtime_store_owner_lock_child_probe") + .arg("--exact") + .env("YOI_TEST_RUNTIME_STORE_LOCK_ROOT", &root) + .output() + .unwrap(); + assert!( + child.status.success(), + "cross-process lock probe failed: {}", + String::from_utf8_lossy(&child.stderr) + ); + assert!(fs::read_dir(&root).unwrap().all(|entry| { + !entry + .unwrap() + .file_name() + .to_string_lossy() + .contains("owner.lock") + })); + + let retained_clone = first.store.clone(); + drop(first); + assert!(matches!( + FsRuntimeStore::open_or_create(root.clone(), "runtime-test").unwrap_err(), + RuntimeError::RuntimeStoreAlreadyOpen { .. } + )); + drop(retained_clone); + fs::remove_dir_all(&root).unwrap(); + FsRuntimeStore::open_or_create(root, "runtime-test").unwrap(); + } + #[test] fn schema_v4_migration_plan_ignores_orphan_worker_directories() { let root = tempfile::tempdir().unwrap(); diff --git a/crates/worker-runtime/src/http_server.rs b/crates/worker-runtime/src/http_server.rs index 313f0876..b067bd68 100644 --- a/crates/worker-runtime/src/http_server.rs +++ b/crates/worker-runtime/src/http_server.rs @@ -2467,6 +2467,7 @@ fn status_for_runtime_error(error: &RuntimeError) -> StatusCode { StatusCode::NOT_FOUND } RuntimeError::RuntimeStopped + | RuntimeError::RuntimeStoreAlreadyOpen { .. } | RuntimeError::WorkerExecutionUnavailable { .. } | RuntimeError::ExecutionBackendUnavailable { .. } | RuntimeError::WorkerExecutionRejected { .. } => StatusCode::CONFLICT, @@ -2488,6 +2489,7 @@ fn status_for_runtime_error(error: &RuntimeError) -> StatusCode { fn code_for_runtime_error(error: &RuntimeError) -> String { match error { RuntimeError::RuntimeStopped => "runtime_stopped".to_string(), + RuntimeError::RuntimeStoreAlreadyOpen { .. } => "runtime_store_already_open".to_string(), RuntimeError::WorkerNotFound { .. } => "worker_not_found".to_string(), RuntimeError::WorkerExecutionUnavailable { .. } => { "worker_execution_unavailable".to_string() diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 8e23c59d..2669f8cb 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -72,6 +72,12 @@ impl RuntimeWorkspaceScope { const SUBSCRIPTION_QUEUE_CAPACITY: usize = 256; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum WorkerRestoreMode { + Explicit, + Automatic, +} + #[derive(Clone, Debug)] pub struct RuntimeSubscriptionUpdate { pub subject_revision: u64, @@ -1295,16 +1301,51 @@ impl Runtime { /// Attach a live execution to a persisted Worker definition. /// - /// Current liveness is never read from disk. If a handle is already - /// present this is idempotent; otherwise the configured backend is tried. + /// Every lifecycle mutation for a Worker is serialized by the same operation + /// lock. A concurrent exact restore waits for the first caller and then + /// converges on its already-installed execution rather than spawning another + /// controller. pub fn restore_worker(&self, worker_ref: &WorkerRef) -> Result { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + self.restore_worker_under_lock(worker_ref, WorkerRestoreMode::Explicit) + } + + fn restore_worker_under_lock( + &self, + worker_ref: &WorkerRef, + mode: WorkerRestoreMode, + ) -> Result { let (backend, request) = { let mut state = self.lock()?; state.ensure_running()?; let (worker_request, previous_working_directory, run_generation) = { let worker = state.worker(worker_ref)?; if worker.execution_handle.is_some() { - return Ok(worker.detail()); + if worker.status.is_active() { + return Ok(worker.detail()); + } + return Err(RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "previous execution cleanup has not completed".to_string(), + }); + } + match mode { + WorkerRestoreMode::Explicit if worker.status != WorkerStatus::Stopped => { + return Err(RuntimeError::InvalidRequest(format!( + "worker {} is not stopped", + worker_ref.worker_id + ))); + } + WorkerRestoreMode::Automatic + if !worker.status.is_active() + || worker.restore_intent != WorkerRestoreIntent::Automatic => + { + return Ok(worker.detail()); + } + _ => {} } ( worker.request.clone(), @@ -1314,7 +1355,7 @@ impl Runtime { }; let backend = state.execution_backend.clone().ok_or_else(|| { RuntimeError::WorkerExecutionUnavailable { - worker_id: worker_ref.worker_id.clone(), + worker_id: worker_ref.worker_id, message: "runtime has no execution backend".to_string(), } })?; @@ -1322,7 +1363,6 @@ impl Runtime { let worker = state.worker_mut(worker_ref)?; worker.run_generation = run_generation; worker.execution_bound = true; - worker.restore_intent = WorkerRestoreIntent::Automatic; } state.persist_worker(&worker_ref.worker_id)?; let workspace_scope = worker_request.workspace_api.as_ref().and_then(|api| { @@ -1350,13 +1390,17 @@ impl Runtime { worker_state, working_directory, } => { - self.commit_restored_worker_execution( + let commit = self.commit_restored_worker_execution( worker_ref, - handle, + handle.clone(), worker_state, WorkerStatus::Idle, working_directory, - )?; + ); + if let Err(error) = commit { + self.cleanup_failed_restore(&backend, worker_ref, &handle)?; + return Err(error); + } self.worker_detail(worker_ref) } WorkerExecutionSpawnResult::Rejected(result) @@ -1367,7 +1411,7 @@ impl Runtime { .record_restore_failure(worker_ref, result.clone())?; } Err(RuntimeError::WorkerExecutionRejected { - worker_id: worker_ref.worker_id.clone(), + worker_id: worker_ref.worker_id, operation: result.operation, outcome: result.outcome, message: result.message_or_default(), @@ -1894,15 +1938,64 @@ impl Runtime { self.stop_worker(worker_ref, reason) } - /// Stop a Worker. Repeated stops are idempotent. + /// Stop a Worker. Repeated stops are idempotent. The per-Worker lifecycle + /// lock remains held until backend cleanup and the terminal catalog commit + /// have both completed. pub fn stop_worker( &self, worker_ref: &WorkerRef, reason: Option, ) -> Result { - self.dispatch_lifecycle_to_backend(worker_ref, WorkerExecutionOperation::Stop)?; + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + + let (backend, handle) = { + let state = self.lock()?; + state.ensure_running()?; + let worker = state.worker(worker_ref)?; + match worker.execution_handle.clone() { + Some(handle) => { + let backend = state.execution_backend.clone().ok_or_else(|| { + RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "runtime has no execution backend".to_string(), + } + })?; + (Some(backend), Some(handle)) + } + None if worker.status == WorkerStatus::Stopped => (None, None), + None => { + return Err(RuntimeError::WorkerExecutionUnavailable { + worker_id: worker_ref.worker_id, + message: "worker has no execution handle".to_string(), + }); + } + } + }; + + if let (Some(backend), Some(handle)) = (backend, handle) { + let result = backend.stop_worker(&handle); + if !result.is_accepted() { + return Err(RuntimeError::WorkerExecutionRejected { + worker_id: worker_ref.worker_id, + operation: result.operation, + outcome: result.outcome, + message: result.message_or_default(), + result, + }); + } + self.transition_worker(worker_ref, WorkerStatus::Stopped)?; + } let _ = reason; - self.transition_worker(worker_ref, WorkerStatus::Stopped) + let state = self.lock()?; + let worker = state.worker(worker_ref)?; + Ok(WorkerLifecycleAck { + worker_ref: worker_ref.clone(), + status: worker.status, + worker_state: worker.worker_state.clone(), + }) } /// Cancel a Worker through a workspace-scoped Runtime authorization context. @@ -2185,10 +2278,26 @@ impl Runtime { #[cfg(feature = "ws-server")] fn execution_context(&self, worker_ref: WorkerRef) -> crate::execution::WorkerExecutionContext { - let runtime = self.clone(); + let runtime = Arc::downgrade(&self.inner); crate::execution::WorkerExecutionContext::new( worker_ref, - Arc::new(move |worker_ref, payload| runtime.observe_worker_event(&worker_ref, payload)), + Arc::new(move |worker_ref, payload| { + let runtime = runtime.upgrade().ok_or(RuntimeError::RuntimeStopped)?; + let mut state = runtime.lock().map_err(|_| RuntimeError::StatePoisoned)?; + state.ensure_worker_ref(&worker_ref)?; + let worker_state_changed = + state.project_protocol_event_to_worker_state(&worker_ref, &payload); + let activity_changed = + state.project_internal_worker_activity(&worker_ref, &payload); + if worker_state_changed || activity_changed { + state.publish_worker_upsert(worker_ref.worker_id)?; + } + if worker_state_changed { + state.persist_runtime_snapshot()?; + state.persist_worker(&worker_ref.worker_id)?; + } + Ok(state.push_worker_observation_event(worker_ref, payload)) + }), ) } @@ -2199,21 +2308,12 @@ impl Runtime { #[cfg(feature = "fs-store")] fn restore_persisted_worker_executions(&self) -> Result<(), RuntimeError> { - #[derive(Clone)] - struct RestoreCandidate { - worker_ref: WorkerRef, - request: CreateWorkerRequest, - run_generation: u64, - previous_working_directory: Option, - config_bundle: Option, - } - let candidates = { - let mut state = self.lock()?; + let state = self.lock()?; if state.execution_backend.is_none() { return Ok(()); } - let worker_ids = state + state .workers .values() .filter(|worker| { @@ -2222,80 +2322,22 @@ impl Runtime { && worker.status.is_active() && worker.restore_intent == WorkerRestoreIntent::Automatic }) - .map(|worker| worker.worker_id) - .collect::>(); - let mut candidates = Vec::with_capacity(worker_ids.len()); - for worker_id in worker_ids { - let (worker_ref, request, previous_working_directory, run_generation) = { - let worker = state - .workers - .get(&worker_id) - .expect("collected Worker exists"); - ( - worker.worker_ref.clone(), - worker.request.clone(), - worker.working_directory.clone(), - worker.run_generation.saturating_add(1).max(1), - ) - }; - state - .workers - .get_mut(&worker_id) - .expect("collected Worker exists") - .run_generation = run_generation; - state.persist_worker(&worker_id)?; - candidates.push(RestoreCandidate { - worker_ref, - request, - run_generation, - previous_working_directory, - config_bundle: None, - }); - } - candidates + .map(|worker| worker.worker_ref.clone()) + .collect::>() }; - for candidate in candidates { - let (backend, workspace_scope) = { - let state = self.lock()?; - let workspace_scope = candidate.request.workspace_api.as_ref().and_then(|api| { - state - .workspace_owners - .get(&api.workspace_id) - .map(|server_id| RuntimeWorkspaceScope::new(&api.workspace_id, server_id)) - }); - (state.execution_backend.clone(), workspace_scope) - }; - let Some(backend) = backend else { - return Ok(()); - }; - let request = WorkerExecutionRestoreRequest { - worker_ref: candidate.worker_ref.clone(), - run_generation: candidate.run_generation, - request: candidate.request, - workspace_scope, - context: self.execution_context(candidate.worker_ref.clone()), - previous_working_directory: candidate.previous_working_directory, - working_directory: None, - config_bundle: candidate.config_bundle, - }; - match backend.restore_worker(request) { - WorkerExecutionSpawnResult::Connected { - handle, - worker_state, - working_directory, - } => self.commit_restored_worker_execution( - &candidate.worker_ref, - handle, - worker_state, - WorkerStatus::Idle, - working_directory, - )?, - WorkerExecutionSpawnResult::Rejected(result) - | WorkerExecutionSpawnResult::Errored(result) => { - let mut state = self.lock()?; - state.record_restore_failure(&candidate.worker_ref, result)?; + for worker_ref in candidates { + let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; + let _operation_guard = operation_lock + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + match self.restore_worker_under_lock(&worker_ref, WorkerRestoreMode::Automatic) { + Ok(_) => {} + Err(RuntimeError::WorkerExecutionRejected { .. }) => { + // The failed restore and fail-closed terminal status were + // persisted by restore_worker_under_lock. } + Err(error) => return Err(error), } } Ok(()) @@ -2313,6 +2355,12 @@ impl Runtime { state.ensure_worker_ref(worker_ref)?; { let worker = state.worker_mut(worker_ref)?; + if worker.execution_handle.is_some() { + return Err(RuntimeError::InvalidRequest(format!( + "worker {} already has a current execution", + worker_ref.worker_id + ))); + } worker.execution_handle = Some(handle); worker.execution_bound = true; worker.status = status; @@ -2326,6 +2374,39 @@ impl Runtime { Ok(()) } + fn cleanup_failed_restore( + &self, + backend: &WorkerExecutionBackendRef, + worker_ref: &WorkerRef, + handle: &WorkerExecutionHandle, + ) -> Result<(), RuntimeError> { + let result = backend.stop_worker(handle); + if !result.is_accepted() { + return Err(RuntimeError::WorkerExecutionRejected { + worker_id: worker_ref.worker_id, + operation: result.operation, + outcome: result.outcome, + message: format!( + "restored execution cleanup failed after catalog commit failure: {}", + result.message_or_default() + ), + result, + }); + } + + let mut state = self.lock()?; + let worker = state.worker_mut(worker_ref)?; + worker.execution_handle = None; + worker.status = WorkerStatus::Stopped; + worker.worker_state = None; + worker.restore_intent = WorkerRestoreIntent::Explicit; + worker.internal_workers.clear(); + state.publish_worker_upsert(worker_ref.worker_id)?; + state.persist_runtime_snapshot()?; + state.persist_worker(&worker_ref.worker_id)?; + Ok(()) + } + /// Bind the Backend registry identity once. Retention evidence fails closed /// until the Runtime host supplies this trusted configuration. pub fn bind_runtime_identity(&self, runtime_id: &str) -> Result<(), RuntimeError> { @@ -3670,6 +3751,7 @@ fn runtime_worker_create_failure_fields( ) -> (&'static str, Option, Option) { match error { RuntimeError::RuntimeStopped => ("runtime_stopped", None, None), + RuntimeError::RuntimeStoreAlreadyOpen { .. } => ("runtime_store_already_open", None, None), RuntimeError::InvalidInitialInputKind { .. } => ("invalid_initial_input_kind", None, None), RuntimeError::WorkerNotFound { .. } => ("worker_not_found", None, None), RuntimeError::WorkerExecutionUnavailable { .. } => { @@ -3854,7 +3936,7 @@ mod tests { use std::collections::BTreeMap; #[cfg(feature = "fs-store")] use std::sync::atomic::{AtomicU64, Ordering}; - use std::sync::{Arc, Mutex}; + use std::sync::{Arc, Barrier, Condvar, Mutex}; #[test] fn worker_create_failure_fields_exclude_raw_error_messages() { @@ -4694,11 +4776,49 @@ mod tests { .with_computed_digest() } + #[derive(Default)] + struct RestoreGate { + entered: Mutex, + entered_changed: Condvar, + released: Mutex, + released_changed: Condvar, + } + + impl RestoreGate { + fn enter_and_wait(&self) { + let mut entered = self.entered.lock().unwrap(); + *entered += 1; + self.entered_changed.notify_all(); + drop(entered); + let mut released = self.released.lock().unwrap(); + while !*released { + released = self.released_changed.wait(released).unwrap(); + } + } + + fn wait_for_entered(&self, expected: u64, timeout: std::time::Duration) -> bool { + let entered = self.entered.lock().unwrap(); + let (entered, _) = self + .entered_changed + .wait_timeout_while(entered, timeout, |entered| *entered < expected) + .unwrap(); + *entered >= expected + } + + fn release(&self) { + *self.released.lock().unwrap() = true; + self.released_changed.notify_all(); + } + } + #[derive(Default)] struct TestExecutionBackend { dispatch_result: Mutex>, stop_result: Mutex>, + stop_gate: Mutex>>, + stop_count: Mutex, restore_result: Mutex>, + restore_gate: Mutex>>, restore_count: Mutex, run_generations: Mutex>, config_bundles: Mutex>>, @@ -4828,6 +4948,10 @@ mod tests { request: WorkerExecutionRestoreRequest, ) -> WorkerExecutionSpawnResult { *self.restore_count.lock().unwrap() += 1; + let restore_gate = self.restore_gate.lock().unwrap().clone(); + if let Some(gate) = restore_gate { + gate.enter_and_wait(); + } self.run_generations .lock() .unwrap() @@ -4888,6 +5012,11 @@ mod tests { } fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult { + *self.stop_count.lock().unwrap() += 1; + let stop_gate = self.stop_gate.lock().unwrap().clone(); + if let Some(gate) = stop_gate { + gate.enter_and_wait(); + } self.stop_result .lock() .unwrap() @@ -5660,6 +5789,131 @@ mod tests { ); } + #[test] + fn failed_stop_cleanup_retains_execution_for_retry() { + let (runtime, backend) = runtime_and_backend(); + runtime.store_config_bundle(test_bundle()).unwrap(); + let created = runtime + .create_worker(task_request("retry stop cleanup")) + .unwrap(); + backend.set_stop_result(WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + "cleanup is still pending", + )); + + let error = runtime.stop_worker(&created.worker_ref, None).unwrap_err(); + assert!(matches!( + error, + RuntimeError::WorkerExecutionRejected { + operation: WorkerExecutionOperation::Stop, + .. + } + )); + assert_eq!( + runtime.worker_detail(&created.worker_ref).unwrap().status, + WorkerStatus::Idle + ); + + let stopped = runtime.stop_worker(&created.worker_ref, None).unwrap(); + assert_eq!(stopped.status, WorkerStatus::Stopped); + assert_eq!(*backend.stop_count.lock().unwrap(), 2); + } + + #[test] + fn restore_waits_for_in_flight_stop_cleanup() { + let backend = Arc::new(TestExecutionBackend::default()); + let runtime = Arc::new( + Runtime::with_execution_backend(RuntimeOptions::default(), backend.clone()).unwrap(), + ); + runtime.store_config_bundle(test_bundle()).unwrap(); + let created = runtime + .create_worker(task_request("stop restore race")) + .unwrap(); + let gate = Arc::new(RestoreGate::default()); + *backend.stop_gate.lock().unwrap() = Some(gate.clone()); + + let stop_runtime = runtime.clone(); + let stop_ref = created.worker_ref.clone(); + let stopping = std::thread::spawn(move || stop_runtime.stop_worker(&stop_ref, None)); + assert!(gate.wait_for_entered(1, std::time::Duration::from_secs(2))); + + let restore_runtime = runtime.clone(); + let restore_ref = created.worker_ref.clone(); + let restoring = std::thread::spawn(move || restore_runtime.restore_worker(&restore_ref)); + std::thread::sleep(std::time::Duration::from_millis(100)); + assert_eq!( + *backend.restore_count.lock().unwrap(), + 0, + "restore reached the backend before stop cleanup completed" + ); + + gate.release(); + stopping.join().unwrap().unwrap(); + let restored = restoring.join().unwrap().unwrap(); + assert_eq!(*backend.stop_count.lock().unwrap(), 1); + assert_eq!(*backend.restore_count.lock().unwrap(), 1); + assert_eq!(restored.status, WorkerStatus::Idle); + assert_eq!( + restored + .worker_state + .as_ref() + .map(|state| state.execution_generation), + Some(2) + ); + } + + #[test] + fn concurrent_exact_restores_share_one_execution() { + let backend = Arc::new(TestExecutionBackend::default()); + let runtime = Arc::new( + Runtime::with_execution_backend(RuntimeOptions::default(), backend.clone()).unwrap(), + ); + runtime.store_config_bundle(test_bundle()).unwrap(); + let created = runtime + .create_worker(task_request("concurrent restore")) + .unwrap(); + runtime.stop_worker(&created.worker_ref, None).unwrap(); + + let gate = Arc::new(RestoreGate::default()); + *backend.restore_gate.lock().unwrap() = Some(gate.clone()); + let start = Arc::new(Barrier::new(3)); + let mut threads = Vec::new(); + for _ in 0..2 { + let runtime = runtime.clone(); + let worker_ref = created.worker_ref.clone(); + let start = start.clone(); + threads.push(std::thread::spawn(move || { + start.wait(); + runtime.restore_worker(&worker_ref) + })); + } + start.wait(); + assert!(gate.wait_for_entered(1, std::time::Duration::from_secs(2))); + let duplicate_entered = gate.wait_for_entered(2, std::time::Duration::from_millis(100)); + gate.release(); + + let restored = threads + .into_iter() + .map(|thread| thread.join().unwrap().unwrap()) + .collect::>(); + assert!( + !duplicate_entered, + "a second backend restore ran concurrently" + ); + assert_eq!(*backend.restore_count.lock().unwrap(), 1); + assert_eq!(restored[0].worker_ref, restored[1].worker_ref); + assert_eq!( + restored[0] + .worker_state + .as_ref() + .map(|state| state.execution_generation), + restored[1] + .worker_state + .as_ref() + .map(|state| state.execution_generation) + ); + } + #[test] fn restored_worker_exposes_the_backend_initial_state_snapshot() { let (runtime, _) = runtime_and_backend(); @@ -6999,6 +7253,7 @@ mod tests { ) .unwrap(); drop(corrupt_runtime); + drop(corrupt_store); let err = Runtime::with_fs_store(crate::fs_store::FsRuntimeStoreOptions { root: corrupt_root.clone(), runtime_id: "test-runtime".to_string(), @@ -7030,6 +7285,7 @@ mod tests { worker_dirs.sort_by_key(|entry| entry.path()); std::fs::remove_file(worker_dirs[0].path().join("worker.json")).unwrap(); drop(missing_runtime); + drop(missing_store); let loaded = Runtime::with_fs_store(crate::fs_store::FsRuntimeStoreOptions { root: missing_root.clone(), runtime_id: "test-runtime".to_string(), diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 758f4188..ace3bf9f 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -10,7 +10,7 @@ use std::collections::HashMap; use std::future::Future; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock, mpsc}; use std::time::Duration; @@ -95,6 +95,7 @@ const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9); pub struct RuntimeWorkerController { pub handle: WorkerHandle, pub shutdown: Arc>>, + pub controller_task: tokio::task::JoinHandle<()>, pub workspace_client: Arc, } @@ -978,7 +979,8 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { run_dir.display() ), })?; - let (handle, shutdown_rx) = (started.handle, started.shutdown); + let (handle, shutdown_rx, controller_task) = + (started.handle, started.shutdown, started.controller_task); if flow_transition_enabled { handle.shared_state.enable_flow_transition(); } @@ -990,6 +992,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { Ok(RuntimeWorkerController { handle, shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))), + controller_task, workspace_client, }) } @@ -1172,7 +1175,8 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { run_dir.display() ), })?; - let (handle, shutdown_rx) = (started.handle, started.shutdown); + let (handle, shutdown_rx, controller_task) = + (started.handle, started.shutdown, started.controller_task); if flow_transition_enabled { handle.shared_state.enable_flow_transition(); } @@ -1184,15 +1188,117 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { Ok(RuntimeWorkerController { handle, shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))), + controller_task, workspace_client, }) } } +#[derive(Clone)] +struct RuntimeExecutionTaskScope { + tasks: Arc>>, + terminal_failure: Arc>>, +} + +struct RuntimeExecutionTask { + name: &'static str, + task: tokio::task::JoinHandle<()>, + abort_before_join: bool, +} + +impl RuntimeExecutionTaskScope { + fn new(controller_task: tokio::task::JoinHandle<()>) -> Self { + Self { + tasks: Arc::new(Mutex::new(vec![RuntimeExecutionTask { + name: "controller", + task: controller_task, + abort_before_join: false, + }])), + terminal_failure: Arc::new(Mutex::new(None)), + } + } + + fn push(&self, name: &'static str, task: tokio::task::JoinHandle<()>, abort_before_join: bool) { + let mut tasks = match self.tasks.lock() { + Ok(tasks) => tasks, + Err(poisoned) => poisoned.into_inner(), + }; + tasks.push(RuntimeExecutionTask { + name, + task, + abort_before_join, + }); + } + + fn abort_all(&self) { + let mut tasks = match self.tasks.lock() { + Ok(tasks) => tasks, + Err(poisoned) => poisoned.into_inner(), + }; + for task in tasks.drain(..) { + task.task.abort(); + } + } + + async fn join(&self) -> Result<(), String> { + if let Some(message) = self + .terminal_failure + .lock() + .map_err(|_| "execution task failure lock is poisoned".to_string())? + .clone() + { + return Err(message); + } + + loop { + let next = { + let mut tasks = self + .tasks + .lock() + .map_err(|_| "execution task registry lock is poisoned".to_string())?; + if tasks.is_empty() { + None + } else { + Some(tasks.remove(0)) + } + }; + let Some(mut task) = next else { + return Ok(()); + }; + let name = task.name; + if task.abort_before_join { + task.task.abort(); + } + match tokio::time::timeout(Duration::from_secs(5), &mut task.task).await { + Ok(Ok(())) => {} + Ok(Err(error)) if task.abort_before_join && error.is_cancelled() => {} + Ok(Err(error)) => { + let message = format!("{name} task failed while stopping Worker: {error}"); + if let Ok(mut failure) = self.terminal_failure.lock() { + *failure = Some(message.clone()); + } + return Err(message); + } + Err(_) => { + self.tasks + .lock() + .map_err(|_| "execution task registry lock is poisoned".to_string())? + .insert(0, task); + return Err(format!( + "{name} task did not stop before timeout; stop remains retryable" + )); + } + } + } + } +} + #[derive(Clone)] struct RuntimeWorkerExecution { handle: WorkerHandle, shutdown: Arc>>, + shutdown_requested: Arc, + tasks: RuntimeExecutionTaskScope, worker_state: Arc>, workspace_client: Option>, } @@ -1261,7 +1367,10 @@ where .map_err(|err| format!("worker adapter task did not complete: {err}"))? } - fn spawn_on_adapter_runtime(&self, task: Fut) -> Result<(), String> + fn spawn_on_adapter_runtime( + &self, + task: Fut, + ) -> Result, String> where Fut: Future + Send + 'static, { @@ -1272,8 +1381,7 @@ where let runtime = runtime .as_ref() .ok_or_else(|| "worker adapter runtime is shutting down".to_string())?; - runtime.spawn(task); - Ok(()) + Ok(runtime.spawn(task)) } fn run_on_adapter_runtime(&self, task: Fut) -> Result @@ -1461,6 +1569,44 @@ where .unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message)) } + fn cleanup_unconnected_controller( + &self, + handle: &WorkerHandle, + shutdown: &Arc>>, + tasks: &RuntimeExecutionTaskScope, + worker_state: &Arc>, + ) -> Result<(), String> { + let command = next_internal_command(worker_state)?; + let handle = handle.clone(); + let shutdown = shutdown.clone(); + let tasks_for_join = tasks.clone(); + let cleanup = self.run_on_adapter_runtime(async move { + handle + .send(Method::Shutdown { command }) + .await + .map_err(|error| format!("failed to request controller cleanup: {error}"))?; + let mut guard = shutdown.lock().await; + if let Some(mut receiver) = guard.take() { + match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await { + Ok(Ok(())) => {} + Ok(Err(_)) => { + return Err("controller cleanup completion channel closed".to_string()); + } + Err(_) => { + *guard = Some(receiver); + return Err("controller cleanup confirmation timed out".to_string()); + } + } + } + drop(guard); + tasks_for_join.join().await + }); + if cleanup.is_err() { + tasks.abort_all(); + } + cleanup + } + fn connect_handle( &self, operation: WorkerExecutionOperation, @@ -1468,17 +1614,19 @@ where bridge_context: crate::execution::WorkerExecutionContext, handle: WorkerHandle, shutdown: Arc>>, + controller_task: tokio::task::JoinHandle<()>, working_directory: Option, workspace_client: Option>, ) -> WorkerExecutionSpawnResult { let worker_state = Arc::new(RwLock::new(handle.shared_state.snapshot())); + let tasks = RuntimeExecutionTaskScope::new(controller_task); #[cfg(feature = "ws-server")] { let streams = subscribe_worker_protocol_session(&handle); let mut events = streams.events; let mut entry_events = streams.log_entries; let bridge_worker_state = worker_state.clone(); - if let Err(message) = self.spawn_on_adapter_runtime(async move { + let bridge_task = match self.spawn_on_adapter_runtime(async move { loop { tokio::select! { event = events.recv() => { @@ -1516,10 +1664,20 @@ where } } }) { - return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored( - operation, message, - )); - } + Ok(task) => task, + Err(message) => { + let cleanup = self + .cleanup_unconnected_controller(&handle, &shutdown, &tasks, &worker_state) + .err() + .map(|cleanup| format!("; controller cleanup failed: {cleanup}")) + .unwrap_or_default(); + return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored( + operation, + format!("{message}{cleanup}"), + )); + } + }; + tasks.push("protocol bridge", bridge_task, true); } #[cfg(not(feature = "ws-server"))] { @@ -1529,12 +1687,29 @@ where let mut workers = match self.workers.lock() { Ok(workers) => workers, Err(_) => { + let cleanup = self + .cleanup_unconnected_controller(&handle, &shutdown, &tasks, &worker_state) + .err() + .map(|cleanup| format!("; controller cleanup failed: {cleanup}")) + .unwrap_or_default(); return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored( operation, - "worker adapter registry lock is poisoned", + format!("worker adapter registry lock is poisoned{cleanup}"), )); } }; + if workers.contains_key(&worker_ref) { + drop(workers); + let cleanup = self + .cleanup_unconnected_controller(&handle, &shutdown, &tasks, &worker_state) + .err() + .map(|cleanup| format!("; controller cleanup failed: {cleanup}")) + .unwrap_or_default(); + return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy( + operation, + format!("Worker is already connected to execution backend{cleanup}"), + )); + } let connected_worker_state = worker_state .read() .unwrap_or_else(std::sync::PoisonError::into_inner) @@ -1544,6 +1719,8 @@ where RuntimeWorkerExecution { handle, shutdown, + shutdown_requested: Arc::new(AtomicBool::new(false)), + tasks, worker_state, workspace_client, }, @@ -1820,6 +1997,7 @@ where bridge_context, controller.handle, controller.shutdown, + controller.controller_task, working_directory, Some(controller.workspace_client), ) @@ -1919,6 +2097,7 @@ where bridge_context, controller.handle, controller.shutdown, + controller.controller_task, working_directory, Some(controller.workspace_client), ) @@ -2103,48 +2282,69 @@ where } }; let Some(execution) = execution else { - return WorkerExecutionResult::rejected( - WorkerExecutionOperation::Stop, - "execution handle does not reference a live Worker", - ); + // The execution backend cleanup may have committed before the + // Runtime catalog commit failed. Treat the retry as converged so + // the Runtime can durably finish its Stopped transition. + return WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop); }; - let artifact_cleanup = execution.handle.clone(); - let shutdown = execution.shutdown.clone(); - let command = match next_internal_command(&execution.worker_state) { - Ok(command) => command, - Err(error) => { - return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, error); - } - }; - let result = self.send_method( - WorkerExecutionOperation::Stop, - execution.handle.clone(), - Method::Shutdown { command }, - ); - if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { - return result; - } - let shutdown_wait = self.run_on_adapter_runtime(async move { - let mut guard = shutdown.lock().await; - let Some(mut receiver) = guard.take() else { - return Ok(()); + + let first_request = !execution.shutdown_requested.swap(true, Ordering::AcqRel); + let result = if first_request { + let command = match next_internal_command(&execution.worker_state) { + Ok(command) => command, + Err(error) => { + execution.shutdown_requested.store(false, Ordering::Release); + return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, error); + } }; - match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await { - Ok(Ok(())) => Ok(()), - Ok(Err(_)) => Err("Worker shutdown completion channel closed".to_string()), - Err(_) => { - *guard = Some(receiver); - Err("Worker shutdown confirmation timed out; stop remains retryable".into()) + let result = self.send_method( + WorkerExecutionOperation::Stop, + execution.handle.clone(), + Method::Shutdown { command }, + ); + if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted { + execution.shutdown_requested.store(false, Ordering::Release); + return result; + } + result + } else { + WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop) + }; + + let shutdown = execution.shutdown.clone(); + let tasks = execution.tasks.clone(); + let shutdown_wait = self.run_on_adapter_runtime(async move { + { + let mut guard = shutdown.lock().await; + if let Some(mut receiver) = guard.take() { + match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await { + Ok(Ok(())) => {} + Ok(Err(_)) => { + return Err("Worker shutdown completion channel closed".to_string()); + } + Err(_) => { + *guard = Some(receiver); + return Err( + "Worker shutdown confirmation timed out; stop remains retryable" + .to_string(), + ); + } + } } } + tasks.join().await }); if let Err(message) = shutdown_wait { return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message); } - let artifact_cleanup_error = artifact_cleanup - .delete_uncommitted_uploaded_files() - .err() - .map(|error| format!("uploaded_file_cleanup_failed: {error}")); + + if let Err(error) = execution.handle.delete_uncommitted_uploaded_files() { + return WorkerExecutionResult::errored( + WorkerExecutionOperation::Stop, + format!("uploaded_file_cleanup_failed: {error}; stop remains retryable"), + ); + } + match self.workers.lock() { Ok(mut workers) => { workers.remove(handle.worker_ref()); @@ -2153,13 +2353,7 @@ where poisoned.into_inner().remove(handle.worker_ref()); } } - if let Some(message) = artifact_cleanup_error { - let mut result = result; - result.message = Some(message); - result - } else { - result - } + result } fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult { @@ -2353,14 +2547,6 @@ mod tests { ); } - fn test_command() -> WorkerCommandEnvelope { - WorkerCommandEnvelope { - command_id: 1, - expected_execution_generation: 1, - expected_worker_state_revision: 0, - } - } - fn adapter_command( backend: &WorkerRuntimeExecutionBackend, worker_ref: &WorkerRef, @@ -2828,6 +3014,7 @@ mod tests { Ok(RuntimeWorkerController { handle, shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))), + controller_task: tokio::spawn(async {}), workspace_client, }) } @@ -3343,7 +3530,7 @@ mod tests { let root = tempfile::tempdir().unwrap(); let long_component = "embedded-workspace-store-segment".repeat(4); let runtime_store_dir = root.path().join(long_component); - let worker_ref = WorkerRef::new(crate::identity::WorkerId::from_legacy_u64(1)); + let worker_ref = WorkerRef::new(crate::identity::WorkerId::now_v7()); let worker_aggregate_dir = runtime_store_dir .join("workers") .join(worker_ref.worker_id.to_string()); @@ -3417,16 +3604,24 @@ mod tests { assert!(!socket_path.exists()); assert!(run_dir.join("worker.out.log").is_file()); assert!(run_dir.join("worker.err.log").is_file()); + let worker_state = Arc::new(RwLock::new(controller.handle.shared_state.snapshot())); controller .handle .send(Method::Shutdown { - command: test_command(), + command: next_internal_command(&worker_state).unwrap(), }) .await .unwrap(); if let Some(receiver) = controller.shutdown.lock().await.take() { - receiver.await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), receiver) + .await + .expect("controller shutdown signal timed out") + .unwrap(); } + tokio::time::timeout(Duration::from_secs(5), controller.controller_task) + .await + .expect("controller task join timed out") + .unwrap(); assert!(!socket_path.exists()); } @@ -3529,25 +3724,8 @@ mod tests { ); assert!(!first_run_socket.exists()); - let (handle, shutdown) = { - let workers = backend.workers.lock().unwrap(); - let execution = workers.get(&worker.worker_ref).unwrap(); - (execution.handle.clone(), execution.shutdown.clone()) - }; - backend - .run_on_adapter_runtime(async move { - handle - .send(Method::Shutdown { - command: test_command(), - }) - .await - .map_err(|error| error.to_string())?; - if let Some(receiver) = shutdown.lock().await.take() { - receiver.await.map_err(|error| error.to_string())?; - } - Ok(()) - }) - .unwrap(); + let handle = WorkerExecutionHandle::new(worker.worker_ref.clone(), backend.backend_id()); + assert!(backend.stop_worker(&handle).is_accepted()); drop(runtime); drop(backend); diff --git a/crates/worker/src/bootstrap.rs b/crates/worker/src/bootstrap.rs index ad292df1..234e53e6 100644 --- a/crates/worker/src/bootstrap.rs +++ b/crates/worker/src/bootstrap.rs @@ -67,6 +67,7 @@ pub struct PreparedWorker { pub struct BootstrappedWorker { pub handle: WorkerHandle, pub shutdown: ShutdownReceiver, + pub controller_task: tokio::task::JoinHandle<()>, } #[derive(Debug, Error)] @@ -227,7 +228,7 @@ where runtime_base, bash_output_dir, } => { - WorkerController::spawn_with_transport( + WorkerController::spawn_with_transport_owned( worker, &runtime_base, &bash_output_dir, @@ -239,7 +240,7 @@ where run_dir, bash_output_dir, } => { - WorkerController::spawn_runtime_managed_run_with_transport( + WorkerController::spawn_runtime_managed_run_with_transport_owned( worker, &run_dir, &bash_output_dir, @@ -250,7 +251,11 @@ where }; match controller { - Ok((handle, shutdown)) => Ok(BootstrappedWorker { handle, shutdown }), + Ok((handle, shutdown, controller_task)) => Ok(BootstrappedWorker { + handle, + shutdown, + controller_task, + }), Err(source) => { let cleanup_failed = match cleanup_session { Some(session) => session.close().await.is_err(), diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index d0fd8293..15f95b41 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -486,7 +486,7 @@ impl WorkerController { C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, { - Self::spawn_inner( + let (handle, shutdown, _task) = Self::spawn_inner( worker, runtime_base, bash_output_dir, @@ -494,7 +494,8 @@ impl WorkerController { None, WorkerControllerTransport::UnixSocket, ) - .await + .await?; + Ok((handle, shutdown)) } /// Spawn a direct Worker while letting an in-process host select the @@ -505,6 +506,22 @@ impl WorkerController { bash_output_dir: &Path, transport: WorkerControllerTransport, ) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error> + where + C: LlmClient + Clone + 'static, + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, + { + let (handle, shutdown, _task) = + Self::spawn_with_transport_owned(worker, runtime_base, bash_output_dir, transport) + .await?; + Ok((handle, shutdown)) + } + + pub(crate) async fn spawn_with_transport_owned( + worker: Worker, + runtime_base: &Path, + bash_output_dir: &Path, + transport: WorkerControllerTransport, + ) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error> where C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, @@ -535,7 +552,7 @@ impl WorkerController { C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, { - Self::spawn_inner( + let (handle, shutdown, _task) = Self::spawn_inner( worker, runtime_base, bash_output_dir, @@ -543,7 +560,8 @@ impl WorkerController { None, WorkerControllerTransport::UnixSocket, ) - .await + .await?; + Ok((handle, shutdown)) } /// Spawn into an exact persistent `runs/` directory. @@ -573,6 +591,26 @@ impl WorkerController { bash_output_dir: &Path, transport: WorkerControllerTransport, ) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error> + where + C: LlmClient + Clone + 'static, + St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, + { + let (handle, shutdown, _task) = Self::spawn_runtime_managed_run_with_transport_owned( + worker, + run_dir, + bash_output_dir, + transport, + ) + .await?; + Ok((handle, shutdown)) + } + + pub(crate) async fn spawn_runtime_managed_run_with_transport_owned( + worker: Worker, + run_dir: &Path, + bash_output_dir: &Path, + transport: WorkerControllerTransport, + ) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error> where C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, @@ -598,7 +636,7 @@ impl WorkerController { runtime_managed: bool, runtime_run: Option<&Path>, transport: WorkerControllerTransport, - ) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error> + ) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error> where C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, @@ -629,7 +667,7 @@ impl WorkerController { runtime_managed: bool, runtime_run: Option<&Path>, transport: WorkerControllerTransport, - ) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error> + ) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error> where C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static, @@ -729,9 +767,9 @@ impl WorkerController { None, ) .await?; - if let Some(session) = fs_for_view.as_ref() { - wire_workdir_command_events(session, &in_flight); - } + let command_observer = fs_for_view + .as_ref() + .and_then(|session| wire_workdir_command_events(session, &in_flight)); // Intake role Workers self-terminate only after a successful // TicketIntakeReady turn has fully settled back to Idle. The request @@ -805,7 +843,7 @@ impl WorkerController { let pause_tx = worker.engine_mut().pause_sender(); let notify_buffer = worker.notify_buffer_handle(); - tokio::spawn(controller_loop( + let controller_task = tokio::spawn(controller_loop( worker, method_rx, working_event_tx, @@ -820,26 +858,27 @@ impl WorkerController { shutdown_tx, socket_server, shutdown_after_idle, + command_observer, )); - Ok((handle, shutdown_rx)) + Ok((handle, shutdown_rx, controller_task)) } } pub(crate) fn wire_workdir_command_events( session: &Arc, in_flight: &InFlightEvents, -) { +) -> Option> { in_flight.replace_command_snapshot(protocol_command_snapshots(session.as_ref())); let Some(mut events) = session.subscribe_command_events() else { - return; + return None; }; // Keep only a weak reference in the observer task. Holding the session // strongly here would keep its broadcast sender alive forever and prevent // the receiver from observing closure during Worker teardown. let session = Arc::downgrade(session); let in_flight = in_flight.clone(); - tokio::spawn(async move { + Some(tokio::spawn(async move { loop { match events.recv().await { Ok(event) => in_flight.publish_command_event(protocol_command_event(event)), @@ -853,7 +892,7 @@ pub(crate) fn wire_workdir_command_events( Err(broadcast::error::RecvError::Closed) => break, } } - }); + })) } fn protocol_command_snapshots(session: &dyn WorkdirSession) -> Vec { @@ -1507,6 +1546,7 @@ async fn controller_loop( shutdown_tx: oneshot::Sender<()>, socket_server: Option, shutdown_after_idle: ShutdownAfterIdleRequest, + mut command_observer: Option>, ) where C: LlmClient + Clone + 'static, St: Store + WorkerMetadataStore + Clone + 'static, @@ -2343,28 +2383,49 @@ async fn controller_loop( } } - drop(_socket_server); - if let Err(error) = runtime_dir.close_socket().await { - tracing::warn!(%error, "Worker runtime socket cleanup failed"); + let had_socket_server = _socket_server.is_some(); + if let Some(socket_server) = _socket_server { + socket_server.shutdown().await; + } + while had_socket_server { + match runtime_dir.close_socket().await { + Ok(()) => break, + Err(error) => { + tracing::warn!(%error, "Worker runtime socket cleanup failed; retrying"); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + } } // Feature callbacks and tasks share the Worker scope. Stop them before // Memory/Workdir teardown so they cannot observe a partially closed Worker. worker.stop_feature_runtime("controller shutdown").await; - let child_cleanup_succeeded = match spawned_registry.shutdown_internal().await { - Ok(()) => true, - Err(error) => { - tracing::warn!(%error, "Internal SubWorker cleanup failed before Workdir shutdown"); - false + loop { + match spawned_registry.shutdown_internal().await { + Ok(()) => break, + Err(error) => { + tracing::warn!(%error, "Internal SubWorker cleanup failed; retrying"); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } } - }; + } - if child_cleanup_succeeded - && let Some(session) = worker.workdir_session() - && let Err(error) = session.close().await - { - tracing::warn!(%error, "Workdir session close failed"); + if let Some(session) = worker.workdir_session() { + loop { + match session.close().await { + Ok(()) => break, + Err(error) => { + tracing::warn!(%error, "Workdir session close failed; retrying"); + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + } + } + } + + if let Some(observer) = command_observer.take() { + observer.abort(); + let _ = observer.await; } // Report upward that this Worker is stopping before the controller diff --git a/crates/worker/src/ipc/server.rs b/crates/worker/src/ipc/server.rs index 88566597..8df81da4 100644 --- a/crates/worker/src/ipc/server.rs +++ b/crates/worker/src/ipc/server.rs @@ -4,7 +4,8 @@ use std::path::PathBuf; use protocol::stream::{JsonLineReader, JsonLineWriter}; use tokio::net::UnixListener; -use tokio::task::JoinHandle; +use tokio::sync::oneshot; +use tokio::task::{JoinHandle, JoinSet}; use crate::controller::WorkerHandle; use crate::ipc::protocol_session::{ @@ -19,7 +20,8 @@ use protocol::{ErrorCode, Event}; /// - Client writes Method lines → forwarded to WorkerController /// - Worker events → written as Event lines to all connected clients pub struct SocketServer { - _accept_task: JoinHandle<()>, + accept_task: Option>, + shutdown: Option>, path: PathBuf, } @@ -33,20 +35,45 @@ impl SocketServer { let listener = UnixListener::bind(&path)?; let handle = handle.clone(); + let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); - let _accept_task = tokio::spawn(async move { + let accept_task = tokio::spawn(async move { + let mut connections = JoinSet::new(); loop { - match listener.accept().await { - Ok((stream, _)) => { - let handle = handle.clone(); - tokio::spawn(handle_connection(stream, handle)); + tokio::select! { + _ = &mut shutdown_rx => break, + accepted = listener.accept() => match accepted { + Ok((stream, _)) => { + let handle = handle.clone(); + connections.spawn(handle_connection(stream, handle)); + } + Err(_) => break, + }, + completed = connections.join_next(), if !connections.is_empty() => { + let _ = completed; } - Err(_) => break, } } + connections.shutdown().await; }); - Ok(Self { _accept_task, path }) + Ok(Self { + accept_task: Some(accept_task), + shutdown: Some(shutdown_tx), + path, + }) + } + + /// Stop accepting connections and join the server task. Dropping its + /// `JoinSet` cancels every active connection task before this returns. + pub async fn shutdown(mut self) { + if let Some(shutdown) = self.shutdown.take() { + let _ = shutdown.send(()); + } + if let Some(task) = self.accept_task.take() { + let _ = task.await; + } + let _ = tokio::fs::remove_file(&self.path).await; } /// The socket file path. @@ -57,6 +84,10 @@ impl SocketServer { impl Drop for SocketServer { fn drop(&mut self) { + if let Some(shutdown) = self.shutdown.take() { + let _ = shutdown.send(()); + } + let _ = self.accept_task.take(); let _ = std::fs::remove_file(&self.path); } } diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 038d99e9..4bcc0f5f 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -2219,6 +2219,33 @@ async fn status_json_reflects_worker_name() { // Socket transport tests // --------------------------------------------------------------------------- +#[tokio::test] +async fn shutdown_joins_socket_server_with_active_connection() { + use tokio::net::UnixStream; + + let worker = make_worker(MockClient::new(simple_text_events())).await; + let runtime_base = tempfile::tempdir().unwrap(); + let bash_output_dir = runtime_base.path().join("bash-output"); + let (handle, shutdown_rx) = + WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir) + .await + .unwrap(); + let socket_path = handle.runtime_dir.socket_path(); + let _connection = UnixStream::connect(&socket_path).await.unwrap(); + + handle + .send(Method::Shutdown { + command: worker_command(&handle), + }) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(5), shutdown_rx) + .await + .expect("controller should join its socket tasks") + .expect("controller shutdown signal should remain open"); + assert!(!socket_path.exists()); +} + #[tokio::test] async fn socket_run_receives_events() { use protocol::stream::{JsonLineReader, JsonLineWriter}; diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index c7c29822..f15845c1 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -5064,6 +5064,11 @@ fn embedded_runtime_diagnostic(error: &EmbeddedRuntimeError) -> RuntimeDiagnosti DiagnosticSeverity::Warning, "Embedded Runtime rejected the request".to_string(), ), + EmbeddedRuntimeError::RuntimeStoreAlreadyOpen { .. } => diagnostic( + "embedded_runtime_store_already_open", + DiagnosticSeverity::Error, + "Embedded Runtime store is already owned by another Runtime process".to_string(), + ), EmbeddedRuntimeError::StoreIo { .. } | EmbeddedRuntimeError::StoreMissing { .. } | EmbeddedRuntimeError::StoreCorrupt { .. } => diagnostic(