synd_api/
shutdown.rs

1use std::{
2    future::Future,
3    io,
4    net::SocketAddr,
5    sync::{Arc, OnceLock},
6    time::Duration,
7};
8
9use axum_server::Handle;
10use tokio_util::sync::CancellationToken;
11use tracing::{error, info};
12
13/// `CancellationToken` wrapper
14#[derive(Clone)]
15pub struct Shutdown {
16    root: CancellationToken,
17    handle: Handle<SocketAddr>,
18    reason: Arc<OnceLock<ShutdownReason>>,
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum ShutdownReason {
23    Manual,
24    Idle,
25    ApiRequest,
26    Signal,
27    SignalError,
28    Unknown,
29}
30
31impl ShutdownReason {
32    #[must_use]
33    pub fn as_str(self) -> &'static str {
34        match self {
35            Self::Manual => "manual",
36            Self::Idle => "idle",
37            Self::ApiRequest => "api_request",
38            Self::Signal => "signal",
39            Self::SignalError => "signal_error",
40            Self::Unknown => "unknown",
41        }
42    }
43}
44
45impl Shutdown {
46    pub fn manual<F>(on_graceful_shutdown: F) -> Self
47    where
48        F: FnOnce() + Send + 'static,
49    {
50        let root = CancellationToken::new();
51        let reason = Arc::new(OnceLock::new());
52
53        let ct = root.clone();
54        let shutdown_reason = Arc::clone(&reason);
55        let handle = axum_server::Handle::new();
56        let notify = handle.clone();
57        tokio::spawn(async move {
58            ct.cancelled().await;
59            let reason = shutdown_reason
60                .get()
61                .copied()
62                .unwrap_or(ShutdownReason::Unknown);
63            info!(reason = reason.as_str(), "Graceful shutdown requested");
64            on_graceful_shutdown();
65            notify.graceful_shutdown(Some(Duration::from_secs(3)));
66        });
67
68        Self {
69            root,
70            handle,
71            reason,
72        }
73    }
74
75    /// When the given signal Future is resolved, call the `cancel` method of the held `CancellationToken`.
76    pub fn watch_signal<Fut, F>(signal: Fut, on_graceful_shutdown: F) -> Self
77    where
78        F: FnOnce() + Send + 'static,
79        Fut: Future<Output = io::Result<()>> + Send + 'static,
80    {
81        let shutdown = Self::manual(on_graceful_shutdown);
82        let notify = shutdown.clone();
83        tokio::spawn(async move {
84            let reason = match signal.await {
85                Ok(()) => ShutdownReason::Signal,
86                Err(err) => {
87                    error!(error = %err, "Failed to handle shutdown signal");
88                    ShutdownReason::SignalError
89                }
90            };
91            notify.shutdown_with_reason(reason);
92        });
93
94        shutdown
95    }
96
97    /// Request shutdown
98    pub fn shutdown(&self) {
99        self.shutdown_with_reason(ShutdownReason::Manual);
100    }
101
102    pub fn shutdown_with_reason(&self, reason: ShutdownReason) {
103        let _ = self.reason.set(reason);
104        self.root.cancel();
105    }
106
107    #[must_use]
108    pub fn reason(&self) -> Option<ShutdownReason> {
109        self.reason.get().copied()
110    }
111
112    pub fn is_shutdown_requested(&self) -> bool {
113        self.root.is_cancelled()
114    }
115
116    pub fn into_handle(self) -> Handle<SocketAddr> {
117        self.handle
118    }
119
120    /// Return `CancellationToken which is cancelled at shutdown`
121    pub fn cancellation_token(&self) -> CancellationToken {
122        self.root.clone()
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use std::{
129        io::ErrorKind,
130        sync::{
131            Arc,
132            atomic::{AtomicBool, Ordering},
133        },
134    };
135
136    use futures_util::future;
137
138    use super::*;
139
140    #[tokio::test(flavor = "multi_thread")]
141    async fn signal_trigger_graceful_shutdown() {
142        for signal_result in [Ok(()), Err(io::Error::from(ErrorKind::Other))] {
143            let called = Arc::new(AtomicBool::new(false));
144            let called_cloned = Arc::clone(&called);
145            let on_graceful_shutdown = move || {
146                called_cloned.store(true, Ordering::Relaxed);
147            };
148            let (tx, rx) = tokio::sync::oneshot::channel::<io::Result<()>>();
149            let s = Shutdown::watch_signal(
150                async move {
151                    rx.await.unwrap().ok();
152                    signal_result
153                },
154                on_graceful_shutdown,
155            );
156            let ct = s.cancellation_token();
157
158            // Mock signal triggered
159            tx.send(Ok(())).unwrap();
160
161            // Check cancellation token is cancelled and axum handler called
162            let mut ok = false;
163            let mut count = 0;
164            loop {
165                count += 1;
166                if count >= 10 {
167                    break;
168                }
169                if s.root.is_cancelled() && ct.is_cancelled() && called.load(Ordering::Relaxed) {
170                    ok = true;
171                    break;
172                }
173                tokio::time::sleep(Duration::from_millis(100)).await;
174            }
175            assert!(ok, "cancelation does not work");
176        }
177    }
178
179    #[tokio::test(flavor = "multi_thread")]
180    async fn shutdown_trigger_graceful_shutdown() {
181        let called = Arc::new(AtomicBool::new(false));
182        let called_cloned = Arc::clone(&called);
183        let on_graceful_shutdown = move || {
184            called_cloned.store(true, Ordering::Relaxed);
185        };
186        let s = Shutdown::watch_signal(future::pending(), on_graceful_shutdown);
187        let ct = s.cancellation_token();
188
189        s.shutdown();
190        assert_eq!(s.reason(), Some(ShutdownReason::Manual));
191
192        // Check cancellation token is cancelled and axum handler called
193        let mut ok = false;
194        let mut count = 0;
195        loop {
196            count += 1;
197            if count >= 10 {
198                break;
199            }
200            if s.root.is_cancelled() && ct.is_cancelled() && called.load(Ordering::Relaxed) {
201                ok = true;
202                break;
203            }
204            tokio::time::sleep(Duration::from_millis(100)).await;
205        }
206        assert!(ok, "cancelation does not work");
207        assert!(s.is_shutdown_requested());
208    }
209}
synd_api/shutdown.rs

synd_api/
shutdown.rs
