diff --git a/tokio-postgres/Cargo.toml b/tokio-postgres/Cargo.toml index 349d4cdc..f058446a 100644 --- a/tokio-postgres/Cargo.toml +++ b/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 } diff --git a/tokio-postgres/src/error/mod.rs b/tokio-postgres/src/error/mod.rs index bc3913a4..1b7fca2c 100644 --- a/tokio-postgres/src/error/mod.rs +++ b/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) + } } diff --git a/tokio-postgres/src/proto/connect_once.rs b/tokio-postgres/src/proto/connect_once.rs index e5035185..4dba5c3c 100644 --- a/tokio-postgres/src/proto/connect_once.rs +++ b/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,12 +42,16 @@ where #[state_machine_future(transitions(Handshaking))] ConnectingUnix { future: tokio_uds::ConnectFuture, + connect_timeout: Option, + timeout: Option, tls_mode: T, params: HashMap, }, #[state_machine_future(transitions(ConnectingTcp))] ResolvingDns { future: CpuFuture, io::Error>, + connect_timeout: Option, + timeout: Option, tls_mode: T, params: HashMap, }, @@ -53,6 +59,8 @@ where ConnectingTcp { future: tokio_tcp::ConnectFuture, addrs: vec::IntoIter, + connect_timeout: Option, + timeout: Option, tls_mode: T, params: HashMap, }, @@ -69,7 +77,20 @@ where T: TlsMode, { fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start>) -> Poll, 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::().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)] { @@ -77,6 +98,8 @@ where 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>, ) -> Poll, 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>, ) -> Poll, 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>, ) -> Poll, 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,