(improvement)TokenAware round robin policy and others - improved query planning.

Optimize TokenAwarePolicy query plan generation

This patch significantly improves the performance of TokenAwarePolicy by
reducing overhead in list materialization and distance calculation.

Key changes:
1. Introduced `make_query_plan_with_exclusion()` to the LoadBalancingPolicy
   interface.
   - This allows a parent policy (like TokenAware) to request a plan from
     a child policy while efficiently skipping a set of already-yielded
     hosts (replicas).
   - Implemented optimized versions in `DCAwareRoundRobinPolicy` and
     `RackAwareRoundRobinPolicy`. These implementations integrate the
     exclusion check directly into their generation loops, avoiding the
     need for inefficient external filtering or full list materialization.

2. Optimized `TokenAwarePolicy.make_query_plan`:
   - Removed list materialization of the child query plan.
   - Replaced multiple passes over replicas (checking `child.distance`
     each time) with a single pass that buckets replicas into local/remote
     lists.
   - Utilizes `make_query_plan_with_exclusion` to yield the remainder
     of the plan.
   - Added `__slots__` to reduce memory overhead and attribute access cost.

Benchmark (100K queries, 45-node/5-DC topology, Python 3.14, median of 5 runs):
Policy                    |  Kops/s | vs master | delta | Mem KB
-----------------------------------------------------------------
DCAware                   |     159 |      +50% |       |    1.5
RackAware                 |     128 |      +88% |       |    2.0
TokenAware(DCAware)       |      81 |     +350% | +305% |    2.9
TokenAware(RackAware)     |      64 |     +276% | +237% |    3.6
Default(DCAware)          |     199 |     +119% |  +62% |    1.6
HostFilter(DCAware)       |      56 |       +6% |       |    1.7

TokenAware throughput increases ~4x: DCAware-wrapped goes from 18 to 81 Kops/s
(+350%), RackAware-wrapped from 17 to 64 Kops/s (+276%).
DefaultPolicy also benefits strongly (+119%) from the new exclusion interface.

Signed-off-by: Yaniv Kaul <yaniv.kaul@scylladb.com>

Author: mykaul

cassandra/policies.py

@@ -14,7 +14,7 @@
 import random
 
 from collections import namedtuple
-from itertools import islice, cycle, groupby, repeat
+from itertools import islice, cycle, groupby, repeat, chain
 import logging
 from random import randint, shuffle
 from threading import Lock
@@ -157,6 +157,20 @@ def make_query_plan(self, working_keyspace=None, query=None):
         """
         raise NotImplementedError()
 
+    def make_query_plan_with_exclusion(
+        self, working_keyspace=None, query=None, excluded=()
+    ):
+        """
+        Same as :meth:`make_query_plan`, but with an additional `excluded` parameter.
+        `excluded` should be a container (set, list, etc.) of hosts to skip.
+
+        The default implementation simply delegates to `make_query_plan` and filters the result.
+        Subclasses may override this for performance.
+        """
+        for host in self.make_query_plan(working_keyspace, query):
+            if host not in excluded:
+                yield host
+
     def check_supported(self):
         """
         This will be called after the cluster Metadata has been initialized.
@@ -199,6 +213,22 @@ def make_query_plan(self, working_keyspace=None, query=None):
         else:
             return []
 
+    def make_query_plan_with_exclusion(
+        self, working_keyspace=None, query=None, excluded=()
+    ):
+        pos = self._position
+        self._position += 1
+
+        hosts = self._live_hosts
+        length = len(hosts)
+        if length:
+            pos %= length
+            for host in islice(cycle(hosts), pos, pos + length):
+                if host not in excluded:
+                    yield host
+        else:
+            return
+
     def on_up(self, host):
         with self._hosts_lock:
             self._live_hosts = self._live_hosts.union((host,))
@@ -297,6 +327,26 @@ def make_query_plan(self, working_keyspace=None, query=None):
         for host in self._remote_hosts:
             yield host
 
+    def make_query_plan_with_exclusion(
+        self, working_keyspace=None, query=None, excluded=()
+    ):
+        # not thread-safe, but we don't care much about lost increments
+        # for the purposes of load balancing
+        pos = self._position
+        self._position += 1
+
+        local_live = self._dc_live_hosts.get(self.local_dc, ())
+        pos = (pos % len(local_live)) if local_live else 0
+        for host in islice(cycle(local_live), pos, pos + len(local_live)):
+            if excluded and host in excluded:
+                continue
+            yield host
+
+        for host in self._remote_hosts:
+            if excluded and host in excluded:
+                continue
+            yield host
+
     def on_up(self, host):
         # not worrying about threads because this will happen during
         # control connection startup/refresh
@@ -449,6 +499,35 @@ def make_query_plan(self, working_keyspace=None, query=None):
         for host in self._remote_hosts:
             yield host
 
+    def make_query_plan_with_exclusion(
+        self, working_keyspace=None, query=None, excluded=()
+    ):
+        pos = self._position
+        self._position += 1
+
+        local_rack_live = self._live_hosts.get((self.local_dc, self.local_rack), ())
+        length = len(local_rack_live)
+        if length:
+            p = pos % length
+            for host in islice(cycle(local_rack_live), p, p + length):
+                if excluded and host in excluded:
+                    continue
+                yield host
+
+        local_non_rack = self._non_local_rack_hosts
+        length = len(local_non_rack)
+        if length:
+            p = pos % length
+            for host in islice(cycle(local_non_rack), p, p + length):
+                if excluded and host in excluded:
+                    continue
+                yield host
+
+        for host in self._remote_hosts:
+            if excluded and host in excluded:
+                continue
+            yield host
+
     def on_up(self, host):
         dc = self._dc(host)
         rack = self._rack(host)
@@ -515,16 +594,12 @@ class TokenAwarePolicy(LoadBalancingPolicy):
     policy's query plan will be used as is.
     """
 
-    _child_policy = None
-    _cluster_metadata = None
-    shuffle_replicas = True
-    """
-    Yield local replicas in a random order.
-    """
+    __slots__ = ("_child_policy", "_cluster_metadata", "shuffle_replicas")
 
     def __init__(self, child_policy, shuffle_replicas=True):
         self._child_policy = child_policy
         self.shuffle_replicas = shuffle_replicas
+        self._cluster_metadata = None
 
     def populate(self, cluster, hosts):
         self._cluster_metadata = cluster.metadata
@@ -548,48 +623,79 @@ def make_query_plan(self, working_keyspace=None, query=None):
 
         child = self._child_policy
         if query is None or query.routing_key is None or keyspace is None:
-            for host in child.make_query_plan(keyspace, query):
-                yield host
+            yield from child.make_query_plan(keyspace, query)
             return
 
+        cluster_metadata = self._cluster_metadata
+        token_map = cluster_metadata.token_map
         replicas = []
-        tablet = self._cluster_metadata._tablets.get_tablet_for_key(
-            keyspace,
-            query.table,
-            self._cluster_metadata.token_map.token_class.from_key(query.routing_key),
-        )
-
-        if tablet is not None:
-            replicas_mapped = set(map(lambda r: r[0], tablet.replicas))
-            child_plan = child.make_query_plan(keyspace, query)
+        if token_map:
+            try:
+                token = token_map.token_class.from_key(query.routing_key)
+                tablet = cluster_metadata._tablets.get_tablet_for_key(
+                    keyspace, query.table, token
+                )
 
-            replicas = [host for host in child_plan if host.host_id in replicas_mapped]
-        else:
-            replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key)
+                if tablet is not None:
+                    replicas_mapped = {r[0] for r in tablet.replicas}
+                    child_plan = child.make_query_plan(keyspace, query)
+                    replicas = [
+                        host for host in child_plan if host.host_id in replicas_mapped
+                    ]
+                else:
+                    try:
+                        replicas = list(token_map.get_replicas(keyspace, token))
+                    except Exception:
+                        log.debug(
+                            "Failed to get replicas from token_map, falling back to cluster metadata"
+                        )
+                        replicas = cluster_metadata.get_replicas(
+                            keyspace, query.routing_key
+                        )
+            except Exception:
+                log.debug(
+                    "Failed to resolve token or tablet for query plan, "
+                    "falling back to child policy",
+                    exc_info=True,
+                )
 
         if self.shuffle_replicas and not query.is_lwt():
             shuffle(replicas)
 
-        def yield_in_order(hosts):
-            for distance in [
-                HostDistance.LOCAL_RACK,
-                HostDistance.LOCAL,
-                HostDistance.REMOTE,
-            ]:
-                for replica in hosts:
-                    if replica.is_up and child.distance(replica) == distance:
-                        yield replica
-
-        # yield replicas: local_rack, local, remote
-        yield from yield_in_order(replicas)
-        # yield rest of the cluster: local_rack, local, remote
-        yield from yield_in_order(
-            [
-                host
-                for host in child.make_query_plan(keyspace, query)
-                if host not in replicas
-            ]
-        )
+        local_rack = []
+        local = []
+        remote = []
+
+        child_distance = child.distance
+
+        for replica in replicas:
+            if replica.is_up:
+                d = child_distance(replica)
+                if d == HostDistance.LOCAL_RACK:
+                    local_rack.append(replica)
+                elif d == HostDistance.LOCAL:
+                    local.append(replica)
+                elif d == HostDistance.REMOTE:
+                    remote.append(replica)
+
+        yielded_sequence = tuple(chain(local_rack, local, remote))
+
+        if yielded_sequence:
+            yield from yielded_sequence
+
+            yielded = set(yielded_sequence)
+
+            # yield rest of the cluster
+            try:
+                yield from child.make_query_plan_with_exclusion(
+                    keyspace, query, yielded
+                )
+            except (AttributeError, TypeError):
+                for host in child.make_query_plan(keyspace, query):
+                    if host not in yielded:
+                        yield host
+        else:
+            yield from child.make_query_plan(keyspace, query)
 
     def on_up(self, *args, **kwargs):
         return self._child_policy.on_up(*args, **kwargs)