1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
//! Graceful-shutdown signal handling. Every serve path uses this
//! one function, so they cannot disagree.
//!
//! ## Why `ctrl_c()` alone is not enough
//!
//! On Unix `tokio::signal::ctrl_c()` waits for SIGINT only. But
//! Kubernetes, `docker stop` and systemd all send SIGTERM, which
//! kills the process by default. A server waiting on `ctrl_c()`
//! therefore drains on a laptop and dies mid-request in
//! production, with exit code 0 and nothing in the log.
/// Waits until the process is asked to stop, by Ctrl-C (SIGINT) or
/// SIGTERM. Off Unix, only Ctrl-C.
///
/// Pass it to `axum::serve(..).with_graceful_shutdown(..)`. The
/// server then stops accepting connections and lets the requests
/// already running finish.
pub async fn shutdown_signal() {
#[cfg(unix)]
{
use tokio::signal::unix::{signal, SignalKind};
// If we cannot install the handler, fall back to Ctrl-C.
// Hanging forever would be worse.
let mut term = match signal(SignalKind::terminate()) {
Ok(s) => s,
Err(e) => {
tracing::warn!(
target: "rustango::shutdown",
error = %e,
"could not install a SIGTERM handler; Ctrl-C only",
);
let _ = tokio::signal::ctrl_c().await;
tracing::info!(target: "rustango::shutdown", signal = "SIGINT", "shutting down");
return;
}
};
let signal_name = tokio::select! {
_ = tokio::signal::ctrl_c() => "SIGINT",
_ = term.recv() => "SIGTERM",
};
tracing::info!(target: "rustango::shutdown", signal = signal_name, "shutting down");
}
#[cfg(not(unix))]
{
let _ = tokio::signal::ctrl_c().await;
tracing::info!(target: "rustango::shutdown", signal = "SIGINT", "shutting down");
}
}
/// How long open connections get after the stop signal, when
/// `server.shutdown_timeout_secs` is unset (#1883).
pub const DEFAULT_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20);
/// The signal a graceful server stops accepting on.
pub type StopSignal = std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>;
/// Serve `app` until [`shutdown_signal`], then give open connections
/// `drain` to finish. SSE and long-poll never finish on their own, so
/// without a deadline SIGTERM waited for SIGKILL (#1883). This is the
/// crate's only `axum::serve`, so no serve path can skip the drain (#1948).
///
/// ```ignore
/// serve_until_drained(listener, app, drain).await?;
/// ```
///
/// # Errors
/// What the server returns.
#[cfg(feature = "_axum")]
pub async fn serve_until_drained(
listener: tokio::net::TcpListener,
app: axum::Router,
drain: std::time::Duration,
) -> std::io::Result<()> {
drain_on(Box::pin(shutdown_signal()), listener, app, drain).await
}
#[cfg(feature = "_axum")]
async fn drain_on(
signal: StopSignal,
listener: tokio::net::TcpListener,
app: axum::Router,
drain: std::time::Duration,
) -> std::io::Result<()> {
let (stopped, on_stop) = tokio::sync::oneshot::channel::<()>();
// `ConnectInfo` feeds `access_log`, IP filters and per-IP limits.
let serve = std::future::IntoFuture::into_future(
axum::serve(
listener,
app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
)
.with_graceful_shutdown(async move {
signal.await;
let _ = stopped.send(());
}),
);
tokio::pin!(serve);
tokio::select! {
r = &mut serve => r,
() = async {
// A dropped sender means the server ended first; that arm wins.
if on_stop.await.is_err() {
std::future::pending::<()>().await;
}
tokio::time::sleep(drain).await;
} => {
tracing::warn!(
target: "rustango::shutdown",
drain_secs = drain.as_secs(),
"drain deadline passed; closing the connections still open",
);
Ok(())
}
}
}
#[cfg(all(test, feature = "_axum"))]
mod tests {
use super::*;
/// A request that never ends (long-poll, SSE) must not hold
/// shutdown past the deadline (#1883).
#[tokio::test]
async fn an_endless_request_does_not_outlive_the_drain_deadline() {
use tokio::io::AsyncWriteExt as _;
let app = axum::Router::new().route(
"/poll",
axum::routing::get(|| async {
std::future::pending::<()>().await;
""
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (stop, on_stop) = tokio::sync::oneshot::channel::<()>();
let server = tokio::spawn(drain_on(
Box::pin(async move {
let _ = on_stop.await;
}),
listener,
app,
std::time::Duration::from_millis(200),
));
let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
client
.write_all(b"GET /poll HTTP/1.1\r\nhost: x\r\n\r\n")
.await
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
stop.send(()).unwrap();
let done = tokio::time::timeout(std::time::Duration::from_secs(3), server).await;
assert!(
done.is_ok(),
"shutdown waited on the open request past its deadline"
);
drop(client);
}
}