ullun/
shutdown.rs

1//! Shutdown signal and background tasks
2//!
3//! When the server is asked to stop (CTRL+C, SIGTERM, or the future given to
4//! `App::serve_with_shutdown` completes) it stops accepting connections, tells
5//! idle keep-alive connections to close, lets in-flight requests finish, and
6//! waits for [background tasks](crate::app::App::background), each at most
7//! `App::shutdown_grace` (default 10 s). Long-lived work such as SSE streams,
8//! WebSocket handlers and periodic jobs should watch [`Shutdown`] to end
9//! promptly instead of running into that limit.
10
11use std::time::Duration;
12use tokio::sync::watch;
13
14/// Resolves when the server begins shutting down
15#[derive(Clone)]
16pub struct Shutdown {
17    rx: watch::Receiver<bool>,
18}
19
20impl Shutdown {
21    pub(crate) fn new(rx: watch::Receiver<bool>) -> Self {
22        Self { rx }
23    }
24
25    /// A signal that never fires (requests built outside a running server)
26    pub fn never() -> Self {
27        let (tx, rx) = watch::channel(false);
28        // keep the sender alive for the process lifetime so `wait` pends instead of erroring
29        std::mem::forget(tx);
30        Self { rx }
31    }
32
33    /// True once shutdown has begun
34    pub fn is_triggered(&self) -> bool {
35        *self.rx.borrow()
36    }
37
38    /// Completes when shutdown begins (immediately if it already has)
39    pub async fn wait(&self) {
40        let mut rx = self.rx.clone();
41        loop {
42            if *rx.borrow_and_update() {
43                return;
44            }
45            if rx.changed().await.is_err() {
46                // sender gone without a signal: the server is being torn down
47                return;
48            }
49        }
50    }
51
52    /// Sleep for `duration`; returns `false` if shutdown began first
53    ///
54    /// Designed for loops: `while shutdown.sleep(Duration::from_secs(30)).await { work().await }`.
55    pub async fn sleep(&self, duration: Duration) -> bool {
56        tokio::select! {
57            _ = tokio::time::sleep(duration) => !self.is_triggered(),
58            _ = self.wait() => false,
59        }
60    }
61}
62
63/// Handed to every task registered with `App::background`; cheap to clone into
64/// sub-tasks
65#[derive(Clone)]
66pub struct Context {
67    pub(crate) state: std::sync::Arc<crate::request::StateMap>,
68    shutdown: Shutdown,
69}
70
71impl Context {
72    pub(crate) fn new(state: std::sync::Arc<crate::request::StateMap>, shutdown: Shutdown) -> Self {
73        Self { state, shutdown }
74    }
75
76    /// Shared state registered with `App::state`
77    pub fn state<T: Send + Sync + 'static>(&self) -> crate::error::Result<&T> {
78        self.state.get::<T>().ok_or_else(|| {
79            crate::error::Error::internal(format!(
80                "State of type {} is not registered",
81                std::any::type_name::<T>()
82            ))
83        })
84    }
85
86    /// The shutdown signal; clone it into spawned sub-tasks
87    pub fn shutdown(&self) -> Shutdown {
88        self.shutdown.clone()
89    }
90
91    /// Shorthand for `shutdown().sleep(duration)`
92    pub async fn sleep(&self, duration: Duration) -> bool {
93        self.shutdown.sleep(duration).await
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[tokio::test]
102    async fn wait_sleep_and_trigger() {
103        let (tx, rx) = watch::channel(false);
104        let s = Shutdown::new(rx);
105        assert!(!s.is_triggered());
106        assert!(s.sleep(Duration::from_millis(5)).await);
107
108        let s2 = s.clone();
109        let waiter = tokio::spawn(async move { s2.wait().await });
110        tx.send(true).unwrap();
111        waiter.await.unwrap();
112        assert!(s.is_triggered());
113        // already triggered: returns at once, sleep reports the interruption
114        s.wait().await;
115        assert!(!s.sleep(Duration::from_secs(60)).await);
116    }
117
118    #[tokio::test]
119    async fn sleep_is_interrupted_by_the_signal() {
120        let (tx, rx) = watch::channel(false);
121        let s = Shutdown::new(rx);
122        let sleeper = tokio::spawn({
123            let s = s.clone();
124            async move { s.sleep(Duration::from_secs(60)).await }
125        });
126        tokio::time::sleep(Duration::from_millis(20)).await;
127        tx.send(true).unwrap();
128        assert!(!tokio::time::timeout(Duration::from_secs(5), sleeper)
129            .await
130            .unwrap()
131            .unwrap());
132    }
133
134    #[tokio::test]
135    async fn never_does_not_fire() {
136        let s = Shutdown::never();
137        assert!(!s.is_triggered());
138        assert!(tokio::time::timeout(Duration::from_millis(30), s.wait())
139            .await
140            .is_err());
141    }
142}
ullun/shutdown.rs

ullun/
shutdown.rs
