blocking client: fix connect and timeout

edgedb/blocking_client.py

@@ -47,21 +47,57 @@ async def connect_addr(self, addr, timeout):
 
         if isinstance(addr, str):
             # UNIX socket
-            sock = socket.socket(socket.AF_UNIX)
+            res_list = [(socket.AF_UNIX, socket.SOCK_STREAM, -1, None, addr)]
         else:
-            sock = socket.socket(socket.AF_INET)
-
-        try:
-            sock.settimeout(timeout)
+            host, port = addr
+            try:
+                # getaddrinfo() doesn't take timeout!!
+                res_list = socket.getaddrinfo(
+                    host, port, socket.AF_UNSPEC, socket.SOCK_STREAM
+                )
+            except socket.gaierror as e:
+                # All name resolution errors are considered temporary
+                err = errors.ClientConnectionFailedTemporarilyError(str(e))
+                raise err from e
 
+        for i, res in enumerate(res_list):
+            af, socktype, proto, _, sa = res
             try:
-                sock.connect(addr)
+                sock = socket.socket(af, socktype, proto)
+            except OSError as e:
+                sock.close()
+                if i < len(res_list) - 1:
+                    continue
+                else:
+                    raise con_utils.wrap_error(e) from e
+            try:
+                await self._connect_addr(sock, addr, sa, deadline)
+            except TimeoutError:
+                raise
+            except Exception:
+                if i < len(res_list) - 1:
+                    continue
+                else:
+                    raise
+            else:
+                break
 
-                if not isinstance(addr, str):
-                    time_left = deadline - time.monotonic()
-                    if time_left <= 0:
-                        raise TimeoutError
+    async def _connect_addr(self, sock, addr, sa, deadline):
+        try:
+            time_left = deadline - time.monotonic()
+            if time_left <= 0:
+                raise TimeoutError
+            try:
+                sock.settimeout(time_left)
+                sock.connect(sa)
+            except OSError as e:
+                raise con_utils.wrap_error(e) from e
 
+            if not isinstance(addr, str):
+                time_left = deadline - time.monotonic()
+                if time_left <= 0:
+                    raise TimeoutError
+                try:
                     # Upgrade to TLS
                     sock.settimeout(time_left)
                     try:
@@ -74,12 +110,8 @@ async def connect_addr(self, addr, timeout):
                         raise con_utils.wrap_error(e) from e
                     else:
                         con_utils.check_alpn_protocol(sock)
-            except socket.gaierror as e:
-                # All name resolution errors are considered temporary
-                err = errors.ClientConnectionFailedTemporarilyError(str(e))
-                raise err from e
-            except OSError as e:
-                raise con_utils.wrap_error(e) from e
+                except OSError as e:
+                    raise con_utils.wrap_error(e) from e
 
             time_left = deadline - time.monotonic()
             if time_left <= 0:
@@ -92,9 +124,9 @@ async def connect_addr(self, addr, timeout):
             proto.set_connection(self)
 
             try:
-                sock.settimeout(time_left)
-                await proto.connect()
-                sock.settimeout(None)
+                await proto.wait_for(proto.connect(), time_left)
+            except TimeoutError:
+                raise
             except OSError as e:
                 raise con_utils.wrap_error(e) from e
 
@@ -133,6 +165,9 @@ async def close(self, timeout=None):
                     await self._protocol.wait_for(
                         self._protocol.wait_for_disconnect(), timeout
                     )
+            except TimeoutError:
+                self.terminate()
+                raise errors.QueryTimeoutError()
             except Exception:
                 self.terminate()
                 raise

edgedb/protocol/blocking_proto.pyx

@@ -63,18 +63,14 @@ cdef class BlockingIOProtocol(protocol.SansIOProtocolBackwardsCompatible):
     async def wait_for_message(self):
         cdef float timeout
         if self.deadline > 0:
-            timeout = self.deadline - time.monotonic()
-            if timeout <= 0:
-                self.abort()
-                raise errors.QueryTimeoutError()
             while not self.buffer.take_message():
+                timeout = self.deadline - time.monotonic()
+                if timeout <= 0:
+                    self.abort()
+                    raise TimeoutError
                 try:
                     self.sock.settimeout(timeout)
                     data = self.sock.recv(RECV_BUF)
-                    timeout = self.deadline - time.monotonic()
-                    if timeout <= 0:
-                        self.abort()
-                        raise TimeoutError
                 except OSError as e:
                     self._disconnect()
                     raise con_utils.wrap_error(e) from e