From 3d73a55878af8aefa08c62a799862c6078054491 Mon Sep 17 00:00:00 2001
From: Liss Heidrich <31625940+Clueliss@users.noreply.github.com>
Date: Fri, 15 Mar 2024 11:09:12 +0100
Subject: [PATCH 1/4] implement Serve::tcp_nodelay

---
 axum/src/serve.rs | 81 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/axum/src/serve.rs b/axum/src/serve.rs
index b253cb5def..ae7693c62d 100644
--- a/axum/src/serve.rs
+++ b/axum/src/serve.rs
@@ -102,6 +102,7 @@ where
     Serve {
         tcp_listener,
         make_service,
+        tcp_nodelay: None,
         _marker: PhantomData,
     }
 }
@@ -112,6 +113,7 @@ where
 pub struct Serve<M, S> {
     tcp_listener: TcpListener,
     make_service: M,
+    tcp_nodelay: Option<bool>,
     _marker: PhantomData<S>,
 }
 
@@ -146,9 +148,35 @@ impl<M, S> Serve<M, S> {
             tcp_listener: self.tcp_listener,
             make_service: self.make_service,
             signal,
+            tcp_nodelay: self.tcp_nodelay,
             _marker: PhantomData,
         }
     }
+
+    /// Instructs the server to set the value of the `TCP_NODELAY` option on every accepted connection.
+    ///
+    /// See also [`TcpStream::set_nodelay`]
+    ///
+    /// # Example
+    /// ```
+    /// use axum::{Router, routing::get};
+    ///
+    /// # async {
+    /// let router = Router::new().route("/", get(|| async { "Hello, World!" }));
+    ///
+    /// let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
+    /// axum::serve(listener, router)
+    ///     .tcp_nodelay(true)
+    ///     .await
+    ///     .unwrap();
+    /// # };
+    /// ```
+    pub fn tcp_nodelay(self, nodelay: bool) -> Self {
+        Self {
+            tcp_nodelay: Some(nodelay),
+            ..self
+        }
+    }
 }
 
 #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
@@ -160,12 +188,14 @@ where
         let Self {
             tcp_listener,
             make_service,
+            tcp_nodelay,
             _marker: _,
         } = self;
 
         f.debug_struct("Serve")
             .field("tcp_listener", tcp_listener)
             .field("make_service", make_service)
+            .field("tcp_nodelay", tcp_nodelay)
             .finish()
     }
 }
@@ -186,6 +216,7 @@ where
             let Self {
                 tcp_listener,
                 mut make_service,
+                tcp_nodelay,
                 _marker: _,
             } = self;
 
@@ -194,6 +225,13 @@ where
                     Some(conn) => conn,
                     None => continue,
                 };
+
+                if let Some(nodelay) = tcp_nodelay {
+                    if let Err(err) = tcp_stream.set_nodelay(nodelay) {
+                        trace!("failed to set TCP_NODELAY on incoming connection: {err:#}");
+                    }
+                }
+
                 let tcp_stream = TokioIo::new(tcp_stream);
 
                 poll_fn(|cx| make_service.poll_ready(cx))
@@ -240,9 +278,42 @@ pub struct WithGracefulShutdown<M, S, F> {
     tcp_listener: TcpListener,
     make_service: M,
     signal: F,
+    tcp_nodelay: Option<bool>,
     _marker: PhantomData<S>,
 }
 
+impl<M, S, F> WithGracefulShutdown<M, S, F> {
+    /// Instructs the server to set the value of the `TCP_NODELAY` option on every accepted connection.
+    ///
+    /// See also [`TcpStream::set_nodelay`]
+    ///
+    /// # Example
+    /// ```
+    /// use axum::{Router, routing::get};
+    ///
+    /// # async {
+    /// let router = Router::new().route("/", get(|| async { "Hello, World!" }));
+    ///
+    /// let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
+    /// axum::serve(listener, router)
+    ///     .with_graceful_shutdown(shutdown_signal())
+    ///     .tcp_nodelay(true)
+    ///     .await
+    ///     .unwrap();
+    /// # };
+    ///
+    /// async fn shutdown_signal() {
+    ///     // ...
+    /// }
+    /// ```
+    pub fn tcp_nodelay(self, nodelay: bool) -> Self {
+        Self {
+            tcp_nodelay: Some(nodelay),
+            ..self
+        }
+    }
+}
+
 #[cfg(all(feature = "tokio", any(feature = "http1", feature = "http2")))]
 impl<M, S, F> Debug for WithGracefulShutdown<M, S, F>
 where
@@ -255,6 +326,7 @@ where
             tcp_listener,
             make_service,
             signal,
+            tcp_nodelay,
             _marker: _,
         } = self;
 
@@ -262,6 +334,7 @@ where
             .field("tcp_listener", tcp_listener)
             .field("make_service", make_service)
             .field("signal", signal)
+            .field("tcp_nodelay", tcp_nodelay)
             .finish()
     }
 }
@@ -283,6 +356,7 @@ where
             tcp_listener,
             mut make_service,
             signal,
+            tcp_nodelay,
             _marker: _,
         } = self;
 
@@ -310,6 +384,13 @@ where
                         break;
                     }
                 };
+
+                if let Some(nodelay) = tcp_nodelay {
+                    if let Err(err) = tcp_stream.set_nodelay(nodelay) {
+                        trace!("failed to set TCP_NODELAY on incoming connection: {err:#}");
+                    }
+                }
+
                 let tcp_stream = TokioIo::new(tcp_stream);
 
                 trace!("connection {remote_addr} accepted");

From dc8dfb727fd69293dd8096f84f1b43115ca44b70 Mon Sep 17 00:00:00 2001
From: Liss Heidrich <31625940+Clueliss@users.noreply.github.com>
Date: Fri, 15 Mar 2024 11:15:21 +0100
Subject: [PATCH 2/4] add compile test

---
 axum/src/serve.rs | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/axum/src/serve.rs b/axum/src/serve.rs
index ae7693c62d..6e33aaf0fd 100644
--- a/axum/src/serve.rs
+++ b/axum/src/serve.rs
@@ -639,6 +639,20 @@ mod tests {
             TcpListener::bind(addr).await.unwrap(),
             handler.into_make_service_with_connect_info::<SocketAddr>(),
         );
+
+        // nodelay
+        serve(
+            TcpListener::bind(addr).await.unwrap(),
+            handler.into_service(),
+        )
+        .tcp_nodelay(true);
+
+        serve(
+            TcpListener::bind(addr).await.unwrap(),
+            handler.into_service(),
+        )
+        .with_graceful_shutdown(async { /*...*/ })
+        .tcp_nodelay(true);
     }
 
     async fn handler() {}

From d127643c32a610ecc1bcd462b2a5770169b54f8e Mon Sep 17 00:00:00 2001
From: Liss Heidrich <31625940+Clueliss@users.noreply.github.com>
Date: Mon, 18 Mar 2024 07:51:00 +0100
Subject: [PATCH 3/4] fix comments

---
 axum/src/serve.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/axum/src/serve.rs b/axum/src/serve.rs
index 6e33aaf0fd..191f7055e7 100644
--- a/axum/src/serve.rs
+++ b/axum/src/serve.rs
@@ -155,7 +155,7 @@ impl<M, S> Serve<M, S> {
 
     /// Instructs the server to set the value of the `TCP_NODELAY` option on every accepted connection.
     ///
-    /// See also [`TcpStream::set_nodelay`]
+    /// See also [`TcpStream::set_nodelay`].
     ///
     /// # Example
     /// ```
@@ -285,7 +285,7 @@ pub struct WithGracefulShutdown<M, S, F> {
 impl<M, S, F> WithGracefulShutdown<M, S, F> {
     /// Instructs the server to set the value of the `TCP_NODELAY` option on every accepted connection.
     ///
-    /// See also [`TcpStream::set_nodelay`]
+    /// See also [`TcpStream::set_nodelay`].
     ///
     /// # Example
     /// ```

From ec04d258b26a856fd7305d1ae4ca732850afd61b Mon Sep 17 00:00:00 2001
From: Liss Heidrich <31625940+Clueliss@users.noreply.github.com>
Date: Mon, 18 Mar 2024 07:51:15 +0100
Subject: [PATCH 4/4] add entry to changelog

---
 axum/CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/axum/CHANGELOG.md b/axum/CHANGELOG.md
index 80ce95c379..b42ce03a94 100644
--- a/axum/CHANGELOG.md
+++ b/axum/CHANGELOG.md
@@ -7,9 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 # Unreleased
 
+- **added:** `axum::serve::Serve::tcp_nodelay` and `axum::serve::WithGracefulShutdown::tcp_nodelay` ([#2653])
 - **fixed:** Fixed layers being cloned when calling `axum::serve` directly with
   a `Router` or `MethodRouter` ([#2586])
 
+[#2653]: https://github.com/tokio-rs/axum/pull/2653
 [#2586]: https://github.com/tokio-rs/axum/pull/2586
 
 # 0.7.4 (13. January, 2024)