diff --git a/crates/dbx-driver-postgres/src/postgres.rs b/crates/dbx-driver-postgres/src/postgres.rs index 87bc1e5807..782cf67586 100644 --- a/crates/dbx-driver-postgres/src/postgres.rs +++ b/crates/dbx-driver-postgres/src/postgres.rs @@ -16,13 +16,15 @@ use std::fmt::Write as _; use std::fs::File; use std::future::Future; use std::io::BufReader; +use std::net::{IpAddr, SocketAddr}; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::AtomicBool; use std::sync::{Arc, Mutex, OnceLock, Weak}; use std::time::{Duration, Instant}; -use tokio::task::JoinHandle; -use tokio_postgres::config::SslMode; +use tokio::net::TcpStream; +use tokio::task::{JoinHandle, JoinSet}; +use tokio_postgres::config::{Host, SslMode}; use tokio_postgres::tls::{MakeTlsConnect, TlsConnect}; use tokio_postgres::types::{FromSql, Kind, Type}; use tokio_postgres::{AsyncMessage, NoTls, Row, SimpleQueryMessage, Socket}; @@ -2561,6 +2563,125 @@ fn postgres_client_keys() -> &'static Mutex { /// (`RAISE NOTICE`/`WARNING`, etc.) into a per-backend buffer instead of /// discarding them. Query execution drains the buffer so notices are /// attached to the `QueryResult` of the statement that raised them. +/// Delay between staggered TCP connection attempts when a hostname resolves +/// to several addresses (Happy Eyeballs, RFC 8305 ยง4): short enough that +/// falling back to a reachable address feels instant, long enough that a +/// healthy preferred address still wins the race. +const HAPPY_EYEBALLS_STAGGER: Duration = Duration::from_millis(250); + +/// Dial `addr` once. Split out so the happy-eyeballs racer below stays readable. +async fn dial_socket_addr_once(addr: SocketAddr) -> std::io::Result { + TcpStream::connect(addr).await +} + +/// Connect to the first reachable address in `addrs`, racing the attempts +/// with a short stagger instead of trying them strictly one after another. +/// See [`happy_eyeballs_dial_with`] for the racing semantics. +async fn happy_eyeballs_dial(addrs: &[SocketAddr], stagger: Duration) -> std::io::Result { + happy_eyeballs_dial_with(addrs, stagger, dial_socket_addr_once).await +} + +/// Racing dial with an injectable per-address dial function, so tests can +/// script a hanging address deterministically (the dial is the external +/// boundary; the racing orchestration is the behavior under test). +/// +/// The first address is dialed immediately; each further address is started +/// after `stagger` unless an earlier attempt already failed, in which case the +/// next address is tried at once. The first successful dial wins and the +/// remaining attempts are aborted. When every attempt fails, the last error +/// is returned. +async fn happy_eyeballs_dial_with( + addrs: &[SocketAddr], + stagger: Duration, + dial: F, +) -> std::io::Result +where + F: Fn(SocketAddr) -> Fut + Send, + Fut: Future> + Send + 'static, +{ + let mut remaining = addrs.iter(); + let mut in_flight: JoinSet> = JoinSet::new(); + let mut last_err: Option = None; + + let Some(first) = remaining.next() else { + return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, "no socket addresses to dial")); + }; + in_flight.spawn(dial(*first)); + + loop { + tokio::select! { + biased; + result = in_flight.join_next(), if !in_flight.is_empty() => { + match result { + Some(Ok(Ok(stream))) => { + in_flight.abort_all(); + return Ok(stream); + } + Some(Ok(Err(err))) => last_err = Some(err), + Some(Err(join_err)) => { + last_err = Some(std::io::Error::other(join_err.to_string())); + } + // Unreachable: this arm is disabled while the set is empty. + None => break, + } + // A finished attempt immediately frees the way for the next address. + if let Some(addr) = remaining.next() { + in_flight.spawn(dial(*addr)); + } else if in_flight.is_empty() { + break; + } + } + _ = tokio::time::sleep(stagger), if remaining.len() > 0 => { + if let Some(addr) = remaining.next() { + in_flight.spawn(dial(*addr)); + } + } + } + } + + Err(last_err + .unwrap_or_else(|| std::io::Error::new(std::io::ErrorKind::InvalidInput, "no socket addresses to dial"))) +} + +/// Resolve a single-host PostgreSQL config ahead of the driver's own dial and +/// pin it to the first reachable resolved address (Happy Eyeballs, RFC 8305). +/// +/// The winning IP is passed as `hostaddr`, so the driver's dial goes straight +/// to the reachable address while TLS still validates against the original +/// hostname. Returns the config unchanged when there is nothing to race: +/// several hosts, an explicit `hostaddr`, a Unix socket, an IP literal, a +/// single resolved address, or an inconclusive probe (the driver's own dial +/// then surfaces the real error). +async fn pin_postgres_hostaddr(pg_config: &tokio_postgres::Config) -> tokio_postgres::Config { + let mut pg_config = pg_config.clone(); + let [Host::Tcp(host)] = pg_config.get_hosts() else { + return pg_config; + }; + if !pg_config.get_hostaddrs().is_empty() || host.parse::().is_ok() { + return pg_config; + } + let port = pg_config.get_ports().first().copied().unwrap_or(5432); + let Ok(resolved) = tokio::net::lookup_host((host.as_str(), port)).await else { + return pg_config; + }; + let mut addrs: Vec = resolved.collect(); + addrs.sort(); + addrs.dedup(); + if addrs.len() < 2 { + return pg_config; + } + // The probe only picks the address; the driver's own dial (with its + // connect_timeout and TLS setup) still opens the real connection. + let probe = + tokio::time::timeout(super::tcp_probe_timeout(), happy_eyeballs_dial(&addrs, HAPPY_EYEBALLS_STAGGER)).await; + if let Ok(Ok(winner)) = probe { + if let Ok(addr) = winner.peer_addr() { + pg_config.hostaddr(addr.ip()); + } + } + pg_config +} + struct NoticeCapturingConnect where T: MakeTlsConnect + Clone + Sync + Send + 'static, @@ -2587,6 +2708,10 @@ where let tls = self.tls.clone(); let pg_config = pg_config.clone(); Box::pin(async move { + // Prefer the first reachable resolved address so an unreachable + // IPv6 address cannot burn the whole connect timeout before IPv4 + // is tried (issue #10955). + let pg_config = pin_postgres_hostaddr(&pg_config).await; let (client, mut connection) = pg_config.connect(tls).await?; // No query can complete before the connection is being driven, so // the notice buffer is handed to the driver task through a slot @@ -15650,4 +15775,94 @@ mod tests { assert!(sql.contains("WHERE n.nspname = $1")); } } + + // Issue #10955: a hostname resolving to both IPv6 and IPv4 must connect + // through the reachable address instead of stalling on the unreachable + // one until the pool's create timeout fires. + + #[tokio::test] + async fn happy_eyeballs_dial_does_not_wait_for_hung_first_attempt() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let reachable = listener.local_addr().unwrap(); + let bad_ip: IpAddr = "192.0.2.1".parse().unwrap(); + let unreachable = SocketAddr::new(bad_ip, reachable.port()); + let started = Instant::now(); + let stream = tokio::time::timeout( + Duration::from_secs(10), + happy_eyeballs_dial_with(&[unreachable, reachable], Duration::from_millis(50), |addr| async move { + if addr.ip() == bad_ip { + // The unreachable IPv6/NAT64 address from #10955: the + // SYN is never answered, so this dial never resolves. + futures::future::pending::>().await + } else { + TcpStream::connect(addr).await + } + }), + ) + .await + .expect("dial must not stall on the unreachable address") + .expect("dial must succeed through the reachable address"); + assert!( + started.elapsed() < Duration::from_secs(5), + "fallback must be near-instant, not wait out the hung attempt" + ); + assert_eq!(stream.peer_addr().unwrap(), reachable); + } + + #[tokio::test] + async fn happy_eyeballs_dial_tries_next_address_on_fast_failure() { + let closed_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let refused = closed_listener.local_addr().unwrap(); + drop(closed_listener); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let reachable = listener.local_addr().unwrap(); + let started = Instant::now(); + let stream = happy_eyeballs_dial(&[refused, reachable], Duration::from_millis(50)) + .await + .expect("dial must fall through to the reachable address"); + assert!(started.elapsed() < Duration::from_secs(5), "a refused address must not cost a full stagger"); + assert_eq!(stream.peer_addr().unwrap(), reachable); + } + + #[tokio::test] + async fn happy_eyeballs_dial_succeeds_with_single_address() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let stream = happy_eyeballs_dial(&[addr], Duration::from_millis(50)) + .await + .expect("a single reachable address must connect"); + assert_eq!(stream.peer_addr().unwrap(), addr); + } + + #[tokio::test] + async fn happy_eyeballs_dial_reports_error_when_all_addresses_refused() { + let first = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let second = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let closed: Vec = [&first, &second].into_iter().map(|l| l.local_addr().unwrap()).collect(); + drop(first); + drop(second); + let started = Instant::now(); + happy_eyeballs_dial(&closed, Duration::from_millis(50)) + .await + .expect_err("all addresses refused: the dial must fail"); + assert!( + started.elapsed() < Duration::from_secs(5), + "refused dials must fail fast instead of waiting out a stagger" + ); + } + + #[tokio::test] + async fn pin_postgres_hostaddr_leaves_ip_literal_alone() { + let config = tokio_postgres::Config::from_str("host=127.0.0.1 port=5432 user=dbx connect_timeout=2").unwrap(); + let pinned = pin_postgres_hostaddr(&config).await; + assert!(pinned.get_hostaddrs().is_empty()); + } + + #[tokio::test] + async fn pin_postgres_hostaddr_keeps_explicit_hostaddr() { + let mut config = tokio_postgres::Config::from_str("host=localhost port=5432 user=dbx").unwrap(); + config.hostaddr("127.0.0.1".parse().unwrap()); + let pinned = pin_postgres_hostaddr(&config).await; + assert_eq!(pinned.get_hostaddrs().first(), Some(&"127.0.0.1".parse().unwrap())); + } }