From ca5fddf89b2f38e54ca8096ada6689404d058897 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 03:48:54 +0900 Subject: [PATCH] fix: fence recursive SubWorker shutdown --- crates/worker/src/controller.rs | 8 +- .../src/feature/builtin/manage_workdir.rs | 65 +++++++- crates/worker/src/spawn/registry.rs | 157 ++++++++++++++++-- crates/worker/src/spawn/tool.rs | 23 ++- 4 files changed, 217 insertions(+), 36 deletions(-) diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index a73a175a..a56c7c59 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1101,14 +1101,16 @@ where "manage Workdir tools require Backend Workspace API authority", )); } - let child_registry = spawned_registry.clone(); + let shutdown_registry = spawned_registry.clone(); + let reopen_registry = spawned_registry.clone(); feature_registry.add_module( - crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_before_workdir_release( + crate::feature::builtin::manage_workdir::ManageWorkdirFeature::with_child_lifecycle( workspace_client, Arc::new(move || { - let child_registry = child_registry.clone(); + let child_registry = shutdown_registry.clone(); Box::pin(async move { child_registry.shutdown_internal().await }) }), + Arc::new(move || reopen_registry.reopen_internal()), ), ); } diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index 24cbc0a9..eb950bbc 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -56,6 +56,7 @@ const ATTACH_DESCRIPTION: &str = "Attach this Worker to one existing Workdir. Th const DETACH_DESCRIPTION: &str = "Detach this Worker from its active Workdir and release Workdir occupancy. Any ephemeral operation session is closed."; pub(crate) type BeforeWorkdirRelease = Arc Pin> + Send>> + Send + Sync>; +pub(crate) type AfterWorkdirAttach = Arc; const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by id through durable Backend Workspace authority. The input includes only the Workdir id and a bounded reason. The result reports removed, retained, or attention_required without exposing operation-table or provider internals."; @@ -63,6 +64,7 @@ const DELETE_DESCRIPTION: &str = "Request removal of one persistent Workdir by i pub struct ManageWorkdirFeature { client: Arc, before_workdir_release: Option, + after_workdir_attach: Option, } impl std::fmt::Debug for ManageWorkdirFeature { @@ -80,16 +82,19 @@ impl ManageWorkdirFeature { Self { client, before_workdir_release: None, + after_workdir_attach: None, } } - pub(crate) fn with_before_workdir_release( + pub(crate) fn with_child_lifecycle( client: Arc, before_workdir_release: BeforeWorkdirRelease, + after_workdir_attach: AfterWorkdirAttach, ) -> Self { Self { client, before_workdir_release: Some(before_workdir_release), + after_workdir_attach: Some(after_workdir_attach), } } } @@ -110,8 +115,10 @@ impl FeatureModule for ManageWorkdirFeature { } fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { - let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()) - .with_before_workdir_release(self.before_workdir_release.clone()); + let backend = WorkspaceHttpWorkdirBackend::new(self.client.clone()).with_child_lifecycle( + self.before_workdir_release.clone(), + self.after_workdir_attach.clone(), + ); for (name, definition) in [ ( LIST_TOOL, @@ -176,6 +183,7 @@ impl FeatureModule for ManageWorkdirFeature { struct WorkspaceHttpWorkdirBackend { client: Arc, before_workdir_release: Option, + after_workdir_attach: Option, } impl std::fmt::Debug for WorkspaceHttpWorkdirBackend { @@ -371,14 +379,17 @@ impl WorkspaceHttpWorkdirBackend { Self { client, before_workdir_release: None, + after_workdir_attach: None, } } - fn with_before_workdir_release( + fn with_child_lifecycle( mut self, before_workdir_release: Option, + after_workdir_attach: Option, ) -> Self { self.before_workdir_release = before_workdir_release; + self.after_workdir_attach = after_workdir_attach; self } @@ -557,9 +568,17 @@ impl Tool for WorkspaceHttpWorkdirTool { parse_input::(input_json)?, ctx.call_id.to_string(), ), - WorkdirOperation::Attach => self - .backend - .attach(parse_input::(input_json)?), + WorkdirOperation::Attach => { + let result = self + .backend + .attach(parse_input::(input_json)?); + if result.is_ok() + && let Some(after_attach) = &self.backend.after_workdir_attach + { + after_attach(); + } + result + } WorkdirOperation::Detach => { let _input = parse_input::(input_json)?; if let Some(before_release) = &self.backend.before_workdir_release { @@ -1338,7 +1357,7 @@ mod tests { }); let tool = WorkspaceHttpWorkdirTool { backend: WorkspaceHttpWorkdirBackend::new(client.clone()) - .with_before_workdir_release(Some(before_release)), + .with_child_lifecycle(Some(before_release), None), operation: WorkdirOperation::Detach, }; @@ -1361,7 +1380,7 @@ mod tests { Arc::new(|| Box::pin(async { Err(std::io::Error::other("child cleanup failed")) })); let tool = WorkspaceHttpWorkdirTool { backend: WorkspaceHttpWorkdirBackend::new(client.clone()) - .with_before_workdir_release(Some(before_release)), + .with_child_lifecycle(Some(before_release), None), operation: WorkdirOperation::Detach, }; @@ -1373,4 +1392,32 @@ mod tests { assert!(error.to_string().contains("stop Internal SubWorkers")); assert!(client.requests().is_empty()); } + + #[tokio::test] + async fn successful_attach_reopens_internal_subworker_admission() { + let client = Arc::new(RecordingWorkspaceClient::new(vec![response(json!({ + "workspace_id": "workspace/test", + "workdir_id": "wd-attached", + "attached": true + }))])); + let reopen_calls = Arc::new(AtomicUsize::new(0)); + let reopen_calls_for_hook = reopen_calls.clone(); + let after_attach: AfterWorkdirAttach = Arc::new(move || { + reopen_calls_for_hook.fetch_add(1, Ordering::SeqCst); + }); + let tool = WorkspaceHttpWorkdirTool { + backend: WorkspaceHttpWorkdirBackend::new(client) + .with_child_lifecycle(None, Some(after_attach)), + operation: WorkdirOperation::Attach, + }; + + tool.execute( + r#"{"workdir_id":"wd-attached"}"#, + ToolExecutionContext::default(), + ) + .await + .unwrap(); + + assert_eq!(reopen_calls.load(Ordering::SeqCst), 1); + } } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index 3db26267..cba26327 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -72,6 +72,7 @@ pub(crate) struct InternalSpawnedWorkerRecord { #[cfg(test)] pub installed_tools: Arc<[String]>, pub session: InternalWorkerSessionHandle, + pub child_registry: Arc, change_tracker: Option, started_at: Instant, stop_lock: Arc>, @@ -89,6 +90,7 @@ impl InternalSpawnedWorkerRecord { workdir_tool_scope: WorkdirScopeLease, #[cfg(test)] installed_tools: Vec, session: InternalWorkerSessionHandle, + child_registry: Arc, change_tracker: Option, ) -> Self { Self { @@ -98,6 +100,7 @@ impl InternalSpawnedWorkerRecord { #[cfg(test)] installed_tools: installed_tools.into(), session, + child_registry, change_tracker, started_at: Instant::now(), stop_lock: Arc::new(tokio::sync::Mutex::new(())), @@ -235,18 +238,39 @@ pub(crate) struct InternalSpawnReservation { } impl InternalSpawnReservation { - pub(crate) fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> { + pub(crate) fn commit( + mut self, + record: InternalSpawnedWorkerRecord, + ) -> Result<(), (io::Error, InternalSpawnedWorkerRecord)> { if record.worker_name != self.worker_name { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "internal SubWorker reservation name does not match record name", + return Err(( + io::Error::new( + io::ErrorKind::InvalidInput, + "internal SubWorker reservation name does not match record name", + ), + record, )); } - self.registry - .internal_records - .lock() - .map_err(|_| io::Error::other("internal spawned-worker registry lock poisoned"))? - .push(record.clone()); + let mut records = match self.registry.internal_records.lock() { + Ok(records) => records, + Err(_) => { + return Err(( + io::Error::other("internal spawned-worker registry lock poisoned"), + record, + )); + } + }; + if self.registry.internal_shutting_down.load(Ordering::Acquire) { + return Err(( + io::Error::new( + io::ErrorKind::Interrupted, + "internal SubWorker registry is shutting down", + ), + record, + )); + } + records.push(record.clone()); + drop(records); self.registry.start_protocol_forwarding(record); self.committed = true; Ok(()) @@ -267,6 +291,7 @@ pub struct SpawnedWorkerRegistry { internal_records: std::sync::Mutex>, service_records: std::sync::Mutex>, internal_names: std::sync::Mutex>, + internal_shutting_down: AtomicBool, parent_scope: Option, parent_protocol: Mutex, String)>>, } @@ -283,6 +308,7 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -294,6 +320,7 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -304,6 +331,7 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), parent_scope: Some(parent_scope), parent_protocol: Mutex::new(None), }) @@ -383,6 +411,7 @@ impl SpawnedWorkerRegistry { internal_records: std::sync::Mutex::new(Vec::new()), service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), + internal_shutting_down: AtomicBool::new(false), parent_scope, parent_protocol: Mutex::new(None), }), @@ -394,6 +423,12 @@ impl SpawnedWorkerRegistry { self: &Arc, worker_name: String, ) -> io::Result { + if self.internal_shutting_down.load(Ordering::Acquire) { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "internal SubWorker registry is shutting down", + )); + } let mut names = self .internal_names .lock() @@ -702,6 +737,7 @@ impl SpawnedWorkerRegistry { let Some(record) = self.get_internal(name) else { return Ok(false); }; + Box::pin(record.child_registry.shutdown_internal()).await?; record .workdir_tool_scope .close() @@ -711,13 +747,17 @@ impl SpawnedWorkerRegistry { } pub(crate) async fn shutdown_internal(&self) -> io::Result<()> { - let names = self - .internal_records - .lock() - .expect("internal Worker registry lock poisoned") - .iter() - .map(|record| record.worker_name.clone()) - .collect::>(); + let names = { + let records = self + .internal_records + .lock() + .map_err(|_| io::Error::other("internal Worker registry lock poisoned"))?; + self.internal_shutting_down.store(true, Ordering::Release); + records + .iter() + .map(|record| record.worker_name.clone()) + .collect::>() + }; let mut first_error = None; for name in names { if let Err(error) = self.remove_internal(&name).await { @@ -727,6 +767,10 @@ impl SpawnedWorkerRegistry { first_error.map_or(Ok(()), Err) } + pub(crate) fn reopen_internal(&self) { + self.internal_shutting_down.store(false, Ordering::Release); + } + /// Stop one direct Internal SubWorker and discard its registry/scope state. /// /// The child actor must acknowledge its stop before the registry is removed. @@ -753,6 +797,7 @@ impl SpawnedWorkerRegistry { .stop() .await .map_err(|error| io::Error::other(error.to_string()))?; + Box::pin(record.child_registry.shutdown_internal()).await?; record .workdir_tool_scope .close() @@ -1021,6 +1066,7 @@ mod tests { delegation, Vec::new(), session, + registry(), None, ), sender, @@ -1276,6 +1322,85 @@ mod tests { assert!(registry.get_internal("second").is_none()); } + #[tokio::test] + async fn shutdown_rejects_new_reservations_until_reopened() { + let registry = registry(); + registry.shutdown_internal().await.unwrap(); + assert!(registry.reserve_internal_name("late-child".into()).is_err()); + + registry.reopen_internal(); + let reservation = registry.reserve_internal_name("late-child".into()).unwrap(); + drop(reservation); + } + + #[tokio::test] + async fn concurrent_commit_and_shutdown_leave_no_live_internal_worker() { + let registry = registry(); + let reservation = registry + .reserve_internal_name("racing-child".into()) + .unwrap(); + let (record, _events) = + record("racing-child", InternalWorkerVisibility::ParentClient).await; + let scope = record.workdir_tool_scope.clone(); + let barrier = Arc::new(std::sync::Barrier::new(2)); + let commit_barrier = barrier.clone(); + let commit = tokio::task::spawn_blocking(move || { + commit_barrier.wait(); + reservation.commit(record) + }); + let shutdown_registry = registry.clone(); + let shutdown = tokio::spawn(async move { + barrier.wait(); + shutdown_registry.shutdown_internal().await + }); + + let commit = commit.await.unwrap(); + shutdown.await.unwrap().unwrap(); + if let Err((_error, record)) = commit { + record.session.stop().await.unwrap(); + record.child_registry.shutdown_internal().await.unwrap(); + record.workdir_tool_scope.close().await.unwrap(); + } + + assert!(registry.list_internal().is_empty()); + assert!(!scope.is_active()); + } + + #[tokio::test] + async fn shutdown_fences_a_reservation_that_has_not_committed() { + let registry = registry(); + let reservation = registry + .reserve_internal_name("racing-child".into()) + .unwrap(); + let (record, _events) = + record("racing-child", InternalWorkerVisibility::ParentClient).await; + + registry.shutdown_internal().await.unwrap(); + let (error, record) = reservation.commit(record).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::Interrupted); + record.session.stop().await.unwrap(); + record.child_registry.shutdown_internal().await.unwrap(); + record.workdir_tool_scope.close().await.unwrap(); + } + + #[tokio::test] + async fn shutdown_recursively_stops_grandchildren_before_parent_scope_release() { + let registry = registry(); + let (child, _child_events) = record("child", InternalWorkerVisibility::ParentClient).await; + let child_registry = child.child_registry.clone(); + let (grandchild, _grandchild_events) = + record("grandchild", InternalWorkerVisibility::ParentClient).await; + let grandchild_scope = grandchild.workdir_tool_scope.clone(); + install_record(&child_registry, grandchild); + install_record(®istry, child); + + registry.shutdown_internal().await.unwrap(); + + assert!(registry.list_internal().is_empty()); + assert!(child_registry.list_internal().is_empty()); + assert!(!grandchild_scope.is_active()); + } + #[tokio::test] async fn running_worker_is_stopped_before_removal() { let registry = registry(); diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 294983b0..547d939d 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -600,15 +600,19 @@ impl Tool for SubWorkerSpawnTool { ), body.to_string(), ); - let response = self - .workspace_context - .client() - .execute(request) - .map_err(|error| { - ToolError::ExecutionFailed(format!("register review capability: {error}")) - })?; + let response = match self.workspace_context.client().execute(request) { + Ok(response) => response, + Err(error) => { + let _ = session.stop().await; + let _ = workdir_scope.close().await; + return Err(ToolError::ExecutionFailed(format!( + "register review capability: {error}" + ))); + } + }; if !response.is_success() { let _ = session.stop().await; + let _ = workdir_scope.close().await; return Err(ToolError::ExecutionFailed(format!( "register review capability failed with status {}: {}", response.status, response.body @@ -623,10 +627,13 @@ impl Tool for SubWorkerSpawnTool { #[cfg(test)] installed_tools, session.clone(), + child_registry, child_change_tracker, ); - if let Err(error) = name_reservation.commit(record) { + if let Err((error, record)) = name_reservation.commit(record) { let _ = session.stop().await; + let _ = record.child_registry.shutdown_internal().await; + let _ = record.workdir_tool_scope.close().await; return Err(ToolError::ExecutionFailed(format!( "register Internal Worker session: {error}" )));