1#[cfg(feature = "async")]
4pub use crate::tasks::{watch, CancellationToken};
5#[cfg(not(feature = "async"))]
6pub mod watch;
7#[cfg(not(feature = "async"))]
8#[derive(Clone)]
9pub struct CancellationToken(std::sync::Arc<watch::Sender<bool>>);
10#[cfg(not(feature = "async"))]
11impl CancellationToken {
12 pub fn new() -> Self {
13 let (sender, _) = watch::channel(false);
14 Self(std::sync::Arc::new(sender))
15 }
16 pub fn is_cancelled(&self) -> bool {
17 *self.0.borrow()
18 }
19 pub fn cancel(&self) {
20 self.0.send_replace(true);
21 }
22 pub async fn cancelled(&self) {
23 let mut receiver = self.0.subscribe();
24 loop {
25 if *receiver.borrow_and_update() {
26 return;
27 }
28 if receiver.changed().await.is_err() {
29 return;
30 }
31 }
32 }
33}
34#[cfg(not(feature = "async"))]
35impl Default for CancellationToken {
36 fn default() -> Self {
37 Self::new()
38 }
39}
40
41#[cfg(all(test, not(feature = "async")))]
42mod tests {
43 use super::*;
44 use std::{
45 future::Future,
46 sync::Arc,
47 task::{Context, Poll, Wake, Waker},
48 };
49 struct Probe(std::sync::atomic::AtomicUsize);
50 impl Wake for Probe {
51 fn wake(self: Arc<Self>) {
52 self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
53 }
54 }
55 #[test]
56 fn dropping_a_cancel_wait_reclaims_its_executor_waker() {
57 let token = CancellationToken::new();
58 let probe = Arc::new(Probe(std::sync::atomic::AtomicUsize::new(0)));
59 let weak = Arc::downgrade(&probe);
60 let waker = Waker::from(probe.clone());
61 let mut context = Context::from_waker(&waker);
62 let mut wait = Box::pin(token.cancelled());
63 assert!(matches!(wait.as_mut().poll(&mut context), Poll::Pending));
64 drop(wait);
65 drop(waker);
66 assert_eq!(
67 Arc::strong_count(&probe),
68 1,
69 "cancelled wait retained its waker"
70 );
71 token.cancel();
72 assert_eq!(probe.0.load(std::sync::atomic::Ordering::SeqCst), 0);
73 drop(probe);
74 assert!(weak.upgrade().is_none());
75 assert!(token.is_cancelled());
76 crate::threads::block_on(token.cancelled());
77 }
78 #[test]
79 fn cancellation_wakes_racing_native_waiters() {
80 for _ in 0..100 {
81 let token = CancellationToken::new();
82 let stopping = token.clone();
83 let worker = std::thread::spawn(move || stopping.cancel());
84 crate::threads::block_on(token.cancelled());
85 worker.join().unwrap();
86 }
87 }
88}