Add a connect timeout

cc rust-postgres#246

Author: sfackler

postgres-shared/src/params/mod.rs

@@ -2,6 +2,7 @@
 use std::error::Error;
 use std::path::PathBuf;
 use std::mem;
+use std::time::Duration;
 
 use params::url::Url;
 
@@ -43,6 +44,7 @@ pub struct ConnectParams {
     user: Option<User>,
     database: Option<String>,
     options: Vec<(String, String)>,
+    connect_timeout: Option<Duration>,
 }
 
 impl ConnectParams {
@@ -79,6 +81,11 @@ impl ConnectParams {
     pub fn options(&self) -> &[(String, String)] {
         &self.options
     }
+
+    /// A timeout to apply to each socket-level connection attempt.
+    pub fn connect_timeout(&self) -> Option<Duration> {
+        self.connect_timeout
+    }
 }
 
 /// A builder for `ConnectParams`.
@@ -87,6 +94,7 @@ pub struct Builder {
     user: Option<User>,
     database: Option<String>,
     options: Vec<(String, String)>,
+    connect_timeout: Option<Duration>,
 }
 
 impl Builder {
@@ -97,6 +105,7 @@ impl Builder {
             user: None,
             database: None,
             options: vec![],
+            connect_timeout: None,
         }
     }
 
@@ -127,6 +136,12 @@ impl Builder {
         self
     }
 
+    /// Sets the connection timeout.
+    pub fn connect_timeout(&mut self, connect_timeout: Option<Duration>) -> &mut Builder {
+        self.connect_timeout = connect_timeout;
+        self
+    }
+
     /// Constructs a `ConnectParams` from the builder.
     pub fn build(&mut self, host: Host) -> ConnectParams {
         ConnectParams {
@@ -135,6 +150,7 @@ impl Builder {
             user: self.user.take(),
             database: self.database.take(),
             options: mem::replace(&mut self.options, vec![]),
+            connect_timeout: self.connect_timeout,
         }
     }
 }
@@ -196,7 +212,16 @@ impl IntoConnectParams for Url {
         }
 
         for (name, value) in options {
-            builder.option(&name, &value);
+            match &*name {
+                "connect_timeout" => {
+                    let timeout = value.parse().map_err(|_| "invalid connect_timeout")?;
+                    let timeout = Duration::from_secs(timeout);
+                    builder.connect_timeout(Some(timeout));
+                }
+                _ => {
+                    builder.option(&name, &value);
+                }
+            }
         }
 
         let maybe_path = url::decode_component(&host)?;
@@ -209,3 +234,29 @@ impl IntoConnectParams for Url {
         Ok(builder.build(host))
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn parse_url() {
+        let params = "postgres://user@host:44/dbname?connect_timeout=10&application_name=foo";
+        let params = params.into_connect_params().unwrap();
+        assert_eq!(
+            params.user(),
+            Some(&User {
+                name: "user".to_string(),
+                password: None,
+            })
+        );
+        assert_eq!(params.host(), &Host::Tcp("host".to_string()));
+        assert_eq!(params.port(), 44);
+        assert_eq!(params.database(), Some("dbname"));
+        assert_eq!(
+            params.options(),
+            &[("application_name".to_string(), "foo".to_string())][..]
+        );
+        assert_eq!(params.connect_timeout(), Some(Duration::from_secs(10)));
+    }
+}

postgres/Cargo.toml

@@ -59,6 +59,7 @@ no-logging = []
 bytes = "0.4"
 fallible-iterator = "0.1.3"
 log = "0.3"
+socket2 = "0.2"
 
 openssl = { version = "0.9.2", optional = true }
 native-tls = { version = "0.1", optional = true }

postgres/src/lib.rs

@@ -77,6 +77,7 @@ extern crate fallible_iterator;
 extern crate log;
 extern crate postgres_protocol;
 extern crate postgres_shared;
+extern crate socket2;
 
 use fallible_iterator::FallibleIterator;
 use std::cell::{Cell, RefCell};

postgres/src/priv_io.rs

@@ -1,17 +1,17 @@
 use std::io::{self, BufWriter, Read, Write};
-use std::fmt;
-use std::net::TcpStream;
+use std::net::{ToSocketAddrs, SocketAddr};
 use std::time::Duration;
 use std::result;
 use bytes::{BufMut, BytesMut};
 #[cfg(unix)]
 use std::os::unix::net::UnixStream;
 #[cfg(unix)]
-use std::os::unix::io::{AsRawFd, RawFd};
+use std::os::unix::io::{AsRawFd, RawFd, FromRawFd, IntoRawFd};
 #[cfg(windows)]
 use std::os::windows::io::{AsRawSocket, RawSocket};
 use postgres_protocol::message::frontend;
 use postgres_protocol::message::backend;
+use socket2::{Socket, SockAddr, Domain, Type};
 
 use {Result, TlsMode};
 use error;
@@ -118,37 +118,20 @@ impl MessageStream {
     }
 
     fn set_read_timeout(&self, timeout: Option<Duration>) -> io::Result<()> {
-        match self.stream.get_ref().get_ref().0 {
-            InternalStream::Tcp(ref s) => s.set_read_timeout(timeout),
-            #[cfg(unix)]
-            InternalStream::Unix(ref s) => s.set_read_timeout(timeout),
-        }
+        self.stream.get_ref().get_ref().0.set_read_timeout(timeout)
     }
 
     fn set_nonblocking(&self, nonblock: bool) -> io::Result<()> {
-        match self.stream.get_ref().get_ref().0 {
-            InternalStream::Tcp(ref s) => s.set_nonblocking(nonblock),
-            #[cfg(unix)]
-            InternalStream::Unix(ref s) => s.set_nonblocking(nonblock),
-        }
+        self.stream.get_ref().get_ref().0.set_nonblocking(nonblock)
     }
 }
 
 /// A connection to the Postgres server.
 ///
 /// It implements `Read`, `Write` and `TlsStream`, as well as `AsRawFd` on
 /// Unix platforms and `AsRawSocket` on Windows platforms.
-pub struct Stream(InternalStream);
-
-impl fmt::Debug for Stream {
-    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
-        match self.0 {
-            InternalStream::Tcp(ref s) => fmt::Debug::fmt(s, fmt),
-            #[cfg(unix)]
-            InternalStream::Unix(ref s) => fmt::Debug::fmt(s, fmt),
-        }
-    }
-}
+#[derive(Debug)]
+pub struct Stream(Socket);
 
 impl Read for Stream {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
@@ -179,69 +162,56 @@ impl TlsStream for Stream {
 #[cfg(unix)]
 impl AsRawFd for Stream {
     fn as_raw_fd(&self) -> RawFd {
-        match self.0 {
-            InternalStream::Tcp(ref s) => s.as_raw_fd(),
-            InternalStream::Unix(ref s) => s.as_raw_fd(),
-        }
+        self.0.as_raw_fd()
     }
 }
 
 #[cfg(windows)]
 impl AsRawSocket for Stream {
     fn as_raw_socket(&self) -> RawSocket {
-        // Unix sockets aren't supported on windows, so no need to match
-        match self.0 {
-            InternalStream::Tcp(ref s) => s.as_raw_socket(),
-        }
-    }
-}
-
-enum InternalStream {
-    Tcp(TcpStream),
-    #[cfg(unix)]
-    Unix(UnixStream),
-}
-
-impl Read for InternalStream {
-    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-        match *self {
-            InternalStream::Tcp(ref mut s) => s.read(buf),
-            #[cfg(unix)]
-            InternalStream::Unix(ref mut s) => s.read(buf),
-        }
+        self.0.as_raw_socket()
     }
 }
 
-impl Write for InternalStream {
-    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
-        match *self {
-            InternalStream::Tcp(ref mut s) => s.write(buf),
-            #[cfg(unix)]
-            InternalStream::Unix(ref mut s) => s.write(buf),
-        }
-    }
-
-    fn flush(&mut self) -> io::Result<()> {
-        match *self {
-            InternalStream::Tcp(ref mut s) => s.flush(),
-            #[cfg(unix)]
-            InternalStream::Unix(ref mut s) => s.flush(),
-        }
-    }
-}
-
-fn open_socket(params: &ConnectParams) -> Result<InternalStream> {
+fn open_socket(params: &ConnectParams) -> Result<Socket> {
     let port = params.port();
     match *params.host() {
         Host::Tcp(ref host) => {
-            Ok(TcpStream::connect(&(&**host, port)).map(
-                InternalStream::Tcp,
-            )?)
+            let mut error = None;
+            for addr in (&**host, port).to_socket_addrs()? {
+                let domain = match addr {
+                    SocketAddr::V4(_) => Domain::ipv4(),
+                    SocketAddr::V6(_) => Domain::ipv6(),
+                };
+                let socket = Socket::new(domain, Type::stream(), None)?;
+                let addr = SockAddr::from(addr);
+                let r = match params.connect_timeout() {
+                    Some(timeout) => socket.connect_timeout(&addr, timeout),
+                    None => socket.connect(&addr),
+                };
+                match r {
+                    Ok(()) => return Ok(socket),
+                    Err(e) => error = Some(e),
+                }
+            }
+
+            Err(
+                error
+                    .unwrap_or_else(|| {
+                        io::Error::new(
+                            io::ErrorKind::InvalidInput,
+                            "could not resolve any addresses",
+                        )
+                    })
+                    .into(),
+            )
         }
         #[cfg(unix)]
         Host::Unix(ref path) => {
             let path = path.join(&format!(".s.PGSQL.{}", port));
-            Ok(UnixStream::connect(&path).map(InternalStream::Unix)?)
+            Ok(UnixStream::connect(&path).map(|s| unsafe {
+                Socket::from_raw_fd(s.into_raw_fd())
+            })?)
         }
         #[cfg(not(unix))]
         Host::Unix(..) => {