Fix connection_acquisition_timeout setting (#675)

The timeout was not respected should the server successfully exchange the
version handshake with the driver but then exceed the timeout by taking too long
to reply to `HELLO`.

Author: robsdedude

neo4j/_async/io/_bolt.py

@@ -23,7 +23,12 @@
 from time import perf_counter
 
 from ..._async_compat.network import AsyncBoltSocket
-from ..._exceptions import BoltHandshakeError
+from ..._async_compat.util import AsyncUtil
+from ..._exceptions import (
+    BoltError,
+    BoltHandshakeError,
+    SocketDeadlineExceeded,
+)
 from ...addressing import Address
 from ...api import (
     ServerInfo,
@@ -32,6 +37,7 @@
 from ...conf import PoolConfig
 from ...exceptions import (
     AuthError,
+    DriverError,
     IncompleteCommit,
     ServiceUnavailable,
     SessionExpired,
@@ -76,6 +82,7 @@ class AsyncBolt:
     idle_since = float("-inf")
 
     # The socket
+    _closing = False
     _closed = False
 
     # The socket
@@ -260,24 +267,42 @@ async def ping(cls, address, *, timeout=None, **config):
 
     @classmethod
     async def open(
-        cls, address, *, auth=None, timeout=None, routing_context=None, **pool_config
+        cls, address, *, auth=None, timeout=None, routing_context=None,
+        **pool_config
     ):
-        """ Open a new Bolt connection to a given server address.
+        """Open a new Bolt connection to a given server address.
 
         :param address:
         :param auth:
         :param timeout: the connection timeout in seconds
         :param routing_context: dict containing routing context
         :param pool_config:
-        :return:
-        :raise BoltHandshakeError: raised if the Bolt Protocol can not negotiate a protocol version.
+
+        :return: connected AsyncBolt instance
+
+        :raise BoltHandshakeError:
+            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()
         pool_config = PoolConfig.consume(pool_config)
+
+        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, pool_config.protocol_version, handshake, data = \
             await AsyncBoltSocket.connect(
                 address,
-                timeout=timeout,
+                timeout=socket_connection_timeout,
                 custom_resolver=pool_config.resolver,
                 ssl_context=pool_config.get_ssl_context(),
                 keep_alive=pool_config.keep_alive,
@@ -308,17 +333,31 @@ async def open(
             AsyncBoltSocket.close_socket(s)
 
             supported_versions = cls.protocol_handlers().keys()
-            raise BoltHandshakeError("The Neo4J server does not support communication with this driver. This driver have support for Bolt Protocols {}".format(supported_versions), address=address, request_data=handshake, response_data=data)
+            raise BoltHandshakeError(
+                "The Neo4J server does not support communication with this "
+                "driver. This driver have support for Bolt Protocols {}"
+                "".format(supported_versions),
+                address=address, request_data=handshake, response_data=data
+            )
 
         connection = bolt_cls(
             address, s, pool_config.max_connection_lifetime, auth=auth,
             user_agent=pool_config.user_agent, routing_context=routing_context
         )
 
         try:
-            await connection.hello()
+            connection.socket.set_deadline(time_remaining())
+            try:
+                await connection.hello()
+            except SocketDeadlineExceeded as e:
+                # connection._defunct = True
+                raise ServiceUnavailable(
+                    "Timeout during initial handshake occurred"
+                ) from e
+            finally:
+                connection.socket.set_deadline(None)
         except Exception:
-            await connection.close()
+            await connection.close_non_blocking()
             raise
 
         return connection
@@ -440,6 +479,11 @@ async def reset(self):
         """
         pass
 
+    @abc.abstractmethod
+    def goodbye(self):
+        """Append a GOODBYE message to the outgoing queued."""
+        pass
+
     def _append(self, signature, fields=(), response=None):
         """ Appends a message to the outgoing queue.
 
@@ -481,7 +525,8 @@ async def send_all(self):
         await self._send_all()
 
     @abc.abstractmethod
-    async def _fetch_message(self):
+    async def _process_message(self, details, summary_signature,
+                               summary_metadata):
         """ Receive at most one message from the server, if available.
 
         :return: 2-tuple of number of detail messages and number of summary
@@ -505,7 +550,12 @@ async def fetch_message(self):
         if not self.responses:
             return 0, 0
 
-        res = await self._fetch_message()
+        # Receive exactly one message
+        details, summary_signature, summary_metadata = \
+            await AsyncUtil.next(self.inbox)
+        res = await self._process_message(
+            details, summary_signature, summary_metadata
+        )
         self.idle_since = perf_counter()
         return res
 
@@ -548,9 +598,13 @@ async def _set_defunct(self, message, error=None, silent=False):
         # connection from the client side, and remove the address
         # from the connection pool.
         self._defunct = True
-        await self.close()
-        if self.pool:
-            await self.pool.deactivate(address=self.unresolved_address)
+        if not self._closing:
+            # If we fail while closing the connection, there is no need to
+            # remove the connection from the pool, nor to try to close the
+            # connection again.
+            await self.close()
+            if self.pool:
+                await self.pool.deactivate(address=self.unresolved_address)
         # Iterate through the outstanding responses, and if any correspond
         # to COMMIT requests then raise an error to signal that we are
         # unable to confirm that the COMMIT completed successfully.
@@ -584,11 +638,37 @@ def stale(self):
     def set_stale(self):
         self._stale = True
 
-    @abc.abstractmethod
     async def close(self):
-        """ Close the connection.
+        """Close the connection."""
+        if self._closed or self._closing:
+            return
+        self._closing = True
+        if not self._defunct:
+            self.goodbye()
+            try:
+                await self._send_all()
+            except (OSError, BoltError, DriverError):
+                pass
+        log.debug("[#%04X]  C: <CLOSE>", self.local_port)
+        try:
+            self.socket.close()
+        except OSError:
+            pass
+        finally:
+            self._closed = True
+
+    async def close_non_blocking(self):
+        """Set the socket to non-blocking and close it.
+
+        This will try to send the `GOODBYE` message (given the socket is not
+        marked as defunct). However, should the write operation require
+        blocking (e.g., a full network buffer), then the socket will be closed
+        immediately (without `GOODBYE` message).
         """
-        pass
+        if self._closed or self._closing:
+            return
+        self.socket.settimeout(0)
+        await self.close()
 
     @abc.abstractmethod
     def closed(self):

neo4j/_async/io/_bolt3.py

@@ -314,16 +314,17 @@ def fail(metadata):
         await self.send_all()
         await self.fetch_all()
 
-    async def _fetch_message(self):
-        """ Receive at most one message from the server, if available.
+    def goodbye(self):
+        log.debug("[#%04X]  C: GOODBYE", self.local_port)
+        self._append(b"\x02", ())
+
+    async def _process_message(self, details, summary_signature,
+                               summary_metadata):
+        """ Process at most one message from the server, if available.
 
         :return: 2-tuple of number of detail messages and number of summary
                  messages fetched
         """
-        # Receive exactly one message
-        details, summary_signature, summary_metadata = \
-            await AsyncUtil.next(self.inbox)
-
         if details:
             log.debug("[#%04X]  S: RECORD * %d", self.local_port, len(details))  # Do not log any data
             await self.responses[0].on_records(details)
@@ -363,25 +364,6 @@ async def _fetch_message(self):
 
         return len(details), 1
 
-    async def close(self):
-        """ Close the connection.
-        """
-        if not self._closed:
-            if not self._defunct:
-                log.debug("[#%04X]  C: GOODBYE", self.local_port)
-                self._append(b"\x02", ())
-                try:
-                    await self._send_all()
-                except (OSError, BoltError, DriverError):
-                    pass
-            log.debug("[#%04X]  C: <CLOSE>", self.local_port)
-            try:
-                self.socket.close()
-            except OSError:
-                pass
-            finally:
-                self._closed = True
-
     def closed(self):
         return self._closed
 

neo4j/_async/io/_bolt4.py

@@ -20,10 +20,7 @@
 from ssl import SSLSocket
 
 from ..._async_compat.util import AsyncUtil
-from ..._exceptions import (
-    BoltError,
-    BoltProtocolError,
-)
+from ..._exceptions import BoltProtocolError
 from ...api import (
     READ_ACCESS,
     SYSTEM_DATABASE,
@@ -32,7 +29,6 @@
 from ...exceptions import (
     ConfigurationError,
     DatabaseUnavailable,
-    DriverError,
     ForbiddenOnReadOnlyDatabase,
     Neo4jError,
     NotALeader,
@@ -265,16 +261,17 @@ def fail(metadata):
         await self.send_all()
         await self.fetch_all()
 
-    async def _fetch_message(self):
-        """ Receive at most one message from the server, if available.
+    def goodbye(self):
+        log.debug("[#%04X]  C: GOODBYE", self.local_port)
+        self._append(b"\x02", ())
+
+    async def _process_message(self, details, summary_signature,
+                               summary_metadata):
+        """ Process at most one message from the server, if available.
 
         :return: 2-tuple of number of detail messages and number of summary
                  messages fetched
         """
-        # Receive exactly one message
-        details, summary_signature, summary_metadata = \
-            await AsyncUtil.next(self.inbox)
-
         if details:
             log.debug("[#%04X]  S: RECORD * %d", self.local_port, len(details))  # Do not log any data
             await self.responses[0].on_records(details)
@@ -315,25 +312,6 @@ async def _fetch_message(self):
 
         return len(details), 1
 
-    async def close(self):
-        """ Close the connection.
-        """
-        if not self._closed:
-            if not self._defunct:
-                log.debug("[#%04X]  C: GOODBYE", self.local_port)
-                self._append(b"\x02", ())
-                try:
-                    await self._send_all()
-                except (OSError, BoltError, DriverError):
-                    pass
-            log.debug("[#%04X]  C: <CLOSE>", self.local_port)
-            try:
-                self.socket.close()
-            except OSError:
-                pass
-            finally:
-                self._closed = True
-
     def closed(self):
         return self._closed
 

neo4j/_async/io/_pool.py

@@ -112,15 +112,14 @@ async def _acquire_from_pool(self, address, health_check):
                 return connection
         return None
 
-    async def _acquire_new(self, address, timeout, health_check):
+    async def _acquire_new(self, address, timeout):
         connections = self.connections[address]
         max_pool_size = self.pool_config.max_connection_pool_size
         infinite_pool_size = (max_pool_size < 0
                               or max_pool_size == float("inf"))
         can_create_new_connection = (infinite_pool_size
                                      or len(connections) < max_pool_size)
         if can_create_new_connection:
-            timeout = min(self.pool_config.connection_timeout, timeout)
             try:
                 connection = await self.opener(address, timeout)
             except ServiceUnavailable:
@@ -170,9 +169,7 @@ async def health_check(connection_):
                 if connection:
                     return connection
                 # all connections in pool are in-use
-                connection = await self._acquire_new(
-                    address, time_remaining(), health_check
-                )
+                connection = await self._acquire_new(address, time_remaining())
                 if connection:
                     return connection