Quelle  graceful.rs   Sprache: Rust

//! Utility to gracefully shutdown a server.
//!
//! This module provides a [`GracefulShutdown`] type,
//! which can be used to gracefully shutdown a server.
//!
//! See <https://github.com/hyperium/hyper-util/blob/master/examples/server_graceful.rs>
//! for an example of how to use this.

use std::{
    fmt::{self, Debug},
    future::Future,
    pin::Pin,
    task::{self, Poll},
};

use pin_project_lite::pin_project;
use tokio::sync::watch;

/// A graceful shutdown utility
// Purposefully not `Clone`, see `watcher()` method for why.
pub struct GracefulShutdown {
    tx: watch::Sender<()>,
}

/// A watcher side of the graceful shutdown.
///
/// This type can only watch a connection, it cannot trigger a shutdown.
///
/// Call [`GracefulShutdown::watcher()`] to construct one of these.
pub struct Watcher {
    rx: watch::Receiver<()>,
}

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> {
        self.watcher().watch(conn)
    }

    /// Create an owned type that can watch a connection.
    ///
    /// This method allows created an owned type that can be sent onto another
    /// task before calling [`Watcher::watch()`].
    // Internal: this function exists because `Clone` allows footguns.
    // If the `tx` were cloned (or the `rx`), race conditions can happens where
    // one task starting a shutdown is scheduled and interwined with a task
    // starting to watch a connection, and the "watch version" is one behind.
    pub fn watcher(&self) -> Watcher {
        let rx = self.tx.subscribe();
        Watcher { 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;
    }

    /// Returns the number of the watching connections.
    pub fn count(&self) -> usize {
        self.tx.receiver_count()
    }
}

impl Debug for GracefulShutdown {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("GracefulShutdown").finish()
    }
}

impl Default for GracefulShutdown {
    fn default() -> Self {
        Self::new()
    }
}

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

impl Debug for Watcher {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("GracefulWatcher").finish()
    }
}

pin_project! {
    struct GracefulConnectionFuture<C, F: Future> {
        #[pin]
        conn: C,
        #[pin]
        cancel: F,
        #[pin]
        // If cancelled, this is held until the inner conn is done.
        cancelled_guard: Option<F::Output>,
    }
}

impl<C, F: Future> GracefulConnectionFuture<C, F> {
    fn new(conn: C, cancel: F) -> Self {
        Self {
            conn,
            cancel,
            cancelled_guard: None,
        }
    }
}

impl<C, F: Future> Debug for GracefulConnectionFuture<C, F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("GracefulConnectionFuture").finish()
    }
}

impl<C, F> Future for GracefulConnectionFuture<C, F>
where
    C: GracefulConnection,
    F: Future,
{
    type Output = C::Output;

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
        let mut this = self.project();
        if this.cancelled_guard.is_none() {
            if let Poll::Ready(guard) = this.cancel.poll(cx) {
                this.cancelled_guard.set(Some(guard));
                this.conn.as_mut().graceful_shutdown();
            }
        }
        this.conn.poll(cx)
    }
}

/// An internal utility trait as an umbrella target for all (hyper) connection
/// types that the [`GracefulShutdown`] can watch.
pub trait GracefulConnection: Future<Output = Result<(), Self::Error>> + private::Sealed {
    /// The error type returned by the connection when used as a future.
    type Error;

    /// Start a graceful shutdown process for this connection.
    fn graceful_shutdown(self: Pin<&mut Self>);
}

#[cfg(feature = "http1")]
impl<I, B, S> GracefulConnection for hyper::server::conn::http1::Connection<I, S>
where
    S: hyper::service::HttpService<hyper::body::Incoming, ResBody = B>,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    I: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
    B: hyper::body::Body + 'static,
    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
{
    type Error = hyper::Error;

    fn graceful_shutdown(self: Pin<&mut Self>) {
        hyper::server::conn::http1::Connection::graceful_shutdown(self);
    }
}

#[cfg(feature = "http2")]
impl<I, B, S, E> GracefulConnection for hyper::server::conn::http2::Connection<I, S, E>
where
    S: hyper::service::HttpService<hyper::body::Incoming, ResBody = B>,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    I: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
    B: hyper::body::Body + 'static,
    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    E: hyper::rt::bounds::Http2ServerConnExec<S::Future, B>,
{
    type Error = hyper::Error;

    fn graceful_shutdown(self: Pin<&mut Self>) {
        hyper::server::conn::http2::Connection::graceful_shutdown(self);
    }
}

#[cfg(feature = "server-auto")]
impl<I, B, S, E> GracefulConnection for crate::server::conn::auto::Connection<'_, I, S, E>
where
    S: hyper::service::Service<http::Request<hyper::body::Incoming>, Response = http::Response<B>>,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    S::Future: 'static,
    I: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
    B: hyper::body::Body + 'static,
    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    E: hyper::rt::bounds::Http2ServerConnExec<S::Future, B>,
{
    type Error = Box<dyn std::error::Error + Send + Sync>;

    fn graceful_shutdown(self: Pin<&mut Self>) {
        crate::server::conn::auto::Connection::graceful_shutdown(self);
    }
}

#[cfg(feature = "server-auto")]
impl<I, B, S, E> GracefulConnection
    for crate::server::conn::auto::UpgradeableConnection<'_, I, S, E>
where
    S: hyper::service::Service<http::Request<hyper::body::Incoming>, Response = http::Response<B>>,
    S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    S::Future: 'static,
    I: hyper::rt::Read + hyper::rt::Write + Unpin + Send + 'static,
    B: hyper::body::Body + 'static,
    B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    E: hyper::rt::bounds::Http2ServerConnExec<S::Future, B>,
{
    type Error = Box<dyn std::error::Error + Send + Sync>;

    fn graceful_shutdown(self: Pin<&mut Self>) {
        crate::server::conn::auto::UpgradeableConnection::graceful_shutdown(self);
    }
}

mod private {
    pub trait Sealed {}

    #[cfg(feature = "http1")]
    impl<I, B, S> Sealed for hyper::server::conn::http1::Connection<I, S>
    where
        S: hyper::service::HttpService<hyper::body::Incoming, ResBody = B>,
        S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        I: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
        B: hyper::body::Body + 'static,
        B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    {
    }

    #[cfg(feature = "http1")]
    impl<I, B, S> Sealed for hyper::server::conn::http1::UpgradeableConnection<I, S>
    where
        S: hyper::service::HttpService<hyper::body::Incoming, ResBody = B>,
        S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        I: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
        B: hyper::body::Body + 'static,
        B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
    {
    }

    #[cfg(feature = "http2")]
    impl<I, B, S, E> Sealed for hyper::server::conn::http2::Connection<I, S, E>
    where
        S: hyper::service::HttpService<hyper::body::Incoming, ResBody = B>,
        S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        I: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
        B: hyper::body::Body + 'static,
        B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        E: hyper::rt::bounds::Http2ServerConnExec<S::Future, B>,
    {
    }

    #[cfg(feature = "server-auto")]
    impl<I, B, S, E> Sealed for crate::server::conn::auto::Connection<'_, I, S, E>
    where
        S: hyper::service::Service<
            http::Request<hyper::body::Incoming>,
            Response = http::Response<B>,
        >,
        S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        S::Future: 'static,
        I: hyper::rt::Read + hyper::rt::Write + Unpin + 'static,
        B: hyper::body::Body + 'static,
        B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        E: hyper::rt::bounds::Http2ServerConnExec<S::Future, B>,
    {
    }

    #[cfg(feature = "server-auto")]
    impl<I, B, S, E> Sealed for crate::server::conn::auto::UpgradeableConnection<'_, I, S, E>
    where
        S: hyper::service::Service<
            http::Request<hyper::body::Incoming>,
            Response = http::Response<B>,
        >,
        S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        S::Future: 'static,
        I: hyper::rt::Read + hyper::rt::Write + Unpin + Send + 'static,
        B: hyper::body::Body + 'static,
        B::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
        E: hyper::rt::bounds::Http2ServerConnExec<S::Future, B>,
    {
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use pin_project_lite::pin_project;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;

    pin_project! {
        #[derive(Debug)]
        struct DummyConnection<F> {
            #[pin]
            future: F,
            shutdown_counter: Arc<AtomicUsize>,
        }
    }

    impl<F> private::Sealed for DummyConnection<F> {}

    impl<F: Future> GracefulConnection for DummyConnection<F> {
        type Error = ();

        fn graceful_shutdown(self: Pin<&mut Self>) {
            self.shutdown_counter.fetch_add(1, Ordering::SeqCst);
        }
    }

    impl<F: Future> Future for DummyConnection<F> {
        type Output = Result<(), ()>;

        fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
            match self.project().future.poll(cx) {
                Poll::Ready(_) => Poll::Ready(Ok(())),
                Poll::Pending => Poll::Pending,
            }
        }
    }

    #[cfg(not(miri))]
    #[tokio::test]
    async fn test_graceful_shutdown_ok() {
        let graceful = GracefulShutdown::new();
        let shutdown_counter = Arc::new(AtomicUsize::new(0));
        let (dummy_tx, _) = tokio::sync::broadcast::channel(1);

        for i in 1..=3 {
            let mut dummy_rx = dummy_tx.subscribe();
            let shutdown_counter = shutdown_counter.clone();

            let future = async move {
                tokio::time::sleep(std::time::Duration::from_millis(i * 10)).await;
                let _ = dummy_rx.recv().await;
            };
            let dummy_conn = DummyConnection {
                future,
                shutdown_counter,
            };
            let conn = graceful.watch(dummy_conn);
            tokio::spawn(async move {
                conn.await.unwrap();
            });
        }

        assert_eq!(shutdown_counter.load(Ordering::SeqCst), 0);
        let _ = dummy_tx.send(());

        tokio::select! {
            _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
                panic!("timeout")
            },
            _ = graceful.shutdown() => {
                assert_eq!(shutdown_counter.load(Ordering::SeqCst), 3);
            }
        }
    }

    #[cfg(not(miri))]
    #[tokio::test]
    async fn test_graceful_shutdown_delayed_ok() {
        let graceful = GracefulShutdown::new();
        let shutdown_counter = Arc::new(AtomicUsize::new(0));

        for i in 1..=3 {
            let shutdown_counter = shutdown_counter.clone();

            //tokio::time::sleep(std::time::Duration::from_millis(i * 5)).await;
            let future = async move {
                tokio::time::sleep(std::time::Duration::from_millis(i * 50)).await;
            };
            let dummy_conn = DummyConnection {
                future,
                shutdown_counter,
            };
            let conn = graceful.watch(dummy_conn);
            tokio::spawn(async move {
                conn.await.unwrap();
            });
        }

        assert_eq!(shutdown_counter.load(Ordering::SeqCst), 0);

        tokio::select! {
            _ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {
                panic!("timeout")
            },
            _ = graceful.shutdown() => {
                assert_eq!(shutdown_counter.load(Ordering::SeqCst), 3);
            }
        }
    }

    #[cfg(not(miri))]
    #[tokio::test]
    async fn test_graceful_shutdown_multi_per_watcher_ok() {
        let graceful = GracefulShutdown::new();
        let shutdown_counter = Arc::new(AtomicUsize::new(0));

        for i in 1..=3 {
            let shutdown_counter = shutdown_counter.clone();

            let mut futures = Vec::new();
            for u in 1..=i {
                let future = tokio::time::sleep(std::time::Duration::from_millis(u * 50));
                let dummy_conn = DummyConnection {
                    future,
                    shutdown_counter: shutdown_counter.clone(),
                };
                let conn = graceful.watch(dummy_conn);
                futures.push(conn);
            }
            tokio::spawn(async move {
                futures_util::future::join_all(futures).await;
            });
        }

        assert_eq!(shutdown_counter.load(Ordering::SeqCst), 0);

        tokio::select! {
            _ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {
                panic!("timeout")
            },
            _ = graceful.shutdown() => {
                assert_eq!(shutdown_counter.load(Ordering::SeqCst), 6);
            }
        }
    }

    #[cfg(not(miri))]
    #[tokio::test]
    async fn test_graceful_shutdown_timeout() {
        let graceful = GracefulShutdown::new();
        let shutdown_counter = Arc::new(AtomicUsize::new(0));

        for i in 1..=3 {
            let shutdown_counter = shutdown_counter.clone();

            let future = async move {
                if i == 1 {
                    std::future::pending::<()>().await
                } else {
                    std::future::ready(()).await
                }
            };
            let dummy_conn = DummyConnection {
                future,
                shutdown_counter,
            };
            let conn = graceful.watch(dummy_conn);
            tokio::spawn(async move {
                conn.await.unwrap();
            });
        }

        assert_eq!(shutdown_counter.load(Ordering::SeqCst), 0);

        tokio::select! {
            _ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
                assert_eq!(shutdown_counter.load(Ordering::SeqCst), 3);
            },
            _ = graceful.shutdown() => {
                panic!("shutdown should not be completed: as not all our conns finish")
            }
        }
    }
}