support connection timeout

Author: Hexilee

src/socket.rs

@@ -3,9 +3,11 @@ use async_std::os::unix::net::UnixStream;
 
 use async_std::io::{self, Read, Write};
 use async_std::net::TcpStream;
+use std::future::Future;
 use std::mem::MaybeUninit;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use std::time::Duration;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_postgres::config::{Config, Host};
 
@@ -81,24 +83,39 @@ pub async fn connect_socket(config: &Config) -> io::Result<Socket> {
     let mut ports = config.get_ports().iter().cloned();
     for host in config.get_hosts() {
         let port = ports.next().unwrap_or(DEFAULT_PORT);
+        let dur = config.get_connect_timeout();
         let result = match host {
             #[cfg(unix)]
             Host::Unix(path) => {
                 let sock = path.join(format!(".s.PGSQL.{}", port));
-                UnixStream::connect(sock).await.map(Into::into)
+                let fut = UnixStream::connect(sock);
+                timeout(dur, fut).await.map(Into::into)
+            }
+            Host::Tcp(tcp) => {
+                let fut = TcpStream::connect((tcp.as_str(), port));
+                timeout(dur, fut).await.map(Into::into)
             }
-            Host::Tcp(tcp) => TcpStream::connect((tcp.as_str(), port))
-                .await
-                .map(Into::into),
             #[cfg(not(unix))]
             Host::Unix(_) => {
                 io::Error::new(io::ErrorKind::Other, "unix domain socket is unsupported")
             }
         };
+
         match result {
             Err(err) => error = err,
             stream => return stream,
         }
     }
     Err(error)
 }
+
+async fn timeout<F, T>(dur: Option<&Duration>, fut: F) -> io::Result<T>
+where
+    F: Future<Output = io::Result<T>>,
+{
+    if let Some(timeout) = dur {
+        io::timeout(timeout.clone(), fut).await
+    } else {
+        fut.await
+    }
+}