From 1c8a74e572c1c66bfd9ee7dc3c5a543e46fb2e2d Mon Sep 17 00:00:00 2001 From: Zyy Date: Mon, 5 Oct 2026 09:27:42 +0800 Subject: [PATCH] fix(core): release dropped cancellation waiters Use event-listener registrations that unregister when cancellation futures are dropped, including waits attached to ancestor tokens. Recheck the cancelled flag after registration to preserve wakeups during races. Add allocation regressions for select churn and child cancellation, plus broadcast and concurrent cancellation coverage. Refs #338, #345 Co-authored-by: Codex --- backend/Cargo.lock | 8 + backend/crates/gateway-core/Cargo.toml | 4 + backend/crates/gateway-core/src/lifecycle.rs | 42 ++-- .../crates/gateway-core/tests/lifecycle.rs | 190 +++++++++++++++++- 4 files changed, 214 insertions(+), 30 deletions(-) diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 49e2bedbf..159646a81 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -31,6 +31,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocation-counter" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "beb9e990c0a33699f1984d85a6abead615ccc72dd8130bf3e15dcabe2ca149c9" + [[package]] name = "allocator-api2" version = "0.2.21" @@ -1728,11 +1734,13 @@ dependencies = [ name = "gateway-core" version = "0.1.0" dependencies = [ + "allocation-counter", "async-trait", "base64 0.23.1", "bytes", "chrono", "chrono-tz", + "event-listener", "futures", "futures-timer", "http", diff --git a/backend/crates/gateway-core/Cargo.toml b/backend/crates/gateway-core/Cargo.toml index 3318ec1ba..520e11ef5 100644 --- a/backend/crates/gateway-core/Cargo.toml +++ b/backend/crates/gateway-core/Cargo.toml @@ -17,6 +17,7 @@ base64 = "0.23.0" bytes = "1.11.1" chrono = { version = "0.4.45", features = ["serde"] } chrono-tz = { version = "0.10", features = ["serde"] } +event-listener = "5.4.2" futures = "0.3.32" http = "1.5.0" http-body = "1.0.1" @@ -31,5 +32,8 @@ tracing = "0.1.44" uuid = { version = "1.23.3", features = ["v7"] } url = "2.5.8" +[dev-dependencies] +allocation-counter = "0.8.1" + [lints] workspace = true diff --git a/backend/crates/gateway-core/src/lifecycle.rs b/backend/crates/gateway-core/src/lifecycle.rs index 6111eec7d..2cd43d6eb 100644 --- a/backend/crates/gateway-core/src/lifecycle.rs +++ b/backend/crates/gateway-core/src/lifecycle.rs @@ -1,19 +1,17 @@ //! 请求与连接的截止、租约、取消和 drain 契约。 use std::fmt; +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime}; -use futures::{ - channel::oneshot, - future::{BoxFuture, pending}, -}; +use event_listener::Event; +use futures::future::{BoxFuture, pending}; use futures_timer::Delay; struct CancellationState { cancelled: AtomicBool, - waiters: Mutex>>, + event: Event, } /// 可克隆的请求、任务与连接取消信号。 @@ -44,7 +42,7 @@ impl CancellationToken { Self { state: Arc::new(CancellationState { cancelled: AtomicBool::new(false), - waiters: Mutex::new(Vec::new()), + event: Event::new(), }), ancestors: Arc::from([]), } @@ -59,7 +57,7 @@ impl CancellationToken { Self { state: Arc::new(CancellationState { cancelled: AtomicBool::new(false), - waiters: Mutex::new(Vec::new()), + event: Event::new(), }), ancestors: ancestors.into(), } @@ -69,13 +67,7 @@ impl CancellationToken { if self.state.cancelled.swap(true, Ordering::AcqRel) { return; } - let waiters = { - let mut guard = lock_unpoisoned(&self.state.waiters); - std::mem::take(&mut *guard) - }; - for waiter in waiters { - let _ = waiter.send(()); - } + self.state.event.notify(usize::MAX); } #[must_use] @@ -103,24 +95,16 @@ fn wait_for_state(state: Arc) -> BoxFuture<'static, ()> { if state.cancelled.load(Ordering::Acquire) { return; } - let (sender, receiver) = oneshot::channel(); - { - let mut waiters = lock_unpoisoned(&state.waiters); - if state.cancelled.load(Ordering::Acquire) { - return; - } - waiters.push(sender); + // listener 随等待 future 丢弃而注销，避免 select 败选分支在自身及祖先状态中积累。 + let listener = state.event.listen(); + // 注册后复查，覆盖取消发生在首次检查与注册之间的竞态。 + if state.cancelled.load(Ordering::Acquire) { + return; } - let _ = receiver.await; + listener.await; }) } -fn lock_unpoisoned(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { - mutex - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) -} - /// 进程已进入 drain，新连接不得再注册。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct ConnectionDraining; diff --git a/backend/crates/gateway-core/tests/lifecycle.rs b/backend/crates/gateway-core/tests/lifecycle.rs index a606b1955..163a44d36 100644 --- a/backend/crates/gateway-core/tests/lifecycle.rs +++ b/backend/crates/gateway-core/tests/lifecycle.rs @@ -1,15 +1,203 @@ -use std::sync::Arc; +use std::future::Future as _; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Barrier}; +use std::task::{Context, Wake, Waker}; use futures::{FutureExt as _, pin_mut}; use gateway_core::lifecycle::CancellationToken; +#[test] +fn discarded_cancellation_branches_release_allocations() { + let token = CancellationToken::new(); + // 预热一次，允许取消信号保留固定大小的通知状态。 + discard_cancellation_branch(&token); + + let allocations = allocation_counter::measure(|| { + for _ in 0..10_000 { + discard_cancellation_branch(&token); + } + }); + + assert_eq!(allocations.count_current, 0, "{allocations:?}"); + assert_eq!(allocations.bytes_current, 0, "{allocations:?}"); + assert!(!token.is_cancelled()); +} + +#[test] +fn dropped_descendant_waits_release_ancestor_allocations() { + let parent = CancellationToken::new(); + let child = parent.child_token(); + let grandchild = child.child_token(); + discard_cancellation_branch(&grandchild); + + let allocations = allocation_counter::measure(|| { + for _ in 0..10_000 { + discard_cancellation_branch(&grandchild); + } + }); + + assert_eq!(allocations.count_current, 0, "{allocations:?}"); + assert_eq!(allocations.bytes_current, 0, "{allocations:?}"); + assert!(!parent.is_cancelled()); +} + +#[test] +fn completing_child_cancellation_releases_ancestor_allocations() { + let parent = CancellationToken::new(); + discard_cancellation_branch(&parent); + + let allocations = allocation_counter::measure(|| { + for _ in 0..1_000 { + let child = parent.child_token(); + let mut waiting = Box::pin(child.cancelled()); + assert!(waiting.as_mut().now_or_never().is_none()); + child.cancel(); + assert!(waiting.now_or_never().is_some()); + } + }); + + assert_eq!(allocations.count_current, 0, "{allocations:?}"); + assert_eq!(allocations.bytes_current, 0, "{allocations:?}"); + assert!(!parent.is_cancelled()); +} + +fn discard_cancellation_branch(token: &CancellationToken) { + // select 先轮询取消分支，再由另一个就绪分支获胜并丢弃取消等待。 + let selected = futures::future::select(Box::pin(token.cancelled()), futures::future::ready(())) + .now_or_never(); + assert!(matches!(selected, Some(futures::future::Either::Right(_)))); +} + +#[derive(Default)] +struct WakeCounter(AtomicUsize); + +impl Wake for WakeCounter { + fn wake(self: Arc) { + self.wake_by_ref(); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +#[test] +fn cancellation_wakes_all_live_waiters_after_other_waits_are_dropped() { + let parent = CancellationToken::new(); + let child = parent.child_token(); + let grandchild = child.child_token(); + let sibling = parent.child_token(); + let tokens = [&parent, &child, &grandchild, &sibling]; + let mut waiters = Vec::new(); + + for token in tokens.into_iter().cycle().take(64) { + let counter = Arc::new(WakeCounter::default()); + let waker = Waker::from(Arc::clone(&counter)); + let mut waiting = Box::pin(token.cancelled()); + assert!( + waiting + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + waiters.push((waiting, counter)); + } + for _ in 0..100 { + discard_cancellation_branch(&grandchild); + } + + parent.clone().cancel(); + parent.cancel(); + + for (waiting, counter) in waiters { + assert!(counter.0.load(Ordering::Relaxed) > 0); + assert!(waiting.now_or_never().is_some()); + assert_eq!(Arc::strong_count(&counter), 1); + } +} + +#[test] +fn cancelling_child_only_wakes_its_descendants() { + let parent = CancellationToken::new(); + let child = parent.child_token(); + let grandchild = child.child_token(); + let sibling = parent.child_token(); + let mut parent_wait = Box::pin(parent.cancelled()); + let mut grandchild_wait = Box::pin(grandchild.cancelled()); + let mut sibling_wait = Box::pin(sibling.cancelled()); + assert!(parent_wait.as_mut().now_or_never().is_none()); + assert!(grandchild_wait.as_mut().now_or_never().is_none()); + assert!(sibling_wait.as_mut().now_or_never().is_none()); + + child.cancel(); + + assert!(grandchild_wait.now_or_never().is_some()); + assert!(parent_wait.now_or_never().is_none()); + assert!(sibling_wait.now_or_never().is_none()); + assert!(!parent.is_cancelled()); + assert!(!sibling.is_cancelled()); +} + +#[test] +fn cancellation_racing_with_wait_registration_and_drop_does_not_lose_wakeup() { + for _ in 0..256 { + let parent = CancellationToken::new(); + let child = parent.child_token(); + let grandchild = child.child_token(); + let mut dropped_wait = Box::pin(grandchild.cancelled()); + assert!(dropped_wait.as_mut().now_or_never().is_none()); + let start = Barrier::new(2); + let finished = Barrier::new(2); + let counter = Arc::new(WakeCounter::default()); + let waker = Waker::from(Arc::clone(&counter)); + + std::thread::scope(|scope| { + scope.spawn(|| { + start.wait(); + parent.cancel(); + finished.wait(); + }); + + start.wait(); + drop(dropped_wait); + let mut waiting = Box::pin(grandchild.cancelled()); + let pending = waiting + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending(); + finished.wait(); + + // 取消完成后再检查，既验证就绪结果，也验证 Pending 路径实际收到唤醒。 + if pending { + assert!(counter.0.load(Ordering::Relaxed) > 0); + assert!(waiting.now_or_never().is_some()); + } + }); + assert!(grandchild.is_cancelled()); + assert_eq!(Arc::strong_count(&counter), 2); + } +} + +#[test] +fn descendants_created_after_cancellation_are_immediately_ready() { + let parent = CancellationToken::new(); + parent.cancel(); + let child = parent.child_token(); + let grandchild = child.child_token(); + drop(parent); + drop(child); + + assert!(grandchild.is_cancelled()); + assert!(grandchild.cancelled().now_or_never().is_some()); +} + #[test] fn cancellation_token_should_wake_current_state() { let token = CancellationToken::new(); token.cancel(); assert!(token.is_cancelled()); + assert!(token.cancelled().now_or_never().is_some()); } #[test]