From 00fb43b650b4af7780a44b07b404e5300d8e7b8e Mon Sep 17 00:00:00 2001
From: "Ngo Iok Ui (Wu Yu Wei)" <wusyong9104@gmail.com>
Date: Thu, 29 Oct 2020 23:23:01 +0800
Subject: [PATCH] Add tcp_keepalive option for ClientBuilder (#1070)

---
 src/async_impl/client.rs | 16 ++++++++++++++++
 src/connect.rs           | 12 ++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/async_impl/client.rs b/src/async_impl/client.rs
index 75610e6..ff103f7 100644
--- a/src/async_impl/client.rs
+++ b/src/async_impl/client.rs
@@ -80,6 +80,7 @@ struct Config {
     connection_verbose: bool,
     pool_idle_timeout: Option<Duration>,
     pool_max_idle_per_host: usize,
+    tcp_keepalive: Option<Duration>,
     #[cfg(feature = "__tls")]
     identity: Option<Identity>,
     proxies: Vec<Proxy>,
@@ -131,6 +132,7 @@ impl ClientBuilder {
                 connection_verbose: false,
                 pool_idle_timeout: Some(Duration::from_secs(90)),
                 pool_max_idle_per_host: std::usize::MAX,
+                tcp_keepalive: Some(Duration::from_secs(60)),
                 proxies: Vec::new(),
                 auto_sys_proxy: true,
                 redirect_policy: redirect::Policy::default(),
@@ -316,6 +318,7 @@ impl ClientBuilder {
 
         builder.pool_idle_timeout(config.pool_idle_timeout);
         builder.pool_max_idle_per_host(config.pool_max_idle_per_host);
+        connector.set_keepalive(config.tcp_keepalive);
 
         if config.http1_title_case_headers {
             builder.http1_title_case_headers(true);
@@ -712,6 +715,19 @@ impl ClientBuilder {
         self
     }
 
+    /// Set that all sockets have `SO_KEEPALIVE` set with the supplied duration.
+    ///
+    /// If `None`, the option will not be set.
+    ///
+    /// Default is 60 seconds.
+    pub fn tcp_keepalive<D>(mut self, val: D) -> ClientBuilder
+        where
+            D: Into<Option<Duration>>,
+    {
+        self.config.tcp_keepalive = val.into();
+        self
+    }
+
     // TLS options
 
     /// Add a custom root certificate.
diff --git a/src/connect.rs b/src/connect.rs
index f52a978..7d5202c 100644
--- a/src/connect.rs
+++ b/src/connect.rs
@@ -71,6 +71,7 @@ impl_http_connector! {
     fn set_local_address(&mut self, addr: Option<IpAddr>);
     fn enforce_http(&mut self, is_enforced: bool);
     fn set_nodelay(&mut self, nodelay: bool);
+    fn set_keepalive(&mut self, dur: Option<Duration>);
 }
 
 impl Service<Uri> for HttpConnector {
@@ -474,6 +475,17 @@ impl Connector {
 
         self.connect_with_maybe_proxy(proxy_dst, true).await
     }
+
+    pub fn set_keepalive(&mut self, dur: Option<Duration>) {
+        match &mut self.inner {
+            #[cfg(feature = "default-tls")]
+            Inner::DefaultTls(http, _tls) => http.set_keepalive(dur),
+            #[cfg(feature = "rustls-tls")]
+            Inner::RustlsTls { http, .. } => http.set_keepalive(dur),
+            #[cfg(not(feature = "__tls"))]
+            Inner::Http(http) => http.set_keepalive(dur),
+        }
+    }
 }
 
 fn into_uri(scheme: Scheme, host: Authority) -> Uri {