Fix bolt handshake not having a timeout (#905)

During the bolt version negotiation the socket was not constrained by any
timeout. This could, in rare occasions cause the driver to get stuck until some
other mechanism (e.g., socket keep alive) would free it.

The handshake is not limited by the `connection_acquisition_timeout`.

Author: robsdedude

src/neo4j/_async/io/_bolt.py

@@ -28,6 +28,7 @@
 from ..._codec.hydration import v1 as hydration_v1
 from ..._codec.packstream import v1 as packstream_v1
 from ..._conf import PoolConfig
+from ..._deadline import Deadline
 from ..._exceptions import (
     BoltError,
     BoltHandshakeError,
@@ -289,17 +290,21 @@ def get_handshake(cls):
         return b"".join(version.to_bytes() for version in offered_versions).ljust(16, b"\x00")
 
     @classmethod
-    async def ping(cls, address, *, timeout=None, pool_config=None):
+    async def ping(cls, address, *, deadline=None, pool_config=None):
         """ Attempt to establish a Bolt connection, returning the
         agreed Bolt protocol version if successful.
         """
         if pool_config is None:
             pool_config = PoolConfig()
+        if deadline is None:
+            deadline = Deadline(None)
+
         try:
             s, protocol_version, handshake, data = \
                 await AsyncBoltSocket.connect(
                     address,
-                    timeout=timeout,
+                    tcp_timeout=pool_config.connection_timeout,
+                    deadline=deadline,
                     custom_resolver=pool_config.resolver,
                     ssl_context=pool_config.get_ssl_context(),
                     keep_alive=pool_config.keep_alive,
@@ -313,14 +318,14 @@ async def ping(cls, address, *, timeout=None, pool_config=None):
     # [bolt-version-bump] search tag when changing bolt version support
     @classmethod
     async def open(
-        cls, address, *, auth=None, timeout=None, routing_context=None,
+        cls, address, *, auth=None, deadline=None, routing_context=None,
         pool_config=None
     ):
         """Open a new Bolt connection to a given server address.
 
         :param address:
         :param auth:
-        :param timeout: the connection timeout in seconds
+        :param deadline: how long to wait for the connection to be established
         :param routing_context: dict containing routing context
         :param pool_config:
 
@@ -330,26 +335,17 @@ async def open(
             raised if the Bolt Protocol can not negotiate a protocol version.
         :raise ServiceUnavailable: raised if there was a connection issue.
         """
-        def time_remaining():
-            if timeout is None:
-                return None
-            t = timeout - (perf_counter() - t0)
-            return t if t > 0 else 0
 
-        t0 = perf_counter()
         if pool_config is None:
             pool_config = PoolConfig()
+        if deadline is None:
+            deadline = Deadline(None)
 
-        socket_connection_timeout = pool_config.connection_timeout
-        if socket_connection_timeout is None:
-            socket_connection_timeout = time_remaining()
-        elif timeout is not None:
-            socket_connection_timeout = min(pool_config.connection_timeout,
-                                            time_remaining())
         s, protocol_version, handshake, data = \
             await AsyncBoltSocket.connect(
                 address,
-                timeout=socket_connection_timeout,
+                tcp_timeout=pool_config.connection_timeout,
+                deadline=deadline,
                 custom_resolver=pool_config.resolver,
                 ssl_context=pool_config.get_ssl_context(),
                 keep_alive=pool_config.keep_alive,
@@ -410,7 +406,7 @@ def time_remaining():
         )
 
         try:
-            connection.socket.set_deadline(time_remaining())
+            connection.socket.set_deadline(deadline)
             try:
                 await connection.hello()
             finally:

src/neo4j/_async/io/_pool.py

@@ -135,9 +135,7 @@ async def connection_creator():
             released_reservation = False
             try:
                 try:
-                    connection = await self.opener(
-                        address, deadline.to_timeout()
-                    )
+                    connection = await self.opener(address, deadline)
                 except ServiceUnavailable:
                     await self.deactivate(address)
                     raise
@@ -382,9 +380,9 @@ def open(cls, address, *, auth, pool_config, workspace_config):
         :returns: BoltPool
         """
 
-        async def opener(addr, timeout):
+        async def opener(addr, deadline):
             return await AsyncBolt.open(
-                addr, auth=auth, timeout=timeout, routing_context=None,
+                addr, auth=auth, deadline=deadline, routing_context=None,
                 pool_config=pool_config
             )
 
@@ -437,9 +435,9 @@ def open(cls, *addresses, auth, pool_config, workspace_config,
             raise ConfigurationError("The key 'address' is reserved for routing context.")
         routing_context["address"] = str(address)
 
-        async def opener(addr, timeout):
+        async def opener(addr, deadline):
             return await AsyncBolt.open(
-                addr, auth=auth, timeout=timeout,
+                addr, auth=auth, deadline=deadline,
                 routing_context=routing_context, pool_config=pool_config
             )
 

src/neo4j/_async_compat/network/_bolt_socket.py

@@ -272,11 +272,11 @@ async def _connect_secure(cls, resolved_address, timeout, keep_alive, ssl):
                 ) from error
             raise
 
-    async def _handshake(self, resolved_address):
+    async def _handshake(self, resolved_address, deadline):
         """
 
-        :param s: Socket
         :param resolved_address:
+        :param deadline: Deadline for handshake
 
         :returns: (socket, version, client_handshake, server_response_data)
         """
@@ -296,47 +296,52 @@ async def _handshake(self, resolved_address):
         log.debug("[#%04X]  C: <HANDSHAKE> %s %s %s %s", local_port,
                   *supported_versions)
 
-        data = self.Bolt.MAGIC_PREAMBLE + self.Bolt.get_handshake()
-        await self.sendall(data)
+        request = self.Bolt.MAGIC_PREAMBLE + self.Bolt.get_handshake()
 
         # Handle the handshake response
         original_timeout = self.gettimeout()
-        if original_timeout is not None:
-            self.settimeout(original_timeout + 1)
+        self.settimeout(deadline.to_timeout())
         try:
-            data = await self.recv(4)
+            await self.sendall(request)
+            response = await self.recv(4)
         except OSError as exc:
             raise ServiceUnavailable(
-                "Failed to read any data from server {!r} "
-                "after connected".format(resolved_address)) from exc
+                f"Failed to read any data from server {resolved_address!r} "
+                f"after connected (deadline {deadline})"
+            ) from exc
         finally:
             self.settimeout(original_timeout)
-        data_size = len(data)
+        data_size = len(response)
         if data_size == 0:
             # If no data is returned after a successful select
             # response, the server has closed the connection
             log.debug("[#%04X]  S: <CLOSE>", local_port)
             await self.close()
             raise ServiceUnavailable(
-                "Connection to {address} closed without handshake response".format(
-                    address=resolved_address))
+                f"Connection to {resolved_address} closed without handshake "
+                "response"
+            )
         if data_size != 4:
             # Some garbled data has been received
             log.debug("[#%04X]  S: @*#!", local_port)
             await self.close()
             raise BoltProtocolError(
-                "Expected four byte Bolt handshake response from %r, received %r instead; check for incorrect port number" % (
-                resolved_address, data), address=resolved_address)
-        elif data == b"HTTP":
+                "Expected four byte Bolt handshake response from "
+                f"{resolved_address!r}, received {response!r} instead; "
+                "check for incorrect port number"
+                , address=resolved_address
+            )
+        elif response == b"HTTP":
             log.debug("[#%04X]  S: <CLOSE>", local_port)
             await self.close()
             raise ServiceUnavailable(
-                "Cannot to connect to Bolt service on {!r} "
-                "(looks like HTTP)".format(resolved_address))
-        agreed_version = data[-1], data[-2]
+                f"Cannot to connect to Bolt service on {resolved_address!r} "
+                "(looks like HTTP)"
+            )
+        agreed_version = response[-1], response[-2]
         log.debug("[#%04X]  S: <HANDSHAKE> 0x%06X%02X", local_port,
                   agreed_version[1], agreed_version[0])
-        return self, agreed_version, handshake, data
+        return self, agreed_version, handshake, response
 
     @classmethod
     async def close_socket(cls, socket_):
@@ -356,8 +361,8 @@ async def close_socket(cls, socket_):
                 pass
 
     @classmethod
-    async def connect(cls, address, *, timeout, custom_resolver, ssl_context,
-                      keep_alive):
+    async def connect(cls, address, *, tcp_timeout, deadline,
+                      custom_resolver, ssl_context, keep_alive):
         """ Connect and perform a handshake and return a valid Connection object,
         assuming a protocol version can be agreed.
         """
@@ -371,12 +376,18 @@ async def connect(cls, address, *, timeout, custom_resolver, ssl_context,
             addressing.Address(address), resolver=custom_resolver
         )
         async for resolved_address in resolved_addresses:
+            deadline_timeout = deadline.to_timeout()
+            if (
+                deadline_timeout is not None
+                and deadline_timeout <= tcp_timeout
+            ):
+                tcp_timeout = deadline_timeout
             s = None
             try:
                 s = await cls._connect_secure(
-                    resolved_address, timeout, keep_alive, ssl_context
+                    resolved_address, tcp_timeout, keep_alive, ssl_context
                 )
-                return await s._handshake(resolved_address)
+                return await s._handshake(resolved_address, deadline)
             except (BoltError, DriverError, OSError) as error:
                 try:
                     local_port = s.getsockname()[1]
@@ -560,11 +571,12 @@ def _secure(cls, s, host, ssl_context):
         return s
 
     @classmethod
-    def _handshake(cls, s, resolved_address):
+    def _handshake(cls, s, resolved_address, deadline):
         """
 
         :param s: Socket
         :param resolved_address:
+        :param deadline:
 
         :returns: (socket, version, client_handshake, server_response_data)
         """
@@ -584,46 +596,52 @@ def _handshake(cls, s, resolved_address):
         log.debug("[#%04X]  C: <HANDSHAKE> %s %s %s %s", local_port,
                   *supported_versions)
 
-        data = cls.Bolt.MAGIC_PREAMBLE + cls.Bolt.get_handshake()
-        s.sendall(data)
+        request = cls.Bolt.MAGIC_PREAMBLE + cls.Bolt.get_handshake()
 
         # Handle the handshake response
-        ready_to_read = False
-        with selectors.DefaultSelector() as selector:
-            selector.register(s, selectors.EVENT_READ)
-            selector.select(1)
+        original_timeout = s.gettimeout()
+        s.settimeout(deadline.to_timeout())
         try:
-            data = s.recv(4)
+            s.sendall(request)
+            response = s.recv(4)
         except OSError as exc:
             raise ServiceUnavailable(
-                "Failed to read any data from server {!r} "
-                "after connected".format(resolved_address)) from exc
-        data_size = len(data)
+                f"Failed to read any data from server {resolved_address!r} "
+                f"after connected (deadline {deadline})"
+            ) from exc
+        finally:
+            s.settimeout(original_timeout)
+        data_size = len(response)
         if data_size == 0:
             # If no data is returned after a successful select
             # response, the server has closed the connection
             log.debug("[#%04X]  S: <CLOSE>", local_port)
             cls.close_socket(s)
             raise ServiceUnavailable(
-                "Connection to {address} closed without handshake response".format(
-                    address=resolved_address))
+                f"Connection to {resolved_address} closed without handshake "
+                "response"
+            )
         if data_size != 4:
             # Some garbled data has been received
             log.debug("[#%04X]  S: @*#!", local_port)
             cls.close_socket(s)
             raise BoltProtocolError(
-                "Expected four byte Bolt handshake response from %r, received %r instead; check for incorrect port number" % (
-                resolved_address, data), address=resolved_address)
-        elif data == b"HTTP":
+                "Expected four byte Bolt handshake response from "
+                f"{resolved_address!r}, received {response!r} instead; "
+                "check for incorrect port number"
+                , address=resolved_address
+            )
+        elif response == b"HTTP":
             log.debug("[#%04X]  S: <CLOSE>", local_port)
             cls.close_socket(s)
             raise ServiceUnavailable(
-                "Cannot to connect to Bolt service on {!r} "
-                "(looks like HTTP)".format(resolved_address))
-        agreed_version = data[-1], data[-2]
+                f"Cannot to connect to Bolt service on {resolved_address!r} "
+                "(looks like HTTP)"
+            )
+        agreed_version = response[-1], response[-2]
         log.debug("[#%04X]  S: <HANDSHAKE> 0x%06X%02X", local_port,
                   agreed_version[1], agreed_version[0])
-        return cls(s), agreed_version, handshake, data
+        return cls(s), agreed_version, handshake, response
 
     @classmethod
     def close_socket(cls, socket_):
@@ -639,8 +657,8 @@ def close_socket(cls, socket_):
             pass
 
     @classmethod
-    def connect(cls, address, *, timeout, custom_resolver, ssl_context,
-                keep_alive):
+    def connect(cls, address, *, tcp_timeout, deadline, custom_resolver,
+                ssl_context, keep_alive):
         """ Connect and perform a handshake and return a valid Connection object,
         assuming a protocol version can be agreed.
         """
@@ -653,12 +671,19 @@ def connect(cls, address, *, timeout, custom_resolver, ssl_context,
             addressing.Address(address), resolver=custom_resolver
         )
         for resolved_address in resolved_addresses:
+            deadline_timeout = deadline.to_timeout()
+            if (
+                deadline_timeout is not None
+                and deadline_timeout <= tcp_timeout
+            ):
+                tcp_timeout = deadline_timeout
             s = None
             try:
-                s = BoltSocket._connect(resolved_address, timeout, keep_alive)
+                s = BoltSocket._connect(resolved_address, tcp_timeout,
+                                        keep_alive)
                 s = BoltSocket._secure(s, resolved_address._host_name,
                                        ssl_context)
-                return BoltSocket._handshake(s, resolved_address)
+                return BoltSocket._handshake(s, resolved_address, deadline)
             except (BoltError, DriverError, OSError) as error:
                 try:
                     local_port = s.getsockname()[1]

src/neo4j/_deadline.py

@@ -72,6 +72,9 @@ def from_timeout_or_deadline(cls, timeout):
             return timeout
         return cls(timeout)
 
+    def __str__(self):
+        return f"Deadline(timeout={self._original_timeout})"
+
 
 merge_deadlines = min