kutil_http/axum/
shutdown.rs

1use {
2    ::axum::{extract::*, http::StatusCode, response::*},
3    axum_server::*,
4    std::{io, time::*},
5    tokio::{signal::*, sync::oneshot::*, task::*, *},
6    tokio_util::sync::*,
7};
8
9//
10// Shutdown
11//
12
13/// Axum server shutdown coordinator.
14///
15/// Clones will retain the same coordination.
16#[derive(Clone, Debug)]
17pub struct Shutdown {
18    /// Axum server handle.
19    ///
20    /// Clones will retain the same coordination.
21    pub handle: Handle,
22
23    /// Grace period. [None](Option::None) means indefinite.
24    pub grace_period: Option<Duration>,
25}
26
27impl Shutdown {
28    /// Constructor.
29    pub fn new(grace_period: Option<Duration>) -> Self {
30        Self { handle: Handle::new(), grace_period }
31    }
32
33    /// Shutdown (graceful).
34    pub fn shutdown(&self) {
35        self.handle.graceful_shutdown(self.grace_period);
36    }
37
38    /// Shutdown *now* (ignore grace period).
39    pub fn shutdown_now(&self) {
40        self.handle.shutdown();
41    }
42
43    /// Get a [CancellationToken].
44    ///
45    /// Call [CancellationToken::cancel] to shutdown.
46    ///
47    /// It's best not to call this function more than once, as it spawns a listener task. If you
48    /// need several tokens, clone the one you get here.
49    ///
50    /// Also returns the [JoinHandle] for the listener task.
51    pub fn cancellation_token(&self) -> (CancellationToken, JoinHandle<()>) {
52        let token = CancellationToken::new();
53
54        let shutdown = self.clone();
55
56        (
57            token.clone(),
58            spawn(async move {
59                tracing::info!("waiting on cancellation token");
60
61                token.cancelled().await;
62                tracing::info!("cancellation token activated");
63                shutdown.shutdown();
64            }),
65        )
66    }
67
68    /// Shutdown on channel send.
69    ///
70    /// Call [Sender::send] to shutdown.
71    ///
72    /// It's best not to call this function more than once, as it spawns a listener task. If you
73    /// need several senders, clone the one you get here.
74    ///
75    /// Also returns the [JoinHandle] for the listener task.
76    pub fn on_channel(&self) -> (Sender<()>, JoinHandle<()>) {
77        let (sender, receiver) = channel();
78        let shutdown = self.clone();
79
80        (
81            sender,
82            spawn(async move {
83                tracing::info!("listening on shutdown channel");
84
85                match receiver.await {
86                    Ok(_) => {
87                        tracing::info!("received shutdown message");
88                    }
89
90                    Err(error) => {
91                        tracing::error!("shutdown channel error: {}", error);
92                    }
93                }
94
95                shutdown.shutdown();
96            }),
97        )
98    }
99
100    /// Shutdown on signals.
101    ///
102    /// Returns the [JoinHandle] for the listener task.
103    pub fn on_signals(&self) -> io::Result<JoinHandle<()>> {
104        #[cfg(all(not(unix), not(windows)))]
105        {
106            tracing::warn!("signals not supported on this platform");
107            return Ok(());
108        }
109
110        let shutdown = self.clone();
111
112        #[cfg(unix)]
113        let mut interrupt = unix::signal(unix::SignalKind::interrupt())?; // ctrl+c
114        #[cfg(unix)]
115        let mut terminate = unix::signal(unix::SignalKind::terminate())?;
116
117        Ok(spawn(async move {
118            tracing::info!("listening for shutdown signals");
119
120            #[cfg(unix)]
121            select! {
122                _ = interrupt.recv() => {},
123                _ = terminate.recv() => {},
124            }
125
126            #[cfg(windows)]
127            select! {
128                _ = windows::ctrl_c() => {},
129                _ = windows::ctrl_break() => {},
130                _ = windows::ctrl_close() => {},
131                _ = windows::ctrl_logoff() => {},
132                _ = windows::ctrl_shutdown() => {},
133            }
134
135            tracing::info!("received shutdown signal");
136            shutdown.shutdown();
137        }))
138    }
139}
140
141/// Axum request handler that calls [Shutdown::shutdown].
142///
143/// Expects the [Shutdown] to be available as state. See
144/// [Router::with_state](::axum::Router::with_state).
145pub async fn shutdown_handler(State(shutdown): State<Shutdown>) -> Response {
146    tracing::info!("shutting down");
147    shutdown.shutdown();
148    StatusCode::NO_CONTENT.into_response()
149}
kutil_http/axum/shutdown.rs

kutil_http/axum/
shutdown.rs
