cluster: fence stale host add futures

Author: dkropachev

cassandra/cluster.py

@@ -1870,8 +1870,11 @@ def _cleanup_failed_on_up_handling(self, host):
 
         self._start_reconnector(host, is_host_addition=False)
 
-    def _cleanup_failed_on_add_handling(self, host):
+    def _cleanup_failed_on_add_handling(self, host, addition_generation=None):
         with host.lock:
+            if (addition_generation is not None and
+                    getattr(host, "_node_addition_generation", None) is not addition_generation):
+                return
             host.set_down()
 
         self.profile_manager.on_down(host)
@@ -1880,7 +1883,10 @@ def _cleanup_failed_on_add_handling(self, host):
             session.remove_pool(host)
 
         with host.lock:
-            host._currently_handling_node_addition = False
+            if (addition_generation is None or
+                    getattr(host, "_node_addition_generation", None) is addition_generation):
+                host._currently_handling_node_addition = False
+                host._node_addition_generation = None
 
         self._start_reconnector(host, is_host_addition=True)
 
@@ -2068,11 +2074,13 @@ def on_add(self, host, refresh_nodes=True):
         log.debug("Handling new host %r and notifying listeners", host)
 
         # Keep refresh-time pool rebuilds from racing this host's pool creation.
+        addition_generation = object()
         with host.lock:
             if getattr(host, "_currently_handling_node_addition", False):
                 log.debug("Another thread is already handling add status of node %s", host)
                 return
             host._currently_handling_node_addition = True
+            host._node_addition_generation = addition_generation
 
         have_future = False
         add_aborted = False
@@ -2113,15 +2121,15 @@ def future_completed(future):
 
                     for exc in [f for f in futures_results if isinstance(f, Exception)]:
                         log.error("Unexpected failure while adding node %s, will not mark up:", host, exc_info=exc)
-                        self._cleanup_failed_on_add_handling(host)
+                        self._cleanup_failed_on_add_handling(host, addition_generation)
                         return
 
                     if not all(futures_results):
                         log.warning("Connection pool could not be created, not marking node %s up", host)
-                        self._cleanup_failed_on_add_handling(host)
+                        self._cleanup_failed_on_add_handling(host, addition_generation)
                         return
 
-                    self._finalize_add(host)
+                    self._finalize_add(host, addition_generation=addition_generation)
 
                 for session in tuple(self.sessions):
                     future = session.add_or_renew_pool(host, is_host_addition=True)
@@ -2137,15 +2145,14 @@ def future_completed(future):
         except Exception:
             add_aborted = True
             for future, session in tuple(futures.items()):
-                if not future.cancel():
-                    future.add_done_callback(lambda f, session=session: session.remove_pool(host))
-            self._cleanup_failed_on_add_handling(host)
+                future.cancel()
+            self._cleanup_failed_on_add_handling(host, addition_generation)
             raise
 
         if finalize_add is not None:
-            self._finalize_add(host, set_up=finalize_add)
+            self._finalize_add(host, set_up=finalize_add, addition_generation=addition_generation)
 
-    def _finalize_add(self, host, set_up=True):
+    def _finalize_add(self, host, set_up=True, addition_generation=None):
         try:
             if set_up:
                 host.set_up()
@@ -2158,7 +2165,10 @@ def _finalize_add(self, host, set_up=True):
                 session.update_created_pools()
         finally:
             with host.lock:
-                host._currently_handling_node_addition = False
+                if (addition_generation is None or
+                        getattr(host, "_node_addition_generation", None) is addition_generation):
+                    host._currently_handling_node_addition = False
+                    host._node_addition_generation = None
 
     def on_remove(self, host):
         if self.is_shutdown:
@@ -3294,15 +3304,27 @@ def add_or_renew_pool(self, host, is_host_addition):
         distance = self._profile_manager.distance(host)
         if distance == HostDistance.IGNORED:
             return None
+        addition_generation = getattr(host, "_node_addition_generation", None) if is_host_addition else None
+
+        def is_stale_addition():
+            return (is_host_addition and
+                    getattr(host, "_node_addition_generation", None) is not addition_generation)
 
         def run_add_or_renew_pool():
+            if is_stale_addition():
+                return False
+
             try:
-               new_pool = HostConnection(host, distance, self)
+                new_pool = HostConnection(host, distance, self)
             except AuthenticationFailed as auth_exc:
+                if is_stale_addition():
+                    return False
                 conn_exc = ConnectionException(str(auth_exc), endpoint=host)
                 self.cluster.signal_connection_failure(host, conn_exc, is_host_addition)
                 return False
             except Exception as conn_exc:
+                if is_stale_addition():
+                    return False
                 log.warning("Failed to create connection pool for new host %s:",
                             host, exc_info=conn_exc)
                 # the host itself will still be marked down, so we need to pass
@@ -3311,6 +3333,10 @@ def run_add_or_renew_pool():
                     host, conn_exc, is_host_addition, expect_host_to_be_down=True)
                 return False
 
+            if is_stale_addition():
+                new_pool.shutdown()
+                return False
+
             previous = self._pools.get(host)
             with self._lock:
                 while new_pool._keyspace != self.keyspace:
@@ -3325,12 +3351,19 @@ def callback(pool, errors):
                     new_pool._set_keyspace_for_all_conns(self.keyspace, callback)
                     set_keyspace_event.wait(self.cluster.connect_timeout)
                     if not set_keyspace_event.is_set() or errors_returned:
+                        if is_stale_addition():
+                            new_pool.shutdown()
+                            self._lock.acquire()
+                            return False
                         log.warning("Failed setting keyspace for pool after keyspace changed during connect: %s", errors_returned)
                         self.cluster.on_down(host, is_host_addition)
                         new_pool.shutdown()
                         self._lock.acquire()
                         return False
                     self._lock.acquire()
+                if is_stale_addition():
+                    new_pool.shutdown()
+                    return False
                 self._pools[host] = new_pool
 
             log.debug("Added pool for host %s to session", host)

tests/unit/test_cluster.py

@@ -137,11 +137,19 @@ class RunningPoolSession(object):
             def __init__(self):
                 self.pool_created = False
                 self.remove_pool_calls = 0
+                self.generation = None
                 self.update_created_pools = Mock(return_value=set())
 
             def add_or_renew_pool(self, add_host, is_host_addition):
+                self.generation = getattr(add_host, "_node_addition_generation", None)
                 return running_future
 
+            def complete_running_pool_creation(self, add_host):
+                current_generation = getattr(add_host, "_node_addition_generation", None)
+                if self.generation is None or current_generation is self.generation:
+                    self.pool_created = True
+                running_future.set_result(True)
+
             def remove_pool(self, remove_host):
                 self.remove_pool_calls += 1
                 if remove_host is host and self.pool_created:
@@ -162,13 +170,103 @@ def shutdown(self):
 
             assert running_session.remove_pool_calls == 1
 
-            running_session.pool_created = True
-            running_future.set_result(True)
+            running_session.complete_running_pool_creation(host)
 
             assert running_session.pool_created is False
         finally:
             cluster.shutdown()
 
+    def test_on_add_aborted_future_does_not_remove_newer_successful_pool(self):
+        cluster = Cluster(protocol_version=4)
+        host = Host("127.0.0.1", SimpleConvictionPolicy, datacenter="dc1", rack="rack1", host_id=uuid.uuid4())
+
+        class FencedPoolSession(object):
+
+            def __init__(self):
+                self.pool = None
+                self.add_calls = 0
+                self.stale_future = Future()
+                self.stale_future.set_running_or_notify_cancel()
+                self.stale_generation = None
+                self.update_created_pools = Mock(return_value=set())
+
+            def add_or_renew_pool(self, add_host, is_host_addition):
+                self.add_calls += 1
+                if self.add_calls == 1:
+                    self.stale_generation = getattr(add_host, "_node_addition_generation", None)
+                    return self.stale_future
+
+                self.pool = "fresh-pool"
+                completed_future = Future()
+                completed_future.set_result(True)
+                return completed_future
+
+            def complete_stale_pool_creation(self, add_host):
+                current_generation = getattr(add_host, "_node_addition_generation", None)
+                if self.stale_generation is None or current_generation is self.stale_generation:
+                    self.pool = "stale-pool"
+                self.stale_future.set_result(True)
+
+            def remove_pool(self, remove_host):
+                if remove_host is host:
+                    self.pool = None
+
+            def shutdown(self):
+                pass
+
+        pool_session = FencedPoolSession()
+        recovered_future = Future()
+        recovered_future.set_result(True)
+        failing_session = Mock()
+        failing_session.add_or_renew_pool.side_effect = [RuntimeError("pool add failed"), recovered_future]
+        failing_session.update_created_pools.return_value = set()
+        cluster.sessions = [pool_session, failing_session]
+
+        try:
+            with pytest.raises(RuntimeError):
+                cluster.on_add(host, refresh_nodes=False)
+
+            cluster.on_add(host, refresh_nodes=False)
+            assert pool_session.pool == "fresh-pool"
+
+            pool_session.complete_stale_pool_creation(host)
+
+            assert pool_session.pool == "fresh-pool"
+        finally:
+            cluster.shutdown()
+
+    def test_add_or_renew_pool_does_not_signal_stale_addition_failure(self):
+        host = Host("127.0.0.1", SimpleConvictionPolicy, datacenter="dc1", rack="rack1", host_id=uuid.uuid4())
+        host._node_addition_generation = object()
+
+        session = object.__new__(Session)
+        session.cluster = Mock(connect_timeout=1)
+        session.cluster.signal_connection_failure = Mock()
+        session._profile_manager = Mock()
+        session._profile_manager.distance.return_value = HostDistance.LOCAL
+        session.keyspace = None
+        session._pools = {}
+
+        def submit(fn, *args, **kwargs):
+            future = Future()
+            try:
+                future.set_result(fn(*args, **kwargs))
+            except Exception as exc:
+                future.set_exception(exc)
+            return future
+
+        session.submit = submit
+
+        def raise_after_add_aborts(*args, **kwargs):
+            host._node_addition_generation = None
+            raise RuntimeError("stale connect failure")
+
+        with patch("cassandra.cluster.HostConnection", side_effect=raise_after_add_aborts):
+            future = session.add_or_renew_pool(host, is_host_addition=True)
+
+        assert future.result() is False
+        session.cluster.signal_connection_failure.assert_not_called()
+
     def test_on_add_excludes_host_from_query_plan_when_pool_future_fails(self):
         cluster = Cluster(protocol_version=4)
         host = Host("127.0.0.1", SimpleConvictionPolicy, datacenter="dc1", rack="rack1", host_id=uuid.uuid4())