cluster: remove late pools after aborted host add

Author: dkropachev

cassandra/cluster.py

@@ -2076,7 +2076,9 @@ def on_add(self, host, refresh_nodes=True):
 
         have_future = False
         add_aborted = False
-        futures = set()
+        futures = {}
+        futures_lock = Lock()
+        futures_results = []
         finalize_add = None
         try:
             self.profile_manager.on_add(host)
@@ -2092,12 +2094,9 @@ def on_add(self, host, refresh_nodes=True):
                           "load balancing policy has marked it as IGNORED", host)
                 finalize_add = False
             else:
-                futures_lock = Lock()
-                futures_results = []
-
                 def future_completed(future):
                     with futures_lock:
-                        futures.discard(future)
+                        futures.pop(future, None)
 
                         if add_aborted:
                             return
@@ -2128,7 +2127,7 @@ def future_completed(future):
                     future = session.add_or_renew_pool(host, is_host_addition=True)
                     if future is not None:
                         have_future = True
-                        futures.add(future)
+                        futures[future] = session
 
                 for future in tuple(futures):
                     future.add_done_callback(future_completed)
@@ -2137,8 +2136,9 @@ def future_completed(future):
                     finalize_add = True
         except Exception:
             add_aborted = True
-            for future in tuple(futures):
-                future.cancel()
+            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)
             raise
 

tests/unit/test_cluster.py

@@ -126,6 +126,49 @@ def test_on_add_clears_in_progress_flag_when_later_session_add_fails(self):
         finally:
             cluster.shutdown()
 
+    def test_on_add_removes_pool_created_by_running_future_after_add_aborts(self):
+        cluster = Cluster(protocol_version=4)
+        host = Host("127.0.0.1", SimpleConvictionPolicy, datacenter="dc1", rack="rack1", host_id=uuid.uuid4())
+        running_future = Future()
+        running_future.set_running_or_notify_cancel()
+
+        class RunningPoolSession(object):
+
+            def __init__(self):
+                self.pool_created = False
+                self.remove_pool_calls = 0
+                self.update_created_pools = Mock(return_value=set())
+
+            def add_or_renew_pool(self, add_host, is_host_addition):
+                return running_future
+
+            def remove_pool(self, remove_host):
+                self.remove_pool_calls += 1
+                if remove_host is host and self.pool_created:
+                    self.pool_created = False
+
+            def shutdown(self):
+                pass
+
+        running_session = RunningPoolSession()
+        failing_session = Mock()
+        failing_session.add_or_renew_pool.side_effect = RuntimeError("pool add failed")
+        failing_session.update_created_pools.return_value = set()
+        cluster.sessions = [running_session, failing_session]
+
+        try:
+            with pytest.raises(RuntimeError):
+                cluster.on_add(host, refresh_nodes=False)
+
+            assert running_session.remove_pool_calls == 1
+
+            running_session.pool_created = True
+            running_future.set_result(True)
+
+            assert running_session.pool_created is False
+        finally:
+            cluster.shutdown()
+
     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())