Routing driver forgets addresses on some errors

lutovich · lutovich · commit b6668840f00d · 2017-08-23T10:31:46.000+02:00
This commit makes routing driver forget addresses and purge their
connections on network and database errors. It also makes driver
forget write server addresses when they fail as writers, connections
remain in the pool because such machines can remain readers/routers.
diff --git a/neo4j/bolt/connection.py b/neo4j/bolt/connection.py
@@ -144,6 +144,9 @@ 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
@@ -237,6 +240,14 @@ def reset(self):
         self.sync()
 
     def send(self):
+        try:
+            self._send()
+        except Exception as error:
+            if self.error_handler is not None:
+                self.error_handler(error)
+            raise error
+
+    def _send(self):
         """ Send all queued messages to the server.
         """
         data = self.output_buffer.view()
@@ -250,6 +261,14 @@ def send(self):
         self.output_buffer.clear()
 
     def fetch(self):
+        try:
+            return self._fetch()
+        except Exception as error:
+            if self.error_handler is not None:
+                self.error_handler(error)
+            raise error
+
+    def _fetch(self):
         """ Receive at least one message from the server, if available.
 
         :return: 2-tuple of number of detail messages and number of summary messages fetched
diff --git a/neo4j/v1/api.py b/neo4j/v1/api.py
@@ -277,7 +277,7 @@ def _disconnect(self, sync):
             if sync:
                 try:
                     self._connection.sync()
-                except ServiceUnavailable:
+                except (SessionError, ServiceUnavailable):
                     pass
             if self._connection:
                 self._connection.in_use = False
diff --git a/neo4j/v1/routing.py b/neo4j/v1/routing.py
@@ -251,6 +251,9 @@ class RoutingConnectionPool(ConnectionPool):
     """ Connection pool with routing table.
     """
 
+    FAILURE_CODES = ("Neo.TransientError.General.DatabaseUnavailable")
+    WRITE_FAILURE_CODES = ("Neo.ClientError.Cluster.NotALeader", "Neo.ClientError.General.ForbiddenOnReadOnlyDatabase")
+
     def __init__(self, connector, initial_address, routing_context, *routers, **config):
         super(RoutingConnectionPool, self).__init__(connector)
         self.initial_address = initial_address
@@ -399,12 +402,21 @@ 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, (SessionExpired, ServiceUnavailable)) or error.code in self.FAILURE_CODES:
+            self.remove(address)
+        elif error.code in self.WRITE_FAILURE_CODES:
+            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.
@@ -416,6 +428,11 @@ def remove(self, address):
         self.routing_table.writers.discard(address)
         super(RoutingConnectionPool, self).remove(address)
 
+    def _remove_writer(self, address):
+        """ Remove a writer address from the routing table, if present.
+        """
+        self.routing_table.writers.discard(address)
+
 
 class RoutingDriver(Driver):
     """ A :class:`.RoutingDriver` is created from a ``bolt+routing`` URI. The
diff --git a/test/stub/scripts/database_unavailable.script b/test/stub/scripts/database_unavailable.script
@@ -0,0 +1,12 @@
+!: AUTO INIT
+!: AUTO RESET
+!: AUTO PULL_ALL
+!: AUTO ACK_FAILURE
+!: AUTO RUN "ROLLBACK" {}
+!: AUTO RUN "BEGIN" {}
+!: AUTO RUN "COMMIT" {}
+
+C: RUN "RETURN 1" {}
+C: PULL_ALL
+S: FAILURE {"code": "Neo.TransientError.General.DatabaseUnavailable", "message": "Database is busy doing store copy"}
+S: IGNORED
diff --git a/test/stub/scripts/forbidden_on_read_only_database.script b/test/stub/scripts/forbidden_on_read_only_database.script
@@ -0,0 +1,12 @@
+!: AUTO INIT
+!: AUTO RESET
+!: AUTO PULL_ALL
+!: AUTO ACK_FAILURE
+!: AUTO RUN "ROLLBACK" {}
+!: AUTO RUN "BEGIN" {}
+!: AUTO RUN "COMMIT" {}
+
+C: RUN "CREATE (n {name:'Bob'})" {}
+C: PULL_ALL
+S: FAILURE {"code": "Neo.ClientError.General.ForbiddenOnReadOnlyDatabase", "message": "Unable to write"}
+S: IGNORED
diff --git a/test/stub/scripts/not_a_leader.script b/test/stub/scripts/not_a_leader.script
@@ -0,0 +1,12 @@
+!: AUTO INIT
+!: AUTO RESET
+!: AUTO PULL_ALL
+!: AUTO ACK_FAILURE
+!: AUTO RUN "ROLLBACK" {}
+!: AUTO RUN "BEGIN" {}
+!: AUTO RUN "COMMIT" {}
+
+C: RUN "CREATE (n {name:'Bob'})" {}
+C: PULL_ALL
+S: FAILURE {"code": "Neo.ClientError.Cluster.NotALeader", "message": "Leader switched has happened"}
+S: IGNORED
diff --git a/test/stub/scripts/rude_reader.script b/test/stub/scripts/rude_reader.script
@@ -0,0 +1,7 @@
+!: AUTO INIT
+!: AUTO RESET
+
+C: RUN "RETURN 1" {}
+   PULL_ALL
+S: <EXIT>
+
diff --git a/test/stub/test_routingdriver.py b/test/stub/test_routingdriver.py
@@ -21,7 +21,7 @@
 
 from neo4j.v1 import GraphDatabase, READ_ACCESS, WRITE_ACCESS, SessionExpired, \
     RoutingDriver, RoutingConnectionPool, LeastConnectedLoadBalancingStrategy, LOAD_BALANCING_STRATEGY_ROUND_ROBIN, \
-    RoundRobinLoadBalancingStrategy
+    RoundRobinLoadBalancingStrategy, TransientError, ClientError
 from neo4j.bolt import ProtocolError, ServiceUnavailable
 
 from test.stub.tools import StubTestCase, StubCluster
@@ -236,8 +236,79 @@ def test_can_select_round_robin_load_balancing_strategy(self):
                 self.assertIsInstance(driver._pool.load_balancing_strategy, RoundRobinLoadBalancingStrategy)
 
     def test_no_other_load_balancing_strategies_are_available(self):
-        with StubCluster({9001: "router.script"}):
+        uri = "bolt+routing://127.0.0.1:9001"
+        with self.assertRaises(ValueError):
+            with GraphDatabase.driver(uri, auth=self.auth_token, encrypted=False, load_balancing_strategy=-1):
+                pass
+
+    def test_forgets_address_on_not_a_leader_error(self):
+        with StubCluster({9001: "router.script", 9006: "not_a_leader.script"}):
             uri = "bolt+routing://127.0.0.1:9001"
-            with self.assertRaises(ValueError):
-                with GraphDatabase.driver(uri, auth=self.auth_token, encrypted=False, load_balancing_strategy=-1):
-                    pass
+            with GraphDatabase.driver(uri, auth=self.auth_token, encrypted=False) as driver:
+                with driver.session(WRITE_ACCESS) as session:
+                    with self.assertRaises(ClientError):
+                        _ = session.run("CREATE (n {name:'Bob'})")
+
+                    pool = driver._pool
+                    table = pool.routing_table
+
+                    # address might still have connections in the pool, failed instance just can't serve writes
+                    assert ('127.0.0.1', 9006) in pool.connections
+                    assert table.routers == {('127.0.0.1', 9001), ('127.0.0.1', 9002), ('127.0.0.1', 9003)}
+                    assert table.readers == {('127.0.0.1', 9004), ('127.0.0.1', 9005)}
+                    # writer 127.0.0.1:9006 should've been forgotten because of an error
+                    assert len(table.writers) == 0
+
+    def test_forgets_address_on_forbidden_on_read_only_database_error(self):
+        with StubCluster({9001: "router.script", 9006: "forbidden_on_read_only_database.script"}):
+            uri = "bolt+routing://127.0.0.1:9001"
+            with GraphDatabase.driver(uri, auth=self.auth_token, encrypted=False) as driver:
+                with driver.session(WRITE_ACCESS) as session:
+                    with self.assertRaises(ClientError):
+                        _ = session.run("CREATE (n {name:'Bob'})")
+
+                    pool = driver._pool
+                    table = pool.routing_table
+
+                    # address might still have connections in the pool, failed instance just can't serve writes
+                    assert ('127.0.0.1', 9006) in pool.connections
+                    assert table.routers == {('127.0.0.1', 9001), ('127.0.0.1', 9002), ('127.0.0.1', 9003)}
+                    assert table.readers == {('127.0.0.1', 9004), ('127.0.0.1', 9005)}
+                    # writer 127.0.0.1:9006 should've been forgotten because of an error
+                    assert len(table.writers) == 0
+
+    def test_forgets_address_on_service_unavailable_error(self):
+        with StubCluster({9001: "router.script", 9004: "rude_reader.script"}):
+            uri = "bolt+routing://127.0.0.1:9001"
+            with GraphDatabase.driver(uri, auth=self.auth_token, encrypted=False) as driver:
+                with driver.session(READ_ACCESS) as session:
+                    with self.assertRaises(SessionExpired):
+                        _ = session.run("RETURN 1")
+
+                    pool = driver._pool
+                    table = pool.routing_table
+
+                    # address should not have connections in the pool, it has failed
+                    assert ('127.0.0.1', 9004) not in pool.connections
+                    assert table.routers == {('127.0.0.1', 9001), ('127.0.0.1', 9002), ('127.0.0.1', 9003)}
+                    # reader 127.0.0.1:9004 should've been forgotten because of an error
+                    assert table.readers == {('127.0.0.1', 9005)}
+                    assert table.writers == {('127.0.0.1', 9006)}
+
+    def test_forgets_address_on_database_unavailable_error(self):
+        with StubCluster({9001: "router.script", 9004: "database_unavailable.script"}):
+            uri = "bolt+routing://127.0.0.1:9001"
+            with GraphDatabase.driver(uri, auth=self.auth_token, encrypted=False) as driver:
+                with driver.session(READ_ACCESS) as session:
+                    with self.assertRaises(TransientError):
+                        _ = session.run("RETURN 1")
+
+                    pool = driver._pool
+                    table = pool.routing_table
+
+                    # address should not have connections in the pool, it has failed
+                    assert ('127.0.0.1', 9004) not in pool.connections
+                    assert table.routers == {('127.0.0.1', 9001), ('127.0.0.1', 9002), ('127.0.0.1', 9003)}
+                    # reader 127.0.0.1:9004 should've been forgotten because of an error
+                    assert table.readers == {('127.0.0.1', 9005)}
+                    assert table.writers == {('127.0.0.1', 9006)}
