How do I properly implement shutdown endpoint in Axum?

[vi]

It should properly reply to the client, serve other ongoing requests (but not accept new requests), then exit from hyper::server::Builder::serve.

Implementing rejection of new requests seems relatively straightforward with an extension layer and Arc<AtomicBool>, but what about moving out from fn main?

Current kludge:

       tokio::spawn(async {
            tokio::time::sleep(Duration::from_millis(50)).await;
            exit(0);
        });

[vi]

Looks like https://docs.rs/hyper/0.14.20/hyper/server/struct.Server.html#method.with_graceful_shutdown is the answer, as hinted by Alexendoo on IRC.

[davidpdrsn]

Example here https://github.com/tokio-rs/axum/blob/main/examples/graceful-shutdown/src/main.rs

[simline]

But hyper has no bind_rustls function, so it is not a great way yet, any suggestions?

[elimirks]

It's a bit more gnarly with axum_server (for rustls), but this is how I'm doing it:

    let handle = axum_server::Handle::new();
    let shutdown_signal_fut = shutdown_signal();
    let server_fut = axum_server::bind_rustls(addr, tls_config)
        .handle(handle.clone())
        .serve(app.into_make_service());
    tokio::select! {
        () = shutdown_signal_fut =>
            handle.graceful_shutdown(Some(Duration::from_secs(30))),
        res = server_fut => res?,
    }
    info!("Server is stopping");

Plus the normal axum setup for the shutdown signal

// https://github.com/tokio-rs/axum/blob/main/examples/graceful-shutdown/src/main.rs
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 => {},
    }
    info!("Terminate signal received");
}