Skip to main content

spawned_rt/sync/
mod.rs

1//! Runtime-independent lifecycle primitives. Default async builds retain Tokio's
2//! primitives; thread-only builds use std synchronization and removable wakers.
3#[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}