feat(server): implement Clone for GracefulShutdown

[tottoto]

Implements Clone for GracefulShutdown. This allows users to use GracefulShutdown in other functions.

[seanmonstar]

in other functions.

Not that I'm against adding this, but knowing the use cases would help.

There's a part of me that worries if having multiple clones could be confusing when shutdown() consumes self...

[tottoto]

hyperium/tonic#1788 is a use case of this change.

[zacknewman]

@seanmonstar, when you need to move GracefulShutdown into a separate task before calling watch, you can't since it doesn't implement Clone. A trivial example is handling TLS connections. The time it takes to accept a TLS connection can take too long to do outside of a dedicated task; thus it should be done in a separate task.

Personally the amount of code it takes to do your own graceful shutdown is very little, so I will likely keep using Sender<()> and Receiver<()> directly since I get to avoid a clone on Sender<()> altogether since I can call subscribe before accepting the TLS connection; however seeing how GracefulShutdown exists, others may be willing to pay the price of an extra clone.

async fn main_loop<S>(
    listener: tokio::net::TcpListener,
    acceptor: tokio_rustls::TlsAcceptor,
    app: S,
    server: hyper::server::conn::http2::Builder<hyper_util::rt::TokioExecutor>,
    sender: &hyper_util::server::graceful::GracefulShutdown,
) -> core::convert::Infallible
where
    S: hyper::service::Service<
            axum::extract::Request<hyper::body::Incoming>,
            Response = axum::response::Response<axum::body::Body>,
            Future = axum::routing::future::RouteFuture<core::convert::Infallible>,
            Error = core::convert::Infallible,
        > + core::clone::Clone
        + core::marker::Send
        + 'static,
{
    loop {
        if let Ok(val) = listener.accept().await {
            // This can take too long; thus should be done inside `tokio::spawn`;
            // however we can't since that `move`s `sender`.
            if let Ok(io) = acceptor.accept(val.0).await {
                let con = sender.watch(server.serve_connection(hyper_util::rt::TokioIo::new(io), app.clone()));
                tokio::spawn(async move {
                    let _r: Result<(), hyper::Error> = con.await;
                });
            }
        }
    }
}

If GracefulShutdown implemented Clone, however, then above could be transformed into below:

async fn main_loop<S>(
    listener: tokio::net::TcpListener,
    acceptor: tokio_rustls::TlsAcceptor,
    app: S,
    server: hyper::server::conn::http2::Builder<hyper_util::rt::TokioExecutor>,
    sender: hyper_util::server::graceful::GracefulShutdown,
) -> core::convert::Infallible
where
    S: hyper::service::Service<
            axum::extract::Request<hyper::body::Incoming>,
            Response = axum::response::Response<axum::body::Body>,
            Future = axum::routing::future::RouteFuture<core::convert::Infallible>,
            Error = core::convert::Infallible,
        > + core::clone::Clone
        + core::marker::Send
        + 'static,
{
    loop {
        if let Ok(val) = listener.accept().await {
            let tls = acceptor.clone();
            let http = app.clone();
            let builder = server.clone();
            let shutdown = sender.clone();
            tokio::spawn(async move {
                if let Ok(io) = tls.accept(val.0).await {
                    let _r: Result<(), hyper::Error> = shutdown.watch(builder.serve_connection(hyper_util::rt::TokioIo::new(io), http)).await;
                }
            });
        }
    }
}

"Hand-rolled" graceful shutdown:

struct ShutdownConnection<S, F>
where
    S: hyper::service::Service<
        axum::extract::Request<hyper::body::Incoming>,
        Response = axum::response::Response<axum::body::Body>,
        Future = axum::routing::future::RouteFuture<core::convert::Infallible>,
        Error = core::convert::Infallible,
    >,
{
    con: hyper::server::conn::http2::Connection<
        hyper_util::rt::TokioIo<tokio_rustls::server::TlsStream<tokio::net::TcpStream>>,
        S,
        hyper_util::rt::TokioExecutor,
    >,
    signal: F,
    shutting_down: bool,
}
impl<S, F> ShutdownConnection<S, F>
where
    S: hyper::service::Service<
        axum::extract::Request<hyper::body::Incoming>,
        Response = axum::response::Response<axum::body::Body>,
        Future = axum::routing::future::RouteFuture<core::convert::Infallible>,
        Error = core::convert::Infallible,
    >,
{
    fn new(
        con: hyper::server::conn::http2::Connection<
            hyper_util::rt::TokioIo<tokio_rustls::server::TlsStream<tokio::net::TcpStream>>,
            S,
            hyper_util::rt::TokioExecutor,
        >,
        signal: F,
    ) -> Self {
        Self {
            con,
            signal,
            shutting_down: false,
        }
    }
}
impl<S, F> core::future::Future for ShutdownConnection<S, F>
where
    S: hyper::service::Service<
        axum::extract::Request<hyper::body::Incoming>,
        Response = axum::response::Response<axum::body::Body>,
        Future = axum::routing::future::RouteFuture<core::convert::Infallible>,
        Error = core::convert::Infallible,
    >,
    F: core::future::Future<Output = Result<(), tokio::sync::watch::error::RecvError>>
        + core::marker::Unpin,
{
    type Output = Result<(), hyper::Error>;
    fn poll(
        mut self: core::pin::Pin<&mut Self>,
        cx: &mut core::task::Context<'_>,
    ) -> core::task::Poll<Self::Output> {
        if !self.shutting_down && core::pin::Pin::new(&mut self.signal).poll(cx).is_ready() {
            self.shutting_down = true;
            core::pin::Pin::new(&mut self.con).graceful_shutdown();
        }
        core::pin::Pin::new(&mut self.con).poll(cx)
    }
}
async fn main_loop<S>(
    listener: tokio::net::TcpListener,
    acceptor: tokio_rustls::TlsAcceptor,
    app: S,
    server: hyper::server::conn::http2::Builder<hyper_util::rt::TokioExecutor>,
    sender: &tokio::sync::watch::Sender<()>,
) -> core::convert::Infallible
where
    S: hyper::service::Service<
            axum::extract::Request<hyper::body::Incoming>,
            Response = axum::response::Response<axum::body::Body>,
            Future = axum::routing::future::RouteFuture<core::convert::Infallible>,
            Error = core::convert::Infallible,
        > + core::clone::Clone
        + core::marker::Send
        + 'static,
{
    loop {
        if let Ok(val) = listener.accept().await {
            let tls = acceptor.clone();
            let http = app.clone();
            let builder = server.clone();
            let mut rx = sender.subscribe();
            tokio::spawn(async move {
                if let Ok(io) = tls.accept(val.0).await {
                    let _r: Result<(), hyper::Error> = ShutdownConnection::new(
                        builder.serve_connection(hyper_util::rt::TokioIo::new(io), http),
                        core::pin::pin!(rx.changed()),
                    )
                    .await;
                }
            });
        }
    }
}