From 1f6c042c0fa33d341205e479a43dde24fdb0ce4a Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Tue, 7 Jul 2026 04:21:11 -0700 Subject: [PATCH 1/2] app-server: claim thread teardown groups atomically --- .../src/request_processors/thread_teardown.rs | 223 ++++++++++++++++-- .../thread_teardown_tests.rs | 114 +++++++++ 2 files changed, 313 insertions(+), 24 deletions(-) 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..84866eafb60a 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,8 @@ use super::*; use std::pin::Pin; use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; type ThreadTeardownFuture = Pin + Send + 'static>>; @@ -18,15 +20,37 @@ struct PendingThreadUnloadsInner { #[derive(Default)] struct PendingThreadUnloadRegistry { closing: bool, - entries: HashMap>, + entries: HashMap>, +} + +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, + /// No requested ID was added. The tracked owner must end and release its existing group + /// before any caller awaits these conflicts, then retry one claim for the full known set. + Pending(PendingThreadUnloadConflicts), + Finished, +} + +pub(super) struct PendingThreadUnloadConflicts { + completions: Vec>, +} + pub(super) enum PendingThreadUnloadStartResult { Started(watch::Receiver), Pending(watch::Receiver), @@ -36,29 +60,60 @@ 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, thread_ids: I) -> PendingThreadUnloadExtendResult + where + I: IntoIterator, + { + self.pending.try_extend(self, thread_ids) + } +} + 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); } } @@ -78,24 +133,41 @@ impl PendingThreadUnloads { self.lock_registry() .entries .get(&thread_id) - .map(watch::Sender::subscribe) + .map(|operation| operation.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, Arc::clone(&operation)); + } 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 +180,50 @@ 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, + 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) { + return PendingThreadUnloadExtendResult::Finished; + } + let conflicts = conflicting_completions(®istry, &thread_ids, Some(&owner.operation)); + if !conflicts.is_empty() { + return PendingThreadUnloadExtendResult::Pending(PendingThreadUnloadConflicts { + completions: conflicts, + }); + } + + let mut owned_thread_ids = owner + .operation + .thread_ids + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for thread_id in thread_ids { + registry + .entries + .entry(thread_id) + .or_insert_with(|| Arc::clone(&owner.operation)); + owned_thread_ids.insert(thread_id); + } + PendingThreadUnloadExtendResult::Extended + } + pub(super) fn try_start( &self, thread_id: ThreadId, @@ -123,24 +236,33 @@ 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); + let thread_ids = operation + .thread_ids + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for thread_id in thread_ids.iter() { + if registry + .entries + .get(thread_id) + .is_some_and(|current| Arc::ptr_eq(current, operation)) + { + registry.entries.remove(thread_id); + } } + operation.finished.store(true, Ordering::Release); + drop(thread_ids); drop(registry); - let _ = completed.send(true); + let _ = operation.completed.send(true); } pub(super) async fn close_and_wait(&self) { @@ -158,6 +280,49 @@ 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(operation) = registry.entries.get(thread_id) else { + continue; + }; + if owner.is_some_and(|owner| Arc::ptr_eq(operation, owner)) + || conflicts + .iter() + .any(|conflict| Arc::ptr_eq(conflict, operation)) + { + continue; + } + conflicts.push(Arc::clone(operation)); + } + 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 +331,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..dd9d409b0a08 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,117 @@ 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 overlapping_multi_claim_extensions_release_before_retrying_without_partial_claims() { + let pending = PendingThreadUnloads::default(); + let root_a = ThreadId::new(); + let root_b = ThreadId::new(); + let late_a = ThreadId::new(); + let late_b = ThreadId::new(); + + let PendingThreadUnloadClaimResult::Claimed(claim_a) = pending.try_claim_many([root_a]) else { + panic!("first root should be claimable"); + }; + let PendingThreadUnloadClaimResult::Claimed(claim_b) = pending.try_claim_many([root_b]) else { + panic!("second root should be claimable"); + }; + let (conflicts_a_tx, conflicts_a_rx) = oneshot::channel(); + let (release_a_tx, release_a_rx) = oneshot::channel(); + let completion_a = claim_a.start_with(move |owner| async move { + let PendingThreadUnloadExtendResult::Pending(conflicts) = + owner.try_extend([root_b, late_a]) + else { + panic!("first extension should conflict with the second owner"); + }; + assert!(conflicts_a_tx.send(conflicts).is_ok()); + let _ = release_a_rx.await; + }); + let (conflicts_b_tx, conflicts_b_rx) = oneshot::channel(); + let (release_b_tx, release_b_rx) = oneshot::channel(); + let completion_b = claim_b.start_with(move |owner| async move { + let PendingThreadUnloadExtendResult::Pending(conflicts) = + owner.try_extend([root_a, late_b]) + else { + panic!("second extension should conflict with the first owner"); + }; + assert!(conflicts_b_tx.send(conflicts).is_ok()); + let _ = release_b_rx.await; + }); + let conflicts_a = conflicts_a_rx.await.expect("first extension result"); + let conflicts_b = conflicts_b_rx.await.expect("second extension result"); + assert_eq!(conflicts_a.completions.len(), 1); + assert_eq!(conflicts_b.completions.len(), 1); + assert!(!pending.contains(late_a)); + assert!(!pending.contains(late_b)); + let unclaimed = ThreadId::new(); + let PendingThreadUnloadClaimResult::Pending(all_conflicts) = + pending.try_claim_many([root_a, root_b, unclaimed]) + else { + panic!("an overlapping multi-claim should report every conflicting owner"); + }; + assert_eq!(all_conflicts.completions.len(), 2); + assert!(!pending.contains(unclaimed)); + + // Neither owner may wait while retaining its partial group: each one is the other's + // conflict. Ending both tracked operations lets every conflict subscription complete. + release_a_tx.send(()).expect("release first owner"); + release_b_tx.send(()).expect("release second owner"); + wait_for_thread_unload(completion_a).await; + wait_for_thread_unload(completion_b).await; + wait_for_thread_unloads(conflicts_a).await; + wait_for_thread_unloads(conflicts_b).await; + wait_for_thread_unloads(all_conflicts).await; + + let PendingThreadUnloadClaimResult::Claimed(full_claim) = + pending.try_claim_many([root_a, root_b, late_a, late_b, late_a]) + else { + panic!("the complete deduplicated set should be claimable after both owners release"); + }; + let same_owner_unclaimed = ThreadId::new(); + let PendingThreadUnloadClaimResult::Pending(same_owner_conflict) = + pending.try_claim_many([root_a, late_a, same_owner_unclaimed]) + else { + panic!("overlap with one owner should conflict"); + }; + assert_eq!(same_owner_conflict.completions.len(), 1); + assert!(!pending.contains(same_owner_unclaimed)); + let full_completion = full_claim.completed.clone(); + drop(full_claim); + wait_for_thread_unload(full_completion).await; + wait_for_thread_unloads(same_owner_conflict).await; + assert!(pending.is_empty()); + + let stale_root = ThreadId::new(); + let closing_late = ThreadId::new(); + let stale_late = ThreadId::new(); + let PendingThreadUnloadClaimResult::Claimed(stale_claim) = pending.try_claim_many([stale_root]) + else { + panic!("stale-handle root should be claimable"); + }; + let (extend_tx, extend_rx) = oneshot::channel(); + let (stale_tx, stale_rx) = oneshot::channel(); + let stale_completion = stale_claim.start_with(move |owner| async move { + let _ = extend_rx.await; + assert!(matches!( + owner.try_extend([closing_late]), + PendingThreadUnloadExtendResult::Extended + )); + assert!(stale_tx.send(owner).is_ok()); + }); + let pending_for_drain = pending.clone(); + let drain = tokio::spawn(async move { pending_for_drain.close_and_wait().await }); + while !pending.lock_registry().closing { + tokio::task::yield_now().await; + } + extend_tx.send(()).expect("extend while coordinator closes"); + let stale_owner = stale_rx.await.expect("stale owner handle"); + wait_for_thread_unload(stale_completion).await; + drain.await.expect("coordinator drain should complete"); + assert!(matches!( + stale_owner.try_extend([stale_late]), + PendingThreadUnloadExtendResult::Finished + )); + assert!(!pending.contains(stale_late)); + assert!(pending.is_empty()); +} From b2c07ba8d823b0ea112e27b78a3a3f5168683662 Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Tue, 7 Jul 2026 04:40:30 -0700 Subject: [PATCH 2/2] app-server: hand off tree teardown claims without gaps --- .../src/request_processors/thread_teardown.rs | 184 +++++++++++++--- .../thread_teardown_tests.rs | 204 ++++++++++-------- 2 files changed, 265 insertions(+), 123 deletions(-) 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 84866eafb60a..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,7 @@ 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; @@ -11,16 +12,31 @@ 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 { @@ -41,8 +57,8 @@ pub(super) enum PendingThreadUnloadClaimResult { )] pub(super) enum PendingThreadUnloadExtendResult { Extended, - /// No requested ID was added. The tracked owner must end and release its existing group - /// before any caller awaits these conflicts, then retry one claim for the full known set. + /// 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, } @@ -51,6 +67,14 @@ 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), @@ -98,11 +122,27 @@ impl PendingThreadUnloadClaimHandle { dead_code, reason = "consumed by the stacked complete-tree teardown layer" )] - pub(super) fn try_extend(&self, thread_ids: I) -> PendingThreadUnloadExtendResult + pub(super) fn try_extend( + &self, + guard: &mut PendingThreadUnloadFreezeGuard, + thread_ids: I, + ) -> PendingThreadUnloadExtendResult where I: IntoIterator, { - self.pending.try_extend(self, thread_ids) + 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 + } + } } } @@ -129,11 +169,23 @@ 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(|operation| operation.completed.subscribe()) + .map(|entry| entry.owner.completed.subscribe()) } pub(super) fn try_claim(&self, thread_id: ThreadId) -> PendingThreadUnloadClaimResult { @@ -163,7 +215,13 @@ impl PendingThreadUnloads { finished: AtomicBool::new(false), }); for thread_id in &thread_ids { - registry.entries.insert(*thread_id, Arc::clone(&operation)); + registry.entries.insert( + *thread_id, + PendingThreadUnloadEntry { + owner: Arc::clone(&operation), + successor: None, + }, + ); } let owner = PendingThreadUnloadOwner { pending: self.clone(), @@ -192,6 +250,7 @@ impl PendingThreadUnloads { fn try_extend( &self, owner: &PendingThreadUnloadClaimHandle, + guard: &mut PendingThreadUnloadFreezeGuard, thread_ids: I, ) -> PendingThreadUnloadExtendResult where @@ -199,27 +258,50 @@ impl PendingThreadUnloads { { let thread_ids = dedupe_thread_ids(thread_ids); let mut registry = self.lock_registry(); - if owner.operation.finished.load(Ordering::Acquire) { + if owner.operation.finished.load(Ordering::Acquire) || !guard.bind(&owner.operation) { return PendingThreadUnloadExtendResult::Finished; } - let conflicts = conflicting_completions(®istry, &thread_ids, Some(&owner.operation)); - if !conflicts.is_empty() { - return PendingThreadUnloadExtendResult::Pending(PendingThreadUnloadConflicts { - completions: conflicts, - }); - } - 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 { - registry - .entries - .entry(thread_id) - .or_insert_with(|| Arc::clone(&owner.operation)); - owned_thread_ids.insert(thread_id); + 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 } @@ -246,20 +328,39 @@ impl PendingThreadUnloads { fn release(&self, operation: &Arc) { let mut registry = self.lock_registry(); + 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() { - if registry - .entries - .get(thread_id) - .is_some_and(|current| Arc::ptr_eq(current, operation)) - { - registry.entries.remove(thread_id); + 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(); } } - operation.finished.store(true, Ordering::Release); drop(thread_ids); drop(registry); let _ = operation.completed.send(true); @@ -269,6 +370,7 @@ impl PendingThreadUnloads { { let mut registry = self.lock_registry(); registry.closing = true; + self.inner.freeze.close(); self.inner.tasks.close(); } self.inner.tasks.wait().await; @@ -305,18 +407,32 @@ fn conflicting_completions( ) -> Vec> { let mut conflicts = Vec::>::new(); for thread_id in thread_ids { - let Some(operation) = registry.entries.get(thread_id) else { + let Some(entry) = registry.entries.get(thread_id) else { continue; }; - if owner.is_some_and(|owner| Arc::ptr_eq(operation, owner)) - || conflicts - .iter() - .any(|conflict| Arc::ptr_eq(conflict, operation)) - { + 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()) 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 dd9d409b0a08..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 @@ -59,115 +59,141 @@ async fn teardown_is_singleflight_and_outlives_its_first_waiter() { } #[tokio::test] -async fn overlapping_multi_claim_extensions_release_before_retrying_without_partial_claims() { +async fn freeze_handoffs_are_gapless_serialized_and_cancellation_safe() { let pending = PendingThreadUnloads::default(); - let root_a = ThreadId::new(); - let root_b = ThreadId::new(); + 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(claim_a) = pending.try_claim_many([root_a]) else { - panic!("first root should be claimable"); + let PendingThreadUnloadClaimResult::Claimed(predecessor_a) = + pending.try_claim_many([late_a, late_a_alias]) + else { + panic!("first predecessor should claim late descendants"); }; - let PendingThreadUnloadClaimResult::Claimed(claim_b) = pending.try_claim_many([root_b]) else { - panic!("second root should be claimable"); + 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 (conflicts_a_tx, conflicts_a_rx) = oneshot::channel(); - let (release_a_tx, release_a_rx) = oneshot::channel(); - let completion_a = claim_a.start_with(move |owner| async move { - let PendingThreadUnloadExtendResult::Pending(conflicts) = - owner.try_extend([root_b, late_a]) - else { - panic!("first extension should conflict with the second owner"); - }; - assert!(conflicts_a_tx.send(conflicts).is_ok()); - let _ = release_a_rx.await; + 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 (conflicts_b_tx, conflicts_b_rx) = oneshot::channel(); - let (release_b_tx, release_b_rx) = oneshot::channel(); - let completion_b = claim_b.start_with(move |owner| async move { - let PendingThreadUnloadExtendResult::Pending(conflicts) = - owner.try_extend([root_a, late_b]) - else { - panic!("second extension should conflict with the first owner"); + 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_b_tx.send(conflicts).is_ok()); - let _ = release_b_rx.await; + 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_a = conflicts_a_rx.await.expect("first extension result"); - let conflicts_b = conflicts_b_rx.await.expect("second extension result"); - assert_eq!(conflicts_a.completions.len(), 1); - assert_eq!(conflicts_b.completions.len(), 1); - assert!(!pending.contains(late_a)); - assert!(!pending.contains(late_b)); - let unclaimed = ThreadId::new(); - let PendingThreadUnloadClaimResult::Pending(all_conflicts) = - pending.try_claim_many([root_a, root_b, unclaimed]) - else { - panic!("an overlapping multi-claim should report every conflicting owner"); - }; - assert_eq!(all_conflicts.completions.len(), 2); - assert!(!pending.contains(unclaimed)); + let conflicts = conflicts_rx.await.expect("extension conflicts"); + assert_eq!(conflicts.completions.len(), 2); + assert!(pending.contains(free_sibling)); - // Neither owner may wait while retaining its partial group: each one is the other's - // conflict. Ending both tracked operations lets every conflict subscription complete. - release_a_tx.send(()).expect("release first owner"); - release_b_tx.send(()).expect("release second owner"); - wait_for_thread_unload(completion_a).await; - wait_for_thread_unload(completion_b).await; - wait_for_thread_unloads(conflicts_a).await; - wait_for_thread_unloads(conflicts_b).await; - wait_for_thread_unloads(all_conflicts).await; + 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()); - let PendingThreadUnloadClaimResult::Claimed(full_claim) = - pending.try_claim_many([root_a, root_b, late_a, late_b, late_a]) - else { - panic!("the complete deduplicated set should be claimable after both owners release"); - }; - let same_owner_unclaimed = ThreadId::new(); - let PendingThreadUnloadClaimResult::Pending(same_owner_conflict) = - pending.try_claim_many([root_a, late_a, same_owner_unclaimed]) + 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!("overlap with one owner should conflict"); + panic!("every ID should transfer to the one successor owner"); }; - assert_eq!(same_owner_conflict.completions.len(), 1); - assert!(!pending.contains(same_owner_unclaimed)); - let full_completion = full_claim.completed.clone(); - drop(full_claim); - wait_for_thread_unload(full_completion).await; - wait_for_thread_unloads(same_owner_conflict).await; + 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 stale_root = ThreadId::new(); - let closing_late = ThreadId::new(); - let stale_late = ThreadId::new(); - let PendingThreadUnloadClaimResult::Claimed(stale_claim) = pending.try_claim_many([stale_root]) + 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!("stale-handle root should be claimable"); + panic!("cancelled successor should claim its root"); }; - let (extend_tx, extend_rx) = oneshot::channel(); - let (stale_tx, stale_rx) = oneshot::channel(); - let stale_completion = stale_claim.start_with(move |owner| async move { - let _ = extend_rx.await; - assert!(matches!( - owner.try_extend([closing_late]), - PendingThreadUnloadExtendResult::Extended - )); - assert!(stale_tx.send(owner).is_ok()); + 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 pending_for_drain = pending.clone(); - let drain = tokio::spawn(async move { pending_for_drain.close_and_wait().await }); - while !pending.lock_registry().closing { - tokio::task::yield_now().await; - } - extend_tx.send(()).expect("extend while coordinator closes"); - let stale_owner = stale_rx.await.expect("stale owner handle"); - wait_for_thread_unload(stale_completion).await; - drain.await.expect("coordinator drain should complete"); + 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([stale_late]), + stale_owner.try_extend(&mut stale_freeze, [cancelled_free]), PendingThreadUnloadExtendResult::Finished )); - assert!(!pending.contains(stale_late)); + 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()); }