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?