mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
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 <noreply@openai.com>
This commit is contained in:
Generated
+8
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<Vec<oneshot::Sender<()>>>,
|
||||
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<CancellationState>) -> 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<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
|
||||
mutex
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
/// 进程已进入 drain,新连接不得再注册。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ConnectionDraining;
|
||||
|
||||
@@ -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>) {
|
||||
self.wake_by_ref();
|
||||
}
|
||||
|
||||
fn wake_by_ref(self: &Arc<Self>) {
|
||||
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]
|
||||
|
||||
Reference in New Issue
Block a user