Graceful shutdown doesn't include upgraded connections

[jeremycline]

Hi,

When using Axum with WebSockets, I noticed having an open WebSocket doesn't stop the server from shutting down gracefully. Here's a demo repository.

My expectation was that the demo would remain running for 30 seconds, printing "beep" every second until the graceful shutdown timeout is reached.

I'm not terribly familiar with Hyper so I don't immediately see a fix, so I figured I'd start with an issue. I'm happy to work on a fix, however.

[programatik29]

The reason is hypers http future is terminated after upgrade response.

Take a look at serve_connection. This future terminates even if there is an upgraded connection.

[jeremycline]

FWIW the way I solved this (and I'm embarrassed it took me a while to think of) is to just stash a channel handle in my shared state structure and each WebSocket takes a clone, just like the graceful shutdown tutorial for tokio. After the server shuts down the state gets dropped so it'll hang out until the WebSockets are done.