From 5028d97efdad07b5b95af89b0727206e8f2af572 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Sun, 31 Dec 2023 11:51:48 +0100
Subject: [PATCH] Remove duplication in serving with and without graceful
 shutdown

---
 axum/src/serve.rs | 55 +++++++----------------------------------------
 1 file changed, 8 insertions(+), 47 deletions(-)

diff --git a/axum/src/serve.rs b/axum/src/serve.rs
index 9850af27..8808ffc5 100644
--- a/axum/src/serve.rs
+++ b/axum/src/serve.rs
@@ -181,54 +181,15 @@ where
     type IntoFuture = private::ServeFuture;
 
     fn into_future(self) -> Self::IntoFuture {
-        private::ServeFuture(Box::pin(async move {
-            let Self {
-                tcp_listener,
-                mut make_service,
-                _marker: _,
-            } = self;
+        let Self {
+            tcp_listener,
+            make_service,
+            _marker: _,
+        } = self;
 
-            loop {
-                let (tcp_stream, remote_addr) = match tcp_accept(&tcp_listener).await {
-                    Some(conn) => conn,
-                    None => continue,
-                };
-                let tcp_stream = TokioIo::new(tcp_stream);
-
-                poll_fn(|cx| make_service.poll_ready(cx))
-                    .await
-                    .unwrap_or_else(|err| match err {});
-
-                let tower_service = make_service
-                    .call(IncomingStream {
-                        tcp_stream: &tcp_stream,
-                        remote_addr,
-                    })
-                    .await
-                    .unwrap_or_else(|err| match err {});
-
-                let hyper_service = TowerToHyperService {
-                    service: tower_service,
-                };
-
-                tokio::spawn(async move {
-                    match Builder::new(TokioExecutor::new())
-                        // upgrades needed for websockets
-                        .serve_connection_with_upgrades(tcp_stream, hyper_service)
-                        .await
-                    {
-                        Ok(()) => {}
-                        Err(_err) => {
-                            // This error only appears when the client doesn't send a request and
-                            // terminate the connection.
-                            //
-                            // If client sends one request then terminate connection whenever, it doesn't
-                            // appear.
-                        }
-                    }
-                });
-            }
-        }))
+        serve(tcp_listener, make_service)
+            .with_graceful_shutdown(std::future::pending())
+            .into_future()
     }
 }