Repository: poem-web/poem (language: rust)

Issue:
Graceful shutdown for established connections
## Expected Behavior

`Poem` should trigger hyper's [graceful showdown method for connection](https://docs.rs/hyper/latest/hyper/server/conn/struct.Connection.html#method.graceful_shutdown) as soon as signal is received.
It handles correct shutdown: `GOAWAY` for `http/2` or disabling `keep-alive` for `http/1`

There is an example in hyper - [link](https://github.com/hyperium/hyper/blob/95a153bbc2717bd4233486e09848622ceb900574/examples/upgrades.rs#L168) (not exhaustive though)

## Actual Behavior

Request can be sent over already established connection after initiating graceful shutdown.
Currently [connection is just awaited ](https://github.com/poem-web/poem/blob/10de9d11daceed3def451999ea8496b636f7d1b2/poem/src/server.rs#L196) without any logic for that.

## Steps to Reproduce the Problem

Just reuse the established connection  for new request after graceful shutdown is initiated.


## Specifications

- Version: up to 1.35 (current main)
- Platform: -
- Subsystem: -


Discussion / hints:
I'll create a PR to address it

Fix patch:
diff --git a/poem/src/server.rs b/poem/src/server.rs
index 310cb0bee4..80db9ed980 100644
--- a/poem/src/server.rs
+++ b/poem/src/server.rs
@@ -14,6 +14,7 @@ use tokio::{
     sync::Notify,
     time::Duration,
 };
+use tokio_util::sync::CancellationToken;
 
 use crate::{
     listener::{Acceptor, AcceptorExt, Listener},
@@ -94,7 +95,8 @@ where
         let name = name.as_deref();
         let alive_connections = Arc::new(AtomicUsize::new(0));
         let notify = Arc::new(Notify::new());
-        let timeout_notify = Arc::new(Notify::new());
+        let timeout_token = CancellationToken::new();
+        let conn_shutdown_token = CancellationToken::new();
 
         let mut acceptor = match listener {
             Either::Listener(listener) => listener.into_acceptor().await?.boxed(),
@@ -111,6 +113,7 @@ where
         loop {
             tokio::select! {
                 _ = &mut signal => {
+                    conn_shutdown_token.cancel();
                     if let Some(timeout) = timeout {
                         tracing::info!(
                             name = name,
@@ -118,10 +121,10 @@ where
                             "initiate graceful shutdown",
                         );
 
-                        let timeout_notify = timeout_notify.clone();
+                        let timeout_token = timeout_token.clone();
                         tokio::spawn(async move {
                             tokio::time::sleep(timeout).await;
-                            timeout_notify.notify_waiters();
+                            timeout_token.cancel();
                         });
                     } else {
                         tracing::info!(name = name, "initiate graceful shutdown");
@@ -135,20 +138,24 @@ where
                         let ep = ep.clone();
                         let alive_connections = alive_connections.clone();
                         let notify = notify.clone();
-                        let timeout_notify = timeout_notify.clone();
+                        let timeout_token = timeout_token.clone();
+                        let conn_shutdown_token = conn_shutdown_token.clone();
 
                         tokio::spawn(async move {
+                            let serve_connection = serve_connection(socket, local_addr, remote_addr, scheme, ep, conn_shutdown_token);
+
                             if timeout.is_some() {
                                 tokio::select! {
-                                    _ = serve_connection(socket, local_addr, remote_addr, scheme, ep) => {}
-                                    _ = timeout_notify.notified() => {}
+                                    _ = serve_connection => {}
+                                    _ = timeout_token.cancelled() => {}
                                 }
                             } else {
-                                serve_connection(socket, local_addr, remote_addr, scheme, ep).await;
+                               serve_connection.await;
                             }
 
                             if alive_connections.fetch_sub(1, Ordering::Acquire) == 1 {
-                                notify.notify_one();
+                                // We have to notify only if there is a registered waiter on shutdown
+                                notify.notify_waiters();
                             }
                         });
                     }
@@ -173,6 +180,7 @@ async fn serve_connection(
     remote_addr: RemoteAddr,
     scheme: Scheme,
     ep: Arc<dyn Endpoint<Output = Response>>,
+    conn_shutdown_token: CancellationToken,
 ) {
     let service = hyper::service::service_fn({
         move |req: hyper::Request<hyper::Body>| {
@@ -190,8 +198,22 @@ async fn serve_connection(
         }
     });
 
-    let conn = Http::new()
+    let mut conn = Http::new()
         .serve_connection(socket, service)
         .with_upgrades();
+
+    tokio::select! {
+        _ = &mut conn => {
+            // Connection completed successfully.
+            return;
+        },
+        _ = conn_shutdown_token.cancelled() => {
+            // Init graceful shutdown for connection (`GOAWAY` for `HTTP/2` or disabling `keep-alive` for `HTTP/1`)
+            let conn = Pin::new(&mut conn);
+            conn.graceful_shutdown();
+        }
+    };
+
+    // Continue awaiting after graceful-shutdown is initiated to handle existed requests.
     let _ = conn.await;
 }


Test patch: