Non-blocking resolution

geldata · Sep 27, 2024 · a060b82 · a060b82
1 parent a96e6e4
commit a060b82
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/edb/server/pgcon/rust_transport.py b/edb/server/pgcon/rust_transport.py
@@ -456,12 +456,16 @@ async def create_postgres_connection(
     )
 
     user_protocol = protocol_factory()
+    host_candidates = await asyncio.get_running_loop().run_in_executor(
+        executor=None, func=lambda: state.config.host_candidates
+    )
+
     protocol, host, port, initial_transport = await _create_connection(
         lambda: PGConnectionProtocol(
             state, user_protocol, ready_future, pg_state
         ),
         connect_timeout,
-        state.config.host_candidates,
+        host_candidates,
     )
 
     upgraded_transport = None

diff --git a/edb/server/pgrust/src/connection/mod.rs b/edb/server/pgrust/src/connection/mod.rs
@@ -96,20 +96,20 @@ pub enum ResolvedTarget {
 
 impl ResolvedTarget {
     /// Resolves the target addresses for a given host.
-    pub fn to_addrs_sync(host: dsn::Host) -> Result<Vec<ResolvedTarget>, std::io::Error> {
+    pub fn to_addrs_sync(host: &dsn::Host) -> Result<Vec<ResolvedTarget>, std::io::Error> {
         use std::net::{SocketAddr, ToSocketAddrs};
 
         let mut resolved_targets = Vec::new();
         let dsn::Host(host_type, port) = host;
         match host_type {
             HostType::Hostname(hostname) => {
-                let socket_addrs = (hostname.as_str(), port).to_socket_addrs()?;
+                let socket_addrs = (hostname.as_str(), *port).to_socket_addrs()?;
                 for addr in socket_addrs {
                     resolved_targets.push(Self::SocketAddr(addr));
                 }
             }
             HostType::IP(ip, None) => {
-                resolved_targets.push(ResolvedTarget::SocketAddr(SocketAddr::new(ip, port)))
+                resolved_targets.push(ResolvedTarget::SocketAddr(SocketAddr::new(*ip, *port)))
             }
             HostType::IP(std::net::IpAddr::V4(_), Some(_)) => {
                 return Err(std::io::Error::new(
@@ -120,7 +120,7 @@ impl ResolvedTarget {
             HostType::IP(std::net::IpAddr::V6(ip), Some(scope_id)) => {
                 if let Ok(scope_id) = str::parse::<u32>(&scope_id) {
                     resolved_targets.push(ResolvedTarget::SocketAddr(
-                        std::net::SocketAddrV6::new(ip, port, 0, scope_id).into(),
+                        std::net::SocketAddrV6::new(*ip, *port, 0, scope_id).into(),
                     ));
                 } else {
                     // TODO: Resolve non-numeric scope IDs

diff --git a/edb/server/pgrust/src/python.rs b/edb/server/pgrust/src/python.rs
@@ -88,13 +88,22 @@ impl PyConnectionParams {
         &self,
         py: Python,
     ) -> PyResult<Vec<(&'static str, Py<PyAny>, String, u16)>> {
+        // As this might be blocking, drop the GIL while we allow for
+        // resolution to take place.
+        let hosts = self.inner.hosts()?;
+        let hosts = py.allow_threads(|| {
+            hosts.into_iter().map(|host| {
+                let addrs = ResolvedTarget::to_addrs_sync(&host);
+                (host, addrs)
+            })
+        });
         let mut errors = Vec::new();
         let mut resolved_hosts = Vec::new();
 
-        for host in self.inner.hosts()? {
+        for (host, resolved) in hosts {
             let hostname = host.0.to_string();
             let port = host.1;
-            match ResolvedTarget::to_addrs_sync(host) {
+            match resolved {
                 Ok(addrs) => {
                     for addr in addrs {
                         match addr {