cluster: preserve successful add after listener errors

Author: dkropachev

cassandra/cluster.py

@@ -2077,6 +2077,7 @@ def on_add(self, host, refresh_nodes=True):
         have_future = False
         add_aborted = False
         futures = set()
+        finalize_add = None
         try:
             self.profile_manager.on_add(host)
             self.control_connection.on_add(host, refresh_nodes)
@@ -2089,59 +2090,61 @@ def on_add(self, host, refresh_nodes=True):
             if distance == HostDistance.IGNORED:
                 log.debug("Not adding connection pool for new host %r because the "
                           "load balancing policy has marked it as IGNORED", host)
-                self._finalize_add(host, set_up=False)
-                return
-
-            futures_lock = Lock()
-            futures_results = []
+                finalize_add = False
+            else:
+                futures_lock = Lock()
+                futures_results = []
 
-            def future_completed(future):
-                with futures_lock:
-                    futures.discard(future)
+                def future_completed(future):
+                    with futures_lock:
+                        futures.discard(future)
 
-                    if add_aborted:
-                        return
+                        if add_aborted:
+                            return
 
-                    try:
-                        futures_results.append(future.result())
-                    except Exception as exc:
-                        futures_results.append(exc)
+                        try:
+                            futures_results.append(future.result())
+                        except Exception as exc:
+                            futures_results.append(exc)
 
-                    if futures:
-                        return
+                        if futures:
+                            return
 
-                log.debug('All futures have completed for added host %s', host)
+                    log.debug('All futures have completed for added host %s', host)
 
-                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)
-                    return
+                    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)
+                        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)
-                    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)
+                        return
 
-                self._finalize_add(host)
+                    self._finalize_add(host)
 
-            for session in tuple(self.sessions):
-                future = session.add_or_renew_pool(host, is_host_addition=True)
-                if future is not None:
-                    have_future = True
-                    futures.add(future)
+                for session in tuple(self.sessions):
+                    future = session.add_or_renew_pool(host, is_host_addition=True)
+                    if future is not None:
+                        have_future = True
+                        futures.add(future)
 
-            for future in tuple(futures):
-                future.add_done_callback(future_completed)
+                for future in tuple(futures):
+                    future.add_done_callback(future_completed)
 
-            if not have_future:
-                self._finalize_add(host)
+                if not have_future:
+                    finalize_add = True
         except Exception:
             add_aborted = True
             for future in tuple(futures):
                 future.cancel()
             self._cleanup_failed_on_add_handling(host)
             raise
 
+        if finalize_add is not None:
+            self._finalize_add(host, set_up=finalize_add)
+
     def _finalize_add(self, host, set_up=True):
         try:
             if set_up:

tests/unit/test_cluster.py

@@ -177,6 +177,25 @@ def reentrant_add_while_cleanup_removes_host(cleanup_host):
         finally:
             cluster.shutdown()
 
+    def test_on_add_listener_failure_does_not_mark_successful_add_down(self):
+        cluster = Cluster(protocol_version=4)
+        host = Host("127.0.0.1", SimpleConvictionPolicy, datacenter="dc1", rack="rack1", host_id=uuid.uuid4())
+        listener = Mock()
+        listener.on_add.side_effect = RuntimeError("listener failed")
+        cluster.register_listener(listener)
+
+        try:
+            with pytest.raises(RuntimeError):
+                cluster.on_add(host, refresh_nodes=False)
+
+            load_balancer = cluster.profile_manager.default.load_balancing_policy
+            assert host.is_up is True
+            assert host in list(load_balancer.make_query_plan())
+            assert not host.is_currently_reconnecting()
+            assert not host._currently_handling_node_addition
+        finally:
+            cluster.shutdown()
+
     def test_on_add_waits_for_all_session_pool_futures_before_marking_host_up(self):
         cluster = Cluster(protocol_version=4)
         host = Host("127.0.0.1", SimpleConvictionPolicy, host_id=uuid.uuid4())