优雅停机

pub struct GracefulShutdown {
    tx: watch::Sender<()>,
}

    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;
    }

    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
        })
    }

    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)
    }

error[E0277]: the trait bound `UpgradeableConnection<TokioIo<TcpStream>, _>: GracefulConnection` is not satisfied
   --> src/main.rs:25:31
    |
 25 |   tokio::spawn(graceful.watch(conn));
    |                         ----- ^^^^ the trait `GracefulConnection` is not implemented for `UpgradeableConnection<TokioIo<TcpStream>, ServiceFn<_, _>>`
    |                         |
    |                         required by a bound introduced by this call

    // 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 }
    }

hyper = { version = "=1.12.0", features = ["client", "server", "http1", "http2"] }
hyper-util = { version = "=0.1.21", features = ["server-auto", "server-graceful", "tokio"] }
http-body-util = "=0.1.5"
bytes = "1"
futures = "0.3"
tokio = { version = "=1.54.1", features = ["rt-multi-thread", "macros", "net", "time", "sync"] }

use std::convert::Infallible;
use std::time::{Duration, Instant};

use bytes::Bytes;
use http_body_util::{combinators::BoxBody, BodyExt, Empty, Full, StreamBody};
use hyper::body::{Frame, Incoming};
use hyper::client::conn::{http1 as client_h1, http2 as client_h2};
use hyper::service::service_fn;
use hyper::{Request, Response};
use hyper_util::rt::{TokioExecutor, TokioIo};
use hyper_util::server::conn::auto;
use hyper_util::server::graceful::GracefulShutdown;
use tokio::net::{TcpListener, TcpStream};

type B = BoxBody<Bytes, Infallible>;

// milliseconds since shutdown() was called
static START: std::sync::OnceLock<Instant> = std::sync::OnceLock::new();
fn ms() -> u128 {
  START.get().map(|t| t.elapsed().as_millis()).unwrap_or(0)
}

async fn handler(req: Request<Incoming>) -> Result<Response<B>, Infallible> {
  match req.uri().path() {
    "/slow" => {
      tokio::time::sleep(Duration::from_millis(300)).await;
      Ok(Response::new(Full::new(Bytes::from("slow done")).boxed()))
    }
    "/stream" => {
      // never ends: one chunk every 100 ms, like SSE
      let s = futures::stream::unfold(0u32, |n| async move {
        tokio::time::sleep(Duration::from_millis(100)).await;
        Some((Ok::<_, Infallible>(Frame::data(Bytes::from(format!("tick {n}\n")))), n + 1))
      });
      Ok(Response::new(StreamBody::new(s).boxed()))
    }
    _ => Ok(Response::new(Full::new(Bytes::from("ok")).boxed())),
  }
}

#[tokio::main]
async fn main() {
  let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
  let addr = listener.local_addr().unwrap();
  let graceful = GracefulShutdown::new();
  let labels = ["idle h1", "busy h1", "idle h2", "silent", "stream h1"];
  let (accepted_tx, mut accepted_rx) = tokio::sync::mpsc::unbounded_channel::<()>();

  // accept exactly five connections, then stop accepting
  let accept = async {
    let builder = auto::Builder::new(TokioExecutor::new());
    for label in labels {
      let (stream, _) = listener.accept().await.unwrap();
      let conn = builder.serve_connection(TokioIo::new(stream), service_fn(handler));
      let conn = graceful.watch(conn.into_owned());
      tokio::spawn(async move {
        let r = conn.await.map_err(|e| e.to_string());
        println!("{:>5} ms  server: {label:<9} ended {r:?}", ms());
      });
      accepted_tx.send(()).unwrap();
    }
  };

  let clients = async {
    // idle h1: one request, then keep the connection open
    let io = TokioIo::new(TcpStream::connect(addr).await.unwrap());
    let (mut idle1, conn) = client_h1::handshake::<_, Empty<Bytes>>(io).await.unwrap();
    tokio::spawn(conn);
    idle1.send_request(Request::get("/").body(Empty::new()).unwrap()).await.unwrap();
    accepted_rx.recv().await;

    // busy h1: /slow is in flight when shutdown starts
    let io = TokioIo::new(TcpStream::connect(addr).await.unwrap());
    let (mut busy, conn) = client_h1::handshake::<_, Empty<Bytes>>(io).await.unwrap();
    tokio::spawn(conn);
    let slow = tokio::spawn(async move {
      let res = busy.send_request(Request::get("/slow").body(Empty::new()).unwrap()).await.unwrap();
      let conn_hdr = res.headers().get("connection").cloned();
      let body = res.into_body().collect().await.unwrap().to_bytes();
      println!("{:>5} ms  client: busy h1 got {body:?}, connection: {conn_hdr:?}", ms());
    });
    accepted_rx.recv().await;

    // idle h2
    let io = TokioIo::new(TcpStream::connect(addr).await.unwrap());
    let (mut idle2, conn) = client_h2::handshake::<_, _, Empty<Bytes>>(TokioExecutor::new(), io).await.unwrap();
    tokio::spawn(async move {
      let r = conn.await;
      println!("{:>5} ms  client: idle h2 conn ended {r:?}", ms());
    });
    idle2.send_request(Request::get("http://x/").body(Empty::new()).unwrap()).await.unwrap();
    accepted_rx.recv().await;

    // silent: TCP connected, not a single byte sent
    let silent = TcpStream::connect(addr).await.unwrap();
    accepted_rx.recv().await;

    // stream h1: response body never ends
    let io = TokioIo::new(TcpStream::connect(addr).await.unwrap());
    let (mut st, conn) = client_h1::handshake::<_, Empty<Bytes>>(io).await.unwrap();
    tokio::spawn(conn);
    let res = st.send_request(Request::get("/stream").body(Empty::new()).unwrap()).await.unwrap();
    accepted_rx.recv().await;
    let reader = tokio::spawn(async move {
      let mut body = res.into_body();
      let mut n = 0;
      while let Some(Ok(_)) = body.frame().await {
        n += 1;
      }
      println!("{:>5} ms  client: stream h1 body ended after {n} chunks", ms());
    });
    (idle1, idle2, silent, slow, reader)
  };

  let ((), keep) = tokio::join!(accept, clients);
  tokio::time::sleep(Duration::from_millis(50)).await;
  println!("count() = {}, calling shutdown()", graceful.count());
  START.set(Instant::now()).unwrap();
  match tokio::time::timeout(Duration::from_secs(1), graceful.shutdown()).await {
    Ok(()) => println!("{:>5} ms  shutdown() returned", ms()),
    Err(_) => println!("{:>5} ms  shutdown() still waiting after 1 s, giving up", ms()),
  }
  drop(keep);
}

count() = 5, calling shutdown()
    0 ms  server: silent    ended Err("Cancelled")
    0 ms  server: idle h1   ended Ok(())
    0 ms  client: idle h2 conn ended Ok(())
    0 ms  server: idle h2   ended Ok(())
  249 ms  server: busy h1   ended Ok(())
  249 ms  client: busy h1 got b"slow done", connection: Some("close")
 1002 ms  shutdown() still waiting after 1 s, giving up

    tokio::select! {
        _ = graceful.shutdown() => {
            eprintln!("Gracefully shutdown!");
        },
        _ = tokio::time::sleep(Duration::from_secs(10)) => {
            eprintln!("Waited 10 seconds for graceful shutdown, aborting...");
        }
    }

server conn ended after 0 ms: Ok(())
client conn ended: Ok(())