rget/
shutdown.rs

1//! Cooperative cancellation (PRD §26).
2//!
3//! Small enough not to justify a dependency: a flag plus a notifier. The
4//! important property is that `cancel()` is observable both by a poll
5//! (`is_cancelled`) on the hot write path and by an await
6//! (`cancelled().await`) inside a `select!`.
7
8use std::sync::Arc;
9use std::sync::atomic::{AtomicBool, Ordering};
10
11use tokio::sync::Notify;
12
13#[derive(Clone, Default)]
14pub struct Cancel {
15    inner: Arc<Inner>,
16}
17
18#[derive(Default)]
19struct Inner {
20    flag: AtomicBool,
21    notify: Notify,
22}
23
24impl Cancel {
25    pub fn new() -> Self {
26        Self::default()
27    }
28
29    pub fn cancel(&self) {
30        self.inner.flag.store(true, Ordering::SeqCst);
31        self.inner.notify.notify_waiters();
32    }
33
34    pub fn is_cancelled(&self) -> bool {
35        self.inner.flag.load(Ordering::SeqCst)
36    }
37
38    /// Resolves as soon as cancellation has been requested, including when it
39    /// was requested before this call.
40    pub async fn cancelled(&self) {
41        if self.is_cancelled() {
42            return;
43        }
44        loop {
45            let notified = self.inner.notify.notified();
46            if self.is_cancelled() {
47                return;
48            }
49            notified.await;
50            if self.is_cancelled() {
51                return;
52            }
53        }
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[tokio::test]
62    async fn resolves_when_cancelled_later() {
63        let c = Cancel::new();
64        let c2 = c.clone();
65        tokio::spawn(async move {
66            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
67            c2.cancel();
68        });
69        c.cancelled().await;
70        assert!(c.is_cancelled());
71    }
72
73    #[tokio::test]
74    async fn resolves_immediately_if_already_cancelled() {
75        let c = Cancel::new();
76        c.cancel();
77        // Would hang if `cancelled()` only watched for future notifications.
78        tokio::time::timeout(std::time::Duration::from_millis(50), c.cancelled())
79            .await
80            .expect("should resolve immediately");
81    }
82
83    #[tokio::test]
84    async fn clones_share_state() {
85        let a = Cancel::new();
86        let b = a.clone();
87        a.cancel();
88        assert!(b.is_cancelled());
89    }
90}
rget/shutdown.rs

rget/
shutdown.rs
