More generally useful doc example for .with_graceful shutdown()

[bittrance]

Motivation

The current doc example for .with_graceful_shutdown() does not quite show how shutdown/cancellation can be implemented. I think most readers will have a piece of code that looks like this and now wonders how to do shutdown.

let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
axum::serve(listener, app).await?;

In this scenario, the main issue is how to handle the fact that ::serve() blocks.

Solution

The easiest and most generally applicable way I can come up with to perform shutdown is to use tokio::sync::Notify to signal shutdown, like so:

let listener = tokio::net::TcpListener::bind("0.0.0.0:8082").await?;
let blocker = Arc::new(Notify::new());
let cancel = Arc::clone(&blocker);
tokio::spawn(
    axum::serve(listener, app)
        .with_graceful_shutdown(async move { blocker.notified().await })
        .into_future(),
);
// ...
cancel.notify_one();

I think this more generally useful pattern since the notify can be passed around and triggered where needed. In contrast, the select! strategy implied by the current doc is difficult to use in e.g. integration tests and where the api is a small component of a larger system.

[mladedav]

Simple blocker.notified() doesn't work here?

[bittrance]

Lifetime issue?

[mladedav]

Nit, can't we use just a oneshot channel for the example? As long as users don't need to see how to cancel from multiple places the whole Arc is unnecessary and it might be cleaner to have the semantics of being able to shutdown only once and having a dedicated side for producing and consuming the notification.

[bittrance]

The oneshot channel is fallible, so it has some noise too, but I still agree that it becomes a bit clearer that way. Updated.

Since rx yields Result<(), _> it would seem we can't just pass rx directly. Is there a way we could skip the async move {...} block and pass a future directly? It would seem we need to import futures::future::FutureExt to "eat" the Result? Would that still be preferable to using an async block? Are Axum users likely to have it in their Cargo.toml already?

[bittrance]

Ping @mladedav Follow-up questions to your comments.

[yanns]

I'm not sure if this example is more "useful|.

I'm using https://lib.rs/crates/vss and it's quite clear how to use it

#[tokio::main]
async fn main() {
axum::serve(tcp, app)
    .with_graceful_shutdown(shutdown_signal())
    .await
    .unwrap();
}

async fn shutdown_signal() {
    vss::shutdown_signal().await;
    info!("signal received, starting graceful shutdown");
}

The current example has this simplicity, not having to introduce oneshot. So it's quite useful for beginners who quickly want something that works.

Your example is still useful if someone wants to send a signal to axum to shutdown. And I agree that your example enables testing.
I really don't know how to choose here.