1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
//! A one-shot "stop what you're doing" signal for work already in flight.
//!
//! Cancelling an agent sets its status, which stops the dispatch systems from
//! starting *new* work - but an inference request or tool batch handed to the
//! async lanes before that keeps running to completion. For a stalled provider
//! call that is up to the job timeout; for a tool batch there is no bound at
//! all, and the batch holds tool-lane capacity the whole time.
//!
//! A [`CancelToken`] is handed to each async job when it is dispatched and kept
//! on the agent, so cancelling the agent drops the in-flight future at its next
//! await point: the HTTP request is aborted, the pool permit and lane capacity
//! are released, and the result is never applied.
//!
//! This is deliberately a few lines rather than a dependency: the whole contract
//! is "set a flag once, wake anyone waiting".
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::sync::Notify;
#[derive(Default)]
struct Inner {
cancelled: AtomicBool,
notify: Notify,
}
/// A cloneable cancellation signal. Every clone observes the same state, and
/// cancelling is idempotent - a token that fires twice (say, a cancel racing a
/// reap) is not an error.
#[derive(Clone, Default)]
pub struct CancelToken {
inner: Arc<Inner>,
}
impl CancelToken {
/// A fresh, un-cancelled token.
pub fn new() -> Self {
Self::default()
}
/// Fire the signal, waking every waiter. Idempotent.
pub fn cancel(&self) {
self.inner.cancelled.store(true, Ordering::SeqCst);
self.inner.notify.notify_waiters();
}
/// Whether this token has already fired.
pub fn is_cancelled(&self) -> bool {
self.inner.cancelled.load(Ordering::SeqCst)
}
/// Resolve once the token fires, immediately if it already has.
///
/// The waiter is armed *before* the flag is re-checked. `notify_waiters`
/// only wakes waiters registered at the time it runs, so checking first and
/// waiting second would lose a cancel that landed in between - and the job
/// would then run to completion despite having been cancelled.
pub async fn cancelled(&self) {
let notified = self.inner.notify.notified();
tokio::pin!(notified);
notified.as_mut().enable();
if self.is_cancelled() {
return;
}
notified.await;
}
}
impl std::fmt::Debug for CancelToken {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CancelToken")
.field("cancelled", &self.is_cancelled())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[tokio::test]
async fn cancelled_resolves_when_the_token_fires() {
let token = CancelToken::new();
assert!(!token.is_cancelled());
let waiter = tokio::spawn({
let token = token.clone();
async move { token.cancelled().await }
});
// Let the waiter arm itself, then fire.
tokio::task::yield_now().await;
token.cancel();
tokio::time::timeout(Duration::from_secs(5), waiter)
.await
.expect("a fired token wakes its waiter")
.unwrap();
assert!(token.is_cancelled());
}
#[tokio::test]
async fn cancelled_returns_immediately_for_an_already_fired_token() {
let token = CancelToken::new();
token.cancel();
tokio::time::timeout(Duration::from_secs(5), token.cancelled())
.await
.expect("no wait for a token that already fired");
}
/// A cancel concurrent with the wait is still observed. This exercises the
/// path but cannot *prove* the arm-then-check ordering: the window it guards
/// is a few instructions with no await point in it, so a test cannot land in
/// it reliably. The ordering is enforced by construction (see `cancelled`),
/// not by this test.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_cancel_concurrent_with_the_wait_is_observed() {
for _ in 0..500 {
let token = CancelToken::new();
let firing = tokio::spawn({
let token = token.clone();
async move { token.cancel() }
});
tokio::time::timeout(Duration::from_secs(5), token.cancelled())
.await
.expect("the cancel was observed");
firing.await.unwrap();
}
}
#[tokio::test]
async fn cancel_is_idempotent_and_wakes_every_clone() {
let token = CancelToken::new();
let waiters: Vec<_> = (0..3)
.map(|_| {
let token = token.clone();
tokio::spawn(async move { token.cancelled().await })
})
.collect();
tokio::task::yield_now().await;
token.cancel();
token.cancel(); // twice is fine
for waiter in waiters {
tokio::time::timeout(Duration::from_secs(5), waiter)
.await
.expect("every clone observes the cancel")
.unwrap();
}
}
#[test]
fn debug_reports_the_state() {
let token = CancelToken::new();
assert!(format!("{token:?}").contains("cancelled: false"));
token.cancel();
assert!(format!("{token:?}").contains("cancelled: true"));
}
}