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]