axum_server/
connection.rs

1//! Connection-level control.
2
3use tokio::sync::watch;
4
5/// How a connection was asked to close.
6#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
7pub(crate) enum ShutdownKind {
8    #[default]
9    None,
10    Graceful,
11    Abort,
12}
13
14/// A handle to the connection that dispatched the current request.
15///
16/// axum-server adds this to the extensions of every request, so a handler can
17/// close the connection it is running on:
18///
19/// ```rust,no_run
20/// use axum::{extract::Extension, routing::get, Router};
21/// use axum_server::ConnectionShutdown;
22///
23/// async fn handle(Extension(shutdown): Extension<ConnectionShutdown>) -> &'static str {
24///     shutdown.graceful_shutdown();
25///     "goodbye"
26/// }
27///
28/// let app: Router = Router::new().route("/", get(handle));
29/// ```
30///
31/// The handle can be moved into another task to close the connection later,
32/// for example to enforce a deadline on a slow client.
33///
34/// An HTTP/2 connection carries many streams, so closing it closes all of
35/// them.
36#[derive(Clone, Debug)]
37pub struct ConnectionShutdown {
38    tx: watch::Sender<ShutdownKind>,
39}
40
41impl ConnectionShutdown {
42    pub(crate) fn new(tx: watch::Sender<ShutdownKind>) -> Self {
43        Self { tx }
44    }
45
46    /// Close the connection immediately.
47    ///
48    /// The in-flight request is cancelled, so its response may not be sent.
49    /// Use [`graceful_shutdown`](Self::graceful_shutdown) to let the response
50    /// finish.
51    pub fn shutdown(&self) {
52        let _ = self.tx.send_replace(ShutdownKind::Abort);
53    }
54
55    /// Close the connection after the in-flight response is sent.
56    ///
57    /// The current response is still delivered and no further requests are
58    /// accepted on the connection.
59    pub fn graceful_shutdown(&self) {
60        let _ = self.tx.send_if_modified(|kind| match kind {
61            ShutdownKind::None => {
62                *kind = ShutdownKind::Graceful;
63                true
64            }
65            ShutdownKind::Graceful | ShutdownKind::Abort => false,
66        });
67    }
68}
axum_server/connection.rs

axum_server/
connection.rs
