Created: August, 24, 2023 Last Modified: August, 25, 2023
Catching Termination and Ctrl C signals with Axum and axum_server
Axum (hyper) The easy one
The default built in server for Axum offers an convenience function "with_graceful_shutdown" that makes this easy. All we have to is provide a future that tells the server to stop when resolved. So all we have to implement is a function that awaits the Ctrl or termination signal, this is easy using tokio::signal.
async fn shutdown_signal() {
let ctrl_c = async {
signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
};
#[cfg(unix)]
let terminate = async {
signal::unix::signal(signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
tracing::info!("signal received, starting graceful shutdown");
}
Now we just call that function and pass what it returns to "with_graceful_shutdown"
Server::bind(&ip)
.serve(router.into_make_service())
.with_graceful_shutdown(shutdown_signal())
.await
.unwrap();
With that you now have a http server that will shut down with Ctrl C and when shutdown by Docker or by the system.
Graceful Shutdown down with axum_server and a HTTP to HTTPs redirect
A common way to enable HTTPs for a Axum based website is using axum_server. axum_server doesn't have a convenient with_graceful_shutdown function, the use of HTTPs also generally requires a http to https redirect. Code for a server like that might look something like this
/// ports order: http, https
async fn redirect_http_to_https<F>(ports: (u16, u16))
{
fn make_https(host: String, uri: Uri, ports: (u16, u16)) -> Result<Uri, BoxError> {
let mut parts = uri.into_parts();
parts.scheme = Some(axum::http::uri::Scheme::HTTPS);
if parts.path_and_query.is_none() {
parts.path_and_query = Some("/".parse().unwrap());
}
let https_host = host.replace(&ports.0.to_string(), &ports.1.to_string());
parts.authority = Some(https_host.parse()?);
Ok(Uri::from_parts(parts)?)
}
let redirect = move |Host(host): Host, uri: Uri| async move {
match make_https(host, uri, ports) {
Ok(uri) => Ok(Redirect::permanent(&uri.to_string())),
Err(error) => {
tracing::warn!(%error, "failed to convert URI to HTTPS");
Err(StatusCode::BAD_REQUEST)
}
}
};
let addr = SocketAddr::from(([0, 0, 0, 0], ports.0));
tracing::debug!("http redirect listening on {}", addr);
axum::Server::bind(&addr)
.serve(redirect.into_make_service())
.await
.unwrap();
}
pub async fn serve(
self,
key_path: impl AsRef<Path>,
cert_path: impl AsRef<Path>,
) -> Result<(), anyhow::Error> {
let http_ip = self.http_ip.clone();
let https_ip = self.https_ip.clone();
let https_router = self.https.clone();
let cert_path = PathBuf::from(cert_path.as_ref());
let key_path = PathBuf::from(key_path.as_ref());
let config = RustlsConfig::from_pem_file(cert_path, key_path)
.await
.expect("Error Loading Certs");
log::info!("Binding HTTP to {http_ip}");
//FIXME: get ports
tokio::spawn(redirect_http_to_https((self.http_port, self.https_port), shutdown_future));
log::info!("Binding HTTPs to {https_ip}");
axum_server::bind_rustls(self.https_ip.parse().unwrap(), config)
.serve(https_router.into_make_service())
.await
.unwrap();
Ok(())
}
In order to enable graceful shutdown a handle must be made and passed to the server, then when they server should be shutdown handle.shutdown should be called. Lets augment our shutdown_signal function to support this use case
async fn shutdown_signal(handle: Option<Handle>) {
let ctrl_c = async {
signal::ctrl_c()
.await
.expect("failed to install Ctrl+C handler");
};
#[cfg(unix)]
let terminate = async {
signal::unix::signal(signal::unix::SignalKind::terminate())
.expect("failed to install signal handler")
.recv()
.await;
};
#[cfg(not(unix))]
let terminate = std::future::pending::<()>();
tokio::select! {
_ = ctrl_c => {},
_ = terminate => {},
}
tracing::info!("Received termination signal shutting down");
if let Some(handle) = handle {
handle.graceful_shutdown(Some(Duration::from_secs(10))); // 10 secs is how long docker will wait
// to force shutdown
}
}
Now shutdown_signal can take in a optional handle (optional to make this compatible with a http only server still). With that handle the server can be signaled to shut down when the signals are caught. Now we need to modify our server setup code to use the shutdown signals. This involves changing our redirect function accept a future, and pass that future to the server it creates. As well as creating and passing a handle to the shutdown signal function and the main https server
/// ports order: http, https
async fn redirect_http_to_https(ports: (u16, u16), signal: impl Future<Output = ()>)
{
fn make_https(host: String, uri: Uri, ports: (u16, u16)) -> Result<Uri, BoxError> {
let mut parts = uri.into_parts();
parts.scheme = Some(axum::http::uri::Scheme::HTTPS);
if parts.path_and_query.is_none() {
parts.path_and_query = Some("/".parse().unwrap());
}
let https_host = host.replace(&ports.0.to_string(), &ports.1.to_string());
parts.authority = Some(https_host.parse()?);
Ok(Uri::from_parts(parts)?)
}
let redirect = move |Host(host): Host, uri: Uri| async move {
match make_https(host, uri, ports) {
Ok(uri) => Ok(Redirect::permanent(&uri.to_string())),
Err(error) => {
tracing::warn!(%error, "failed to convert URI to HTTPS");
Err(StatusCode::BAD_REQUEST)
}
}
};
let addr = SocketAddr::from(([0, 0, 0, 0], ports.0));
tracing::debug!("http redirect listening on {}", addr);
axum::Server::bind(&addr)
.serve(redirect.into_make_service())
.with_graceful_shutdown(signal)
.await
.unwrap();
}
#[tracing::instrument(level = "debug", skip(self, key_path, cert_path))]
pub async fn serve(
self,
key_path: impl AsRef<Path>,
cert_path: impl AsRef<Path>,
) -> Result<(), anyhow::Error> {
let http_ip = self.http_ip.clone();
let https_ip = self.https_ip.clone();
let https_router = self.https.clone();
let cert_path = PathBuf::from(cert_path.as_ref());
let key_path = PathBuf::from(key_path.as_ref());
let config = RustlsConfig::from_pem_file(cert_path, key_path)
.await
.expect("Error Loading Certs");
let handle = Handle::new();
let shutdown_future = shutdown_signal(Some(handle.clone()));
log::info!("Binding HTTP to {http_ip}");
//FIXME: get ports
tokio::spawn(redirect_http_to_https((self.http_port, self.https_port), shutdown_future));
log::info!("Binding HTTPs to {https_ip}");
axum_server::bind_rustls(self.https_ip.parse().unwrap(), config)
.handle(handle)
.serve(https_router.into_make_service())
.await
.unwrap();
Ok(())
}
With that change you now have a HTTPS/HTTP capable server that can catch and handle shutdown signals!