From d6bef5d1c7e64315cb48ff5c8413fd994ddaf6a6 Mon Sep 17 00:00:00 2001 From: Hare Date: Thu, 17 Sep 2026 00:36:08 +0900 Subject: [PATCH] fix: address lifecycle review findings --- crates/worker-runtime/src/fs_store.rs | 26 ++- crates/worker-runtime/src/runtime.rs | 190 ++++++++++++++++++-- crates/worker-runtime/src/worker_backend.rs | 90 +++++++--- crates/worker/src/controller.rs | 4 + crates/worker/tests/controller_test.rs | 37 ++++ 5 files changed, 307 insertions(+), 40 deletions(-) diff --git a/crates/worker-runtime/src/fs_store.rs b/crates/worker-runtime/src/fs_store.rs index 393696b2..a2706ba4 100644 --- a/crates/worker-runtime/src/fs_store.rs +++ b/crates/worker-runtime/src/fs_store.rs @@ -1725,12 +1725,36 @@ mod tests { let Some(root) = std::env::var_os("YOI_TEST_RUNTIME_STORE_LOCK_ROOT") else { return; }; + let root = PathBuf::from(root); + if std::env::var_os("YOI_TEST_RUNTIME_STORE_LOCK_EXIT_WITHOUT_DROP").is_some() { + let _store = FsRuntimeStore::open_or_create(root, "runtime-test").unwrap(); + std::process::exit(0); + } assert!(matches!( - FsRuntimeStore::open_or_create(PathBuf::from(root), "runtime-test").unwrap_err(), + FsRuntimeStore::open_or_create(root, "runtime-test").unwrap_err(), RuntimeError::RuntimeStoreAlreadyOpen { .. } )); } + #[test] + fn runtime_store_owner_lock_is_released_when_process_exits_without_drop() { + let parent = tempfile::tempdir().unwrap(); + let root = parent.path().join("crashed-runtime-store"); + 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) + .env("YOI_TEST_RUNTIME_STORE_LOCK_EXIT_WITHOUT_DROP", "1") + .output() + .unwrap(); + assert!( + child.status.success(), + "process-exit lock probe failed: {}", + String::from_utf8_lossy(&child.stderr) + ); + acquire_runtime_store_owner_lock(&root).unwrap(); + } + #[test] fn second_runtime_store_open_conflicts_before_store_mutation() { let parent = tempfile::tempdir().unwrap(); diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 2669f8cb..a8f32b28 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -162,6 +162,36 @@ pub struct Runtime { worker_operations: Arc>>>>, } +struct WorkerOperationLease { + worker_id: WorkerId, + operation_lock: Arc>, + registry: Arc>>>>, +} + +impl std::ops::Deref for WorkerOperationLease { + type Target = Mutex<()>; + + fn deref(&self) -> &Self::Target { + self.operation_lock.as_ref() + } +} + +impl Drop for WorkerOperationLease { + fn drop(&mut self) { + let mut registry = match self.registry.lock() { + Ok(registry) => registry, + Err(poisoned) => poisoned.into_inner(), + }; + if Arc::strong_count(&self.operation_lock) == 2 + && registry + .get(&self.worker_id) + .is_some_and(|current| Arc::ptr_eq(current, &self.operation_lock)) + { + registry.remove(&self.worker_id); + } + } +} + impl Runtime { /// Create a memory-backed Runtime with generated identity. pub fn new_memory() -> Self { @@ -1295,8 +1325,7 @@ impl Runtime { scope: &RuntimeWorkspaceScope, worker_ref: &WorkerRef, ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.restore_worker(worker_ref) + self.restore_worker_with_scope(worker_ref, Some(scope)) } /// Attach a live execution to a persisted Worker definition. @@ -1306,10 +1335,21 @@ impl Runtime { /// converges on its already-installed execution rather than spawning another /// controller. pub fn restore_worker(&self, worker_ref: &WorkerRef) -> Result { + self.restore_worker_with_scope(worker_ref, None) + } + + fn restore_worker_with_scope( + &self, + worker_ref: &WorkerRef, + scope: Option<&RuntimeWorkspaceScope>, + ) -> Result { let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; let _operation_guard = operation_lock .lock() .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } self.restore_worker_under_lock(worker_ref, WorkerRestoreMode::Explicit) } @@ -1934,8 +1974,7 @@ impl Runtime { worker_ref: &WorkerRef, reason: Option, ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.stop_worker(worker_ref, reason) + self.stop_worker_with_scope(worker_ref, reason, Some(scope)) } /// Stop a Worker. Repeated stops are idempotent. The per-Worker lifecycle @@ -1945,11 +1984,23 @@ impl Runtime { &self, worker_ref: &WorkerRef, reason: Option, + ) -> Result { + self.stop_worker_with_scope(worker_ref, reason, None) + } + + fn stop_worker_with_scope( + &self, + worker_ref: &WorkerRef, + reason: Option, + scope: Option<&RuntimeWorkspaceScope>, ) -> Result { let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; let _operation_guard = operation_lock .lock() .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } let (backend, handle) = { let state = self.lock()?; @@ -2043,19 +2094,29 @@ impl Runtime { scope: &RuntimeWorkspaceScope, worker_ref: &WorkerRef, ) -> Result { - self.ensure_worker_in_workspace(scope, worker_ref)?; - self.delete_worker(worker_ref) + self.delete_worker_with_scope(worker_ref, Some(scope)) } /// Delete a non-running Worker from Runtime state and persisted Worker storage. pub fn delete_worker( &self, worker_ref: &WorkerRef, + ) -> Result { + self.delete_worker_with_scope(worker_ref, None) + } + + fn delete_worker_with_scope( + &self, + worker_ref: &WorkerRef, + scope: Option<&RuntimeWorkspaceScope>, ) -> Result { let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?; let _operation_guard = operation_lock .lock() .map_err(|_| RuntimeError::StatePoisoned)?; + if let Some(scope) = scope { + self.ensure_worker_in_workspace(scope, worker_ref)?; + } let (backend, execution_handle) = { let state = self.lock()?; state.ensure_running()?; @@ -2540,15 +2601,25 @@ impl Runtime { Ok(result) } - fn worker_operation_lock(&self, worker_id: WorkerId) -> Result>, RuntimeError> { - let mut operations = self - .worker_operations - .lock() - .map_err(|_| RuntimeError::StatePoisoned)?; - Ok(operations - .entry(worker_id) - .or_insert_with(|| Arc::new(Mutex::new(()))) - .clone()) + fn worker_operation_lock( + &self, + worker_id: WorkerId, + ) -> Result { + let operation_lock = { + let mut operations = self + .worker_operations + .lock() + .map_err(|_| RuntimeError::StatePoisoned)?; + operations + .entry(worker_id) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone() + }; + Ok(WorkerOperationLease { + worker_id, + operation_lock, + registry: self.worker_operations.clone(), + }) } fn lock(&self) -> Result, RuntimeError> { @@ -5789,6 +5860,43 @@ mod tests { ); } + #[test] + fn remove_waits_for_in_flight_restore_and_cannot_delete_live_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("restore remove race")) + .unwrap(); + runtime.stop_worker(&created.worker_ref, None).unwrap(); + let gate = Arc::new(RestoreGate::default()); + *backend.restore_gate.lock().unwrap() = Some(gate.clone()); + + let restore_runtime = runtime.clone(); + let restore_ref = created.worker_ref.clone(); + let restoring = std::thread::spawn(move || restore_runtime.restore_worker(&restore_ref)); + assert!(gate.wait_for_entered(1, std::time::Duration::from_secs(2))); + let remove_runtime = runtime.clone(); + let remove_ref = created.worker_ref.clone(); + let removing = std::thread::spawn(move || remove_runtime.delete_worker(&remove_ref)); + std::thread::sleep(std::time::Duration::from_millis(100)); + assert_eq!(runtime.list_workers().unwrap().len(), 1); + + gate.release(); + assert_eq!( + restoring.join().unwrap().unwrap().status, + WorkerStatus::Idle + ); + assert!(matches!( + removing.join().unwrap().unwrap_err(), + RuntimeError::InvalidRequest(message) if message.contains("must be stopped") + )); + assert_eq!(runtime.list_workers().unwrap().len(), 1); + assert!(runtime.worker_operations.lock().unwrap().is_empty()); + } + #[test] fn failed_stop_cleanup_retains_execution_for_retry() { let (runtime, backend) = runtime_and_backend(); @@ -5912,6 +6020,58 @@ mod tests { .as_ref() .map(|state| state.execution_generation) ); + assert!(runtime.worker_operations.lock().unwrap().is_empty()); + } + + #[test] + fn operation_lock_registry_reclaims_idle_worker_entries() { + let runtime = runtime_with_backend(); + for _ in 0..256 { + let lease = runtime.worker_operation_lock(WorkerId::now_v7()).unwrap(); + let guard = lease.lock().unwrap(); + drop(guard); + drop(lease); + } + assert!(runtime.worker_operations.lock().unwrap().is_empty()); + } + + #[test] + fn scoped_stop_revalidates_workspace_after_waiting_for_worker_lock() { + let (runtime, backend) = runtime_and_backend(); + let runtime = Arc::new(runtime); + runtime.store_config_bundle(test_bundle()).unwrap(); + let workspace_scope = scope("workspace-a", "server-a"); + let worker = runtime + .create_worker_scoped( + &workspace_scope, + scoped_task_request("scoped stop race", "workspace-a"), + ) + .unwrap(); + let lease = runtime + .worker_operation_lock(worker.worker_ref.worker_id) + .unwrap(); + let guard = lease.lock().unwrap(); + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let stopping_runtime = runtime.clone(); + let stopping_ref = worker.worker_ref.clone(); + let stopping_scope = workspace_scope.clone(); + let stopping = std::thread::spawn(move || { + started_tx.send(()).unwrap(); + stopping_runtime.stop_worker_scoped(&stopping_scope, &stopping_ref, None) + }); + started_rx.recv().unwrap(); + { + let mut state = runtime.lock().unwrap(); + state.worker_mut(&worker.worker_ref).unwrap().workspace_id = + Some("workspace-b".to_string()); + } + drop(guard); + drop(lease); + + let error = stopping.join().unwrap().unwrap_err(); + assert!(matches!(error, RuntimeError::WorkerNotFound { .. })); + assert_eq!(*backend.stop_count.lock().unwrap(), 0); + assert!(runtime.worker_operations.lock().unwrap().is_empty()); } #[test] diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index ace3bf9f..cc405c64 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -1197,7 +1197,6 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { #[derive(Clone)] struct RuntimeExecutionTaskScope { tasks: Arc>>, - terminal_failure: Arc>>, } struct RuntimeExecutionTask { @@ -1214,7 +1213,6 @@ impl RuntimeExecutionTaskScope { task: controller_task, abort_before_join: false, }])), - terminal_failure: Arc::new(Mutex::new(None)), } } @@ -1241,15 +1239,8 @@ impl RuntimeExecutionTaskScope { } 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); - } - + let mut first_failure = None; + let mut pending = Vec::new(); loop { let next = { let mut tasks = self @@ -1263,7 +1254,7 @@ impl RuntimeExecutionTaskScope { } }; let Some(mut task) = next else { - return Ok(()); + break; }; let name = task.name; if task.abort_before_join { @@ -1273,23 +1264,29 @@ impl RuntimeExecutionTaskScope { 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); + first_failure.get_or_insert_with(|| { + format!("{name} task failed while stopping Worker: {error}") + }); } 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" - )); + first_failure.get_or_insert_with(|| { + format!("{name} task did not stop before timeout; stop remains retryable") + }); + pending.push(task); } } } + if !pending.is_empty() { + let mut tasks = match self.tasks.lock() { + Ok(tasks) => tasks, + Err(poisoned) => poisoned.into_inner(), + }; + tasks.extend(pending); + } + match first_failure { + Some(message) => Err(message), + None => Ok(()), + } } } @@ -2547,6 +2544,22 @@ mod tests { ); } + #[tokio::test] + async fn execution_scope_drains_remaining_tasks_after_join_failure_and_allows_retry() { + let failed = tokio::spawn(async { panic!("injected controller failure") }); + let scope = RuntimeExecutionTaskScope::new(failed); + scope.push( + "protocol bridge", + tokio::spawn(std::future::pending()), + true, + ); + + let error = scope.join().await.unwrap_err(); + assert!(error.contains("controller task failed")); + assert!(scope.tasks.lock().unwrap().is_empty()); + scope.join().await.unwrap(); + } + fn adapter_command( backend: &WorkerRuntimeExecutionBackend, worker_ref: &WorkerRef, @@ -4197,7 +4210,36 @@ mod tests { let detail = runtime .create_worker(create_request("restore-after-stop")) .unwrap(); + let failed_task = backend + .spawn_on_adapter_runtime(async { panic!("injected owned task failure") }) + .unwrap(); + backend + .workers + .lock() + .unwrap() + .get(&detail.worker_ref) + .unwrap() + .tasks + .push("injected failure", failed_task, false); + let first_stop = runtime.stop_worker(&detail.worker_ref, None).unwrap_err(); + assert!( + first_stop + .to_string() + .contains("injected failure task failed") + ); + assert_eq!( + runtime.worker_detail(&detail.worker_ref).unwrap().status, + crate::catalog::WorkerStatus::Idle + ); + assert!( + backend + .workers + .lock() + .unwrap() + .contains_key(&detail.worker_ref), + "failed cleanup must retain retry authority" + ); runtime.stop_worker(&detail.worker_ref, None).unwrap(); assert_eq!( runtime.worker_detail(&detail.worker_ref).unwrap().status, diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 15f95b41..b22bd998 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -2383,6 +2383,10 @@ async fn controller_loop( } } + // Close method admission before any fallible child/session cleanup. Existing + // senders then fail instead of accepting work that this execution can no + // longer process. + drop(method_rx); let had_socket_server = _socket_server.is_some(); if let Some(socket_server) = _socket_server { socket_server.shutdown().await; diff --git a/crates/worker/tests/controller_test.rs b/crates/worker/tests/controller_test.rs index 4bcc0f5f..cba7156e 100644 --- a/crates/worker/tests/controller_test.rs +++ b/crates/worker/tests/controller_test.rs @@ -2219,6 +2219,43 @@ async fn status_json_reflects_worker_name() { // Socket transport tests // --------------------------------------------------------------------------- +#[tokio::test] +async fn shutdown_closes_method_admission_before_terminal_confirmation() { + 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, mut shutdown_rx) = + WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir) + .await + .unwrap(); + handle + .send(Method::Shutdown { + command: worker_command(&handle), + }) + .await + .unwrap(); + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + tokio::select! { + biased; + result = handle.send(Method::ListRewindTargets) => { + if result.is_err() { + break; + } + } + result = &mut shutdown_rx => { + result.expect("controller shutdown signal should remain open"); + panic!("method admission remained open until terminal confirmation"); + } + } + } + }) + .await + .expect("method admission did not close during shutdown"); + shutdown_rx.await.unwrap(); +} + #[tokio::test] async fn shutdown_joins_socket_server_with_active_connection() { use tokio::net::UnixStream;