Skip to content

Commit

Permalink
feat(transport): add connect timeout to Endpoint (#662)
Browse files Browse the repository at this point in the history
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 #498

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

Co-authored-by: Lucio Franco <[email protected]>
  • Loading branch information
davidpdrsn and LucioFranco authored Jul 1, 2021
1 parent b90bb7b commit 2b60a00
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 3 deletions.
2 changes: 2 additions & 0 deletions tonic/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ transport = [
"tracing-futures",
"tokio/macros",
"tokio/time",
"hyper-timeout",
]
tls = ["transport", "tokio-rustls"]
tls-roots-common = ["tls"]
Expand Down Expand Up @@ -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 }
Expand Down
45 changes: 42 additions & 3 deletions tonic/src/transport/channel/endpoint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>,
}

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand All @@ -272,14 +296,22 @@ 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.
///
/// 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,
Expand All @@ -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.
Expand Down Expand Up @@ -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,
}
}
Expand Down

0 comments on commit 2b60a00

Please sign in to comment.