Improve connection error handling infra

Connection now takes a dedicated error handler object in constructor
instead of having one dynamically attached after creation.

Author: lutovich

neo4j/bolt/connection.py

@@ -119,6 +119,26 @@ def supports_bytes(self):
         return self.version_info() >= (3, 2)
 
 
+class ConnectionErrorHandler(object):
+    """ A handler for send and receive errors.
+    """
+
+    def __init__(self, handlers_by_error_class=None):
+        if handlers_by_error_class is None:
+            handlers_by_error_class = {}
+
+        self.handlers_by_error_class = handlers_by_error_class
+        self.known_errors = tuple(handlers_by_error_class.keys())
+
+    def handle(self, error, address):
+        try:
+            error_class = error.__class__
+            handler = self.handlers_by_error_class[error_class]
+            handler(address)
+        except KeyError:
+            pass
+
+
 class Connection(object):
     """ Server connection for Bolt protocol v1.
 
@@ -144,15 +164,14 @@ class Connection(object):
     #: Error class used for raising connection errors
     Error = ServiceUnavailable
 
-    #: The function to handle send and receive errors
-    error_handler = None
-
     _supports_statement_reuse = False
 
     _last_run_statement = None
 
-    def __init__(self, sock, **config):
+    def __init__(self, address, sock, error_handler, **config):
+        self.address = address
         self.socket = sock
+        self.error_handler = error_handler
         self.server = ServerInfo(SocketAddress.from_socket(sock))
         self.input_buffer = ChunkedInputBuffer()
         self.output_buffer = ChunkedOutputBuffer()
@@ -242,9 +261,8 @@ def reset(self):
     def send(self):
         try:
             self._send()
-        except Exception as error:
-            if self.error_handler is not None:
-                self.error_handler(error)
+        except self.error_handler.known_errors as error:
+            self.error_handler.handle(error, self.address)
             raise error
 
     def _send(self):
@@ -263,9 +281,8 @@ def _send(self):
     def fetch(self):
         try:
             return self._fetch()
-        except Exception as error:
-            if self.error_handler is not None:
-                self.error_handler(error)
+        except self.error_handler.known_errors as error:
+            self.error_handler.handle(error, self.address)
             raise error
 
     def _fetch(self):
@@ -379,8 +396,9 @@ class ConnectionPool(object):
 
     _closed = False
 
-    def __init__(self, connector):
+    def __init__(self, connector, connection_error_handler):
         self.connector = connector
+        self.connection_error_handler = connection_error_handler
         self.connections = {}
         self.lock = RLock()
 
@@ -414,7 +432,7 @@ def acquire_direct(self, address):
                     connection.in_use = True
                     return connection
             try:
-                connection = self.connector(address)
+                connection = self.connector(address, self.connection_error_handler)
             except ServiceUnavailable:
                 self.remove(address)
                 raise
@@ -476,7 +494,7 @@ def closed(self):
             return self._closed
 
 
-def connect(address, ssl_context=None, **config):
+def connect(address, ssl_context=None, error_handler=None, **config):
     """ Connect and perform a handshake and return a valid Connection object, assuming
     a protocol version can be agreed.
     """
@@ -582,7 +600,8 @@ def connect(address, ssl_context=None, **config):
         s.shutdown(SHUT_RDWR)
         s.close()
     elif agreed_version == 1:
-        return Connection(s, der_encoded_server_certificate=der_encoded_server_certificate, **config)
+        return Connection(address, s, der_encoded_server_certificate=der_encoded_server_certificate,
+                          error_handler=error_handler, **config)
     elif agreed_version == 0x48545450:
         log_error("S: [CLOSE]")
         s.close()

neo4j/v1/direct.py

@@ -20,17 +20,25 @@
 
 
 from neo4j.addressing import SocketAddress, resolve
-from neo4j.bolt import DEFAULT_PORT, ConnectionPool, connect
+from neo4j.bolt import DEFAULT_PORT, ConnectionPool, connect, ConnectionErrorHandler
 from neo4j.exceptions import ServiceUnavailable
 from neo4j.v1.api import Driver
 from neo4j.v1.security import SecurityPlan
 from neo4j.v1.session import BoltSession
 
 
+class DirectConnectionErrorHandler(ConnectionErrorHandler):
+    """ Handler for errors in direct driver connections.
+    """
+
+    def __init__(self):
+        super(DirectConnectionErrorHandler, self).__init__({}) # does not need to handle errors
+
+
 class DirectConnectionPool(ConnectionPool):
 
     def __init__(self, connector, address):
-        super(DirectConnectionPool, self).__init__(connector)
+        super(DirectConnectionPool, self).__init__(connector, DirectConnectionErrorHandler())
         self.address = address
 
     def acquire(self, access_mode=None):
@@ -61,7 +69,11 @@ def __init__(self, uri, **config):
         self.address = SocketAddress.from_uri(uri, DEFAULT_PORT)
         self.security_plan = security_plan = SecurityPlan.build(**config)
         self.encrypted = security_plan.encrypted
-        pool = DirectConnectionPool(lambda a: connect(a, security_plan.ssl_context, **config), self.address)
+
+        def connector(address, error_handler):
+            return connect(address, security_plan.ssl_context, error_handler, **config)
+
+        pool = DirectConnectionPool(connector, self.address)
         pool.release(pool.acquire())
         Driver.__init__(self, pool, **config)
 

neo4j/v1/routing.py

@@ -23,7 +23,7 @@
 from time import clock
 
 from neo4j.addressing import SocketAddress, resolve
-from neo4j.bolt import ConnectionPool, ServiceUnavailable, ProtocolError, DEFAULT_PORT, connect
+from neo4j.bolt import ConnectionPool, ServiceUnavailable, ProtocolError, DEFAULT_PORT, connect, ConnectionErrorHandler
 from neo4j.compat.collections import MutableSet, OrderedDict
 from neo4j.exceptions import CypherError, DatabaseUnavailableError, NotALeaderError, ForbiddenOnReadOnlyDatabaseError
 from neo4j.util import ServerVersion
@@ -32,7 +32,6 @@
 from neo4j.v1.security import SecurityPlan
 from neo4j.v1.session import BoltSession
 
-
 LOAD_BALANCING_STRATEGY_LEAST_CONNECTED = 0
 LOAD_BALANCING_STRATEGY_ROUND_ROBIN = 1
 LOAD_BALANCING_STRATEGY_DEFAULT = LOAD_BALANCING_STRATEGY_LEAST_CONNECTED
@@ -247,15 +246,26 @@ def _select(self, offset, addresses):
                 return least_connected_address
 
 
+class RoutingConnectionErrorHandler(ConnectionErrorHandler):
+    """ Handler for errors in routing driver connections.
+    """
+
+    def __init__(self, pool):
+        super(RoutingConnectionErrorHandler, self).__init__({
+            SessionExpired: lambda address: pool.remove(address),
+            ServiceUnavailable: lambda address: pool.remove(address),
+            DatabaseUnavailableError: lambda address: pool.remove(address),
+            NotALeaderError: lambda address: pool.remove_writer(address),
+            ForbiddenOnReadOnlyDatabaseError: lambda address: pool.remove_writer(address)
+        })
+
+
 class RoutingConnectionPool(ConnectionPool):
     """ Connection pool with routing table.
     """
 
-    CLUSTER_MEMBER_FAILURE_ERRORS = (ServiceUnavailable, SessionExpired, DatabaseUnavailableError)
-    WRITE_FAILURE_ERRORS = (NotALeaderError, ForbiddenOnReadOnlyDatabaseError)
-
     def __init__(self, connector, initial_address, routing_context, *routers, **config):
-        super(RoutingConnectionPool, self).__init__(connector)
+        super(RoutingConnectionPool, self).__init__(connector, RoutingConnectionErrorHandler(self))
         self.initial_address = initial_address
         self.routing_context = routing_context
         self.routing_table = RoutingTable(routers)
@@ -402,21 +412,12 @@ def acquire(self, access_mode=None):
             try:
                 connection = self.acquire_direct(address)  # should always be a resolved address
                 connection.Error = SessionExpired
-                connection.error_handler = lambda error: self._handle_connection_error(address, error)
             except ServiceUnavailable:
                 self.remove(address)
             else:
                 return connection
         raise SessionExpired("Failed to obtain connection towards '%s' server." % access_mode)
 
-    def _handle_connection_error(self, address, error):
-        """ Handle routing connection send or receive error.
-        """
-        if isinstance(error, self.CLUSTER_MEMBER_FAILURE_ERRORS):
-            self.remove(address)
-        elif isinstance(error, self.WRITE_FAILURE_ERRORS):
-            self._remove_writer(address)
-
     def remove(self, address):
         """ Remove an address from the connection pool, if present, closing
         all connections to that address. Also remove from the routing table.
@@ -428,7 +429,7 @@ def remove(self, address):
         self.routing_table.writers.discard(address)
         super(RoutingConnectionPool, self).remove(address)
 
-    def _remove_writer(self, address):
+    def remove_writer(self, address):
         """ Remove a writer address from the routing table, if present.
         """
         self.routing_table.writers.discard(address)
@@ -450,8 +451,8 @@ def __init__(self, uri, **config):
             # scenario right now
             raise ValueError("TRUST_ON_FIRST_USE is not compatible with routing")
 
-        def connector(a):
-            return connect(a, security_plan.ssl_context, **config)
+        def connector(address, error_handler):
+            return connect(address, security_plan.ssl_context, error_handler, **config)
 
         pool = RoutingConnectionPool(connector, initial_address, routing_context, *resolve(initial_address), **config)
         try:

test/integration/test_connection.py

@@ -21,7 +21,7 @@
 
 from socket import create_connection
 
-from neo4j.v1 import ConnectionPool, ServiceUnavailable
+from neo4j.v1 import ConnectionPool, ServiceUnavailable, DirectConnectionErrorHandler
 
 from test.integration.tools import IntegrationTestCase
 
@@ -45,10 +45,14 @@ def defunct(self):
         return False
 
 
+def connector(address, _):
+    return QuickConnection(create_connection(address))
+
+
 class ConnectionPoolTestCase(IntegrationTestCase):
 
     def setUp(self):
-        self.pool = ConnectionPool(lambda a: QuickConnection(create_connection(a)))
+        self.pool = ConnectionPool(connector, DirectConnectionErrorHandler())
 
     def tearDown(self):
         self.pool.close()
@@ -104,7 +108,7 @@ def test_releasing_twice(self):
         self.assert_pool_size(address, 0, 1)
 
     def test_cannot_acquire_after_close(self):
-        with ConnectionPool(lambda a: QuickConnection(create_connection(a))) as pool:
+        with ConnectionPool(lambda a: QuickConnection(create_connection(a)), DirectConnectionErrorHandler()) as pool:
             pool.close()
             with self.assertRaises(ServiceUnavailable):
                 _ = pool.acquire_direct("X")

test/stub/test_routing.py

@@ -50,8 +50,8 @@
 UNREACHABLE_ADDRESS = ("127.0.0.1", 8080)
 
 
-def connector(address):
-    return connect(address, auth=basic_auth("neotest", "neotest"))
+def connector(address, error_handler):
+    return connect(address, error_handler=error_handler, auth=basic_auth("neotest", "neotest"))
 
 
 def RoutingPool(*routers):

test/unit/test_routing.py

@@ -52,8 +52,8 @@
 }
 
 
-def connector(address):
-    return connect(address, auth=basic_auth("neotest", "neotest"))
+def connector(address, error_handler):
+    return connect(address, error_handler=error_handler, auth=basic_auth("neotest", "neotest"))
 
 
 class RoundRobinSetTestCase(TestCase):