ados_protocol/
shutdown.rs

1//! A latching, clone-able shutdown signal for a service's run loops.
2//!
3//! `tokio::sync::Notify::notify_waiters` stores no permit: it wakes only the
4//! futures already waiting at the instant it runs. A loop that is inside its
5//! body at that instant (building a snapshot, writing a frame) re-enters its
6//! `select!` with a fresh `notified()` that nothing will ever wake, so a
7//! shutdown that joins its tasks hangs until the service manager kills it.
8//!
9//! [`Shutdown`] latches instead: [`Shutdown::trigger`] sets a flag and wakes
10//! every current waiter, and [`Shutdown::wait`] returns at once when the flag is
11//! already set, so a loop that reaches its `select!` after the trigger still
12//! stops.
13
14use std::sync::atomic::{AtomicBool, Ordering};
15use std::sync::Arc;
16
17use tokio::sync::Notify;
18
19/// A cancellation handle every run loop awaits and the service triggers on
20/// SIGTERM / SIGINT. Cheap to clone (shared `Arc`s).
21#[derive(Clone, Debug)]
22pub struct Shutdown {
23    fired: Arc<AtomicBool>,
24    notify: Arc<Notify>,
25}
26
27impl Shutdown {
28    /// A fresh, un-fired shutdown handle.
29    pub fn new() -> Self {
30        Self {
31            fired: Arc::new(AtomicBool::new(false)),
32            notify: Arc::new(Notify::new()),
33        }
34    }
35
36    /// Fire the signal. Idempotent; wakes every current and future waiter.
37    pub fn trigger(&self) {
38        self.fired.store(true, Ordering::SeqCst);
39        self.notify.notify_waiters();
40    }
41
42    /// Whether the signal has fired.
43    pub fn is_fired(&self) -> bool {
44        self.fired.load(Ordering::SeqCst)
45    }
46
47    /// Resolve once the signal has fired. Returns immediately when it fired
48    /// before this call.
49    pub async fn wait(&self) {
50        if self.is_fired() {
51            return;
52        }
53        // Register for the wakeup, then re-check, closing the race where
54        // trigger() ran between the flag check and the registration.
55        let notified = self.notify.notified();
56        if self.is_fired() {
57            return;
58        }
59        notified.await;
60    }
61}
62
63impl Default for Shutdown {
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72    use std::time::Duration;
73
74    #[tokio::test]
75    async fn a_parked_waiter_wakes_on_trigger() {
76        let s = Shutdown::new();
77        let waiter = s.clone();
78        let handle = tokio::spawn(async move { waiter.wait().await });
79        tokio::time::sleep(Duration::from_millis(20)).await;
80        s.trigger();
81        tokio::time::timeout(Duration::from_secs(1), handle)
82            .await
83            .expect("a parked waiter must wake on trigger")
84            .unwrap();
85    }
86
87    /// The property `Notify::notify_waiters` lacks: a loop that was busy when
88    /// the signal fired, and only reaches its wait afterwards, still stops.
89    #[tokio::test]
90    async fn a_waiter_that_arrives_after_the_trigger_returns_at_once() {
91        let s = Shutdown::new();
92        s.trigger();
93        tokio::time::timeout(Duration::from_millis(100), s.wait())
94            .await
95            .expect("a fired signal must not block a later waiter");
96        assert!(s.is_fired());
97    }
98}
ados_protocol/shutdown.rs

ados_protocol/
shutdown.rs
