From 2b60a00614c5c4260ce0acaaa599da89bebfd267 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Thu, 1 Jul 2021 16:23:22 +0200
Subject: [PATCH] feat(transport): add connect timeout to `Endpoint` (#662)

Pulls in [hyper-timeout] which has a connector that can have a timeout
applied. I did consider vendoring hyper-timeout since its a fairly small
crate. I guess we can always do that later since its not exposed
publicly.

I would also like to add a test but I'm not sure about the best way of
testing this.

Fixes https://github.com/hyperium/tonic/issues/498

[hyper-timeout]: https://github.com/hjr3/hyper-timeout

Co-authored-by: Lucio Franco <luciofranco14@gmail.com>
---
 tonic/Cargo.toml                        |  2 ++
 tonic/src/transport/channel/endpoint.rs | 45 +++++++++++++++++++++++--
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/tonic/Cargo.toml b/tonic/Cargo.toml
index 989b8f8c0..a28372700 100644
--- a/tonic/Cargo.toml
+++ b/tonic/Cargo.toml
@@ -33,6 +33,7 @@ transport = [
     "tracing-futures",
     "tokio/macros",
     "tokio/time",
+    "hyper-timeout",
 ]
 tls = ["transport", "tokio-rustls"]
 tls-roots-common = ["tls"]
@@ -74,6 +75,7 @@ tokio = { version = "1.0.1", features = ["net"], optional = true }
 tokio-stream = "0.1"
 tower = { version = "0.4.7", features = ["balance", "buffer", "discover", "limit", "load", "make", "timeout", "util"], optional = true }
 tracing-futures = { version = "0.2", optional = true }
+hyper-timeout = { version = "0.4", optional = true }
 
 # rustls
 tokio-rustls = { version = "0.22", optional = true }
diff --git a/tonic/src/transport/channel/endpoint.rs b/tonic/src/transport/channel/endpoint.rs
index 204ab7bc9..8cd3e9d9e 100644
--- a/tonic/src/transport/channel/endpoint.rs
+++ b/tonic/src/transport/channel/endpoint.rs
@@ -38,6 +38,7 @@ pub struct Endpoint {
     pub(crate) http2_keep_alive_interval: Option<Duration>,
     pub(crate) http2_keep_alive_timeout: Option<Duration>,
     pub(crate) http2_keep_alive_while_idle: Option<bool>,
+    pub(crate) connect_timeout: Option<Duration>,
     pub(crate) http2_adaptive_window: Option<bool>,
 }
 
@@ -119,6 +120,23 @@ impl Endpoint {
         }
     }
 
+    /// Apply a timeout to connecting to the uri.
+    ///
+    /// Defaults to no timeout.
+    ///
+    /// ```
+    /// # use tonic::transport::Endpoint;
+    /// # use std::time::Duration;
+    /// # let mut builder = Endpoint::from_static("https://example.com");
+    /// builder.connect_timeout(Duration::from_secs(5));
+    /// ```
+    pub fn connect_timeout(self, dur: Duration) -> Self {
+        Endpoint {
+            connect_timeout: Some(dur),
+            ..self
+        }
+    }
+
     /// Set whether TCP keepalive messages are enabled on accepted connections.
     ///
     /// If `None` is specified, keepalive is disabled, otherwise the duration
@@ -253,7 +271,13 @@ impl Endpoint {
         #[cfg(not(feature = "tls"))]
         let connector = service::connector(http);
 
-        Channel::connect(connector, self.clone()).await
+        if let Some(connect_timeout) = self.connect_timeout {
+            let mut connector = hyper_timeout::TimeoutConnector::new(connector);
+            connector.set_connect_timeout(Some(connect_timeout));
+            Channel::connect(connector, self.clone()).await
+        } else {
+            Channel::connect(connector, self.clone()).await
+        }
     }
 
     /// Create a channel from this config.
@@ -272,7 +296,13 @@ impl Endpoint {
         #[cfg(not(feature = "tls"))]
         let connector = service::connector(http);
 
-        Ok(Channel::new(connector, self.clone()))
+        if let Some(connect_timeout) = self.connect_timeout {
+            let mut connector = hyper_timeout::TimeoutConnector::new(connector);
+            connector.set_connect_timeout(Some(connect_timeout));
+            Ok(Channel::new(connector, self.clone()))
+        } else {
+            Ok(Channel::new(connector, self.clone()))
+        }
     }
 
     /// Connect with a custom connector.
@@ -280,6 +310,8 @@ impl Endpoint {
     /// This allows you to build a [Channel](struct.Channel.html) that uses a non-HTTP transport.
     /// See the `uds` example for an example on how to use this function to build channel that
     /// uses a Unix socket transport.
+    ///
+    /// The [`connect_timeout`](Endpoint::connect_timeout) will still be applied.
     pub async fn connect_with_connector<C>(&self, connector: C) -> Result<Channel, Error>
     where
         C: MakeConnection<Uri> + Send + 'static,
@@ -293,7 +325,13 @@ impl Endpoint {
         #[cfg(not(feature = "tls"))]
         let connector = service::connector(connector);
 
-        Channel::connect(connector, self.clone()).await
+        if let Some(connect_timeout) = self.connect_timeout {
+            let mut connector = hyper_timeout::TimeoutConnector::new(connector);
+            connector.set_connect_timeout(Some(connect_timeout));
+            Channel::connect(connector, self.clone()).await
+        } else {
+            Channel::connect(connector, self.clone()).await
+        }
     }
 
     /// Get the endpoint uri.
@@ -328,6 +366,7 @@ impl From<Uri> for Endpoint {
             http2_keep_alive_interval: None,
             http2_keep_alive_timeout: None,
             http2_keep_alive_while_idle: None,
+            connect_timeout: None,
             http2_adaptive_window: None,
         }
     }