Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix: fence recursive SubWorker shutdown

Browse Source
This commit is contained in:
Hare
2026-09-06 03:48:54 +09:00
parent 7b1cf854f2
commit ca5fddf89b
4 changed files with 217 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/controller.rs
+5 -3
View File
@@ -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()),
),
);
}
crates/worker/src/feature/builtin/manage_workdir.rs
+55 -8
View File
@@ -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<dyn Fn() -> Pin<Box<dyn Future<Output = std::io::Result<()>> + Send>> + Send + Sync>;
pub(crate) type AfterWorkdirAttach = Arc<dyn Fn() + Send + Sync>;
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<dyn WorkspaceClient>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
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<dyn WorkspaceClient>,
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<dyn WorkspaceClient>,
before_workdir_release: Option<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
}
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<BeforeWorkdirRelease>,
after_workdir_attach: Option<AfterWorkdirAttach>,
) -> 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::<WorkdirCreateInput>(input_json)?,
ctx.call_id.to_string(),
),
WorkdirOperation::Attach => self
WorkdirOperation::Attach => {
let result = self
.backend
.attach(parse_input::<WorkdirAttachInput>(input_json)?),
.attach(parse_input::<WorkdirAttachInput>(input_json)?);
if result.is_ok()
&& let Some(after_attach) = &self.backend.after_workdir_attach
{
after_attach();
}
result
}
WorkdirOperation::Detach => {
let _input = parse_input::<WorkdirDetachInput>(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);
}
}
crates/worker/src/spawn/registry.rs
+135 -10
View File
@@ -72,6 +72,7 @@ pub(crate) struct InternalSpawnedWorkerRecord {
#[cfg(test)]
pub installed_tools: Arc<[String]>,
pub session: InternalWorkerSessionHandle,
pub child_registry: Arc<SpawnedWorkerRegistry>,
change_tracker: Option<tools::Tracker>,
started_at: Instant,
stop_lock: Arc<tokio::sync::Mutex<()>>,
@@ -89,6 +90,7 @@ impl InternalSpawnedWorkerRecord {
workdir_tool_scope: WorkdirScopeLease,
#[cfg(test)] installed_tools: Vec<String>,
session: InternalWorkerSessionHandle,
child_registry: Arc<SpawnedWorkerRegistry>,
change_tracker: Option<tools::Tracker>,
) -> 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(
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<Vec<InternalSpawnedWorkerRecord>>,
service_records: std::sync::Mutex<Vec<InternalServiceWorkerRecord>>,
internal_names: std::sync::Mutex<HashSet<String>>,
internal_shutting_down: AtomicBool,
parent_scope: Option<SharedScope>,
parent_protocol: Mutex<Option<(broadcast::Sender<Event>, 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<Self>,
worker_name: String,
) -> io::Result<InternalSpawnReservation> {
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
let names = {
let records = self
.internal_records
.lock()
.expect("internal Worker registry lock poisoned")
.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::<Vec<_>>();
.collect::<Vec<_>>()
};
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(&registry, 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();
crates/worker/src/spawn/tool.rs
+15 -8
View File
@@ -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}"
)));