control-connection: run topology removals through cluster

cassandra/cluster.py

@@ -2207,7 +2207,7 @@ def _finalize_add(self, host, set_up=True):
         for session in tuple(self.sessions):
             session.update_created_pools()
 
-    def on_remove(self, host):
+    def on_remove(self, host, refresh_nodes=True):
         if self.is_shutdown:
             return
 
@@ -2218,7 +2218,7 @@ def on_remove(self, host):
             session.on_remove(host)
         for listener in self.listeners:
             listener.on_remove(host)
-        self.control_connection.on_remove(host)
+        self.control_connection.on_remove(host, refresh_nodes=refresh_nodes)
 
         reconnection_handler = host.get_and_set_reconnection_handler(None)
         if reconnection_handler:
@@ -2248,14 +2248,31 @@ def add_host(self, endpoint, datacenter=None, rack=None, signal=True, refresh_no
 
         return host, new
 
-    def remove_host(self, host):
+    def remove_host(self, host, refresh_nodes=True):
         """
         Called when the control connection observes that a node has left the
         ring.  Intended for internal use only.
         """
         if host and self.metadata.remove_host(host):
             log.info("Cassandra host %s removed", host)
-            self.on_remove(host)
+            self.on_remove(host, refresh_nodes=refresh_nodes)
+
+    def remove_host_by_host_id(self, host_id, endpoint=None,
+                               refresh_nodes=True):
+        """Remove the host stored under a specific metadata key."""
+        host = self.metadata.get_host_by_host_id(host_id)
+        if not host or not self.metadata.remove_host_by_host_id(
+                host_id, endpoint):
+            return
+
+        # A refresh can reindex the same Host under a new host ID before
+        # cleaning up its stale old key. In that case only the alias was
+        # removed; the Host itself is still part of the cluster.
+        if self.metadata._is_host_registered(host):
+            return
+
+        log.info("Cassandra host %s removed", host)
+        self.on_remove(host, refresh_nodes=refresh_nodes)
 
     def register_listener(self, listener):
         """
@@ -3914,6 +3931,7 @@ def __init__(self, cluster, timeout,
         self._schema_meta_page_size = schema_meta_page_size
 
         self._lock = RLock()
+        self._refresh_nodes_lock = RLock()
         self._schema_agreement_lock = Lock()
 
         self._reconnection_handler = None
@@ -4192,6 +4210,13 @@ def refresh_node_list_and_token_map(self, force_token_rebuild=False):
 
     def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
                                          force_token_rebuild=False):
+        with self._refresh_nodes_lock:
+            return self._refresh_node_list_and_token_map_locked(
+                connection, preloaded_results, force_token_rebuild)
+
+    def _refresh_node_list_and_token_map_locked(
+            self, connection, preloaded_results=None,
+            force_token_rebuild=False):
         if preloaded_results:
             log.debug("[control connection] Refreshing node list and token map using preloaded results")
             peers_result = preloaded_results[0]
@@ -4318,7 +4343,8 @@ def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
             if old_host_id not in found_host_ids:
                 should_rebuild_token_map = True
                 log.debug("[control connection] Removing host not found in peers metadata: %r", old_host)
-                self._cluster.metadata.remove_host_by_host_id(old_host_id, old_host.endpoint)
+                self._cluster.remove_host_by_host_id(
+                    old_host_id, old_host.endpoint, refresh_nodes=False)
 
         log.debug("[control connection] Finished fetching ring info")
         if partitioner and should_rebuild_token_map:
@@ -4735,13 +4761,13 @@ def on_add(self, host, refresh_nodes=True):
         if refresh_nodes:
             self.refresh_node_list_and_token_map(force_token_rebuild=True)
 
-    def on_remove(self, host):
+    def on_remove(self, host, refresh_nodes=True):
         c = self._connection
         if self._connection_matches_host(c, host):
             log.debug("[control connection] Control connection host (%s) is being removed. Reconnecting", host)
             # refresh will be done on reconnect
             self.reconnect()
-        else:
+        elif refresh_nodes:
             self.refresh_node_list_and_token_map(force_token_rebuild=True)
 
     def get_connections(self):

cassandra/metadata.py

@@ -359,7 +359,8 @@ def remove_host(self, host):
     def remove_host_by_host_id(self, host_id, endpoint=None):
         self._tablets.drop_tablets_by_host_id(host_id)
         with self._hosts_lock:
-            if endpoint and self._host_id_by_endpoint[endpoint] == host_id:
+            if (endpoint and
+                    self._host_id_by_endpoint.get(endpoint) == host_id):
                 self._host_id_by_endpoint.pop(endpoint, False)
             return bool(self._hosts.pop(host_id, False))
 

tests/unit/test_cluster.py

@@ -151,6 +151,46 @@ def test_backward_compat_positional(self):
 
 class ClusterTest(unittest.TestCase):
 
+    def test_remove_host_by_id_runs_lifecycle_without_refresh(self):
+        cluster = Cluster()
+        self.addCleanup(cluster.shutdown)
+        host = Host(
+            "127.0.0.1", SimpleConvictionPolicy, host_id=uuid.uuid4())
+        cluster.metadata.add_or_return_host(host)
+        cluster.profile_manager.on_remove = Mock()
+        cluster.control_connection.on_remove = Mock()
+        session = Mock()
+        cluster.sessions.add(session)
+
+        cluster.remove_host_by_host_id(
+            host.host_id, host.endpoint, refresh_nodes=False)
+
+        cluster.profile_manager.on_remove.assert_called_once_with(host)
+        session.on_remove.assert_called_once_with(host)
+        cluster.control_connection.on_remove.assert_called_once_with(
+            host, refresh_nodes=False)
+
+    def test_remove_host_by_stale_id_preserves_reindexed_host(self):
+        cluster = Cluster()
+        self.addCleanup(cluster.shutdown)
+        old_host_id = uuid.uuid4()
+        new_host_id = uuid.uuid4()
+        host = Host(
+            "127.0.0.1", SimpleConvictionPolicy, host_id=old_host_id)
+        host.set_up()
+        cluster.metadata.add_or_return_host(host)
+        host.host_id = new_host_id
+        cluster.metadata.update_host(host, old_endpoint=host.endpoint)
+        cluster.on_remove = Mock()
+
+        cluster.remove_host_by_host_id(old_host_id, host.endpoint)
+
+        assert cluster.metadata.get_host_by_host_id(old_host_id) is None
+        assert cluster.metadata.get_host_by_host_id(new_host_id) is host
+        assert cluster.metadata.get_host(host.endpoint) is host
+        assert host.is_up
+        cluster.on_remove.assert_not_called()
+
     def test_tuple_for_contact_points(self):
         cluster = Cluster(contact_points=[('localhost', 9045), ('127.0.0.2', 9046), '127.0.0.3'], port=9999)
         self.addCleanup(cluster.shutdown)

tests/unit/test_control_connection.py

@@ -15,6 +15,7 @@
 import unittest
 
 from concurrent.futures import ThreadPoolExecutor
+from threading import Event, Lock, Thread
 from unittest.mock import Mock, ANY, call, patch
 
 from cassandra import OperationTimedOut, SchemaTargetType, SchemaChangeType
@@ -90,10 +91,12 @@ def all_hosts_items(self):
         return list(self.hosts.items())
 
     def remove_host_by_host_id(self, host_id, endpoint=None):
-        if endpoint and self._host_id_by_endpoint[endpoint] == host_id:
+        if endpoint and self._host_id_by_endpoint.get(endpoint) == host_id:
             self._host_id_by_endpoint.pop(endpoint, False)
-        self.removed_hosts.append(self.hosts.pop(host_id, False))
-        return bool(self.hosts.pop(host_id, False))
+        removed_host = self.hosts.pop(host_id, None)
+        if removed_host:
+            self.removed_hosts.append(removed_host)
+        return bool(removed_host)
 
 
 class MockCluster(object):
@@ -109,6 +112,8 @@ class MockCluster(object):
     def __init__(self):
         self.metadata = MockMetadata()
         self.added_hosts = []
+        self.removed_host = None
+        self.removed_host_refresh_nodes = None
         self.scheduler = Mock(spec=_Scheduler)
         self.executor = Mock(spec=ThreadPoolExecutor)
         self.profile_manager.profiles[EXEC_PROFILE_DEFAULT] = ExecutionProfile(RoundRobinPolicy())
@@ -121,8 +126,21 @@ def add_host(self, endpoint, datacenter, rack, signal=False, refresh_nodes=True,
         self.added_hosts.append(host)
         return host, True
 
-    def remove_host(self, host):
-        pass
+    def remove_host(self, host, refresh_nodes=True):
+        removed = self.metadata.remove_host_by_host_id(
+            host.host_id, host.endpoint)
+        if removed:
+            self.removed_host = host
+            self.removed_host_refresh_nodes = refresh_nodes
+
+    def remove_host_by_host_id(self, host_id, endpoint=None,
+                               refresh_nodes=True):
+        host = self.metadata.get_host_by_host_id(host_id)
+        removed = self.metadata.remove_host_by_host_id(host_id, endpoint)
+        if (removed and
+                self.metadata.get_host_by_host_id(host.host_id) is not host):
+            self.removed_host = host
+            self.removed_host_refresh_nodes = refresh_nodes
 
     def on_up(self, host):
         pass
@@ -615,6 +633,14 @@ def test_remove_matches_control_connection_by_host_id(self):
         self.cluster.executor.submit.assert_called_once_with(
             self.control_connection._reconnect)
 
+    def test_remove_non_control_host_can_skip_refresh(self):
+        host = self.cluster.metadata.get_host_by_host_id('uuid2')
+        self.control_connection.refresh_node_list_and_token_map = Mock()
+
+        self.control_connection.on_remove(host, refresh_nodes=False)
+
+        self.control_connection.refresh_node_list_and_token_map.assert_not_called()
+
     def test_down_matches_replacement_at_stale_control_endpoint(self):
         self.control_connection.refresh_node_list_and_token_map()
         old_host = self.cluster.metadata.get_host_by_host_id('uuid1')
@@ -769,6 +795,84 @@ def test_refresh_nodes_and_tokens_remove_host(self):
         self.control_connection.refresh_node_list_and_token_map()
         assert 1 == len(self.cluster.metadata.removed_hosts)
         assert self.cluster.metadata.removed_hosts[0].address == "192.168.1.2"
+        assert self.cluster.removed_host is \
+            self.cluster.metadata.removed_hosts[0]
+        assert self.cluster.removed_host_refresh_nodes is False
+
+    def test_refresh_nodes_and_tokens_preserves_reindexed_host(self):
+        old_host = self.cluster.metadata.get_host_by_host_id('uuid2')
+        endpoint = old_host.endpoint
+        self.connection.peer_results[1][0][-1] = 'replacement-id'
+
+        self.control_connection.refresh_node_list_and_token_map()
+        self.control_connection.refresh_node_list_and_token_map()
+
+        assert self.cluster.metadata.get_host_by_host_id('uuid2') is None
+        assert self.cluster.metadata.get_host_by_host_id(
+            'replacement-id') is old_host
+        assert self.cluster.metadata.get_host(endpoint) is old_host
+        assert len(self.cluster.metadata.all_hosts()) == 3
+        assert old_host.is_up
+        assert self.cluster.removed_host is None
+
+    def test_overlapping_refreshes_serialize_host_reindex(self):
+        old_host = self.cluster.metadata.get_host_by_host_id('uuid2')
+        endpoint = old_host.endpoint
+        self.connection.peer_results[1][0][-1] = 'replacement-id'
+
+        original_refresh = (
+            self.control_connection._refresh_node_list_and_token_map_locked)
+        first_entered = Event()
+        release_first = Event()
+        state_lock = Lock()
+        active = [0]
+        max_active = [0]
+
+        def controlled_refresh(*args, **kwargs):
+            with state_lock:
+                active[0] += 1
+                max_active[0] = max(max_active[0], active[0])
+                is_first = active[0] == 1 and not first_entered.is_set()
+            try:
+                if is_first:
+                    first_entered.set()
+                    assert release_first.wait(2)
+                return original_refresh(*args, **kwargs)
+            finally:
+                with state_lock:
+                    active[0] -= 1
+
+        self.control_connection._refresh_node_list_and_token_map_locked = \
+            controlled_refresh
+        results = []
+
+        def refresh():
+            results.append(
+                self.control_connection.refresh_node_list_and_token_map())
+
+        first = Thread(target=refresh)
+        second = Thread(target=refresh)
+        first.start()
+        assert first_entered.wait(2)
+        acquired = self.control_connection._refresh_nodes_lock.acquire(False)
+        if acquired:
+            self.control_connection._refresh_nodes_lock.release()
+        assert not acquired
+        second.start()
+        release_first.set()
+        first.join(2)
+        second.join(2)
+
+        assert not first.is_alive()
+        assert not second.is_alive()
+        assert results == [True, True]
+        assert max_active[0] == 1
+        assert self.cluster.metadata.get_host_by_host_id('uuid2') is None
+        assert self.cluster.metadata.get_host_by_host_id(
+            'replacement-id') is old_host
+        assert self.cluster.metadata.get_host(endpoint) is old_host
+        assert old_host.is_up
+        assert self.cluster.removed_host is None
 
     def test_refresh_nodes_and_tokens_timeout(self):