From ca2c6b3bbaf1c54f0f30640eab4ea16cfbe90863 Mon Sep 17 00:00:00 2001
From: Finn Bear <finnbearone@gmail.com>
Date: Thu, 31 Mar 2022 07:22:06 -0700
Subject: [PATCH] Add AddrIncomingConfig (#38)

---
 examples/configure_addr_incoming.rs | 27 ++++++++++++++
 src/addr_incoming_config.rs         | 55 +++++++++++++++++++++++++++++
 src/lib.rs                          |  2 ++
 src/server.rs                       | 14 ++++++++
 4 files changed, 98 insertions(+)
 create mode 100644 examples/configure_addr_incoming.rs
 create mode 100644 src/addr_incoming_config.rs

diff --git a/examples/configure_addr_incoming.rs b/examples/configure_addr_incoming.rs
new file mode 100644
index 00000000..bc09229d
--- /dev/null
+++ b/examples/configure_addr_incoming.rs
@@ -0,0 +1,27 @@
+//! Run with `cargo run --example configure_http` command.
+//!
+//! To connect through browser, navigate to "http://localhost:3000" url.
+
+use axum::{routing::get, Router};
+use axum_server::AddrIncomingConfig;
+use std::net::SocketAddr;
+use std::time::Duration;
+
+#[tokio::main]
+async fn main() {
+    let app = Router::new().route("/", get(|| async { "Hello, world!" }));
+
+    let config = AddrIncomingConfig::new()
+        .tcp_nodelay(true)
+        .tcp_sleep_on_accept_errors(true)
+        .tcp_keepalive(Some(Duration::from_secs(32)))
+        .build();
+
+    let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
+    println!("listening on {}", addr);
+    axum_server::bind(addr)
+        .addr_incoming_config(config)
+        .serve(app.into_make_service())
+        .await
+        .unwrap();
+}
diff --git a/src/addr_incoming_config.rs b/src/addr_incoming_config.rs
new file mode 100644
index 00000000..d57741fe
--- /dev/null
+++ b/src/addr_incoming_config.rs
@@ -0,0 +1,55 @@
+use std::time::Duration;
+
+/// A configuration for [`AddrIncoming`].
+#[derive(Debug, Clone)]
+pub struct AddrIncomingConfig {
+    pub(crate) tcp_sleep_on_accept_errors: bool,
+    pub(crate) tcp_keepalive: Option<Duration>,
+    pub(crate) tcp_nodelay: bool,
+}
+
+impl Default for AddrIncomingConfig {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl AddrIncomingConfig {
+    /// Creates a default [`AddrIncoming`] config.
+    pub fn new() -> AddrIncomingConfig {
+        Self {
+            tcp_sleep_on_accept_errors: true,
+            tcp_keepalive: None,
+            tcp_nodelay: false,
+        }
+    }
+
+    /// Builds the config, creating an owned version of it.
+    pub fn build(&mut self) -> Self {
+        self.clone()
+    }
+
+    /// Set whether to sleep on accept errors, to avoid exhausting file descriptor limits.
+    ///
+    /// Default is `true`.
+    pub fn tcp_sleep_on_accept_errors(&mut self, val: bool) -> &mut Self {
+        self.tcp_sleep_on_accept_errors = val;
+        self
+    }
+
+    /// Set how often to send TCP keepalive probes.
+    ///
+    /// Default is `false`.
+    pub fn tcp_keepalive(&mut self, val: Option<Duration>) -> &mut Self {
+        self.tcp_keepalive = val;
+        self
+    }
+
+    /// Set the value of `TCP_NODELAY` option for accepted connections.
+    ///
+    /// Default is `false`.
+    pub fn tcp_nodelay(&mut self, val: bool) -> &mut Self {
+        self.tcp_nodelay = val;
+        self
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 2f248d0e..dd921ef9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -87,6 +87,7 @@
 )]
 #![cfg_attr(docsrs, feature(doc_cfg))]
 
+mod addr_incoming_config;
 mod handle;
 mod http_config;
 mod notify_once;
@@ -96,6 +97,7 @@ pub mod accept;
 pub mod service;
 
 pub use self::{
+    addr_incoming_config::AddrIncomingConfig,
     handle::Handle,
     http_config::HttpConfig,
     server::{bind, Server},
diff --git a/src/server.rs b/src/server.rs
index 2f44e0df..c374fa70 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,3 +1,4 @@
+use crate::addr_incoming_config::AddrIncomingConfig;
 use crate::{
     accept::{Accept, DefaultAcceptor},
     handle::Handle,
@@ -25,6 +26,7 @@ use tokio::{
 pub struct Server<A = DefaultAcceptor> {
     acceptor: A,
     addr: SocketAddr,
+    addr_incoming_conf: AddrIncomingConfig,
     handle: Handle,
     http_conf: HttpConfig,
 }
@@ -43,6 +45,7 @@ impl Server {
         Self {
             acceptor,
             addr,
+            addr_incoming_conf: AddrIncomingConfig::default(),
             handle,
             http_conf: HttpConfig::default(),
         }
@@ -55,6 +58,7 @@ impl<A> Server<A> {
         Server {
             acceptor,
             addr: self.addr,
+            addr_incoming_conf: self.addr_incoming_conf,
             handle: self.handle,
             http_conf: self.http_conf,
         }
@@ -72,6 +76,12 @@ impl<A> Server<A> {
         self
     }
 
+    /// Overwrite addr incoming configuration.
+    pub fn addr_incoming_config(mut self, config: AddrIncomingConfig) -> Self {
+        self.addr_incoming_conf = config;
+        self
+    }
+
     /// Serve provided [`MakeService`].
     ///
     /// # Errors
@@ -93,11 +103,15 @@ impl<A> Server<A> {
         A::Future: Send,
     {
         let acceptor = self.acceptor;
+        let addr_incoming_conf = self.addr_incoming_conf;
         let handle = self.handle;
         let http_conf = self.http_conf;
 
         let listener = TcpListener::bind(self.addr).await?;
         let mut incoming = AddrIncoming::from_listener(listener).map_err(io_other)?;
+        incoming.set_sleep_on_errors(addr_incoming_conf.tcp_sleep_on_accept_errors);
+        incoming.set_keepalive(addr_incoming_conf.tcp_keepalive);
+        incoming.set_nodelay(addr_incoming_conf.tcp_nodelay);
 
         handle.notify_listening(incoming.local_addr());