How to enable graceful shutdown for HTTPS requests?

[pingiun]

Hi, I'm trying to follow the examples for hyper-rustls and the guides and examples for hyper. Specifically I'm trying to combine this example from hyper-rustls and the graceful shutdown guide from hyper.

The problem I'm currently facing is that the rustls example accepts the TLS connection in a tokio::task::spawn, like this:

loop {
    let (tcp_stream, _remote_addr) = incoming.accept().await?;

    let tls_acceptor = tls_acceptor.clone();
    tokio::spawn(async move {
        let tls_stream = match tls_acceptor.accept(tcp_stream).await {
            Ok(tls_stream) => tls_stream,
            Err(err) => {
                eprintln!("failed to perform tls handshake: {err:#}");
                return;
            }
        };
        if let Err(err) = Builder::new(TokioExecutor::new())
            .serve_connection(TokioIo::new(tls_stream), service)
            .await
        {
            eprintln!("failed to serve connection: {err:#}");
        }
    });
}

So, I wanted to add a graceful watch to this connection, and I tried doing it like this:

let listener = TcpListener::bind(addr2).await.unwrap();
let graceful = hyper_util::server::graceful::GracefulShutdown::new();
loop {
    let tls_acceptor = tls_acceptor.clone();
    tokio::select! {
        Ok((tcp_stream, _addr)) = listener.accept() => {
            tokio::task::spawn(async move {
                let tls_stream = match tls_acceptor.accept(tcp_stream).into_fallible().await {
                    Ok(tls_stream) => tls_stream,
                    Err((err, mut io)) => {
                        eprintln!("failed to perform tls handshake: {err:?}");
                        let _ = io.write_all(PLAIN_TO_HTTPS_ERROR).await; // Ignore errors when sending the wrong protocl response
                        return;
                    }
                };
                let io = TokioIo::new(tls_stream);
                let builder = auto::Builder::new(TokioExecutor::new());
                let conn = builder.serve_connection(io, service_fn(index2));
                let fut = graceful.watch(conn.into_owned());

                if let Err(err) = fut.await
                {
                    info!("Error serving connection: {:?}", err);
                }
            });
        }

        Ok(msg) = info_receiver.recv() => {
            match msg {
                WorkerMessage::Shutdown => {
                    info!("Stopping https_l");
                    break;
                },
            }
        }
    }
}
tokio::select! {
    _ = graceful.shutdown() => {
        info!("all connections gracefully closed");
    },
    _ = tokio::time::sleep(std::time::Duration::from_secs(10)) => {
        error!("timed out wait for all connections to close");
    }
}

(The info_receiver is connected to a sender which sends after ctrl+c)

The problem here is that the graceful value cannot be moved into the tokio::spawn call, because then it cannot be used on the next loop iteration. This is the diagnostic I get from rustc:

error[E0382]: use of moved value: `graceful`
   --> src/main.rs:187:40
    |
182 |         let graceful = hyper_util::server::graceful::GracefulShutdown::new();
    |             -------- move occurs because `graceful` has type `GracefulShutdown`, which does not implement the `Copy` trait
183 |         loop {
    |         ---- inside of this loop
...
187 |                     tokio::task::spawn(async move {
    |                                        ^^^^^^^^^^ value moved here, in previous iteration of loop
...
199 |                         let fut = graceful.watch(conn.into_owned());
    |                                   -------- use occurs due to use in coroutine

For more information about this error, try `rustc --explain E0382`.

This can be fixed by moving the tokio::spawn around the connection serving only, but I'm worried that this will cause the server to not accept new requests while a TLS connection is being handshaked. So this will limit the throughput of connection handling.

Is there a good solution to this problem that I'm missing? I want to both support graceful stopping and HTTPS in my server!

[pingiun]

I've found a possible workaround, I copied graceful.rs to my own project and changed it like this:

impl GracefulShutdown {
    /// Create a new graceful shutdown helper.
    pub fn new() -> Self {
        let (tx, _) = watch::channel(());
        Self { tx }
    }

    /// Wrap a future for graceful shutdown watching.
    pub fn watch<C: GracefulConnection>(&self, conn: C) -> impl Future<Output = C::Output> {
        let mut rx = self.tx.subscribe();
        GracefulConnectionFuture::new(conn, async move {
            let _ = rx.changed().await;
            // hold onto the rx until the watched future is completed
            rx
        })
    }

    pub fn get_rx(&self) -> watch::Receiver<()> {
        self.tx.subscribe()
    }

    pub fn from_rx<C: GracefulConnection>(
        mut rx: watch::Receiver<()>,
        conn: C,
    ) -> impl Future<Output = C::Output> {
        GracefulConnectionFuture::new(conn, async move {
            let _ = rx.changed().await;
            // hold onto the rx until the watched future is completed
            rx
        })
    }

    /// Signal shutdown for all watched connections.
    ///
    /// This returns a `Future` which will complete once all watched
    /// connections have shutdown.
    pub async fn shutdown(self) {
        let Self { tx } = self;

        // signal all the watched futures about the change
        let _ = tx.send(());
        // and then wait for all of them to complete
        tx.closed().await;
    }
}

So this allows me to get an rx value out of the GracefulShutdown struct, which I can then move to my handler. The handler code now looks like this:

let listener = TcpListener::bind(addr2).await.unwrap();
let graceful = GracefulShutdown::new();
loop {
    let tls_acceptor = tls_acceptor.clone();
    tokio::select! {
        Ok((tcp_stream, _addr)) = listener.accept() => {
            let graceful_rx = graceful.get_rx();
            tokio::task::spawn(async move {
                let tls_stream = match tls_acceptor.accept(tcp_stream).into_fallible().await {
                    Ok(tls_stream) => tls_stream,
                    Err((err, mut io)) => {
                        eprintln!("failed to perform tls handshake: {err:?}");
                        let _ = io.write_all(PLAIN_TO_HTTPS_ERROR).await; // Ignore errors when sending the wrong protocl response
                        return;
                    }
                };
                let io = TokioIo::new(tls_stream);
                let builder = auto::Builder::new(TokioExecutor::new());
                let conn = builder.serve_connection(io, service_fn(index2));
                let fut = GracefulShutdown::from_rx(graceful_rx, conn.into_owned());

                if let Err(err) = fut.await
                {
                    info!("Error serving connection: {:?}", err);
                }
            });
        }

        Ok(msg) = info_receiver.recv() => {
            match msg {
                WorkerMessage::Shutdown => {
                    info!("Stopping https_l");
                    break;
                },
            }
        }
    }
}
tokio::select! {
    _ = graceful.shutdown() => {
        info!("all connections gracefully closed");
    },
    _ = tokio::time::sleep(std::time::Duration::from_secs(10)) => {
        error!("timed out wait for all connections to close");
    }
}

Maybe an idea to add an API like this? Or am I missing another way to do it?