token-map: rebuild on zero-token transitions

Author: dkropachev

cassandra/cluster.py

@@ -3919,8 +3919,10 @@ def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
             host.dse_workloads = row.get("workloads")
 
             tokens = row.get("tokens", None)
-            if partitioner and tokens and self._token_meta_enabled:
-                token_map[host] = tokens
+            if partitioner and self._token_meta_enabled:
+                should_rebuild_token_map |= self._tokens_changed(host, tokens)
+                if tokens:
+                    token_map[host] = tokens
             self._cluster.metadata.update_host(host, old_endpoint=endpoint)
 
         for old_host_id, old_host in self._cluster.metadata.all_hosts_items():
@@ -3965,12 +3967,41 @@ def _is_valid_peer(row):
 
         if "tokens" in row and not row.get("tokens"):
             log.debug(
-                "Found a zero-token node - tokens is None (broadcast_rpc: %s, host_id: %s). Ignoring host." %
+                "Found a zero-token node (broadcast_rpc: %s, host_id: %s). "
+                "Keeping host for load balancing, but omitting it from the token map." %
                 (broadcast_rpc, host_id))
-            return False
 
         return True
 
+    def _tokens_changed(self, host, tokens):
+        current_token_map = self._cluster.metadata.token_map
+        if current_token_map is None:
+            return False
+
+        if isinstance(current_token_map, dict):
+            current_tokens = current_token_map.get(host)
+            return set(current_tokens or ()) != set(tokens or ())
+
+        token_to_host_owner = getattr(current_token_map, 'token_to_host_owner', None)
+        token_class = getattr(current_token_map, 'token_class', None)
+        if token_to_host_owner is None or token_class is None:
+            return False
+
+        current_tokens = set(
+            token for token, owner in token_to_host_owner.items()
+            if owner == host)
+        if not tokens:
+            return bool(current_tokens)
+
+        try:
+            refreshed_tokens = set(token_class.from_string(token) for token in tokens)
+        except Exception:
+            log.debug("[control connection] Unable to compare refreshed tokens for %s",
+                      host, exc_info=True)
+            return True
+
+        return current_tokens != refreshed_tokens
+
     def _update_location_info(self, host, datacenter, rack):
         if host.datacenter == datacenter and host.rack == rack:
             return False

cassandra/policies.py

@@ -464,6 +464,11 @@ class TokenAwarePolicy(LoadBalancingPolicy):
 
     If no :attr:`~.Statement.routing_key` is set on the query, the child
     policy's query plan will be used as is.
+
+    Token awareness only applies to token-owning nodes. Zero-token CQL
+    proxy or front-layer nodes are not replicas. When those nodes are the
+    intended query front layer, configure a round-robin policy such as
+    :class:`.RoundRobinPolicy` or :class:`.DCAwareRoundRobinPolicy` directly.
     """
 
     _child_policy = None

tests/unit/test_control_connection.py

@@ -406,7 +406,6 @@ def refresh_and_validate_added_hosts():
              [None, None, "a", "dc1", "rack1", ["1", "101", "201"], 'uuid1'],
              ["192.168.1.7", "10.0.0.1", "a", None, "rack1", ["1", "101", "201"], 'uuid2'],
              ["192.168.1.6", "10.0.0.1", "a", "dc1", None, ["1", "101", "201"], 'uuid3'],
-             ["192.168.1.5", "10.0.0.1", "a", "dc1", "rack1", None, 'uuid4'],
              ["192.168.1.4", "10.0.0.1", "a", "dc1", "rack1", ["1", "101", "201"], None]]])
         refresh_and_validate_added_hosts()
 
@@ -420,10 +419,57 @@ def refresh_and_validate_added_hosts():
              [None, 9042, None, 7040, "a", "dc1", "rack1", ["2", "102", "202"], "uuid2"],
              ["192.168.1.5", 9042, "10.0.0.2", 7040, "a", None, "rack1", ["2", "102", "202"], "uuid2"],
              ["192.168.1.5", 9042, "10.0.0.2", 7040, "a", "dc1", None, ["2", "102", "202"], "uuid2"],
-             ["192.168.1.5", 9042, "10.0.0.2", 7040, "a", "dc1", "rack1", None, "uuid2"],
              ["192.168.1.5", 9042, "10.0.0.2", 7040, "a", "dc1", "rack1", ["2", "102", "202"], None]]])
         refresh_and_validate_added_hosts()
 
+    def test_refresh_nodes_and_tokens_keeps_zero_token_peer_for_load_balancing(self):
+        self.connection.peer_results[1].append(
+            ["192.168.1.3", "10.0.0.3", "a", "dc1", "rack1", None, "uuid4"]
+        )
+        self.cluster.scheduler.schedule = lambda delay, f, *args, **kwargs: f(*args, **kwargs)
+
+        self.control_connection.refresh_node_list_and_token_map()
+
+        zero_token_host = self.cluster.metadata.get_host_by_host_id("uuid4")
+        assert zero_token_host is not None
+        assert zero_token_host in self.cluster.metadata.all_hosts()
+        assert zero_token_host not in self.cluster.metadata.token_map
+
+    def test_refresh_nodes_and_tokens_keeps_zero_token_local_host_for_load_balancing(self):
+        self.connection.local_results[1][0][7] = None
+
+        self.control_connection.refresh_node_list_and_token_map()
+
+        local_host = self.cluster.metadata.get_host_by_host_id("uuid1")
+        assert local_host is not None
+        assert local_host in self.cluster.metadata.all_hosts()
+        assert local_host not in self.cluster.metadata.token_map
+
+    def test_refresh_nodes_and_tokens_rebuilds_token_map_when_existing_host_loses_tokens(self):
+        self.control_connection.refresh_node_list_and_token_map()
+
+        zero_token_host = self.cluster.metadata.get_host_by_host_id("uuid2")
+        assert zero_token_host in self.cluster.metadata.token_map
+
+        self.connection.peer_results[1][0][5] = None
+        self.control_connection.refresh_node_list_and_token_map()
+
+        assert zero_token_host in self.cluster.metadata.all_hosts()
+        assert zero_token_host not in self.cluster.metadata.token_map
+
+    def test_refresh_nodes_and_tokens_rebuilds_token_map_when_existing_host_gains_tokens(self):
+        self.connection.peer_results[1][0][5] = None
+        self.control_connection.refresh_node_list_and_token_map()
+
+        token_host = self.cluster.metadata.get_host_by_host_id("uuid2")
+        assert token_host in self.cluster.metadata.all_hosts()
+        assert token_host not in self.cluster.metadata.token_map
+
+        self.connection.peer_results[1][0][5] = ["1", "101", "201"]
+        self.control_connection.refresh_node_list_and_token_map()
+
+        assert token_host in self.cluster.metadata.token_map
+
     def test_change_ip(self):
         """
         Tests node IPs are updated while the nodes themselves are not