Support connect_timeout

Author: sfackler

tokio-postgres/Cargo.toml

@@ -28,7 +28,7 @@ circle-ci = { repository = "sfackler/rust-postgres" }
 
 [features]
 default = ["runtime"]
-runtime = ["tokio-tcp", "tokio-uds", "futures-cpupool", "lazy_static"]
+runtime = ["tokio-tcp", "tokio-timer", "tokio-uds", "futures-cpupool", "lazy_static"]
 
 "with-bit-vec-0.5" = ["bit-vec-05"]
 "with-chrono-0.4" = ["chrono-04"]
@@ -53,6 +53,7 @@ void = "1.0"
 tokio-tcp = { version = "0.1", optional = true }
 futures-cpupool = { version = "0.1", optional = true }
 lazy_static = { version = "1.0", optional = true }
+tokio-timer = { version = "0.2", optional = true }
 
 bit-vec-05 = { version = "0.5", package = "bit-vec", optional = true }
 chrono-04 = { version = "0.4", package = "chrono", optional = true }

tokio-postgres/src/error/mod.rs

@@ -356,6 +356,12 @@ enum Kind {
     InvalidPort,
     #[cfg(feature = "runtime")]
     InvalidPortCount,
+    #[cfg(feature = "runtime")]
+    InvalidConnectTimeout,
+    #[cfg(feature = "runtime")]
+    Timer,
+    #[cfg(feature = "runtime")]
+    ConnectTimeout,
 }
 
 struct ErrorInner {
@@ -401,6 +407,12 @@ impl fmt::Display for Error {
             Kind::InvalidPort => "invalid port",
             #[cfg(feature = "runtime")]
             Kind::InvalidPortCount => "wrong number of ports provided",
+            #[cfg(feature = "runtime")]
+            Kind::InvalidConnectTimeout => "invalid connect_timeout",
+            #[cfg(feature = "runtime")]
+            Kind::Timer => "timer error",
+            #[cfg(feature = "runtime")]
+            Kind::ConnectTimeout => "timed out connecting to server",
         };
         fmt.write_str(s)?;
         if let Some(ref cause) = self.0.cause {
@@ -523,4 +535,19 @@ impl Error {
     pub(crate) fn invalid_port_count() -> Error {
         Error::new(Kind::InvalidPortCount, None)
     }
+
+    #[cfg(feature = "runtime")]
+    pub(crate) fn invalid_connect_timeout(e: ParseIntError) -> Error {
+        Error::new(Kind::InvalidConnectTimeout, Some(Box::new(e)))
+    }
+
+    #[cfg(feature = "runtime")]
+    pub(crate) fn timer(e: tokio_timer::Error) -> Error {
+        Error::new(Kind::Timer, Some(Box::new(e)))
+    }
+
+    #[cfg(feature = "runtime")]
+    pub(crate) fn connect_timeout() -> Error {
+        Error::new(Kind::ConnectTimeout, None)
+    }
 }

tokio-postgres/src/proto/connect_once.rs

@@ -7,8 +7,10 @@ use std::io;
 use std::net::{SocketAddr, ToSocketAddrs};
 #[cfg(unix)]
 use std::path::Path;
+use std::time::{Duration, Instant};
 use std::vec;
 use tokio_tcp::TcpStream;
+use tokio_timer::Delay;
 #[cfg(unix)]
 use tokio_uds::UnixStream;
 
@@ -40,19 +42,25 @@ where
     #[state_machine_future(transitions(Handshaking))]
     ConnectingUnix {
         future: tokio_uds::ConnectFuture,
+        connect_timeout: Option<Duration>,
+        timeout: Option<Delay>,
         tls_mode: T,
         params: HashMap<String, String>,
     },
     #[state_machine_future(transitions(ConnectingTcp))]
     ResolvingDns {
         future: CpuFuture<vec::IntoIter<SocketAddr>, io::Error>,
+        connect_timeout: Option<Duration>,
+        timeout: Option<Delay>,
         tls_mode: T,
         params: HashMap<String, String>,
     },
     #[state_machine_future(transitions(Handshaking))]
     ConnectingTcp {
         future: tokio_tcp::ConnectFuture,
         addrs: vec::IntoIter<SocketAddr>,
+        connect_timeout: Option<Duration>,
+        timeout: Option<Delay>,
         tls_mode: T,
         params: HashMap<String, String>,
     },
@@ -69,14 +77,29 @@ where
     T: TlsMode<Socket>,
 {
     fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<T>>) -> Poll<AfterStart<T>, Error> {
-        let state = state.take();
+        let mut state = state.take();
+
+        let connect_timeout = match state.params.remove("connect_timeout") {
+            Some(s) => {
+                let seconds = s.parse::<i64>().map_err(Error::invalid_connect_timeout)?;
+                if seconds <= 0 {
+                    None
+                } else {
+                    Some(Duration::from_secs(seconds as u64))
+                }
+            }
+            None => None,
+        };
+        let timeout = connect_timeout.map(|d| Delay::new(Instant::now() + d));
 
         #[cfg(unix)]
         {
             if state.host.starts_with('/') {
                 let path = Path::new(&state.host).join(format!(".s.PGSQL.{}", state.port));
                 transition!(ConnectingUnix {
                     future: UnixStream::connect(path),
+                    connect_timeout,
+                    timeout,
                     tls_mode: state.tls_mode,
                     params: state.params,
                 })
@@ -87,6 +110,8 @@ where
         let port = state.port;
         transition!(ResolvingDns {
             future: DNS_POOL.spawn_fn(move || (&*host, port).to_socket_addrs()),
+            connect_timeout,
+            timeout,
             tls_mode: state.tls_mode,
             params: state.params,
         })
@@ -96,6 +121,14 @@ where
     fn poll_connecting_unix<'a>(
         state: &'a mut RentToOwn<'a, ConnectingUnix<T>>,
     ) -> Poll<AfterConnectingUnix<T>, Error> {
+        if let Some(timeout) = &mut state.timeout {
+            match timeout.poll() {
+                Ok(Async::Ready(())) => return Err(Error::connect_timeout()),
+                Ok(Async::NotReady) => {}
+                Err(e) => return Err(Error::timer(e)),
+            }
+        }
+
         let stream = try_ready!(state.future.poll().map_err(Error::connect));
         let stream = Socket::new_unix(stream);
         let state = state.take();
@@ -108,6 +141,14 @@ where
     fn poll_resolving_dns<'a>(
         state: &'a mut RentToOwn<'a, ResolvingDns<T>>,
     ) -> Poll<AfterResolvingDns<T>, Error> {
+        if let Some(timeout) = &mut state.timeout {
+            match timeout.poll() {
+                Ok(Async::Ready(())) => return Err(Error::connect_timeout()),
+                Ok(Async::NotReady) => {}
+                Err(e) => return Err(Error::timer(e)),
+            }
+        }
+
         let mut addrs = try_ready!(state.future.poll().map_err(Error::connect));
         let state = state.take();
 
@@ -124,6 +165,8 @@ where
         transition!(ConnectingTcp {
             future: TcpStream::connect(&addr),
             addrs,
+            connect_timeout: state.connect_timeout,
+            timeout: state.timeout,
             tls_mode: state.tls_mode,
             params: state.params,
         })
@@ -132,6 +175,14 @@ where
     fn poll_connecting_tcp<'a>(
         state: &'a mut RentToOwn<'a, ConnectingTcp<T>>,
     ) -> Poll<AfterConnectingTcp<T>, Error> {
+        if let Some(timeout) = &mut state.timeout {
+            match timeout.poll() {
+                Ok(Async::Ready(())) => return Err(Error::connect_timeout()),
+                Ok(Async::NotReady) => {}
+                Err(e) => return Err(Error::timer(e)),
+            }
+        }
+
         let stream = loop {
             match state.future.poll() {
                 Ok(Async::Ready(stream)) => break stream,