From f48636916adb8c81700f1688e53018abbfd7ad32 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20Ml=C3=A1dek?= <david.mladek.cz@gmail.com>
Date: Thu, 7 Mar 2024 23:32:33 +0100
Subject: [PATCH] tests: add teest for graceful shutdown

---
 src/server/conn/auto.rs | 41 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/src/server/conn/auto.rs b/src/server/conn/auto.rs
index 9ddd4066..89004c4b 100644
--- a/src/server/conn/auto.rs
+++ b/src/server/conn/auto.rs
@@ -816,8 +816,11 @@ mod tests {
     use http_body::Body;
     use http_body_util::{BodyExt, Empty, Full};
     use hyper::{body, body::Bytes, client, service::service_fn};
-    use std::{convert::Infallible, error::Error as StdError, net::SocketAddr};
-    use tokio::net::{TcpListener, TcpStream};
+    use std::{convert::Infallible, error::Error as StdError, net::SocketAddr, time::Duration};
+    use tokio::{
+        net::{TcpListener, TcpStream},
+        pin,
+    };
 
     const BODY: &[u8] = b"Hello, world!";
 
@@ -871,6 +874,40 @@ mod tests {
         assert_eq!(body, BODY);
     }
 
+    #[cfg(not(miri))]
+    #[tokio::test]
+    async fn graceful_shutdown() {
+        let listener = TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
+            .await
+            .unwrap();
+
+        let listener_addr = listener.local_addr().unwrap();
+
+        // Spawn the task in background so that we can connect there
+        let listen_task = tokio::spawn(async move { listener.accept().await.unwrap() });
+        // Only connect a stream, do not send headers or anything
+        let _stream = TcpStream::connect(listener_addr).await.unwrap();
+
+        let (stream, _) = listen_task.await.unwrap();
+        let stream = TokioIo::new(stream);
+        let builder = auto::Builder::new(TokioExecutor::new());
+        let connection = builder.serve_connection(stream, service_fn(hello));
+
+        pin!(connection);
+
+        connection.as_mut().graceful_shutdown();
+
+        let connection_error = tokio::time::timeout(Duration::from_millis(200), connection)
+            .await
+            .expect("Connection should have finished in a timely manner after graceful shutdown.")
+            .expect_err("Connection should have been interrupted.");
+
+        let connection_error = connection_error
+            .downcast_ref::<std::io::Error>()
+            .expect("The error should have been `std::io::Error`.");
+        assert_eq!(connection_error.kind(), std::io::ErrorKind::Interrupted);
+    }
+
     async fn connect_h1<B>(addr: SocketAddr) -> client::conn::http1::SendRequest<B>
     where
         B: Body + Send + 'static,