feat(client): Add connect timeout to HttpConnector

This takes the same strategy as golang, where the timeout value is
divided equally between the candidate socket addresses.

If happy eyeballs is enabled, the division takes place "below" the
IPv4/IPv6 partitioning.

Backported to 0.12 from master.

Author: chewi

src/client/connect/dns.rs

@@ -261,6 +261,10 @@ impl IpAddrs {
     pub(super) fn is_empty(&self) -> bool {
         self.iter.as_slice().is_empty()
     }
+
+    pub(super) fn len(&self) -> usize {
+        self.iter.as_slice().len()
+    }
 }
 
 impl Iterator for IpAddrs {

src/client/connect/http.rs

@@ -12,7 +12,7 @@ use http::uri::Scheme;
 use net2::TcpBuilder;
 use tokio_reactor::Handle;
 use tokio_tcp::{TcpStream, ConnectFuture};
-use tokio_timer::Delay;
+use tokio_timer::{Delay, Timeout};
 
 use super::{Connect, Connected, Destination};
 use super::dns::{self, GaiResolver, Resolve, TokioThreadpoolGaiResolver};
@@ -29,6 +29,7 @@ use super::dns::{self, GaiResolver, Resolve, TokioThreadpoolGaiResolver};
 pub struct HttpConnector<R = GaiResolver> {
     enforce_http: bool,
     handle: Option<Handle>,
+    connect_timeout: Option<Duration>,
     happy_eyeballs_timeout: Option<Duration>,
     keep_alive_timeout: Option<Duration>,
     local_address: Option<IpAddr>,
@@ -120,6 +121,7 @@ impl<R> HttpConnector<R> {
         HttpConnector {
             enforce_http: true,
             handle: None,
+            connect_timeout: None,
             happy_eyeballs_timeout: Some(Duration::from_millis(300)),
             keep_alive_timeout: None,
             local_address: None,
@@ -187,6 +189,17 @@ impl<R> HttpConnector<R> {
         self.local_address = addr;
     }
 
+    /// Set the connect timeout.
+    ///
+    /// If a domain resolves to multiple IP addresses, the timeout will be
+    /// evenly divided across them.
+    ///
+    /// Default is `None`.
+    #[inline]
+    pub fn set_connect_timeout(&mut self, dur: Option<Duration>) {
+        self.connect_timeout = dur;
+    }
+
     /// Set timeout for [RFC 6555 (Happy Eyeballs)][RFC 6555] algorithm.
     ///
     /// If hostname resolves to both IPv4 and IPv6 addresses and connection
@@ -259,6 +272,7 @@ where
         HttpConnecting {
             state: State::Lazy(self.resolver.clone(), host.into(), self.local_address),
             handle: self.handle.clone(),
+            connect_timeout: self.connect_timeout,
             happy_eyeballs_timeout: self.happy_eyeballs_timeout,
             keep_alive_timeout: self.keep_alive_timeout,
             nodelay: self.nodelay,
@@ -285,6 +299,7 @@ fn invalid_url<R: Resolve>(err: InvalidUrl, handle: &Option<Handle>) -> HttpConn
         keep_alive_timeout: None,
         nodelay: false,
         port: 0,
+        connect_timeout: None,
         happy_eyeballs_timeout: None,
         reuse_address: false,
         send_buffer_size: None,
@@ -319,6 +334,7 @@ impl StdError for InvalidUrl {
 pub struct HttpConnecting<R: Resolve = GaiResolver> {
     state: State<R>,
     handle: Option<Handle>,
+    connect_timeout: Option<Duration>,
     happy_eyeballs_timeout: Option<Duration>,
     keep_alive_timeout: Option<Duration>,
     nodelay: bool,
@@ -348,7 +364,7 @@ impl<R: Resolve> Future for HttpConnecting<R> {
                     // skip resolving the dns and start connecting right away.
                     if let Some(addrs) = dns::IpAddrs::try_parse(host, self.port) {
                         state = State::Connecting(ConnectingTcp::new(
-                            local_addr, addrs, self.happy_eyeballs_timeout, self.reuse_address));
+                            local_addr, addrs, self.connect_timeout, self.happy_eyeballs_timeout, self.reuse_address));
                     } else {
                         let name = dns::Name::new(mem::replace(host, String::new()));
                         state = State::Resolving(resolver.resolve(name), local_addr);
@@ -364,7 +380,7 @@ impl<R: Resolve> Future for HttpConnecting<R> {
                                 .collect();
                             let addrs = dns::IpAddrs::new(addrs);
                             state = State::Connecting(ConnectingTcp::new(
-                                local_addr, addrs, self.happy_eyeballs_timeout, self.reuse_address));
+                                local_addr, addrs, self.connect_timeout, self.happy_eyeballs_timeout, self.reuse_address));
                         }
                     };
                 },
@@ -417,6 +433,7 @@ impl ConnectingTcp {
     fn new(
         local_addr: Option<IpAddr>,
         remote_addrs: dns::IpAddrs,
+        connect_timeout: Option<Duration>,
         fallback_timeout: Option<Duration>,
         reuse_address: bool,
     ) -> ConnectingTcp {
@@ -425,25 +442,25 @@ impl ConnectingTcp {
             if fallback_addrs.is_empty() {
                 return ConnectingTcp {
                     local_addr,
-                    preferred: ConnectingTcpRemote::new(preferred_addrs),
+                    preferred: ConnectingTcpRemote::new(preferred_addrs, connect_timeout),
                     fallback: None,
                     reuse_address,
                 };
             }
 
             ConnectingTcp {
                 local_addr,
-                preferred: ConnectingTcpRemote::new(preferred_addrs),
+                preferred: ConnectingTcpRemote::new(preferred_addrs, connect_timeout),
                 fallback: Some(ConnectingTcpFallback {
                     delay: Delay::new(Instant::now() + fallback_timeout),
-                    remote: ConnectingTcpRemote::new(fallback_addrs),
+                    remote: ConnectingTcpRemote::new(fallback_addrs, connect_timeout),
                 }),
                 reuse_address,
             }
         } else {
             ConnectingTcp {
                 local_addr,
-                preferred: ConnectingTcpRemote::new(remote_addrs),
+                preferred: ConnectingTcpRemote::new(remote_addrs, connect_timeout),
                 fallback: None,
                 reuse_address,
             }
@@ -458,13 +475,17 @@ struct ConnectingTcpFallback {
 
 struct ConnectingTcpRemote {
     addrs: dns::IpAddrs,
-    current: Option<ConnectFuture>,
+    connect_timeout: Option<Duration>,
+    current: Option<MaybeTimedConnectFuture>,
 }
 
 impl ConnectingTcpRemote {
-    fn new(addrs: dns::IpAddrs) -> Self {
+    fn new(addrs: dns::IpAddrs, connect_timeout: Option<Duration>) -> Self {
+        let connect_timeout = connect_timeout.map(|t| t / (addrs.len() as u32));
+
         Self {
             addrs,
+            connect_timeout,
             current: None,
         }
     }
@@ -481,7 +502,18 @@ impl ConnectingTcpRemote {
         let mut err = None;
         loop {
             if let Some(ref mut current) = self.current {
-                match current.poll() {
+                let poll: Poll<TcpStream, io::Error> = match current {
+                    MaybeTimedConnectFuture::Timed(future) => match future.poll() {
+                        Ok(tcp) => Ok(tcp),
+                        Err(err) => if err.is_inner() {
+                            Err(err.into_inner().unwrap())
+                        } else {
+                            Err(io::Error::new(io::ErrorKind::TimedOut, err.description()))
+                        }
+                    },
+                    MaybeTimedConnectFuture::Untimed(future) => future.poll(),
+                };
+                match poll {
                     Ok(Async::Ready(tcp)) => {
                         debug!("connected to {:?}", tcp.peer_addr().ok());
                         return Ok(Async::Ready(tcp));
@@ -492,14 +524,14 @@ impl ConnectingTcpRemote {
                         err = Some(e);
                         if let Some(addr) = self.addrs.next() {
                             debug!("connecting to {}", addr);
-                            *current = connect(&addr, local_addr, handle, reuse_address)?;
+                            *current = connect(&addr, local_addr, handle, reuse_address, self.connect_timeout)?;
                             continue;
                         }
                     }
                 }
             } else if let Some(addr) = self.addrs.next() {
                 debug!("connecting to {}", addr);
-                self.current = Some(connect(&addr, local_addr, handle, reuse_address)?);
+                self.current = Some(connect(&addr, local_addr, handle, reuse_address, self.connect_timeout)?);
                 continue;
             }
 
@@ -508,7 +540,12 @@ impl ConnectingTcpRemote {
     }
 }
 
-fn connect(addr: &SocketAddr, local_addr: &Option<IpAddr>, handle: &Option<Handle>, reuse_address: bool) -> io::Result<ConnectFuture> {
+enum MaybeTimedConnectFuture {
+    Timed(Timeout<ConnectFuture>),
+    Untimed(ConnectFuture),
+}
+
+fn connect(addr: &SocketAddr, local_addr: &Option<IpAddr>, handle: &Option<Handle>, reuse_address: bool, connect_timeout: Option<Duration>) -> io::Result<MaybeTimedConnectFuture> {
     let builder = match addr {
         &SocketAddr::V4(_) => TcpBuilder::new_v4()?,
         &SocketAddr::V6(_) => TcpBuilder::new_v6()?,
@@ -540,7 +577,13 @@ fn connect(addr: &SocketAddr, local_addr: &Option<IpAddr>, handle: &Option<Handl
         None => Cow::Owned(Handle::default()),
     };
 
-    Ok(TcpStream::connect_std(builder.to_tcp_stream()?, addr, &handle))
+    let stream = TcpStream::connect_std(builder.to_tcp_stream()?, addr, &handle);
+
+    if let Some(timeout) = connect_timeout {
+        Ok(MaybeTimedConnectFuture::Timed(Timeout::new(stream, timeout)))
+    } else {
+        Ok(MaybeTimedConnectFuture::Untimed(stream))
+    }
 }
 
 impl ConnectingTcp {
@@ -706,7 +749,7 @@ mod tests {
             }
 
             let addrs = hosts.iter().map(|host| (host.clone(), addr.port()).into()).collect();
-            let connecting_tcp = ConnectingTcp::new(None, dns::IpAddrs::new(addrs), Some(fallback_timeout), false);
+            let connecting_tcp = ConnectingTcp::new(None, dns::IpAddrs::new(addrs), None, Some(fallback_timeout), false);
             let fut = ConnectingTcpFuture(connecting_tcp);
 
             let start = Instant::now();