diff --git a/codex-rs/app-server/src/request_processors/thread_teardown.rs b/codex-rs/app-server/src/request_processors/thread_teardown.rs index e7fa235e22ee..69460daa684d 100644 --- a/codex-rs/app-server/src/request_processors/thread_teardown.rs +++ b/codex-rs/app-server/src/request_processors/thread_teardown.rs @@ -1,6 +1,9 @@ use super::*; use std::pin::Pin; use std::sync::Mutex as StdMutex; +use std::sync::Weak; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; type ThreadTeardownFuture = Pin + Send + 'static>>; @@ -9,24 +12,69 @@ pub(crate) struct PendingThreadUnloads { inner: Arc, } -#[derive(Default)] struct PendingThreadUnloadsInner { registry: StdMutex, tasks: TaskTracker, + freeze: Arc, +} + +impl Default for PendingThreadUnloadsInner { + fn default() -> Self { + Self { + registry: StdMutex::new(PendingThreadUnloadRegistry::default()), + tasks: TaskTracker::new(), + freeze: Arc::new(Semaphore::new(1)), + } + } } #[derive(Default)] struct PendingThreadUnloadRegistry { closing: bool, - entries: HashMap>, + entries: HashMap, +} + +struct PendingThreadUnloadEntry { + owner: Arc, + successor: Option>, +} + +struct PendingThreadUnloadOperation { + completed: watch::Sender, + thread_ids: StdMutex>, + finished: AtomicBool, } pub(super) enum PendingThreadUnloadClaimResult { Claimed(PendingThreadUnloadClaim), - Pending(watch::Receiver), + Pending(PendingThreadUnloadConflicts), Closing, } +#[allow( + dead_code, + reason = "consumed by the stacked complete-tree teardown layer" +)] +pub(super) enum PendingThreadUnloadExtendResult { + Extended, + /// Free IDs are owned and contested IDs have exact successor handoffs registered. The caller + /// retains both its owner and freeze guard while awaiting these predecessors, then retries. + Pending(PendingThreadUnloadConflicts), + Finished, +} + +pub(super) struct PendingThreadUnloadConflicts { + completions: Vec>, +} + +/// Coordinator-wide affine token for one tree-freeze extension phase. Acquire it before the +/// global thread-list permit and retain it across any temporary permit drops until the tree is +/// stable; no other operation may register successor handoffs while it is held. +pub(super) struct PendingThreadUnloadFreezeGuard { + _permit: tokio::sync::OwnedSemaphorePermit, + operation: Option>, +} + pub(super) enum PendingThreadUnloadStartResult { Started(watch::Receiver), Pending(watch::Receiver), @@ -36,29 +84,76 @@ pub(super) enum PendingThreadUnloadStartResult { pub(super) struct PendingThreadUnloadClaim { start_tx: Option>, completed: watch::Receiver, + pending: PendingThreadUnloads, + operation: Arc, } impl PendingThreadUnloadClaim { - pub(super) fn start(mut self, teardown: F) -> watch::Receiver + pub(super) fn start(self, teardown: F) -> watch::Receiver where F: std::future::Future + Send + 'static, + { + self.start_with(|_| teardown) + } + + pub(super) fn start_with(mut self, teardown: F) -> watch::Receiver + where + F: FnOnce(PendingThreadUnloadClaimHandle) -> Fut, + Fut: std::future::Future + Send + 'static, { if let Some(start_tx) = self.start_tx.take() { - let _ = start_tx.send(Box::pin(teardown)); + let handle = PendingThreadUnloadClaimHandle { + pending: self.pending.clone(), + operation: Arc::clone(&self.operation), + }; + let _ = start_tx.send(Box::pin(teardown(handle))); } self.completed.clone() } } +pub(super) struct PendingThreadUnloadClaimHandle { + pending: PendingThreadUnloads, + operation: Arc, +} + +impl PendingThreadUnloadClaimHandle { + #[allow( + dead_code, + reason = "consumed by the stacked complete-tree teardown layer" + )] + pub(super) fn try_extend( + &self, + guard: &mut PendingThreadUnloadFreezeGuard, + thread_ids: I, + ) -> PendingThreadUnloadExtendResult + where + I: IntoIterator, + { + self.pending.try_extend(self, guard, thread_ids) + } +} + +impl PendingThreadUnloadFreezeGuard { + fn bind(&mut self, operation: &Arc) -> bool { + match self.operation.as_ref().and_then(Weak::upgrade) { + Some(bound) => Arc::ptr_eq(&bound, operation), + None => { + self.operation = Some(Arc::downgrade(operation)); + true + } + } + } +} + struct PendingThreadUnloadOwner { pending: PendingThreadUnloads, - thread_id: ThreadId, - completed: watch::Sender, + operation: Arc, } impl Drop for PendingThreadUnloadOwner { fn drop(&mut self) { - self.pending.release(self.thread_id, &self.completed); + self.pending.release(&self.operation); } } @@ -74,28 +169,63 @@ impl PendingThreadUnloads { self.lock_registry().entries.contains_key(&thread_id) } + #[allow( + dead_code, + reason = "consumed by the stacked complete-tree teardown layer" + )] + pub(super) async fn acquire_freeze_guard(&self) -> Option { + let permit = self.inner.freeze.clone().acquire_owned().await.ok()?; + Some(PendingThreadUnloadFreezeGuard { + _permit: permit, + operation: None, + }) + } + pub(super) fn subscribe(&self, thread_id: ThreadId) -> Option> { self.lock_registry() .entries .get(&thread_id) - .map(watch::Sender::subscribe) + .map(|entry| entry.owner.completed.subscribe()) } pub(super) fn try_claim(&self, thread_id: ThreadId) -> PendingThreadUnloadClaimResult { + self.try_claim_many([thread_id]) + } + + pub(super) fn try_claim_many(&self, thread_ids: I) -> PendingThreadUnloadClaimResult + where + I: IntoIterator, + { + let thread_ids = dedupe_thread_ids(thread_ids); let mut registry = self.lock_registry(); if registry.closing { return PendingThreadUnloadClaimResult::Closing; } - if let Some(completed) = registry.entries.get(&thread_id) { - return PendingThreadUnloadClaimResult::Pending(completed.subscribe()); + let conflicts = conflicting_completions(®istry, &thread_ids, /*owner*/ None); + if !conflicts.is_empty() { + return PendingThreadUnloadClaimResult::Pending(PendingThreadUnloadConflicts { + completions: conflicts, + }); } let (completed, completed_rx) = watch::channel(false); - registry.entries.insert(thread_id, completed.clone()); + let operation = Arc::new(PendingThreadUnloadOperation { + completed, + thread_ids: StdMutex::new(thread_ids.iter().copied().collect()), + finished: AtomicBool::new(false), + }); + for thread_id in &thread_ids { + registry.entries.insert( + *thread_id, + PendingThreadUnloadEntry { + owner: Arc::clone(&operation), + successor: None, + }, + ); + } let owner = PendingThreadUnloadOwner { pending: self.clone(), - thread_id, - completed, + operation: Arc::clone(&operation), }; let (start_tx, start_rx) = oneshot::channel::(); let task = self.inner.tasks.spawn(async move { @@ -108,9 +238,74 @@ impl PendingThreadUnloads { PendingThreadUnloadClaimResult::Claimed(PendingThreadUnloadClaim { start_tx: Some(start_tx), completed: completed_rx, + pending: self.clone(), + operation, }) } + #[allow( + dead_code, + reason = "consumed by the stacked complete-tree teardown layer" + )] + fn try_extend( + &self, + owner: &PendingThreadUnloadClaimHandle, + guard: &mut PendingThreadUnloadFreezeGuard, + thread_ids: I, + ) -> PendingThreadUnloadExtendResult + where + I: IntoIterator, + { + let thread_ids = dedupe_thread_ids(thread_ids); + let mut registry = self.lock_registry(); + if owner.operation.finished.load(Ordering::Acquire) || !guard.bind(&owner.operation) { + return PendingThreadUnloadExtendResult::Finished; + } + let mut owned_thread_ids = owner + .operation + .thread_ids + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut conflicts = Vec::new(); + for thread_id in thread_ids { + match registry.entries.entry(thread_id) { + std::collections::hash_map::Entry::Vacant(entry) => { + entry.insert(PendingThreadUnloadEntry { + owner: Arc::clone(&owner.operation), + successor: None, + }); + owned_thread_ids.insert(thread_id); + } + std::collections::hash_map::Entry::Occupied(mut entry) => { + let entry = entry.get_mut(); + if Arc::ptr_eq(&entry.owner, &owner.operation) { + owned_thread_ids.insert(thread_id); + continue; + } + let successor = entry.successor.as_ref().and_then(Weak::upgrade); + if successor + .as_ref() + .is_none_or(|successor| successor.finished.load(Ordering::Acquire)) + { + entry.successor = Some(Arc::downgrade(&owner.operation)); + } else if let Some(successor) = successor.as_ref() + && !Arc::ptr_eq(successor, &owner.operation) + { + push_conflict(&mut conflicts, successor); + } + push_conflict(&mut conflicts, &entry.owner); + } + } + } + drop(owned_thread_ids); + if !conflicts.is_empty() { + return PendingThreadUnloadExtendResult::Pending(PendingThreadUnloadConflicts { + completions: conflict_receivers(conflicts), + }); + } + PendingThreadUnloadExtendResult::Extended + } + pub(super) fn try_start( &self, thread_id: ThreadId, @@ -123,30 +318,59 @@ impl PendingThreadUnloads { PendingThreadUnloadClaimResult::Claimed(claim) => { PendingThreadUnloadStartResult::Started(claim.start(teardown)) } - PendingThreadUnloadClaimResult::Pending(completed) => { - PendingThreadUnloadStartResult::Pending(completed) - } + PendingThreadUnloadClaimResult::Pending(conflicts) => match conflicts.into_single() { + Some(completed) => PendingThreadUnloadStartResult::Pending(completed), + None => PendingThreadUnloadStartResult::Closing, + }, PendingThreadUnloadClaimResult::Closing => PendingThreadUnloadStartResult::Closing, } } - fn release(&self, thread_id: ThreadId, completed: &watch::Sender) { + fn release(&self, operation: &Arc) { let mut registry = self.lock_registry(); - if registry - .entries - .get(&thread_id) - .is_some_and(|current| current.same_channel(completed)) - { - registry.entries.remove(&thread_id); + operation.finished.store(true, Ordering::Release); + let thread_ids = operation + .thread_ids + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for thread_id in thread_ids.iter() { + let std::collections::hash_map::Entry::Occupied(mut entry) = + registry.entries.entry(*thread_id) + else { + continue; + }; + if !Arc::ptr_eq(&entry.get().owner, operation) { + continue; + } + let successor = entry + .get() + .successor + .as_ref() + .and_then(Weak::upgrade) + .filter(|successor| !successor.finished.load(Ordering::Acquire)); + if let Some(successor) = successor { + let mut successor_thread_ids = successor + .thread_ids + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let entry = entry.get_mut(); + entry.owner = Arc::clone(&successor); + entry.successor = None; + successor_thread_ids.insert(*thread_id); + } else { + entry.remove(); + } } + drop(thread_ids); drop(registry); - let _ = completed.send(true); + let _ = operation.completed.send(true); } pub(super) async fn close_and_wait(&self) { { let mut registry = self.lock_registry(); registry.closing = true; + self.inner.freeze.close(); self.inner.tasks.close(); } self.inner.tasks.wait().await; @@ -158,6 +382,63 @@ impl PendingThreadUnloads { } } +impl PendingThreadUnloadConflicts { + fn into_single(mut self) -> Option> { + debug_assert_eq!(self.completions.len(), 1); + self.completions.pop() + } +} + +fn dedupe_thread_ids(thread_ids: I) -> Vec +where + I: IntoIterator, +{ + let mut seen = HashSet::new(); + thread_ids + .into_iter() + .filter(|thread_id| seen.insert(*thread_id)) + .collect() +} + +fn conflicting_completions( + registry: &PendingThreadUnloadRegistry, + thread_ids: &[ThreadId], + owner: Option<&Arc>, +) -> Vec> { + let mut conflicts = Vec::>::new(); + for thread_id in thread_ids { + let Some(entry) = registry.entries.get(thread_id) else { + continue; + }; + if owner.is_some_and(|owner| Arc::ptr_eq(&entry.owner, owner)) { + continue; + } + push_conflict(&mut conflicts, &entry.owner); + } + conflict_receivers(conflicts) +} + +fn push_conflict( + conflicts: &mut Vec>, + operation: &Arc, +) { + if !conflicts + .iter() + .any(|conflict| Arc::ptr_eq(conflict, operation)) + { + conflicts.push(Arc::clone(operation)); + } +} + +fn conflict_receivers( + conflicts: Vec>, +) -> Vec> { + conflicts + .into_iter() + .map(|operation| operation.completed.subscribe()) + .collect() +} + pub(super) async fn wait_for_thread_unload(mut completed: watch::Receiver) { while !*completed.borrow_and_update() { if completed.changed().await.is_err() { @@ -166,6 +447,16 @@ pub(super) async fn wait_for_thread_unload(mut completed: watch::Receiver) } } +#[allow( + dead_code, + reason = "consumed by the stacked complete-tree teardown layer" +)] +pub(super) async fn wait_for_thread_unloads(conflicts: PendingThreadUnloadConflicts) { + for completed in conflicts.completions { + wait_for_thread_unload(completed).await; + } +} + pub(super) fn start_thread_teardown( pending: PendingThreadUnloads, thread_id: ThreadId, diff --git a/codex-rs/app-server/src/request_processors/thread_teardown_tests.rs b/codex-rs/app-server/src/request_processors/thread_teardown_tests.rs index 7e2f466121de..40f974d91ae5 100644 --- a/codex-rs/app-server/src/request_processors/thread_teardown_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_teardown_tests.rs @@ -57,3 +57,143 @@ async fn teardown_is_singleflight_and_outlives_its_first_waiter() { assert_eq!(calls.load(Ordering::SeqCst), 1); assert!(pending.is_empty()); } + +#[tokio::test] +async fn freeze_handoffs_are_gapless_serialized_and_cancellation_safe() { + let pending = PendingThreadUnloads::default(); + let mut freeze = pending + .acquire_freeze_guard() + .await + .expect("freeze guard should be available"); + let root = ThreadId::new(); + let free_sibling = ThreadId::new(); + let late_a = ThreadId::new(); + let late_a_alias = ThreadId::new(); + let late_b = ThreadId::new(); + + let PendingThreadUnloadClaimResult::Claimed(predecessor_a) = + pending.try_claim_many([late_a, late_a_alias]) + else { + panic!("first predecessor should claim late descendants"); + }; + let (release_predecessor_a_tx, release_predecessor_a_rx) = oneshot::channel(); + let predecessor_a_completion = predecessor_a.start(async move { + let _ = release_predecessor_a_rx.await; + }); + let PendingThreadUnloadClaimResult::Claimed(predecessor_b) = pending.try_claim(late_b) else { + panic!("second predecessor should claim its late descendant"); + }; + let (release_predecessor_b_tx, release_predecessor_b_rx) = oneshot::channel(); + let predecessor_b_completion = predecessor_b.start(async move { + let _ = release_predecessor_b_rx.await; + }); + let PendingThreadUnloadClaimResult::Claimed(successor) = pending.try_claim(root) else { + panic!("successor should claim its root"); + }; + let (conflicts_tx, conflicts_rx) = oneshot::channel(); + let (retry_tx, retry_rx) = oneshot::channel(); + let (stable_tx, stable_rx) = oneshot::channel(); + let (release_successor_tx, release_successor_rx) = oneshot::channel(); + let successor_completion = successor.start_with(move |owner| async move { + let PendingThreadUnloadExtendResult::Pending(conflicts) = owner.try_extend( + &mut freeze, + [root, free_sibling, late_a, late_a_alias, late_b], + ) else { + panic!("late descendants should conflict with their predecessor"); + }; + assert!(conflicts_tx.send(conflicts).is_ok()); + let _ = retry_rx.await; + assert!(matches!( + owner.try_extend( + &mut freeze, + [root, free_sibling, late_a, late_a_alias, late_b] + ), + PendingThreadUnloadExtendResult::Extended + )); + let _ = stable_tx.send(()); + let _ = release_successor_rx.await; + }); + let conflicts = conflicts_rx.await.expect("extension conflicts"); + assert_eq!(conflicts.completions.len(), 2); + assert!(pending.contains(free_sibling)); + + let pending_for_freeze = pending.clone(); + let next_freeze = tokio::spawn(async move { pending_for_freeze.acquire_freeze_guard().await }); + tokio::task::yield_now().await; + assert!(!next_freeze.is_finished()); + + release_predecessor_a_tx + .send(()) + .expect("release first predecessor"); + release_predecessor_b_tx + .send(()) + .expect("release second predecessor"); + wait_for_thread_unload(predecessor_a_completion).await; + wait_for_thread_unload(predecessor_b_completion).await; + wait_for_thread_unloads(conflicts).await; + let PendingThreadUnloadClaimResult::Pending(successor_conflict) = + pending.try_claim_many([root, free_sibling, late_a, late_a_alias, late_b]) + else { + panic!("every ID should transfer to the one successor owner"); + }; + assert_eq!(successor_conflict.completions.len(), 1); + retry_tx.send(()).expect("retry extension after handoff"); + stable_rx.await.expect("successor should become stable"); + release_successor_tx.send(()).expect("release successor"); + wait_for_thread_unload(successor_completion).await; + wait_for_thread_unloads(successor_conflict).await; + let mut freeze = next_freeze + .await + .expect("freeze waiter should not fail") + .expect("next freeze guard should become available"); + assert!(pending.is_empty()); + + let cancelled_root = ThreadId::new(); + let cancelled_free = ThreadId::new(); + let contested = ThreadId::new(); + let PendingThreadUnloadClaimResult::Claimed(predecessor) = pending.try_claim(contested) else { + panic!("cancellation predecessor should be claimable"); + }; + let (release_predecessor_tx, release_predecessor_rx) = oneshot::channel(); + let predecessor_completion = predecessor.start(async move { + let _ = release_predecessor_rx.await; + }); + let PendingThreadUnloadClaimResult::Claimed(cancelled) = pending.try_claim(cancelled_root) + else { + panic!("cancelled successor should claim its root"); + }; + let (cancelled_conflicts_tx, cancelled_conflicts_rx) = oneshot::channel(); + let (stale_owner_tx, stale_owner_rx) = oneshot::channel(); + let cancelled_completion = cancelled.start_with(move |owner| async move { + let PendingThreadUnloadExtendResult::Pending(conflicts) = + owner.try_extend(&mut freeze, [cancelled_root, cancelled_free, contested]) + else { + panic!("cancelled successor should register a handoff"); + }; + assert!(cancelled_conflicts_tx.send(conflicts).is_ok()); + assert!(stale_owner_tx.send(owner).is_ok()); + }); + let cancelled_conflicts = cancelled_conflicts_rx + .await + .expect("cancelled extension conflicts"); + let stale_owner = stale_owner_rx.await.expect("cancelled owner handle"); + wait_for_thread_unload(cancelled_completion).await; + assert!(!pending.contains(cancelled_root)); + assert!(!pending.contains(cancelled_free)); + assert!(pending.contains(contested)); + let mut stale_freeze = pending + .acquire_freeze_guard() + .await + .expect("cancelled freeze guard should be released"); + assert!(matches!( + stale_owner.try_extend(&mut stale_freeze, [cancelled_free]), + PendingThreadUnloadExtendResult::Finished + )); + drop(stale_freeze); + release_predecessor_tx + .send(()) + .expect("release cancellation predecessor"); + wait_for_thread_unload(predecessor_completion).await; + wait_for_thread_unloads(cancelled_conflicts).await; + assert!(pending.is_empty()); +}