From 052d60bd7d2f0f9cbe3da19ced380e3ce6af82d9 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 6 Sep 2026 04:10:23 +0900 Subject: [PATCH] fix: fence late command events and spawn rollback --- crates/workdir/src/scope.rs | 34 ++++- crates/worker/src/spawn/registry.rs | 188 +++++++++++++++++++++------- crates/worker/src/spawn/tool.rs | 5 +- 3 files changed, 173 insertions(+), 54 deletions(-) diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index e730d9f0..131ae6cf 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -75,6 +75,7 @@ impl WorkdirToolBroker { owned_commands: Arc::new(Mutex::new(HashSet::new())), pending_command_events: Arc::new(Mutex::new(HashMap::new())), starting_tool_calls: Arc::new(Mutex::new(HashSet::new())), + forwarded_starts: Arc::new(Mutex::new(HashSet::new())), forwarded_terminals: Arc::new(Mutex::new(HashSet::new())), command_events, closes_source: true, @@ -314,6 +315,7 @@ struct ScopedWorkdirSession { owned_commands: Arc>>, pending_command_events: Arc>>>, starting_tool_calls: Arc>>, + forwarded_starts: Arc>>, forwarded_terminals: Arc>>, command_events: broadcast::Sender, closes_source: bool, @@ -675,6 +677,7 @@ impl ScopedWorkdirSession { let owned_commands = Arc::new(Mutex::new(HashSet::new())); let pending_command_events = Arc::new(Mutex::new(HashMap::new())); let starting_tool_calls = Arc::new(Mutex::new(HashSet::new())); + let forwarded_starts = Arc::new(Mutex::new(HashSet::new())); let forwarded_terminals = Arc::new(Mutex::new(HashSet::new())); let (command_events, _) = broadcast::channel(64); let event_forwarder = forward_owned_command_events( @@ -682,6 +685,7 @@ impl ScopedWorkdirSession { owned_commands.clone(), pending_command_events.clone(), starting_tool_calls.clone(), + forwarded_starts.clone(), forwarded_terminals.clone(), command_events.clone(), ) @@ -699,6 +703,7 @@ impl ScopedWorkdirSession { owned_commands, pending_command_events, starting_tool_calls, + forwarded_starts, forwarded_terminals, command_events, closes_source: false, @@ -862,6 +867,7 @@ impl WorkdirSession for ScopedWorkdirSession { { publish_owned_command_event( &self.command_events, + &self.forwarded_starts, &self.forwarded_terminals, CommandEvent::Started { command_id: handle.0.clone(), @@ -871,7 +877,12 @@ impl WorkdirSession for ScopedWorkdirSession { ); } for event in pending { - publish_owned_command_event(&self.command_events, &self.forwarded_terminals, event); + publish_owned_command_event( + &self.command_events, + &self.forwarded_starts, + &self.forwarded_terminals, + event, + ); } Ok(handle) } @@ -1033,6 +1044,7 @@ fn forward_owned_command_events( owned_commands: Arc>>, pending_command_events: Arc>>>, starting_tool_calls: Arc>>, + forwarded_starts: Arc>>, forwarded_terminals: Arc>>, sender: broadcast::Sender, ) -> Option> { @@ -1082,7 +1094,7 @@ fn forward_owned_command_events( } drop(pending); drop(owned); - publish_owned_command_event(&sender, &forwarded_terminals, event); + publish_owned_command_event(&sender, &forwarded_starts, &forwarded_terminals, event); } })) } @@ -1097,6 +1109,7 @@ fn command_event_id(event: &CommandEvent) -> &str { fn publish_owned_command_event( sender: &broadcast::Sender, + forwarded_starts: &Mutex>, forwarded_terminals: &Mutex>, event: CommandEvent, ) { @@ -1106,11 +1119,16 @@ fn publish_owned_command_event( .expect("forwarded terminal command mutex poisoned"); match &event { CommandEvent::Terminal { .. } if !terminals.insert(command_id.to_string()) => return, - CommandEvent::Started { .. } | CommandEvent::Output { .. } - if terminals.contains(command_id) => + CommandEvent::Started { .. } if terminals.contains(command_id) => return, + CommandEvent::Started { .. } + if !forwarded_starts + .lock() + .expect("forwarded command start mutex poisoned") + .insert(command_id.to_string()) => { return; } + CommandEvent::Output { .. } if terminals.contains(command_id) => return, _ => {} } drop(terminals); @@ -1657,8 +1675,16 @@ mod tests { } assert_eq!(kinds.first(), Some(&"started")); assert_eq!(kinds.last(), Some(&"terminal")); + assert_eq!(kinds.iter().filter(|kind| **kind == "started").count(), 1); assert!(kinds.contains(&"output")); assert!(streamed.contains("fast-output")); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(100), events.recv()) + .await + .is_err(), + "no provider event may follow the terminal event" + ); + assert!(child.command_snapshot().is_empty()); child.close().await.unwrap(); } diff --git a/crates/worker/src/spawn/registry.rs b/crates/worker/src/spawn/registry.rs index cba26327..fd946e54 100644 --- a/crates/worker/src/spawn/registry.rs +++ b/crates/worker/src/spawn/registry.rs @@ -12,7 +12,7 @@ use std::collections::{BTreeMap, HashSet}; use std::io; use std::sync::{ Arc, Mutex, - atomic::{AtomicBool, AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, }; use std::time::Instant; @@ -23,7 +23,7 @@ use protocol::{Event, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnaps use session_store::{ LoggedItem, WorkerMetadataStore, WorkerReclaimedChild, WorkerSpawnedChild, WorkerStoreError, }; -use tokio::sync::broadcast; +use tokio::sync::{Notify, broadcast}; use tracing::warn; use workdir::WorkdirScopeLease; @@ -238,39 +238,56 @@ pub(crate) struct InternalSpawnReservation { } impl InternalSpawnReservation { - 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", - ), - record, - )); - } - 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, - )); + pub(crate) async fn commit(mut self, record: InternalSpawnedWorkerRecord) -> io::Result<()> { + let rejection = if record.worker_name != self.worker_name { + Some(io::Error::new( + io::ErrorKind::InvalidInput, + "internal SubWorker reservation name does not match record name", + )) + } else { + match self.registry.internal_records.lock() { + Ok(mut records) => { + if self.registry.internal_shutting_down.load(Ordering::Acquire) { + Some(io::Error::new( + io::ErrorKind::Interrupted, + "internal SubWorker registry is shutting down", + )) + } else { + records.push(record.clone()); + None + } + } + Err(_) => Some(io::Error::other( + "internal spawned-worker registry lock poisoned", + )), } }; - if self.registry.internal_shutting_down.load(Ordering::Acquire) { - return Err(( - io::Error::new( - io::ErrorKind::Interrupted, - "internal SubWorker registry is shutting down", - ), - record, - )); + if let Some(error) = rejection { + let mut cleanup_failures = Vec::new(); + if let Err(cleanup) = record.session.stop().await { + cleanup_failures.push(format!("stop rejected Internal SubWorker: {cleanup}")); + } + if let Err(cleanup) = Box::pin(record.child_registry.shutdown_internal()).await { + cleanup_failures.push(format!( + "stop rejected Internal SubWorker descendants: {cleanup}" + )); + } + if let Err(cleanup) = record.workdir_tool_scope.close().await { + cleanup_failures.push(format!( + "close rejected Internal SubWorker Workdir tools: {cleanup}" + )); + } + if cleanup_failures.is_empty() { + return Err(error); + } + self.registry + .internal_spawn_cleanup_failed + .store(true, Ordering::Release); + return Err(io::Error::other(format!( + "{error}; {}", + cleanup_failures.join("; ") + ))); } - records.push(record.clone()); - drop(records); self.registry.start_protocol_forwarding(record); self.committed = true; Ok(()) @@ -284,6 +301,10 @@ impl Drop for InternalSpawnReservation { names.remove(&self.worker_name); } } + self.registry + .pending_internal_spawns + .fetch_sub(1, Ordering::AcqRel); + self.registry.pending_internal_notify.notify_waiters(); } } @@ -292,6 +313,9 @@ pub struct SpawnedWorkerRegistry { service_records: std::sync::Mutex>, internal_names: std::sync::Mutex>, internal_shutting_down: AtomicBool, + pending_internal_spawns: AtomicUsize, + pending_internal_notify: Notify, + internal_spawn_cleanup_failed: AtomicBool, parent_scope: Option, parent_protocol: Mutex, String)>>, } @@ -309,6 +333,9 @@ impl SpawnedWorkerRegistry { service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -321,6 +348,9 @@ impl SpawnedWorkerRegistry { service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope: None, parent_protocol: Mutex::new(None), }) @@ -332,6 +362,9 @@ impl SpawnedWorkerRegistry { service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope: Some(parent_scope), parent_protocol: Mutex::new(None), }) @@ -412,6 +445,9 @@ impl SpawnedWorkerRegistry { service_records: std::sync::Mutex::new(Vec::new()), internal_names: std::sync::Mutex::new(HashSet::new()), internal_shutting_down: AtomicBool::new(false), + pending_internal_spawns: AtomicUsize::new(0), + pending_internal_notify: Notify::new(), + internal_spawn_cleanup_failed: AtomicBool::new(false), parent_scope, parent_protocol: Mutex::new(None), }), @@ -423,6 +459,10 @@ impl SpawnedWorkerRegistry { self: &Arc, worker_name: String, ) -> io::Result { + let records = self + .internal_records + .lock() + .map_err(|_| io::Error::other("internal Worker registry lock poisoned"))?; if self.internal_shutting_down.load(Ordering::Acquire) { return Err(io::Error::new( io::ErrorKind::Interrupted, @@ -439,7 +479,9 @@ impl SpawnedWorkerRegistry { format!("spawned worker `{worker_name}` is already registered"), )); } + self.pending_internal_spawns.fetch_add(1, Ordering::AcqRel); drop(names); + drop(records); Ok(InternalSpawnReservation { registry: Arc::clone(self), worker_name, @@ -758,17 +800,31 @@ impl SpawnedWorkerRegistry { .map(|record| record.worker_name.clone()) .collect::>() }; + loop { + let notified = self.pending_internal_notify.notified(); + if self.pending_internal_spawns.load(Ordering::Acquire) == 0 { + break; + } + notified.await; + } let mut first_error = None; for name in names { if let Err(error) = self.remove_internal(&name).await { first_error.get_or_insert(error); } } + if first_error.is_none() && self.internal_spawn_cleanup_failed.load(Ordering::Acquire) { + first_error = Some(io::Error::other( + "an in-flight Internal SubWorker failed cleanup during shutdown", + )); + } first_error.map_or(Ok(()), Err) } pub(crate) fn reopen_internal(&self) { self.internal_shutting_down.store(false, Ordering::Release); + self.internal_spawn_cleanup_failed + .store(false, Ordering::Release); } /// Stop one direct Internal SubWorker and discard its registry/scope state. @@ -1342,24 +1398,22 @@ mod tests { 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 barrier = Arc::new(tokio::sync::Barrier::new(2)); let commit_barrier = barrier.clone(); - let commit = tokio::task::spawn_blocking(move || { - commit_barrier.wait(); - reservation.commit(record) + let commit = tokio::spawn(async move { + commit_barrier.wait().await; + reservation.commit(record).await }); let shutdown_registry = registry.clone(); let shutdown = tokio::spawn(async move { - barrier.wait(); + barrier.wait().await; 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(); + if let Err(error) = commit { + assert_eq!(error.kind(), io::ErrorKind::Interrupted); } assert!(registry.list_internal().is_empty()); @@ -1375,12 +1429,54 @@ mod tests { let (record, _events) = record("racing-child", InternalWorkerVisibility::ParentClient).await; - registry.shutdown_internal().await.unwrap(); - let (error, record) = reservation.commit(record).unwrap_err(); + let mut shutdown = { + let registry = registry.clone(); + tokio::spawn(async move { registry.shutdown_internal().await }) + }; + while !registry.internal_shutting_down.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), &mut shutdown) + .await + .is_err(), + "shutdown must wait for the pending spawn to roll back" + ); + let error = reservation.commit(record).await.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(); + shutdown.await.unwrap().unwrap(); + } + + #[tokio::test] + async fn rejected_spawn_cleanup_failure_keeps_shutdown_failed_closed() { + let registry = registry(); + let reservation = registry + .reserve_internal_name("cleanup-failure".into()) + .unwrap(); + let (record, _events) = + record("cleanup-failure", InternalWorkerVisibility::ParentClient).await; + record.session.force_stop_failure(); + let shutdown = { + let registry = registry.clone(); + tokio::spawn(async move { registry.shutdown_internal().await }) + }; + while !registry.internal_shutting_down.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + + let error = reservation.commit(record).await.unwrap_err(); + assert!( + error + .to_string() + .contains("stop rejected Internal SubWorker") + ); + let shutdown_error = shutdown.await.unwrap().unwrap_err(); + assert!( + shutdown_error + .to_string() + .contains("failed cleanup during shutdown") + ); + assert!(registry.internal_shutting_down.load(Ordering::Acquire)); } #[tokio::test] diff --git a/crates/worker/src/spawn/tool.rs b/crates/worker/src/spawn/tool.rs index 547d939d..6370653f 100644 --- a/crates/worker/src/spawn/tool.rs +++ b/crates/worker/src/spawn/tool.rs @@ -630,10 +630,7 @@ impl Tool for SubWorkerSpawnTool { child_registry, child_change_tracker, ); - 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; + if let Err(error) = name_reservation.commit(record).await { return Err(ToolError::ExecutionFailed(format!( "register Internal Worker session: {error}" )));