From a12b338aac4e730c4117877ae1e5bba89f5cdbad Mon Sep 17 00:00:00 2001
From: Conrad Ludgate <conradludgate@gmail.com>
Date: Tue, 23 Apr 2024 18:18:03 +0100
Subject: [PATCH] parse ip eagerly

---
 proxy/src/compute.rs | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/proxy/src/compute.rs b/proxy/src/compute.rs
index 3af561d232..6e14982844 100644
--- a/proxy/src/compute.rs
+++ b/proxy/src/compute.rs
@@ -11,7 +11,12 @@ use crate::{
 use futures::TryFutureExt;
 use itertools::Itertools;
 use pq_proto::StartupMessageParams;
-use std::{io, net::SocketAddr, time::Duration};
+use std::{
+    io,
+    net::{IpAddr, SocketAddr},
+    str::FromStr,
+    time::Duration,
+};
 use thiserror::Error;
 use tokio::net::TcpStream;
 use tokio_postgres::{
@@ -182,11 +187,15 @@ impl ConnCfg {
 
         // wrap TcpStream::connect with timeout
         let connect_with_timeout = |host, port| async move {
-            let addrs = dns
-                .resolve(host)
-                .await
-                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
-                .collect::<Vec<_>>();
+            let addrs = if let Ok(ip) = IpAddr::from_str(host) {
+                vec![ip]
+            } else {
+                dns.resolve(host)
+                    .await
+                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?
+                    .collect()
+            };
+
             let timeout = timeout / addrs.len() as u32;
 
             let mut last_err = None;