Integrate upstream hostaddr/load balancing support

Author: benesch

postgres/src/config.rs

@@ -6,6 +6,7 @@ use crate::connection::Connection;
 use crate::Client;
 use log::info;
 use std::fmt;
+use std::net::IpAddr;
 use std::path::Path;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -43,6 +44,19 @@ use tokio_postgres::{Error, Socket};
 ///     path to the directory containing Unix domain sockets. Otherwise, it is treated as a hostname. Multiple hosts
 ///     can be specified, separated by commas. Each host will be tried in turn when connecting. Required if connecting
 ///     with the `connect` method.
+/// * `hostaddr` - Numeric IP address of host to connect to. This should be in the standard IPv4 address format,
+///     e.g., 172.28.40.9. If your machine supports IPv6, you can also use those addresses.
+///     If this parameter is not specified, the value of `host` will be looked up to find the corresponding IP address,
+///     - or if host specifies an IP address, that value will be used directly.
+///     Using `hostaddr` allows the application to avoid a host name look-up, which might be important in applications
+///     with time constraints. However, a host name is required for verify-full SSL certificate verification.
+///     Specifically:
+///         * If `hostaddr` is specified without `host`, the value for `hostaddr` gives the server network address.
+///             The connection attempt will fail if the authentication method requires a host name;
+///         * If `host` is specified without `hostaddr`, a host name lookup occurs;
+///         * If both `host` and `hostaddr` are specified, the value for `hostaddr` gives the server network address.
+///             The value for `host` is ignored unless the authentication method requires it,
+///             in which case it will be used as the host name.
 /// * `port` - The port to connect to. Multiple ports can be specified, separated by commas. The number of ports must be
 ///     either 1, in which case it will be used for all hosts, or the same as the number of hosts. Defaults to 5432 if
 ///     omitted or the empty string.
@@ -74,6 +88,10 @@ use tokio_postgres::{Error, Socket};
 /// ```
 ///
 /// ```not_rust
+/// host=host1,host2,host3 port=1234,,5678 hostaddr=127.0.0.1,127.0.0.2,127.0.0.3 user=postgres target_session_attrs=read-write
+/// ```
+///
+/// ```not_rust
 /// host=host1,host2,host3 port=1234,,5678 user=postgres target_session_attrs=read-write
 /// ```
 ///
@@ -250,6 +268,7 @@ impl Config {
     ///
     /// Multiple hosts can be specified by calling this method multiple times, and each will be tried in order. On Unix
     /// systems, a host starting with a `/` is interpreted as a path to a directory containing Unix domain sockets.
+    /// There must be either no hosts, or the same number of hosts as hostaddrs.
     pub fn host(&mut self, host: &str) -> &mut Config {
         self.config.host(host);
         self
@@ -260,6 +279,11 @@ impl Config {
         self.config.get_hosts()
     }
 
+    /// Gets the hostaddrs that have been added to the configuration with `hostaddr`.
+    pub fn get_hostaddrs(&self) -> &[IpAddr] {
+        self.config.get_hostaddrs()
+    }
+
     /// Adds a Unix socket host to the configuration.
     ///
     /// Unlike `host`, this method allows non-UTF8 paths.
@@ -272,6 +296,15 @@ impl Config {
         self
     }
 
+    /// Adds a hostaddr to the configuration.
+    ///
+    /// Multiple hostaddrs can be specified by calling this method multiple times, and each will be tried in order.
+    /// There must be either no hostaddrs, or the same number of hostaddrs as hosts.
+    pub fn hostaddr(&mut self, hostaddr: IpAddr) -> &mut Config {
+        self.config.hostaddr(hostaddr);
+        self
+    }
+
     /// Adds a port to the configuration.
     ///
     /// Multiple ports can be specified by calling this method multiple times. There must either be no ports, in which

tokio-postgres/Cargo.toml

@@ -59,6 +59,7 @@ serde = { version = "1.0", optional = true }
 socket2 = { version = "0.5", features = ["all"] }
 tokio = { version = "1.27", features = ["io-util"] }
 tokio-util = { version = "0.7", features = ["codec"] }
+rand = "0.8.5"
 
 [dev-dependencies]
 futures-executor = "0.3"

tokio-postgres/src/cancel_query.rs

@@ -1,5 +1,5 @@
 use crate::client::SocketConfig;
-use crate::config::{Host, SslMode};
+use crate::config::SslMode;
 use crate::tls::MakeTlsConnect;
 use crate::{cancel_query_raw, connect_socket, Error, Socket};
 use std::io;
@@ -24,24 +24,20 @@ where
         }
     };
 
-    let hostname = match &config.host {
-        Host::Tcp(host) => &**host,
-        // postgres doesn't support TLS over unix sockets, so the choice here doesn't matter
-        #[cfg(unix)]
-        Host::Unix(_) => "",
-    };
     let tls = tls
-        .make_tls_connect(hostname)
+        .make_tls_connect(config.hostname.as_deref().unwrap_or(""))
         .map_err(|e| Error::tls(e.into()))?;
+    let has_hostname = config.hostname.is_some();
 
     let socket = connect_socket::connect_socket(
-        &config.host,
+        &config.addr,
         config.port,
         config.connect_timeout,
         config.tcp_user_timeout,
         config.keepalive.as_ref(),
     )
     .await?;
 
-    cancel_query_raw::cancel_query_raw(socket, ssl_mode, tls, process_id, secret_key).await
+    cancel_query_raw::cancel_query_raw(socket, ssl_mode, tls, has_hostname, process_id, secret_key)
+        .await
 }

tokio-postgres/src/cancel_query_raw.rs

@@ -9,14 +9,15 @@ pub async fn cancel_query_raw<S, T>(
     stream: S,
     mode: SslMode,
     tls: T,
+    has_hostname: bool,
     process_id: i32,
     secret_key: i32,
 ) -> Result<(), Error>
 where
     S: AsyncRead + AsyncWrite + Unpin,
     T: TlsConnect<S>,
 {
-    let mut stream = connect_tls::connect_tls(stream, mode, tls).await?;
+    let mut stream = connect_tls::connect_tls(stream, mode, tls, has_hostname).await?;
 
     let mut buf = BytesMut::new();
     frontend::cancel_request(process_id, secret_key, &mut buf);

tokio-postgres/src/cancel_token.rs

@@ -55,6 +55,7 @@ impl CancelToken {
             stream,
             self.ssl_mode,
             tls,
+            true,
             self.process_id,
             self.secret_key,
         )

tokio-postgres/src/client.rs

@@ -1,6 +1,4 @@
 use crate::codec::{BackendMessages, FrontendMessage};
-#[cfg(feature = "runtime")]
-use crate::config::Host;
 use crate::config::SslMode;
 use crate::connection::{Request, RequestMessages};
 use crate::copy_both::CopyBothDuplex;
@@ -29,6 +27,8 @@ use postgres_protocol::message::{backend::Message, frontend};
 use postgres_types::BorrowToSql;
 use std::collections::HashMap;
 use std::fmt;
+use std::net::IpAddr;
+use std::path::PathBuf;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 #[cfg(feature = "runtime")]
@@ -155,13 +155,22 @@ impl InnerClient {
 #[cfg(feature = "runtime")]
 #[derive(Clone)]
 pub(crate) struct SocketConfig {
-    pub host: Host,
+    pub addr: Addr,
+    pub hostname: Option<String>,
     pub port: u16,
     pub connect_timeout: Option<Duration>,
     pub tcp_user_timeout: Option<Duration>,
     pub keepalive: Option<KeepaliveConfig>,
 }
 
+#[cfg(feature = "runtime")]
+#[derive(Clone)]
+pub(crate) enum Addr {
+    Tcp(IpAddr),
+    #[cfg(unix)]
+    Unix(PathBuf),
+}
+
 /// An asynchronous PostgreSQL client.
 ///
 /// The client is one half of what is returned when a connection is established. Users interact with the database