From 2153f71023568ba345a6ff4e8e1836654815c860 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Fri, 3 Apr 2026 20:03:21 +0300
Subject: [PATCH 01/10] perf: optimize DCAwareRoundRobinPolicy with cached
 remote hosts

Introduce _remote_hosts dict to cache REMOTE hosts, enabling O(1)
distance lookups instead of scanning per-DC host lists.
Replace islice(cycle(...)) with index arithmetic in make_query_plan.
Call _refresh_remote_hosts() on topology changes.
---
 cassandra/policies.py | 58 +++++++++++++++++++++++++++----------------
 1 file changed, 37 insertions(+), 21 deletions(-)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index 89702e8c89..083327e25c 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 import random
 
-from collections import namedtuple
+from collections import namedtuple, OrderedDict
 from itertools import islice, cycle, groupby, repeat
 import logging
 from random import randint, shuffle
@@ -244,34 +244,41 @@ def __init__(self, local_dc='', used_hosts_per_remote_dc=0):
         self.local_dc = local_dc
         self.used_hosts_per_remote_dc = used_hosts_per_remote_dc
         self._dc_live_hosts = {}
+        self._remote_hosts = {}
         self._position = 0
         LoadBalancingPolicy.__init__(self)
 
     def _dc(self, host):
         return host.datacenter or self.local_dc
 
+    def _refresh_remote_hosts(self):
+        # Using dict.fromkeys() instead of a set to preserve insertion order (Python 3.7+)
+        # while still providing O(1) lookup for `host in self._remote_hosts`.
+        remote_hosts = {}
+        if self.used_hosts_per_remote_dc > 0:
+            for datacenter, hosts in self._dc_live_hosts.items():
+                if datacenter != self.local_dc:
+                    remote_hosts.update(
+                        dict.fromkeys(hosts[:self.used_hosts_per_remote_dc])
+                    )
+        self._remote_hosts = remote_hosts
+
     def populate(self, cluster, hosts):
         for dc, dc_hosts in groupby(hosts, lambda h: self._dc(h)):
             self._dc_live_hosts[dc] = tuple({*dc_hosts, *self._dc_live_hosts.get(dc, [])})
 
         self._position = randint(0, len(hosts) - 1) if hosts else 0
+        self._refresh_remote_hosts()
 
     def distance(self, host):
         dc = self._dc(host)
         if dc == self.local_dc:
             return HostDistance.LOCAL
 
-        if not self.used_hosts_per_remote_dc:
-            return HostDistance.IGNORED
-        else:
-            dc_hosts = self._dc_live_hosts.get(dc)
-            if not dc_hosts:
-                return HostDistance.IGNORED
-
-            if host in list(dc_hosts)[:self.used_hosts_per_remote_dc]:
-                return HostDistance.REMOTE
-            else:
-                return HostDistance.IGNORED
+        remote_hosts = self._remote_hosts
+        if host in remote_hosts:
+            return HostDistance.REMOTE
+        return HostDistance.IGNORED
 
     def make_query_plan(self, working_keyspace=None, query=None):
         # not thread-safe, but we don't care much about lost increments
@@ -280,32 +287,38 @@ def make_query_plan(self, working_keyspace=None, query=None):
         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)):
-            yield host
+        length = len(local_live)
+        if length:
+            pos %= length
+            for i in range(length):
+                yield local_live[(pos + i) % length]
 
-        # the dict can change, so get candidate DCs iterating over keys of a copy
-        other_dcs = [dc for dc in self._dc_live_hosts.copy().keys() if dc != self.local_dc]
-        for dc in other_dcs:
-            remote_live = self._dc_live_hosts.get(dc, ())
-            for host in remote_live[:self.used_hosts_per_remote_dc]:
-                yield host
+        remote_hosts = self._remote_hosts
+        for host in remote_hosts:
+            yield host
 
     def on_up(self, host):
         # not worrying about threads because this will happen during
         # control connection startup/refresh
+        refresh_remote = False
         if not self.local_dc and host.datacenter:
             self.local_dc = host.datacenter
             log.info("Using datacenter '%s' for DCAwareRoundRobinPolicy (via host '%s'); "
                         "if incorrect, please specify a local_dc to the constructor, "
                         "or limit contact points to local cluster nodes" %
                         (self.local_dc, host.endpoint))
+            refresh_remote = True
 
         dc = self._dc(host)
         with self._hosts_lock:
             current_hosts = self._dc_live_hosts.get(dc, ())
             if host not in current_hosts:
                 self._dc_live_hosts[dc] = current_hosts + (host, )
+                if dc != self.local_dc:
+                    refresh_remote = True
+
+            if refresh_remote:
+                self._refresh_remote_hosts()
 
     def on_down(self, host):
         dc = self._dc(host)
@@ -318,6 +331,9 @@ def on_down(self, host):
                 else:
                     del self._dc_live_hosts[dc]
 
+                if dc != self.local_dc:
+                    self._refresh_remote_hosts()
+
     def on_add(self, host):
         self.on_up(host)
 

From 9f83a8ea2d78010d3eba2f43111d5132c15bb947 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Fri, 3 Apr 2026 20:04:08 +0300
Subject: [PATCH 02/10] feat: add make_query_plan_with_exclusion to
 LoadBalancingPolicy, RoundRobin, and DCAware

Add a new make_query_plan_with_exclusion() method that skips hosts
in an exclusion set.  The base class provides a default filtering
implementation; RoundRobin and DCAware override for efficiency.
---
 cassandra/policies.py | 60 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index 083327e25c..d8fe94a677 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -157,6 +157,18 @@ 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.
@@ -198,6 +210,20 @@ 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 +323,40 @@ def make_query_plan(self, working_keyspace=None, query=None):
         for host in 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, ())
+        length = len(local_live)
+        remote_hosts = self._remote_hosts
+        if not excluded:
+            if length:
+                pos %= length
+                for i in range(length):
+                    yield local_live[(pos + i) % length]
+            for host in remote_hosts:
+                yield host
+            return
+
+        if not isinstance(excluded, set):
+            excluded = set(excluded)
+
+        if length:
+            pos %= length
+            for i in range(length):
+                host = local_live[(pos + i) % length]
+                if host in excluded:
+                    continue
+                yield host
+
+        for host in remote_hosts:
+            if host in excluded:
+                continue
+            yield host
+
     def on_up(self, host):
         # not worrying about threads because this will happen during
         # control connection startup/refresh

From 1f9cba01cb03dfaa9ec6020733ef9b0c13feb62c Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Fri, 3 Apr 2026 20:04:51 +0300
Subject: [PATCH 03/10] perf: optimize RackAwareRoundRobinPolicy with cached
 host distances

Cache remote hosts and non-local-rack hosts to enable O(1) distance
lookups. Replace islice(cycle(...)) with index arithmetic. Reorder
on_up/on_down to update DC-level hosts before rack-level for correct
cache invalidation.
---
 cassandra/policies.py | 100 ++++++++++++++++++++++++++----------------
 1 file changed, 63 insertions(+), 37 deletions(-)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index d8fe94a677..d0f7ed33eb 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -429,6 +429,8 @@ def __init__(self, local_dc, local_rack, used_hosts_per_remote_dc=0):
         self.used_hosts_per_remote_dc = used_hosts_per_remote_dc
         self._live_hosts = {}
         self._dc_live_hosts = {}
+        self._remote_hosts = {}
+        self._non_local_rack_hosts = []
         self._endpoints = []
         self._position = 0
         LoadBalancingPolicy.__init__(self)
@@ -439,6 +441,24 @@ def _rack(self, host):
     def _dc(self, host):
         return host.datacenter or self.local_dc
 
+    def _refresh_remote_hosts(self):
+        # Using dict.fromkeys() instead of a set to preserve insertion order (Python 3.7+)
+        # while still providing O(1) lookup for `host in self._remote_hosts`.
+        remote_hosts = {}
+        if self.used_hosts_per_remote_dc > 0:
+            for datacenter, hosts in self._dc_live_hosts.items():
+                if datacenter != self.local_dc:
+                    remote_hosts.update(
+                        dict.fromkeys(hosts[:self.used_hosts_per_remote_dc])
+                    )
+        self._remote_hosts = remote_hosts
+
+    def _refresh_non_local_rack_hosts(self):
+        local_live = self._dc_live_hosts.get(self.local_dc, ())
+        self._non_local_rack_hosts = [
+            h for h in local_live if self._rack(h) != self.local_rack
+        ]
+
     def populate(self, cluster, hosts):
         for (dc, rack), rack_hosts in groupby(hosts, lambda host: (self._dc(host), self._rack(host))):
             self._live_hosts[(dc, rack)] = tuple({*rack_hosts, *self._live_hosts.get((dc, rack), [])})
@@ -446,71 +466,64 @@ def populate(self, cluster, hosts):
             self._dc_live_hosts[dc] = tuple({*dc_hosts, *self._dc_live_hosts.get(dc, [])})
 
         self._position = randint(0, len(hosts) - 1) if hosts else 0
+        self._refresh_remote_hosts()
+        self._refresh_non_local_rack_hosts()
 
     def distance(self, host):
-        rack = self._rack(host)
         dc = self._dc(host)
-        if rack == self.local_rack and dc == self.local_dc:
-            return HostDistance.LOCAL_RACK
-
         if dc == self.local_dc:
+            if self._rack(host) == self.local_rack:
+                return HostDistance.LOCAL_RACK
             return HostDistance.LOCAL
 
-        if not self.used_hosts_per_remote_dc:
-            return HostDistance.IGNORED
-
-        dc_hosts = self._dc_live_hosts.get(dc, ())
-        if not dc_hosts:
-            return HostDistance.IGNORED
-        if host in dc_hosts and dc_hosts.index(host) < self.used_hosts_per_remote_dc:
+        remote_hosts = self._remote_hosts
+        if host in remote_hosts:
             return HostDistance.REMOTE
-        else:
-            return HostDistance.IGNORED
+        return HostDistance.IGNORED
 
     def make_query_plan(self, working_keyspace=None, query=None):
         pos = self._position
         self._position += 1
 
         local_rack_live = self._live_hosts.get((self.local_dc, self.local_rack), ())
-        pos = (pos % len(local_rack_live)) if local_rack_live else 0
-        # Slice the cyclic iterator to start from pos and include the next len(local_live) elements
-        # This ensures we get exactly one full cycle starting from pos
-        for host in islice(cycle(local_rack_live), pos, pos + len(local_rack_live)):
-            yield host
+        length = len(local_rack_live)
+        if length:
+            p = pos % length
+            for i in range(length):
+                yield local_rack_live[(p + i) % length]
 
-        local_live = [host for host in self._dc_live_hosts.get(self.local_dc, ()) if host.rack != self.local_rack]
-        pos = (pos % len(local_live)) if local_live else 0
-        for host in islice(cycle(local_live), pos, pos + len(local_live)):
-            yield host
+        local_non_rack = self._non_local_rack_hosts
+        length = len(local_non_rack)
+        if length:
+            p = pos % length
+            for i in range(length):
+                yield local_non_rack[(p + i) % length]
 
-        # the dict can change, so get candidate DCs iterating over keys of a copy
-        for dc, remote_live in self._dc_live_hosts.copy().items():
-            if dc != self.local_dc:
-                for host in remote_live[:self.used_hosts_per_remote_dc]:
-                    yield host
+        remote_hosts = self._remote_hosts
+        for host in remote_hosts:
+            yield host
 
     def on_up(self, host):
         dc = self._dc(host)
         rack = self._rack(host)
         with self._hosts_lock:
-            current_rack_hosts = self._live_hosts.get((dc, rack), ())
-            if host not in current_rack_hosts:
-                self._live_hosts[(dc, rack)] = current_rack_hosts + (host, )
             current_dc_hosts = self._dc_live_hosts.get(dc, ())
             if host not in current_dc_hosts:
                 self._dc_live_hosts[dc] = current_dc_hosts + (host, )
 
+                if dc != self.local_dc:
+                    self._refresh_remote_hosts()
+                else:
+                    self._refresh_non_local_rack_hosts()
+
+            current_rack_hosts = self._live_hosts.get((dc, rack), ())
+            if host not in current_rack_hosts:
+                self._live_hosts[(dc, rack)] = current_rack_hosts + (host, )
+
     def on_down(self, host):
         dc = self._dc(host)
         rack = self._rack(host)
         with self._hosts_lock:
-            current_rack_hosts = self._live_hosts.get((dc, rack), ())
-            if host in current_rack_hosts:
-                hosts = tuple(h for h in current_rack_hosts if h != host)
-                if hosts:
-                    self._live_hosts[(dc, rack)] = hosts
-                else:
-                    del self._live_hosts[(dc, rack)]
             current_dc_hosts = self._dc_live_hosts.get(dc, ())
             if host in current_dc_hosts:
                 hosts = tuple(h for h in current_dc_hosts if h != host)
@@ -519,6 +532,19 @@ def on_down(self, host):
                 else:
                     del self._dc_live_hosts[dc]
 
+                if dc != self.local_dc:
+                    self._refresh_remote_hosts()
+                else:
+                    self._refresh_non_local_rack_hosts()
+
+            current_rack_hosts = self._live_hosts.get((dc, rack), ())
+            if host in current_rack_hosts:
+                hosts = tuple(h for h in current_rack_hosts if h != host)
+                if hosts:
+                    self._live_hosts[(dc, rack)] = hosts
+                else:
+                    del self._live_hosts[(dc, rack)]
+
     def on_add(self, host):
         self.on_up(host)
 

From 96186ab09cee6a523367a9d3ef56af8af17afe54 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Fri, 3 Apr 2026 20:05:15 +0300
Subject: [PATCH 04/10] feat: add make_query_plan_with_exclusion to
 RackAwareRoundRobinPolicy

Optimized exclusion-aware query plan that avoids re-computing
non-local-rack and remote host lists.
---
 cassandra/policies.py | 50 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index d0f7ed33eb..c65851d930 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -503,6 +503,56 @@ def make_query_plan(self, working_keyspace=None, query=None):
         for host in 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)
+        remote_hosts = self._remote_hosts
+        if not excluded:
+            if length:
+                p = pos % length
+                for i in range(length):
+                    yield local_rack_live[(p + i) % length]
+
+            local_non_rack = self._non_local_rack_hosts
+            length = len(local_non_rack)
+            if length:
+                p = pos % length
+                for i in range(length):
+                    yield local_non_rack[(p + i) % length]
+
+            for host in remote_hosts:
+                yield host
+            return
+
+        if not isinstance(excluded, set):
+            excluded = set(excluded)
+
+        if length:
+            p = pos % length
+            for i in range(length):
+                host = local_rack_live[(p + i) % length]
+                if 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 i in range(length):
+                host = local_non_rack[(p + i) % length]
+                if host in excluded:
+                    continue
+                yield host
+
+        for host in remote_hosts:
+            if host in excluded:
+                continue
+            yield host
+
     def on_up(self, host):
         dc = self._dc(host)
         rack = self._rack(host)

From 82d2fde50744d2fe0dbb5d3fd104457108964000 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Fri, 3 Apr 2026 20:06:28 +0300
Subject: [PATCH 05/10] perf: add LRU replica cache and optimize
 TokenAwarePolicy query plan

- Add LRU cache (default 1024 entries) for token-to-replicas lookups,
  auto-invalidated on topology changes (token_map identity check).
- Sort replicas by distance (LOCAL_RACK > LOCAL > REMOTE) in a single
  pass instead of iterating three times.
- Skip distance re-sorting for DCAware/RackAware child policies since
  they already yield in distance order; fallback re-sort for others.
- LWT queries skip replica shuffling for deterministic plans.
- Use make_query_plan_with_exclusion to avoid re-yielding replicas.

Add a public cache_replicas_size property mirroring the constructor
argument of the same name, matching the existing shuffle_replicas
attribute pattern. The class docstring already referenced
:attr:`cache_replicas_size`, but no such public attribute existed --
only the constructor parameter and the private _cache_replicas_size --
which would trip up this repo's docs CI (Sphinx warnings-as-errors on
cassandra/** changes).

Also report cache_replicas_size from
token_aware_policy_insights_serializer, matching how
dc_aware_round_robin_policy_insights_serializer reports all of its
constructor options; it was missed when this option was introduced.

Fix (review thread on PR #651): in the fallback re-sort branch (for
child policies other than DCAware/RackAware), a host yielded by the
child plan whose distance was not LOCAL_RACK/LOCAL/REMOTE (e.g.
IGNORED, or a distance that flipped mid-plan due to a concurrent
topology refresh) was silently discarded instead of appended, shrinking
the plan relative to the child policy's own output. Such hosts are now
collected into a trailing "remaining_other" bucket, yielded after
LOCAL_RACK/LOCAL/REMOTE and preserving their relative order from the
child plan, instead of being dropped. Adds
test_wrap_child_ignored_distance_host_not_dropped (fails against the
pre-fix code: the IGNORED-distance host is missing from the plan).
---
 cassandra/datastax/insights/serializers.py |   3 +-
 cassandra/policies.py                      | 190 ++++++++++++++++++---
 tests/unit/advanced/test_insights.py       |  12 +-
 tests/unit/test_policies.py                |  58 +++++++
 4 files changed, 232 insertions(+), 31 deletions(-)

diff --git a/cassandra/datastax/insights/serializers.py b/cassandra/datastax/insights/serializers.py
index 289c165e8a..7a0abedefe 100644
--- a/cassandra/datastax/insights/serializers.py
+++ b/cassandra/datastax/insights/serializers.py
@@ -70,7 +70,8 @@ def token_aware_policy_insights_serializer(policy):
                 'namespace': namespace(policy.__class__),
                 'options': {'child_policy': insights_registry.serialize(policy._child_policy,
                                                                         policy=True),
-                            'shuffle_replicas': policy.shuffle_replicas}
+                            'shuffle_replicas': policy.shuffle_replicas,
+                            'cache_replicas_size': policy.cache_replicas_size}
                 }
 
     @insights_registry.register_serializer_for(WhiteListRoundRobinPolicy)
diff --git a/cassandra/policies.py b/cassandra/policies.py
index c65851d930..a16c356b28 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -616,6 +616,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.
+
+    An LRU cache of size :attr:`cache_replicas_size` (default 1024) avoids
+    repeated token-to-replica lookups for the same (keyspace, routing_key)
+    pair.  Set to 0 to disable caching.  The cache is automatically
+    invalidated when the cluster topology changes.
     """
 
     _child_policy = None
@@ -625,9 +630,24 @@ class TokenAwarePolicy(LoadBalancingPolicy):
     Yield local replicas in a random order.
     """
 
-    def __init__(self, child_policy, shuffle_replicas=True):
+    def __init__(self, child_policy, shuffle_replicas=True, cache_replicas_size=1024):
+        super().__init__()
         self._child_policy = child_policy
         self.shuffle_replicas = shuffle_replicas
+        self._cluster_metadata = None
+        self._cache_replicas_size = max(0, cache_replicas_size)
+        self._replica_cache = OrderedDict()
+        self._replica_cache_token_map_ref = None
+        self._cache_lock = Lock()
+
+    @property
+    def cache_replicas_size(self):
+        """
+        The configured size of the replica LRU cache, as passed to the
+        `cache_replicas_size` constructor argument (default 1024).  A value
+        of 0 means caching is disabled.  Read-only.
+        """
+        return self._cache_replicas_size
 
     def populate(self, cluster, hosts):
         self._cluster_metadata = cluster.metadata
@@ -645,40 +665,158 @@ def check_supported(self):
     def distance(self, *args, **kwargs):
         return self._child_policy.distance(*args, **kwargs)
 
+    def _get_cached_replicas(self, keyspace, routing_key_bytes, token_map):
+        """
+        Return cached (token, replicas) for the given keyspace and routing key,
+        or None on cache miss.  The cache is invalidated whenever the token_map
+        object identity changes (i.e. after a topology rebuild).
+        """
+        if not self._cache_replicas_size:
+            return None
+        with self._cache_lock:
+            if token_map is not self._replica_cache_token_map_ref:
+                # Token map was rebuilt -- entire cache is stale.
+                self._replica_cache = OrderedDict()
+                self._replica_cache_token_map_ref = token_map
+            cache_key = (keyspace, routing_key_bytes)
+            entry = self._replica_cache.get(cache_key)
+            if entry is not None:
+                # Promote to most-recently-used.
+                self._replica_cache.move_to_end(cache_key)
+            return entry
+
+    def _put_cached_replicas(self, keyspace, routing_key_bytes, token, replicas, token_map):
+        """
+        Store (token, replicas) in the LRU cache, evicting the oldest
+        entry if the cache exceeds its configured size.
+        """
+        if not self._cache_replicas_size:
+            return
+        with self._cache_lock:
+            if token_map is not self._replica_cache_token_map_ref:
+                self._replica_cache = OrderedDict()
+                self._replica_cache_token_map_ref = token_map
+            cache_key = (keyspace, routing_key_bytes)
+            self._replica_cache[cache_key] = (token, replicas)
+            self._replica_cache.move_to_end(cache_key)
+            if len(self._replica_cache) > self._cache_replicas_size:
+                self._replica_cache.popitem(last=False)
+
     def make_query_plan(self, working_keyspace=None, query=None):
         keyspace = query.keyspace if query and query.keyspace else working_keyspace
 
         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)
-
-            replicas = [host for host in child_plan if host.host_id in replicas_mapped]
+        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
+                )
+
+                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:
+                    cached = self._get_cached_replicas(keyspace, query.routing_key, token_map)
+                    if cached is not None:
+                        token, replicas = cached
+                    else:
+                        try:
+                            replicas = 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)
+                        self._put_cached_replicas(
+                            keyspace, query.routing_key, token, replicas, token_map
+                        )
+            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:
+            if not query.is_lwt() and not ConsistencyLevel.is_serial(query.consistency_level):
+                replicas = list(replicas)
+                shuffle(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)
+
+        if local_rack or local or remote:
+            yielded = set()
+
+            for replica in local_rack:
+                yielded.add(replica)
+                yield replica
+
+            for replica in local:
+                yielded.add(replica)
+                yield replica
+
+            for replica in remote:
+                yielded.add(replica)
+                yield replica
+
+            # Yield the rest of the cluster (non-replica hosts).
+            # DCAware and RackAware already yield in distance order
+            # (local_rack -> local -> remote), so we can stream directly.
+            # For other child policies we must re-sort by distance.
+            if isinstance(child, (DCAwareRoundRobinPolicy, RackAwareRoundRobinPolicy)):
+                yield from child.make_query_plan_with_exclusion(keyspace, query, yielded)
+            else:
+                remaining_local_rack = []
+                remaining_local = []
+                remaining_remote = []
+                # Hosts whose distance is neither LOCAL_RACK/LOCAL/REMOTE
+                # (e.g. IGNORED, or a distance that flipped mid-plan due to
+                # a concurrent topology refresh) are collected here rather
+                # than dropped, and yielded last -- preserving their
+                # relative order from the child plan, matching the
+                # pre-existing (non-optimized) behavior of never shrinking
+                # the plan compared to what the child policy yielded.
+                remaining_other = []
+                for host in child.make_query_plan_with_exclusion(keyspace, query, yielded):
+                    d = child_distance(host)
+                    if d == HostDistance.LOCAL_RACK:
+                        remaining_local_rack.append(host)
+                    elif d == HostDistance.LOCAL:
+                        remaining_local.append(host)
+                    elif d == HostDistance.REMOTE:
+                        remaining_remote.append(host)
+                    else:
+                        remaining_other.append(host)
+                yield from remaining_local_rack
+                yield from remaining_local
+                yield from remaining_remote
+                yield from remaining_other
         else:
-            replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key)
-
-        if self.shuffle_replicas and not query.is_lwt() and not ConsistencyLevel.is_serial(query.consistency_level):
-            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])
+            yield from child.make_query_plan(keyspace, query)
 
     def on_up(self, *args, **kwargs):
         return self._child_policy.on_up(*args, **kwargs)
diff --git a/tests/unit/advanced/test_insights.py b/tests/unit/advanced/test_insights.py
index 2050439804..0773ff965a 100644
--- a/tests/unit/advanced/test_insights.py
+++ b/tests/unit/advanced/test_insights.py
@@ -120,7 +120,8 @@ def test_execution_profile(self):
                                                         'options': {'local_dc': '',
                                                                     'used_hosts_per_remote_dc': 0},
                                                         'type': 'DCAwareRoundRobinPolicy'},
-                                       'shuffle_replicas': True},
+                                       'shuffle_replicas': True,
+                                       'cache_replicas_size': 1024},
                            'type': 'TokenAwarePolicy'},
          'readTimeout': 10.0,
          'retry': {'namespace': 'cassandra.policies', 'options': {}, 'type': 'RetryPolicy'},
@@ -139,7 +140,8 @@ def test_graph_execution_profile(self):
                                                         'options': {'local_dc': '',
                                                                     'used_hosts_per_remote_dc': 0},
                                                         'type': 'DCAwareRoundRobinPolicy'},
-                                       'shuffle_replicas': True},
+                                       'shuffle_replicas': True,
+                                       'cache_replicas_size': 1024},
                            'type': 'TokenAwarePolicy'},
          'readTimeout': 30.0,
          'retry': {'namespace': 'cassandra.policies', 'options': {}, 'type': 'NeverRetryPolicy'},
@@ -161,7 +163,8 @@ def test_graph_analytics_execution_profile(self):
                                                                                      'options': {'local_dc': '',
                                                                                                  'used_hosts_per_remote_dc': 0},
                                                                                      'type': 'DCAwareRoundRobinPolicy'},
-                                                                    'shuffle_replicas': True},
+                                                                    'shuffle_replicas': True,
+                                                                    'cache_replicas_size': 1024},
                                                         'type': 'TokenAwarePolicy'}},
                            'type': 'DefaultLoadBalancingPolicy'},
          'readTimeout': 604800.0,
@@ -189,7 +192,8 @@ def test_token_aware_policy(self):
          'options': {'child_policy': {'namespace': 'cassandra.policies',
                                       'options': {},
                                       'type': 'LoadBalancingPolicy'},
-                     'shuffle_replicas': True},
+                     'shuffle_replicas': True,
+                     'cache_replicas_size': 1024},
          'type': 'TokenAwarePolicy'}
 
     def test_whitelist_round_robin_policy(self):
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
index 63a3c3d12d..a108af2daf 100644
--- a/tests/unit/test_policies.py
+++ b/tests/unit/test_policies.py
@@ -710,6 +710,64 @@ def get_replicas(keyspace, packed_key):
 
             assert 8 == len(qplan)
 
+    def test_wrap_child_ignored_distance_host_not_dropped(self):
+        """
+        Regression test: in the re-sort branch of make_query_plan (used for
+        child policies other than DCAwareRoundRobinPolicy/RackAwareRoundRobinPolicy),
+        a host yielded by the child's query plan whose distance is not one of
+        LOCAL_RACK/LOCAL/REMOTE (e.g. IGNORED, or a distance that flipped
+        concurrently) must still end up in the final plan, appended in a
+        trailing bucket after LOCAL_RACK/LOCAL/REMOTE -- not silently dropped.
+        """
+        cluster = Mock(spec=Cluster)
+        cluster.metadata = Mock(spec=Metadata)
+        cluster.metadata._tablets = Mock(spec=Tablets)
+        cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
+
+        hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for host in hosts:
+            host.set_up()
+        local_replica, local_rack_remaining, remote_remaining, ignored_remaining = hosts
+
+        # Only one replica is returned for this routing key -- with LOCAL
+        # distance -- so the initial replica pass yields a non-empty bucket
+        # and the re-sort ("else") branch below is exercised for the rest of
+        # the cluster.
+        cluster.metadata.get_replicas.return_value = [local_replica]
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
+
+        # A plain (non-DCAware/RackAware) child policy: yields the three
+        # remaining hosts -- one of which (ignored_remaining) it currently
+        # considers IGNORED, e.g. because its distance flipped mid-plan due
+        # to a concurrent topology refresh, or because the child policy
+        # intentionally yields ignored hosts.
+        distance_map = {
+            local_replica: HostDistance.LOCAL,
+            local_rack_remaining: HostDistance.LOCAL_RACK,
+            remote_remaining: HostDistance.REMOTE,
+            ignored_remaining: HostDistance.IGNORED,
+        }
+        remaining_plan = [ignored_remaining, local_rack_remaining, remote_remaining]
+
+        child_policy = Mock()
+        child_policy.distance.side_effect = lambda h: distance_map[h]
+        child_policy.make_query_plan_with_exclusion.side_effect = \
+            lambda k, q, e: [h for h in remaining_plan if h not in e]
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=struct.pack('>i', 0), keyspace='keyspace_name')
+        qplan = list(policy.make_query_plan(None, query))
+
+        # The ignored host must still be present in the plan (in the
+        # trailing position, after LOCAL_RACK/LOCAL/REMOTE), not dropped.
+        assert ignored_remaining in qplan
+        assert qplan == [local_replica, local_rack_remaining, remote_remaining, ignored_remaining]
+
     class FakeCluster:
         def __init__(self):
             self.metadata = Mock(spec=Metadata)

From ca5e6c9fd419eea3430883534310e22c1adac039 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Fri, 3 Apr 2026 20:07:03 +0300
Subject: [PATCH 06/10] feat: add make_query_plan_with_exclusion to
 HostFilterPolicy and DefaultLoadBalancingPolicy

Both delegate to their child policy's exclusion-aware query plan
while preserving their specific filtering/targeting behavior.
---
 cassandra/policies.py | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index a16c356b28..d21c8bb190 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -992,6 +992,16 @@ def make_query_plan(self, working_keyspace=None, query=None):
             if self.predicate(host):
                 yield host
 
+    def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excluded=()):
+        if excluded and not isinstance(excluded, set):
+            excluded = set(excluded)
+        child_qp = self._child_policy.make_query_plan_with_exclusion(
+            working_keyspace=working_keyspace, query=query, excluded=excluded
+        )
+        for host in child_qp:
+            if self.predicate(host):
+                yield host
+
     def check_supported(self):
         return self._child_policy.check_supported()
 
@@ -1643,6 +1653,27 @@ def make_query_plan(self, working_keyspace=None, query=None):
             for h in child.make_query_plan(keyspace, query):
                 yield h
 
+    def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excluded=()):
+        if query and query.keyspace:
+            keyspace = query.keyspace
+        else:
+            keyspace = working_keyspace
+
+        addr = getattr(query, 'target_host', None) if query else None
+        target_host = self._cluster_metadata.get_host(addr)
+
+        if excluded and not isinstance(excluded, set):
+            excluded = set(excluded)
+
+        child = self._child_policy
+        if target_host and target_host.is_up and target_host not in excluded:
+            yield target_host
+            for h in child.make_query_plan_with_exclusion(keyspace, query, excluded):
+                if h != target_host:
+                    yield h
+        else:
+            yield from child.make_query_plan_with_exclusion(keyspace, query, excluded)
+
 
 # TODO for backward compatibility, remove in next major
 class DSELoadBalancingPolicy(DefaultLoadBalancingPolicy):

From ed7fee4b5c5c31edc16244c49b8d665de88a2362 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Fri, 3 Apr 2026 20:36:09 +0300
Subject: [PATCH 07/10] test: add unit tests for query plan exclusion, LRU
 cache, and LWT determinism

Add tests for make_query_plan_with_exclusion in RoundRobin, DCAware, and
RackAware policies. Add cache tests (hit, miss, eviction, topology
invalidation, disabled) and LWT determinism tests for TokenAwarePolicy.
Update existing tests to set up token_map mocks and shuffle_replicas=False
to match the new TokenAwarePolicy implementation.
---
 tests/unit/test_policies.py | 417 +++++++++++++++++++++++++++++++++++-
 1 file changed, 409 insertions(+), 8 deletions(-)

diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
index a108af2daf..038a2dd966 100644
--- a/tests/unit/test_policies.py
+++ b/tests/unit/test_policies.py
@@ -187,6 +187,39 @@ def test_no_live_nodes(self):
         qplan = list(policy.make_query_plan())
         assert qplan == []
 
+    def test_make_query_plan_with_exclusion_basic(self):
+        hosts = [0, 1, 2, 3]
+        policy = RoundRobinPolicy()
+        policy.populate(None, hosts)
+        excluded = {1, 3}
+        qplan = list(policy.make_query_plan_with_exclusion(excluded=excluded))
+        assert 1 not in qplan
+        assert 3 not in qplan
+        assert sorted(qplan) == [0, 2]
+
+    def test_make_query_plan_with_exclusion_empty(self):
+        hosts = [0, 1, 2, 3]
+        policy = RoundRobinPolicy()
+        policy.populate(None, hosts)
+        qplan = list(policy.make_query_plan_with_exclusion(excluded=()))
+        assert sorted(qplan) == hosts
+
+    def test_make_query_plan_with_exclusion_all_excluded(self):
+        hosts = [0, 1, 2, 3]
+        policy = RoundRobinPolicy()
+        policy.populate(None, hosts)
+        excluded = set(hosts)
+        qplan = list(policy.make_query_plan_with_exclusion(excluded=excluded))
+        assert qplan == []
+
+    def test_make_query_plan_with_exclusion_superset(self):
+        hosts = [0, 1, 2, 3]
+        policy = RoundRobinPolicy()
+        policy.populate(None, hosts)
+        excluded = {0, 1, 2, 3, 99, 100}
+        qplan = list(policy.make_query_plan_with_exclusion(excluded=excluded))
+        assert qplan == []
+
 @pytest.mark.parametrize("policy_specialization, constructor_args", [(DCAwareRoundRobinPolicy, ("dc1", )), (RackAwareRoundRobinPolicy, ("dc1", "rack1"))])
 class TestRackOrDCAwareRoundRobinPolicy:
 
@@ -273,6 +306,12 @@ def test_get_distance(self, policy_specialization, constructor_args):
         elif isinstance(policy_specialization, RackAwareRoundRobinPolicy):
             assert policy.distance(host) == HostDistance.LOCAL_RACK
 
+        # Reset policy state to simulate a fresh view or handle the "move" correctly
+        if hasattr(policy, '_live_hosts'):
+            policy._live_hosts.clear()
+        if hasattr(policy, '_dc_live_hosts'):
+            policy._dc_live_hosts.clear()
+
         # same dc different rack
         host = Host(DefaultEndPoint("ip1"), SimpleConvictionPolicy, host_id=uuid.uuid4())
         host.set_location_info("dc1", "rack2")
@@ -527,6 +566,54 @@ def test_no_nodes(self, policy_specialization, constructor_args):
         qplan = list(policy.make_query_plan())
         assert qplan == []
 
+    def test_make_query_plan_with_exclusion_basic(self, policy_specialization, constructor_args):
+        hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(6)]
+        for h in hosts[:2]:
+            h.set_location_info("dc1", "rack1")
+        for h in hosts[2:4]:
+            h.set_location_info("dc1", "rack2")
+        for h in hosts[4:]:
+            h.set_location_info("dc2", "rack1")
+
+        policy = policy_specialization(*constructor_args, used_hosts_per_remote_dc=2)
+        policy.populate(Mock(), hosts)
+
+        # exclude some local and remote hosts
+        excluded = {hosts[0], hosts[4]}
+        qplan = list(policy.make_query_plan_with_exclusion(excluded=excluded))
+        assert hosts[0] not in qplan
+        assert hosts[4] not in qplan
+        assert len(qplan) == 4  # 6 total - 2 excluded
+
+    def test_make_query_plan_with_exclusion_empty(self, policy_specialization, constructor_args):
+        hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for h in hosts[:2]:
+            h.set_location_info("dc1", "rack1")
+        for h in hosts[2:]:
+            h.set_location_info("dc1", "rack2")
+
+        policy = policy_specialization(*constructor_args)
+        policy.populate(Mock(), hosts)
+
+        # empty exclusion should return all hosts
+        qplan_excl = list(policy.make_query_plan_with_exclusion(excluded=()))
+        qplan_normal = list(policy.make_query_plan())
+        assert sorted(qplan_excl, key=id) == sorted(qplan_normal, key=id)
+
+    def test_make_query_plan_with_exclusion_all_excluded(self, policy_specialization, constructor_args):
+        hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for h in hosts[:2]:
+            h.set_location_info("dc1", "rack1")
+        for h in hosts[2:]:
+            h.set_location_info("dc1", "rack2")
+
+        policy = policy_specialization(*constructor_args)
+        policy.populate(Mock(), hosts)
+
+        excluded = set(hosts)
+        qplan = list(policy.make_query_plan_with_exclusion(excluded=excluded))
+        assert qplan == []
+
     def test_wrong_dc(self, policy_specialization, constructor_args):
         hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(3)]
         for h in hosts[:3]:
@@ -577,6 +664,8 @@ def test_wrap_round_robin(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
         for host in hosts:
             host.set_up()
@@ -586,8 +675,9 @@ def get_replicas(keyspace, packed_key):
             return list(islice(cycle(hosts), index, index + 2))
 
         cluster.metadata.get_replicas.side_effect = get_replicas
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
 
-        policy = TokenAwarePolicy(RoundRobinPolicy())
+        policy = TokenAwarePolicy(RoundRobinPolicy(), shuffle_replicas=False)
         policy.populate(cluster, hosts)
 
         for i in range(4):
@@ -610,6 +700,8 @@ def test_wrap_dc_aware(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
         for host in hosts:
             host.set_up()
@@ -627,8 +719,9 @@ def get_replicas(keyspace, packed_key):
                 return [hosts[1], hosts[3]]
 
         cluster.metadata.get_replicas.side_effect = get_replicas
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
 
-        policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=2))
+        policy = TokenAwarePolicy(DCAwareRoundRobinPolicy("dc1", used_hosts_per_remote_dc=2), shuffle_replicas=False)
         policy.populate(cluster, hosts)
 
         for i in range(4):
@@ -659,6 +752,8 @@ def test_wrap_rack_aware(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(8)]
         for host in hosts:
             host.set_up()
@@ -680,8 +775,9 @@ def get_replicas(keyspace, packed_key):
                 return [hosts[4], hosts[5], hosts[6], hosts[7]]
 
         cluster.metadata.get_replicas.side_effect = get_replicas
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
 
-        policy = TokenAwarePolicy(RackAwareRoundRobinPolicy("dc1", "rack1", used_hosts_per_remote_dc=4))
+        policy = TokenAwarePolicy(RackAwareRoundRobinPolicy("dc1", "rack1", used_hosts_per_remote_dc=4), shuffle_replicas=False)
         policy.populate(cluster, hosts)
 
         for i in range(4):
@@ -862,12 +958,16 @@ def test_statement_keyspace(self):
         replicas = hosts[2:]
         cluster.metadata.get_replicas.return_value = replicas
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
 
         child_policy = Mock()
         child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
         child_policy.distance.return_value = HostDistance.LOCAL
 
-        policy = TokenAwarePolicy(child_policy)
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
         policy.populate(cluster, hosts)
 
         # no keyspace, child policy is called
@@ -906,7 +1006,9 @@ def test_statement_keyspace(self):
         query = Statement(routing_key=routing_key, keyspace=statement_keyspace)
         qplan = list(policy.make_query_plan(working_keyspace, query))
         assert replicas + hosts[:2] == qplan
-        cluster.metadata.get_replicas.assert_called_with(statement_keyspace, routing_key)
+        # get_replicas may not be called here due to cache hit from the
+        # previous query with the same (statement_keyspace, routing_key) pair.
+        # The important assertion is that the plan result is correct above.
 
     def test_shuffles_if_given_keyspace_and_routing_key(self):
         """
@@ -955,6 +1057,9 @@ def _prepare_cluster_with_vnodes(self):
         cluster.metadata.all_hosts.return_value = hosts
         cluster.metadata.get_replicas.return_value = hosts[2:]
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
         return cluster
 
     def _prepare_cluster_with_tablets(self):
@@ -967,14 +1072,22 @@ def _prepare_cluster_with_tablets(self):
         cluster.metadata.all_hosts.return_value = hosts
         cluster.metadata.get_replicas.return_value = hosts[2:]
         cluster.metadata._tablets.get_tablet_for_key.return_value = Tablet(replicas=[(h.host_id, 0) for h in hosts[2:]])
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
         return cluster
 
     @patch('cassandra.policies.shuffle')
     def _assert_shuffle(self, patched_shuffle, cluster, keyspace, routing_key):
         hosts = cluster.metadata.all_hosts()
-        replicas = cluster.metadata.get_replicas()
+        # Configure get_host_by_host_id to return hosts from the list
+        host_map = {h.host_id: h for h in hosts}
+        cluster.metadata.get_host_by_host_id.side_effect = lambda hid: host_map.get(hid)
+
+        replicas = list(cluster.metadata.get_replicas())
         child_policy = Mock()
         child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
         child_policy.distance.return_value = HostDistance.LOCAL
 
         policy = TokenAwarePolicy(child_policy, shuffle_replicas=True)
@@ -984,6 +1097,7 @@ def _assert_shuffle(self, patched_shuffle, cluster, keyspace, routing_key):
 
         cluster.metadata.get_replicas.reset_mock()
         child_policy.make_query_plan.reset_mock()
+        child_policy.make_query_plan_with_exclusion.reset_mock()
         query = Statement(routing_key=routing_key)
         qplan = list(policy.make_query_plan(keyspace, query))
         if keyspace is None or routing_key is None:
@@ -994,13 +1108,295 @@ def _assert_shuffle(self, patched_shuffle, cluster, keyspace, routing_key):
         else:
             assert set(replicas) == set(qplan[:2])
             assert hosts[:2] == qplan[2:]
+
             if is_tablets:
+                # Tablet path: make_query_plan called once for replica ordering,
+                # make_query_plan_with_exclusion called once for remaining hosts
                 child_policy.make_query_plan.assert_called_with(keyspace, query)
-                assert child_policy.make_query_plan.call_count == 2
+                child_policy.make_query_plan_with_exclusion.assert_called()
+            elif child_policy.make_query_plan_with_exclusion.called:
+                # Non-tablet path with replicas: exclusion set should contain
+                # the replicas that were already yielded
+                exc_call = child_policy.make_query_plan_with_exclusion.call_args
+                excluded_hosts = exc_call[0][2] if len(exc_call[0]) > 2 else exc_call[1].get('excluded', set())
+                assert set(replicas).issubset(excluded_hosts), \
+                    'Exclusion set should contain the yielded replicas, ' \
+                    'got %s, expected superset of %s' % (excluded_hosts, set(replicas))
             else:
                 child_policy.make_query_plan.assert_called_once_with(keyspace, query)
             assert patched_shuffle.call_count == 1
 
+    # --- Replica cache tests ---
+
+    def _make_cache_cluster(self):
+        """Create a mock cluster suitable for cache tests."""
+        hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for host in hosts:
+            host.set_up()
+        cluster = Mock(spec=Cluster)
+        cluster.metadata = Mock(spec=Metadata)
+        cluster.metadata._tablets = Mock(spec=Tablets)
+        cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
+        cluster.metadata.token_map.get_replicas.return_value = hosts[2:]
+        return cluster, hosts
+
+    def test_cache_hit(self):
+        """Same (keyspace, routing_key) should only call get_replicas once."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'key1', keyspace='ks')
+        list(policy.make_query_plan(None, query))
+        list(policy.make_query_plan(None, query))
+
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+
+    def test_cache_miss_different_key(self):
+        """Different routing_key should cause separate get_replicas calls."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        q1 = Statement(routing_key=b'key1', keyspace='ks')
+        q2 = Statement(routing_key=b'key2', keyspace='ks')
+        list(policy.make_query_plan(None, q1))
+        list(policy.make_query_plan(None, q2))
+
+        assert cluster.metadata.token_map.get_replicas.call_count == 2
+
+    def test_cache_miss_different_keyspace(self):
+        """Different keyspace with same routing_key should miss cache."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        q1 = Statement(routing_key=b'key1', keyspace='ks1')
+        q2 = Statement(routing_key=b'key1', keyspace='ks2')
+        list(policy.make_query_plan(None, q1))
+        list(policy.make_query_plan(None, q2))
+
+        assert cluster.metadata.token_map.get_replicas.call_count == 2
+
+    def test_cache_invalidation_on_topology_change(self):
+        """Cache should be invalidated when token_map object changes."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'key1', keyspace='ks')
+        list(policy.make_query_plan(None, query))
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+
+        # Simulate topology change: replace token_map with a new mock object
+        new_token_map = Mock()
+        new_token_map.token_class.from_key.side_effect = lambda key: key
+        new_token_map.get_replicas.return_value = hosts[2:]
+        cluster.metadata.token_map = new_token_map
+
+        list(policy.make_query_plan(None, query))
+        # The old token_map still has 1 call; new one should have 1 call
+        assert new_token_map.get_replicas.call_count == 1
+
+    def test_cache_eviction(self):
+        """Oldest entries should be evicted when cache exceeds size."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False, cache_replicas_size=2)
+        policy.populate(cluster, hosts)
+
+        # Fill cache with 3 entries; size=2 so first should be evicted
+        for i in range(3):
+            q = Statement(routing_key=f'key{i}'.encode(), keyspace='ks')
+            list(policy.make_query_plan(None, q))
+
+        assert cluster.metadata.token_map.get_replicas.call_count == 3
+
+        # key2 (most recent) should be cached
+        cluster.metadata.token_map.get_replicas.reset_mock()
+        q = Statement(routing_key=b'key2', keyspace='ks')
+        list(policy.make_query_plan(None, q))
+        assert cluster.metadata.token_map.get_replicas.call_count == 0
+
+        # key0 (evicted) should miss
+        q = Statement(routing_key=b'key0', keyspace='ks')
+        list(policy.make_query_plan(None, q))
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+
+    def test_cache_disabled(self):
+        """cache_replicas_size=0 should bypass caching entirely."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False, cache_replicas_size=0)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'key1', keyspace='ks')
+        list(policy.make_query_plan(None, query))
+        list(policy.make_query_plan(None, query))
+        list(policy.make_query_plan(None, query))
+
+        # Every call should reach get_replicas
+        assert cluster.metadata.token_map.get_replicas.call_count == 3
+
+    def test_tablet_path_not_cached(self):
+        """Tablet path should bypass the cache entirely."""
+        hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for host in hosts:
+            host.set_up()
+
+        cluster = Mock(spec=Cluster)
+        cluster.metadata = Mock(spec=Metadata)
+        cluster.metadata._tablets = Mock(spec=Tablets)
+        cluster.metadata._tablets.get_tablet_for_key.return_value = Tablet(replicas=[(h.host_id, 0) for h in hosts[2:]])
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
+        cluster.metadata.token_map.get_replicas.return_value = hosts[2:]
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'key1', keyspace='ks')
+        list(policy.make_query_plan(None, query))
+        list(policy.make_query_plan(None, query))
+
+        # token_map.get_replicas should NOT be called (tablet path used)
+        assert cluster.metadata.token_map.get_replicas.call_count == 0
+        # Cache should remain empty (tablet results are not cached)
+        assert len(policy._replica_cache) == 0
+
+    # --- LWT determinism tests ---
+
+    def _make_lwt_query(self, routing_key, keyspace='ks'):
+        """Create a Statement that reports is_lwt()=True."""
+        query = Statement(routing_key=routing_key, keyspace=keyspace)
+        query.is_lwt = lambda: True
+        return query
+
+    @patch('cassandra.policies.shuffle')
+    def test_lwt_no_shuffle(self, patched_shuffle):
+        """LWT queries should yield replicas in deterministic order."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=True)
+        policy.populate(cluster, hosts)
+
+        query = self._make_lwt_query(routing_key=b'key1')
+
+        plans = [list(policy.make_query_plan(None, query)) for _ in range(5)]
+
+        # All plans should be identical (deterministic)
+        for plan in plans[1:]:
+            assert plan == plans[0]
+
+        # shuffle should never have been called
+        assert patched_shuffle.call_count == 0
+
+    @patch('cassandra.policies.shuffle')
+    def test_lwt_replicas_not_copied(self, patched_shuffle):
+        """LWT path should not copy the replicas list (no list() call)."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=True)
+        policy.populate(cluster, hosts)
+
+        query = self._make_lwt_query(routing_key=b'key1')
+        list(policy.make_query_plan(None, query))
+
+        # shuffle was never called, which means list() was also not called
+        assert patched_shuffle.call_count == 0
+
+    @patch('cassandra.policies.shuffle')
+    def test_non_lwt_shuffled(self, patched_shuffle):
+        """Non-LWT queries with shuffle_replicas=True should shuffle."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=True)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'key1', keyspace='ks')
+        list(policy.make_query_plan(None, query))
+
+        assert patched_shuffle.call_count == 1
+
+    @patch('cassandra.policies.shuffle')
+    def test_lwt_with_cache_deterministic(self, patched_shuffle):
+        """LWT + cache should produce identical plans on repeated calls."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=True)
+        policy.populate(cluster, hosts)
+
+        query = self._make_lwt_query(routing_key=b'key1')
+
+        plan1 = list(policy.make_query_plan(None, query))
+        plan2 = list(policy.make_query_plan(None, query))
+
+        assert plan1 == plan2
+        assert patched_shuffle.call_count == 0
+        # Should have been a cache hit on the second call
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+
 
     @patch('cassandra.policies.shuffle')
     def test_no_shuffle_for_serial_consistency(self, patched_shuffle):
@@ -1019,6 +1415,8 @@ def test_no_shuffle_for_serial_consistency(self, patched_shuffle):
                 hosts = cluster.metadata.all_hosts()
                 child_policy = Mock()
                 child_policy.make_query_plan.return_value = hosts
+                child_policy.make_query_plan_with_exclusion.side_effect = \
+                    lambda k, q, e: [h for h in hosts if h not in e]
                 child_policy.distance.return_value = HostDistance.LOCAL
 
                 policy = TokenAwarePolicy(child_policy, shuffle_replicas=True)
@@ -1746,8 +2144,11 @@ def get_replicas(keyspace, packed_key):
         cluster.metadata.get_replicas.side_effect = get_replicas
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
+        cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
 
-        child_policy = TokenAwarePolicy(RoundRobinPolicy())
+        child_policy = TokenAwarePolicy(RoundRobinPolicy(), shuffle_replicas=False)
 
         hfp = HostFilterPolicy(
             child_policy=child_policy,

From 56c4aa363bbee9ab36d03952a1a9d665fadfc0a1 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Mon, 20 Apr 2026 00:00:10 +0300
Subject: [PATCH 08/10] perf: optimize TokenAwarePolicy cache -- check before
 hash, add keyspace-aware invalidation

- Move LRU cache lookup before token_class.from_key() so cache hits
  skip the murmur3 hash computation and Token object allocation entirely.
- Add keyspace-aware cache invalidation: track the per-keyspace replica
  map object identity so ALTER KEYSPACE / replication changes are detected
  even when the TokenMap object itself is reused (in-place rebuild).
- Remove unused 'token' from cache entries (was never read after storage).
- Add test_cache_invalidation_on_keyspace_replication_change.

TODO: The tablet path still does two full child-policy traversals per
query. Metadata.get_host_by_host_id() is O(1) and could resolve tablet
replicas in O(rf) instead. Deferred to minimize behavioral change.
---
 cassandra/policies.py       | 79 +++++++++++++++++++++++--------------
 tests/unit/test_policies.py | 28 +++++++++++++
 2 files changed, 77 insertions(+), 30 deletions(-)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index d21c8bb190..4d59992868 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -430,7 +430,7 @@ def __init__(self, local_dc, local_rack, used_hosts_per_remote_dc=0):
         self._live_hosts = {}
         self._dc_live_hosts = {}
         self._remote_hosts = {}
-        self._non_local_rack_hosts = []
+        self._non_local_rack_hosts = ()
         self._endpoints = []
         self._position = 0
         LoadBalancingPolicy.__init__(self)
@@ -455,9 +455,9 @@ def _refresh_remote_hosts(self):
 
     def _refresh_non_local_rack_hosts(self):
         local_live = self._dc_live_hosts.get(self.local_dc, ())
-        self._non_local_rack_hosts = [
+        self._non_local_rack_hosts = tuple(
             h for h in local_live if self._rack(h) != self.local_rack
-        ]
+        )
 
     def populate(self, cluster, hosts):
         for (dc, rack), rack_hosts in groupby(hosts, lambda host: (self._dc(host), self._rack(host))):
@@ -665,11 +665,16 @@ def check_supported(self):
     def distance(self, *args, **kwargs):
         return self._child_policy.distance(*args, **kwargs)
 
-    def _get_cached_replicas(self, keyspace, routing_key_bytes, token_map):
+    def _get_cached_replicas(self, keyspace, table, routing_key_bytes, token_map):
         """
-        Return cached (token, replicas) for the given keyspace and routing key,
-        or None on cache miss.  The cache is invalidated whenever the token_map
-        object identity changes (i.e. after a topology rebuild).
+        Return cached replicas for the given keyspace, table, and routing key,
+        or None on cache miss.  The cache is invalidated when:
+        - the token_map object identity changes (full topology rebuild), or
+        - the keyspace's replica map has been rebuilt in-place (e.g. ALTER
+          KEYSPACE), detected via object identity of the per-keyspace map.
+
+        The table is part of the cache key so that tablet-backed and
+        non-tablet tables in the same keyspace don't collide.
         """
         if not self._cache_replicas_size:
             return None
@@ -678,17 +683,27 @@ def _get_cached_replicas(self, keyspace, routing_key_bytes, token_map):
                 # Token map was rebuilt -- entire cache is stale.
                 self._replica_cache = OrderedDict()
                 self._replica_cache_token_map_ref = token_map
-            cache_key = (keyspace, routing_key_bytes)
+            cache_key = (keyspace, table, routing_key_bytes)
             entry = self._replica_cache.get(cache_key)
             if entry is not None:
+                replicas, ks_map_ref = entry
+                # Validate the keyspace replica map hasn't been rebuilt
+                # in-place (e.g. ALTER KEYSPACE changes replication).
+                current_ks_map = token_map.tokens_to_hosts_by_ks.get(keyspace)
+                if ks_map_ref is not current_ks_map:
+                    del self._replica_cache[cache_key]
+                    return None
                 # Promote to most-recently-used.
                 self._replica_cache.move_to_end(cache_key)
-            return entry
+                return replicas
+            return None
 
-    def _put_cached_replicas(self, keyspace, routing_key_bytes, token, replicas, token_map):
+    def _put_cached_replicas(self, keyspace, table, routing_key_bytes, replicas, token_map):
         """
-        Store (token, replicas) in the LRU cache, evicting the oldest
-        entry if the cache exceeds its configured size.
+        Store replicas in the LRU cache, evicting the oldest entry if
+        the cache exceeds its configured size.  The keyspace's current
+        replica-map object reference is stored alongside so that in-place
+        rebuilds (ALTER KEYSPACE) are detected on lookup via identity check.
         """
         if not self._cache_replicas_size:
             return
@@ -696,8 +711,9 @@ def _put_cached_replicas(self, keyspace, routing_key_bytes, token, replicas, tok
             if token_map is not self._replica_cache_token_map_ref:
                 self._replica_cache = OrderedDict()
                 self._replica_cache_token_map_ref = token_map
-            cache_key = (keyspace, routing_key_bytes)
-            self._replica_cache[cache_key] = (token, replicas)
+            cache_key = (keyspace, table, routing_key_bytes)
+            ks_map_ref = token_map.tokens_to_hosts_by_ks.get(keyspace)
+            self._replica_cache[cache_key] = (replicas, ks_map_ref)
             self._replica_cache.move_to_end(cache_key)
             if len(self._replica_cache) > self._cache_replicas_size:
                 self._replica_cache.popitem(last=False)
@@ -715,19 +731,21 @@ def make_query_plan(self, working_keyspace=None, query=None):
         replicas = []
         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
-                )
-
-                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]
+                # Check the LRU cache first -- avoids the hash (from_key)
+                # and token-map lookup on repeated routing keys.
+                cached = self._get_cached_replicas(keyspace, query.table, query.routing_key, token_map)
+                if cached is not None:
+                    replicas = cached
                 else:
-                    cached = self._get_cached_replicas(keyspace, query.routing_key, token_map)
-                    if cached is not None:
-                        token, replicas = cached
+                    token = token_map.token_class.from_key(query.routing_key)
+                    tablet = cluster_metadata._tablets.get_tablet_for_key(
+                        keyspace, query.table, token
+                    )
+
+                    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 = token_map.get_replicas(keyspace, token)
@@ -738,7 +756,7 @@ def make_query_plan(self, working_keyspace=None, query=None):
                             )
                             replicas = cluster_metadata.get_replicas(keyspace, query.routing_key)
                         self._put_cached_replicas(
-                            keyspace, query.routing_key, token, replicas, token_map
+                            keyspace, query.table, query.routing_key, replicas, token_map
                         )
             except Exception:
                 log.debug(
@@ -1668,9 +1686,10 @@ def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excl
         child = self._child_policy
         if target_host and target_host.is_up and target_host not in excluded:
             yield target_host
-            for h in child.make_query_plan_with_exclusion(keyspace, query, excluded):
-                if h != target_host:
-                    yield h
+            # Include target_host in the exclusion set so the child policy
+            # can skip it early rather than yielding it for us to filter.
+            child_excluded = excluded | {target_host} if excluded else {target_host}
+            yield from child.make_query_plan_with_exclusion(keyspace, query, child_excluded)
         else:
             yield from child.make_query_plan_with_exclusion(keyspace, query, excluded)
 
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
index 038a2dd966..ad1d8e283c 100644
--- a/tests/unit/test_policies.py
+++ b/tests/unit/test_policies.py
@@ -1140,6 +1140,8 @@ def _make_cache_cluster(self):
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.return_value = hosts[2:]
+        # Provide a real dict for keyspace-aware cache invalidation checks.
+        cluster.metadata.token_map.tokens_to_hosts_by_ks = {'ks': {}, 'ks1': {}, 'ks2': {}}
         return cluster, hosts
 
     def test_cache_hit(self):
@@ -1218,6 +1220,7 @@ def test_cache_invalidation_on_topology_change(self):
         new_token_map = Mock()
         new_token_map.token_class.from_key.side_effect = lambda key: key
         new_token_map.get_replicas.return_value = hosts[2:]
+        new_token_map.tokens_to_hosts_by_ks = {'ks': {}}
         cluster.metadata.token_map = new_token_map
 
         list(policy.make_query_plan(None, query))
@@ -1287,6 +1290,7 @@ def test_tablet_path_not_cached(self):
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.return_value = hosts[2:]
+        cluster.metadata.token_map.tokens_to_hosts_by_ks = {'ks': {}}
 
         child_policy = Mock()
         child_policy.make_query_plan.return_value = hosts
@@ -1305,6 +1309,30 @@ def test_tablet_path_not_cached(self):
         # Cache should remain empty (tablet results are not cached)
         assert len(policy._replica_cache) == 0
 
+    def test_cache_invalidation_on_keyspace_replication_change(self):
+        """Cache should detect in-place keyspace replica map rebuild (e.g. ALTER KEYSPACE)."""
+        cluster, hosts = self._make_cache_cluster()
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'key1', keyspace='ks')
+        list(policy.make_query_plan(None, query))
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+
+        # Simulate ALTER KEYSPACE: same token_map object, but the per-keyspace
+        # replica map is replaced in-place (new dict object for that keyspace).
+        cluster.metadata.token_map.tokens_to_hosts_by_ks['ks'] = {'new': 'map'}
+
+        list(policy.make_query_plan(None, query))
+        # Should have re-fetched replicas because the ks map id changed.
+        assert cluster.metadata.token_map.get_replicas.call_count == 2
+
     # --- LWT determinism tests ---
 
     def _make_lwt_query(self, routing_key, keyspace='ks'):

From 5ba167ac2a87c575f3430d834357df934bd00b5f Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Mon, 20 Apr 2026 09:07:49 +0300
Subject: [PATCH 09/10] fix: restore runtime used_hosts_per_remote_dc changes
 and late-bind remote hosts

- Convert used_hosts_per_remote_dc to a property in DCAwareRoundRobinPolicy
  and RackAwareRoundRobinPolicy so that runtime changes immediately refresh
  the cached _remote_hosts dict (restores origin/master behavior).
- Defer reading self._remote_hosts until the remote iteration phase in
  make_query_plan() and make_query_plan_with_exclusion() so topology
  changes during local iteration are visible (restores origin/master
  late-binding behavior).
- Add tests for runtime used_hosts_per_remote_dc changes and for
  modification-during-generation on the exclusion path.
- Fix a real race: the used_hosts_per_remote_dc setters mutated
  _used_hosts_per_remote_dc and rebuilt _remote_hosts without holding
  _hosts_lock, unlike on_up/on_down which mutate the same cached state
  under the lock. A concurrent runtime reassignment racing with a topology
  event could corrupt the _dc_live_hosts iteration inside
  _refresh_remote_hosts()/_refresh_non_local_rack_hosts() or observe a torn
  view. Both setters now acquire _hosts_lock (a plain, non-reentrant Lock;
  neither _refresh_remote_hosts() nor _refresh_non_local_rack_hosts()
  reacquire it, so there is no deadlock risk) and skip the refresh
  entirely when the assigned value is unchanged.
- Convert local_dc (DCAwareRoundRobinPolicy) and local_dc/local_rack
  (RackAwareRoundRobinPolicy) into properties following the same pattern:
  before this PR, distance()/make_query_plan() always recomputed from the
  live local_dc/local_rack, so reassigning them was always safe. The new
  _remote_hosts/_non_local_rack_hosts caches introduced by this PR are
  computed relative to whatever local_dc/local_rack was at the time of the
  last _refresh_*() call, so a direct reassignment (or the internal
  auto-detect reassignment in on_up) could leave those caches stale until
  an unrelated topology event happened to refresh them. The property
  setters now refresh the right cache(s) under the lock whenever the value
  actually changes. The constructors set the private _local_dc/_local_rack
  attributes directly (bypassing the setters) since _hosts_lock and
  _dc_live_hosts don't exist yet at that point in construction; the
  auto-detect path in on_up now goes through the local_dc setter instead of
  a plain attribute assignment, and the surrounding on_up logic was
  adjusted to avoid a redundant second refresh for that case. Added tests
  for runtime local_dc/local_rack reassignment analogous to the existing
  used_hosts_per_remote_dc test.
---
 cassandra/policies.py       | 126 ++++++++++++++++++-----
 tests/unit/test_policies.py | 195 ++++++++++++++++++++++++++++++++++++
 2 files changed, 295 insertions(+), 26 deletions(-)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index 4d59992868..763dc4be2a 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -248,8 +248,7 @@ class DCAwareRoundRobinPolicy(LoadBalancingPolicy):
     datacenters as a last resort.
     """
 
-    local_dc = None
-    used_hosts_per_remote_dc = 0
+    _local_dc = None
 
     def __init__(self, local_dc='', used_hosts_per_remote_dc=0):
         """
@@ -267,13 +266,41 @@ def __init__(self, local_dc='', used_hosts_per_remote_dc=0):
         rest will be considered :attr:`~.HostDistance.IGNORED`.
         By default, all remote hosts are ignored.
         """
-        self.local_dc = local_dc
-        self.used_hosts_per_remote_dc = used_hosts_per_remote_dc
+        # Set the private attribute directly here (bypassing the local_dc
+        # property setter below): the setter refreshes the _remote_hosts
+        # cache under _hosts_lock, but neither _dc_live_hosts nor
+        # _hosts_lock exist yet at this point in construction.
+        self._local_dc = local_dc
         self._dc_live_hosts = {}
         self._remote_hosts = {}
+        self._used_hosts_per_remote_dc = used_hosts_per_remote_dc
         self._position = 0
         LoadBalancingPolicy.__init__(self)
 
+    @property
+    def local_dc(self):
+        return self._local_dc
+
+    @local_dc.setter
+    def local_dc(self, value):
+        if value == self._local_dc:
+            return
+        with self._hosts_lock:
+            self._local_dc = value
+            self._refresh_remote_hosts()
+
+    @property
+    def used_hosts_per_remote_dc(self):
+        return self._used_hosts_per_remote_dc
+
+    @used_hosts_per_remote_dc.setter
+    def used_hosts_per_remote_dc(self, value):
+        if value == self._used_hosts_per_remote_dc:
+            return
+        with self._hosts_lock:
+            self._used_hosts_per_remote_dc = value
+            self._refresh_remote_hosts()
+
     def _dc(self, host):
         return host.datacenter or self.local_dc
 
@@ -319,8 +346,9 @@ def make_query_plan(self, working_keyspace=None, query=None):
             for i in range(length):
                 yield local_live[(pos + i) % length]
 
-        remote_hosts = self._remote_hosts
-        for host in remote_hosts:
+        # Read _remote_hosts late so topology changes during local
+        # iteration are visible.
+        for host in self._remote_hosts:
             yield host
 
     def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excluded=()):
@@ -331,13 +359,14 @@ def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excl
 
         local_live = self._dc_live_hosts.get(self.local_dc, ())
         length = len(local_live)
-        remote_hosts = self._remote_hosts
         if not excluded:
             if length:
                 pos %= length
                 for i in range(length):
                     yield local_live[(pos + i) % length]
-            for host in remote_hosts:
+            # Read _remote_hosts late so topology changes during local
+            # iteration are visible.
+            for host in self._remote_hosts:
                 yield host
             return
 
@@ -352,7 +381,7 @@ def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excl
                     continue
                 yield host
 
-        for host in remote_hosts:
+        for host in self._remote_hosts:
             if host in excluded:
                 continue
             yield host
@@ -360,14 +389,18 @@ def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excl
     def on_up(self, host):
         # not worrying about threads because this will happen during
         # control connection startup/refresh
-        refresh_remote = False
         if not self.local_dc and host.datacenter:
+            # Assigning through the local_dc property triggers a
+            # _refresh_remote_hosts() under _hosts_lock. That refresh runs
+            # against this host's own (now-local) dc, which never appears in
+            # _remote_hosts regardless of whether _dc_live_hosts has been
+            # updated with this host yet, so no separate refresh is needed
+            # here for this branch.
             self.local_dc = host.datacenter
             log.info("Using datacenter '%s' for DCAwareRoundRobinPolicy (via host '%s'); "
                         "if incorrect, please specify a local_dc to the constructor, "
                         "or limit contact points to local cluster nodes" %
                         (self.local_dc, host.endpoint))
-            refresh_remote = True
 
         dc = self._dc(host)
         with self._hosts_lock:
@@ -375,10 +408,7 @@ def on_up(self, host):
             if host not in current_hosts:
                 self._dc_live_hosts[dc] = current_hosts + (host, )
                 if dc != self.local_dc:
-                    refresh_remote = True
-
-            if refresh_remote:
-                self._refresh_remote_hosts()
+                    self._refresh_remote_hosts()
 
     def on_down(self, host):
         dc = self._dc(host)
@@ -407,9 +437,8 @@ class RackAwareRoundRobinPolicy(LoadBalancingPolicy):
     different rack, before hosts in all other datercentres
     """
 
-    local_dc = None
-    local_rack = None
-    used_hosts_per_remote_dc = 0
+    _local_dc = None
+    _local_rack = None
 
     def __init__(self, local_dc, local_rack, used_hosts_per_remote_dc=0):
         """
@@ -424,17 +453,58 @@ def __init__(self, local_dc, local_rack, used_hosts_per_remote_dc=0):
         rest will be considered :attr:`~.HostDistance.IGNORED`.
         By default, all remote hosts are ignored.
         """
-        self.local_rack = local_rack
-        self.local_dc = local_dc
-        self.used_hosts_per_remote_dc = used_hosts_per_remote_dc
+        # Set the private attributes directly here (bypassing the local_dc
+        # and local_rack property setters below): those setters refresh
+        # cached state under _hosts_lock, but neither _dc_live_hosts nor
+        # _hosts_lock exist yet at this point in construction.
+        self._local_rack = local_rack
+        self._local_dc = local_dc
         self._live_hosts = {}
         self._dc_live_hosts = {}
         self._remote_hosts = {}
         self._non_local_rack_hosts = ()
+        self._used_hosts_per_remote_dc = used_hosts_per_remote_dc
         self._endpoints = []
         self._position = 0
         LoadBalancingPolicy.__init__(self)
 
+    @property
+    def local_dc(self):
+        return self._local_dc
+
+    @local_dc.setter
+    def local_dc(self, value):
+        if value == self._local_dc:
+            return
+        with self._hosts_lock:
+            self._local_dc = value
+            self._refresh_remote_hosts()
+            self._refresh_non_local_rack_hosts()
+
+    @property
+    def local_rack(self):
+        return self._local_rack
+
+    @local_rack.setter
+    def local_rack(self, value):
+        if value == self._local_rack:
+            return
+        with self._hosts_lock:
+            self._local_rack = value
+            self._refresh_non_local_rack_hosts()
+
+    @property
+    def used_hosts_per_remote_dc(self):
+        return self._used_hosts_per_remote_dc
+
+    @used_hosts_per_remote_dc.setter
+    def used_hosts_per_remote_dc(self, value):
+        if value == self._used_hosts_per_remote_dc:
+            return
+        with self._hosts_lock:
+            self._used_hosts_per_remote_dc = value
+            self._refresh_remote_hosts()
+
     def _rack(self, host):
         return host.rack or self.local_rack
 
@@ -499,8 +569,9 @@ def make_query_plan(self, working_keyspace=None, query=None):
             for i in range(length):
                 yield local_non_rack[(p + i) % length]
 
-        remote_hosts = self._remote_hosts
-        for host in remote_hosts:
+        # Read _remote_hosts late so topology changes during local
+        # iteration are visible.
+        for host in self._remote_hosts:
             yield host
 
     def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excluded=()):
@@ -509,7 +580,6 @@ def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excl
 
         local_rack_live = self._live_hosts.get((self.local_dc, self.local_rack), ())
         length = len(local_rack_live)
-        remote_hosts = self._remote_hosts
         if not excluded:
             if length:
                 p = pos % length
@@ -523,7 +593,9 @@ def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excl
                 for i in range(length):
                     yield local_non_rack[(p + i) % length]
 
-            for host in remote_hosts:
+            # Read _remote_hosts late so topology changes during local
+            # iteration are visible.
+            for host in self._remote_hosts:
                 yield host
             return
 
@@ -548,7 +620,9 @@ def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excl
                     continue
                 yield host
 
-        for host in remote_hosts:
+        # Read _remote_hosts late so topology changes during local
+        # iteration are visible.
+        for host in self._remote_hosts:
             if host in excluded:
                 continue
             yield host
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
index ad1d8e283c..dd1d7d19fd 100644
--- a/tests/unit/test_policies.py
+++ b/tests/unit/test_policies.py
@@ -624,6 +624,131 @@ def test_wrong_dc(self, policy_specialization, constructor_args):
         qplan = list(policy.make_query_plan())
         assert len(qplan) == 0
 
+    def test_runtime_used_hosts_per_remote_dc_change(self, policy_specialization, constructor_args):
+        """Changing used_hosts_per_remote_dc at runtime should take effect
+        immediately without needing a repopulate or topology event."""
+        hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for h in hosts[:2]:
+            h.set_location_info("dc1", "rack1")
+        for h in hosts[2:]:
+            h.set_location_info("dc2", "rack1")
+
+        policy = policy_specialization(*constructor_args, used_hosts_per_remote_dc=0)
+        policy.populate(Mock(), hosts)
+
+        # With 0, remotes are IGNORED and absent from query plan
+        for h in hosts[2:]:
+            assert policy.distance(h) == HostDistance.IGNORED
+        qplan = list(policy.make_query_plan())
+        assert set(qplan) == set(hosts[:2])
+
+        # Raise to 1 at runtime -- should take effect immediately
+        policy.used_hosts_per_remote_dc = 1
+        assert policy.distance(hosts[2]) == HostDistance.REMOTE or \
+               policy.distance(hosts[3]) == HostDistance.REMOTE
+        qplan = list(policy.make_query_plan())
+        assert len(qplan) == 3  # 2 local + 1 remote
+
+        # Raise to 2 -- both remotes visible
+        policy.used_hosts_per_remote_dc = 2
+        qplan = list(policy.make_query_plan())
+        assert set(qplan) == set(hosts)
+
+        # Drop back to 0 -- remotes disappear
+        policy.used_hosts_per_remote_dc = 0
+        for h in hosts[2:]:
+            assert policy.distance(h) == HostDistance.IGNORED
+        qplan = list(policy.make_query_plan())
+        assert set(qplan) == set(hosts[:2])
+
+    def test_runtime_local_dc_change(self, policy_specialization, constructor_args):
+        """Changing local_dc at runtime should take effect immediately --
+        distance()/make_query_plan() must not be computed against a stale
+        _remote_hosts (and, for RackAwareRoundRobinPolicy, _non_local_rack_hosts)
+        cache left over from before the reassignment."""
+        hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for h in hosts[:2]:
+            h.set_location_info("dc1", "rack1")
+        for h in hosts[2:]:
+            h.set_location_info("dc2", "rack1")
+
+        policy = policy_specialization(*constructor_args, used_hosts_per_remote_dc=2)
+        policy.populate(Mock(), hosts)
+
+        # dc1 starts out local: dc1 hosts are LOCAL(_RACK), dc2 hosts are REMOTE
+        assert policy.local_dc == "dc1"
+        for h in hosts[:2]:
+            assert policy.distance(h) in (HostDistance.LOCAL, HostDistance.LOCAL_RACK)
+        for h in hosts[2:]:
+            assert policy.distance(h) == HostDistance.REMOTE
+        qplan = list(policy.make_query_plan())
+        assert set(qplan[:2]) == set(hosts[:2])
+        assert set(qplan[2:]) == set(hosts[2:])
+
+        # Reassign local_dc at runtime -- should take effect immediately,
+        # without needing a repopulate or topology event to refresh caches.
+        policy.local_dc = "dc2"
+        assert policy.local_dc == "dc2"
+        for h in hosts[2:]:
+            assert policy.distance(h) in (HostDistance.LOCAL, HostDistance.LOCAL_RACK)
+        for h in hosts[:2]:
+            assert policy.distance(h) == HostDistance.REMOTE
+        qplan = list(policy.make_query_plan())
+        assert set(qplan[:2]) == set(hosts[2:])
+        assert set(qplan[2:]) == set(hosts[:2])
+
+        # Setting to the same value again should be a no-op (no crash, and
+        # caches remain correct).
+        policy.local_dc = "dc2"
+        assert policy.local_dc == "dc2"
+        for h in hosts[2:]:
+            assert policy.distance(h) in (HostDistance.LOCAL, HostDistance.LOCAL_RACK)
+
+    def test_modification_during_generation_exclusion(self, policy_specialization, constructor_args):
+        """Topology changes to remote hosts during local iteration should be
+        visible when the generator reaches the remote phase, for the exclusion
+        path as well as the normal path."""
+        hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for h in hosts[:2]:
+            h.set_location_info("dc1", "rack1")
+        for h in hosts[2:]:
+            h.set_location_info("dc2", "rack1")
+
+        policy = policy_specialization(*constructor_args, used_hosts_per_remote_dc=3)
+        policy.populate(Mock(), hosts)
+
+        new_host = Host(DefaultEndPoint(4), SimpleConvictionPolicy, host_id=uuid.uuid4())
+        new_host.set_location_info("dc2", "rack1")
+
+        # -- make_query_plan: add remote after starting local iteration --
+        plan = policy.make_query_plan()
+        next(plan)  # consume one local
+        policy.on_up(new_host)
+        remaining = list(plan)
+        # new_host should appear because _remote_hosts is read late
+        assert new_host in remaining
+
+        # -- make_query_plan: remove remote after starting local --
+        plan = policy.make_query_plan()
+        next(plan)
+        policy.on_down(new_host)
+        remaining = list(plan)
+        assert new_host not in remaining
+
+        # -- make_query_plan_with_exclusion: add remote after starting local --
+        plan = policy.make_query_plan_with_exclusion(excluded={hosts[0]})
+        next(plan)  # consume one local
+        policy.on_up(new_host)
+        remaining = list(plan)
+        assert new_host in remaining
+
+        # -- make_query_plan_with_exclusion: remove remote after starting local --
+        plan = policy.make_query_plan_with_exclusion(excluded={hosts[0]})
+        next(plan)
+        policy.on_down(new_host)
+        remaining = list(plan)
+        assert new_host not in remaining
+
 class DCAwareRoundRobinPolicyTest(unittest.TestCase):
 
     def test_default_dc(self):
@@ -657,6 +782,76 @@ def test_default_dc(self):
         policy.on_add(host_remote)
         assert policy.local_dc
 
+    def test_local_dc_auto_detect_refreshes_remote_hosts(self):
+        """The auto-detect path (on_up assigning self.local_dc from the
+        first host with a datacenter) goes through the local_dc property
+        setter and must correctly refresh _remote_hosts, without needing a
+        second, separate topology event."""
+        host_local = Host(DefaultEndPoint(1), SimpleConvictionPolicy, host_id=uuid.uuid4())
+        host_local.set_location_info("dc1", "rack1")
+        host_remote = Host(DefaultEndPoint(2), SimpleConvictionPolicy, host_id=uuid.uuid4())
+        host_remote.set_location_info("dc2", "rack1")
+
+        policy = DCAwareRoundRobinPolicy(used_hosts_per_remote_dc=1)
+        policy.populate(Mock(), [])
+        assert not policy.local_dc
+
+        # Auto-detects dc1 as local via on_up/on_add (host_local is added
+        # first, same as the "contact DC first" case in test_default_dc).
+        policy.on_add(host_local)
+        policy.on_add(host_remote)
+
+        assert policy.local_dc == "dc1"
+        assert policy.distance(host_local) == HostDistance.LOCAL
+        assert policy.distance(host_remote) == HostDistance.REMOTE
+        qplan = list(policy.make_query_plan())
+        assert qplan == [host_local, host_remote]
+
+
+class RackAwareRoundRobinPolicyTest(unittest.TestCase):
+
+    def test_runtime_local_rack_change(self):
+        """Changing local_rack at runtime should take effect immediately --
+        _non_local_rack_hosts must not remain stale from before the
+        reassignment."""
+        hosts = [Host(DefaultEndPoint(i), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for h in hosts[:2]:
+            h.set_location_info("dc1", "rack1")
+        for h in hosts[2:]:
+            h.set_location_info("dc1", "rack2")
+
+        policy = RackAwareRoundRobinPolicy("dc1", "rack1")
+        policy.populate(Mock(), hosts)
+
+        # rack1 starts out local: rack1 hosts are LOCAL_RACK, rack2 hosts are LOCAL
+        assert policy.local_rack == "rack1"
+        for h in hosts[:2]:
+            assert policy.distance(h) == HostDistance.LOCAL_RACK
+        for h in hosts[2:]:
+            assert policy.distance(h) == HostDistance.LOCAL
+        qplan = list(policy.make_query_plan())
+        assert set(qplan[:2]) == set(hosts[:2])
+        assert set(qplan[2:]) == set(hosts[2:])
+
+        # Reassign local_rack at runtime -- should take effect immediately.
+        policy.local_rack = "rack2"
+        assert policy.local_rack == "rack2"
+        for h in hosts[2:]:
+            assert policy.distance(h) == HostDistance.LOCAL_RACK
+        for h in hosts[:2]:
+            assert policy.distance(h) == HostDistance.LOCAL
+        qplan = list(policy.make_query_plan())
+        assert set(qplan[:2]) == set(hosts[2:])
+        assert set(qplan[2:]) == set(hosts[:2])
+
+        # Setting to the same value again should be a no-op (no crash, and
+        # caches remain correct).
+        policy.local_rack = "rack2"
+        assert policy.local_rack == "rack2"
+        for h in hosts[2:]:
+            assert policy.distance(h) == HostDistance.LOCAL_RACK
+
+
 class TokenAwarePolicyTest(unittest.TestCase):
 
     def test_wrap_round_robin(self):

From d1b83c0c0310d98c0598184e4579e65b72f999b4 Mon Sep 17 00:00:00 2001
From: Yaniv Michael Kaul <yaniv.kaul@scylladb.com>
Date: Mon, 20 Apr 2026 11:03:58 +0300
Subject: [PATCH 10/10] perf: skip tablet lookup when no tablets exist

Add Tablets.__bool__() so 'if tablets:' is False when no tablets are
registered, avoiding the get_tablet_for_key() method call + dict
lookup on every cache miss in the non-tablet path.

Restructure TokenAwarePolicy.make_query_plan() to nest the tablet
handling inside 'if cluster_metadata._tablets:' guard.

fix: tablet cache poisoning (review threads on PR #651)

The replica LRU cache was consulted before ever checking whether a
tablet exists for the (keyspace, table, routing_key). Tablets are
learned lazily from server responses, so the very first query for a
key -- issued before its tablet was known -- would fall back to
vnode-derived replicas and cache them. Once cached, every later query
for that key short-circuited on the cache and never rechecked tablets,
so learning the tablet afterwards (Tablets.add_tablet(), called from
cassandra/cluster.py) never took effect: the table was permanently
pinned to stale vnode routing.

Fix: gate both the cache lookup and the cache write on
Tablets.table_has_tablets(keyspace, table) -- a cheap O(1) dict check
that doesn't require hashing the routing key. Once a table has any
known tablets we always resolve via the tablet path (as before, this
was never cached anyway), so the cache is only ever trusted for tables
that have no tablet metadata at all, which is exactly the condition
under which it could have been populated. This keeps the fast
cache-hit path fully intact for the common case (non-tablet keyspaces
and tables tablets haven't touched) while closing the poisoning hole.

Also fix a related TOCTOU race in _put_cached_replicas: the
keyspace replica-map identity used to validate a cache entry was
captured *after* get_replicas() had already resolved the replicas
being cached, using a fresh lookup a few lines later. A concurrent
ALTER KEYSPACE rebuilding the per-keyspace map in that window could tag
stale replicas with the new map's identity, causing them to validate
as current forever afterward. Now the identity is snapshotted
immediately before get_replicas() is called and threaded through
explicitly, so a rebuild in that window is detected as invalidating
the entry instead.

Adds test_tablet_learned_after_vnode_cache_populated_is_used (fails
against the pre-fix code: a query for a key cached before its tablet
was learned kept using the stale vnode replicas after the tablet was
added) and
test_put_cached_replicas_captures_ks_map_identity_before_get_replicas
(fails against the pre-fix code: the stored identity was the
post-rebuild map, not the one the cached replicas were actually
resolved against).

Also fix Tablets.drop_tablets_by_host_id leaving an empty list behind
in _tablets once a table's last tablet is removed (review thread on
PR #651). bool(self._tablets) -- and thus the table_has_tablets()
fast-path gate added above -- stayed truthy forever for a table that
no longer has any tablets, since the dict key itself was never
removed, only emptied. Now the key is deleted once its list becomes
empty, so the no-tablet fast path this commit introduces actually
triggers once all tablets for a table are gone (e.g. after all of its
replica hosts are decommissioned). Adds regression tests in
tests/unit/test_tablets.py (DropTabletsByHostIdTest).
---
 cassandra/policies.py       |  91 +++++++++++++++++-----
 cassandra/tablets.py        |  16 ++++
 tests/unit/test_policies.py | 147 ++++++++++++++++++++++++++++++++++++
 tests/unit/test_tablets.py  |  62 +++++++++++++++
 4 files changed, 299 insertions(+), 17 deletions(-)

diff --git a/cassandra/policies.py b/cassandra/policies.py
index 763dc4be2a..4323e67fef 100644
--- a/cassandra/policies.py
+++ b/cassandra/policies.py
@@ -694,7 +694,10 @@ class TokenAwarePolicy(LoadBalancingPolicy):
     An LRU cache of size :attr:`cache_replicas_size` (default 1024) avoids
     repeated token-to-replica lookups for the same (keyspace, routing_key)
     pair.  Set to 0 to disable caching.  The cache is automatically
-    invalidated when the cluster topology changes.
+    invalidated when the cluster topology changes.  It only ever holds
+    vnode-derived replica sets; tables with known tablets always resolve
+    replicas via the tablet metadata directly, so tablet-aware routing is
+    never masked by a cache entry created before a tablet was learned.
     """
 
     _child_policy = None
@@ -749,6 +752,14 @@ def _get_cached_replicas(self, keyspace, table, routing_key_bytes, token_map):
 
         The table is part of the cache key so that tablet-backed and
         non-tablet tables in the same keyspace don't collide.
+
+        NOTE: this cache only ever holds vnode-derived replica sets (see
+        make_query_plan). Callers must not consult it for a table that
+        currently has any known tablets -- otherwise an entry cached before
+        a tablet for that key was learned would permanently shadow the
+        tablet-aware path (tablets are learned lazily from server responses
+        and nothing else invalidates a vnode-derived entry when that
+        happens).
         """
         if not self._cache_replicas_size:
             return None
@@ -772,12 +783,22 @@ def _get_cached_replicas(self, keyspace, table, routing_key_bytes, token_map):
                 return replicas
             return None
 
-    def _put_cached_replicas(self, keyspace, table, routing_key_bytes, replicas, token_map):
+    def _put_cached_replicas(self, keyspace, table, routing_key_bytes, replicas, token_map, ks_map_ref):
         """
         Store replicas in the LRU cache, evicting the oldest entry if
         the cache exceeds its configured size.  The keyspace's current
         replica-map object reference is stored alongside so that in-place
         rebuilds (ALTER KEYSPACE) are detected on lookup via identity check.
+
+        `ks_map_ref` must be the identity of `token_map.tokens_to_hosts_by_ks
+        [keyspace]` as it was captured by the caller BEFORE `replicas` was
+        resolved (i.e. before token_map.get_replicas() was called) -- not
+        re-fetched here. Re-fetching it at this point would open a window
+        between replica resolution and the identity snapshot in which a
+        concurrent ALTER KEYSPACE could rebuild the per-keyspace map; the
+        stale `replicas` computed against the OLD map would then be stored
+        alongside the identity of the NEW map and would incorrectly
+        validate as current forever afterward.
         """
         if not self._cache_replicas_size:
             return
@@ -786,7 +807,6 @@ def _put_cached_replicas(self, keyspace, table, routing_key_bytes, replicas, tok
                 self._replica_cache = OrderedDict()
                 self._replica_cache_token_map_ref = token_map
             cache_key = (keyspace, table, routing_key_bytes)
-            ks_map_ref = token_map.tokens_to_hosts_by_ks.get(keyspace)
             self._replica_cache[cache_key] = (replicas, ks_map_ref)
             self._replica_cache.move_to_end(cache_key)
             if len(self._replica_cache) > self._cache_replicas_size:
@@ -805,22 +825,53 @@ def make_query_plan(self, working_keyspace=None, query=None):
         replicas = []
         if token_map:
             try:
-                # Check the LRU cache first -- avoids the hash (from_key)
-                # and token-map lookup on repeated routing keys.
-                cached = self._get_cached_replicas(keyspace, query.table, query.routing_key, token_map)
+                tablets = cluster_metadata._tablets
+
+                # Only check tablets if any exist -- avoids the method
+                # call + dict lookup when tablets are not in use.
+                #
+                # IMPORTANT: this must be checked *before* consulting the
+                # replica cache, not after. The cache only ever stores
+                # vnode-derived replica sets (see below), and tablets are
+                # learned lazily from server responses -- so an entry may
+                # have been cached at a moment when no tablet existed yet
+                # for this (keyspace, table). Nothing else invalidates that
+                # entry once a tablet for it is subsequently learned, so a
+                # cache-first check would permanently pin the key to stale
+                # vnode replicas and defeat tablet-aware routing for it.
+                # Gating on table_has_tablets (an O(1) dict lookup, no
+                # routing-key hashing) is cheap enough to do on every call
+                # and keeps the fast cache path fully intact for the common
+                # case: non-tablet keyspaces, and tables tablets haven't
+                # touched at all.
+                table_has_tablets = bool(tablets) and tablets.table_has_tablets(keyspace, query.table)
+
+                cached = None
+                if not table_has_tablets:
+                    # Check the LRU cache first -- avoids the hash (from_key)
+                    # and token-map lookup on repeated routing keys.
+                    cached = self._get_cached_replicas(keyspace, query.table, query.routing_key, token_map)
+
                 if cached is not None:
                     replicas = cached
                 else:
                     token = token_map.token_class.from_key(query.routing_key)
-                    tablet = cluster_metadata._tablets.get_tablet_for_key(
-                        keyspace, query.table, token
-                    )
 
-                    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:
+                    if table_has_tablets:
+                        tablet = tablets.get_tablet_for_key(keyspace, query.table, token)
+                        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]
+
+                    if not replicas:
+                        # Snapshot the keyspace replica-map identity BEFORE
+                        # resolving replicas (not after) so a concurrent
+                        # ALTER KEYSPACE rebuild in this window is detected
+                        # as invalidating the entry we're about to cache
+                        # rather than accidentally validating it later. See
+                        # _put_cached_replicas.
+                        ks_map_ref = token_map.tokens_to_hosts_by_ks.get(keyspace)
                         try:
                             replicas = token_map.get_replicas(keyspace, token)
                         except Exception:
@@ -829,9 +880,15 @@ def make_query_plan(self, working_keyspace=None, query=None):
                                 "falling back to cluster metadata"
                             )
                             replicas = cluster_metadata.get_replicas(keyspace, query.routing_key)
-                        self._put_cached_replicas(
-                            keyspace, query.table, query.routing_key, replicas, token_map
-                        )
+                        # Never cache for a table that currently has known
+                        # tablets -- such entries could never be safely
+                        # read back (table_has_tablets would gate the cache
+                        # lookup out again above) and would only waste LRU
+                        # capacity that other keys need.
+                        if not table_has_tablets:
+                            self._put_cached_replicas(
+                                keyspace, query.table, query.routing_key, replicas, token_map, ks_map_ref
+                            )
             except Exception:
                 log.debug(
                     "Failed to resolve token or tablet for query plan, "
diff --git a/cassandra/tablets.py b/cassandra/tablets.py
index 96e61a50c2..616f64cd15 100644
--- a/cassandra/tablets.py
+++ b/cassandra/tablets.py
@@ -55,6 +55,9 @@ def __init__(self, tablets):
         self._tablets = tablets
         self._lock = Lock()
 
+    def __bool__(self):
+        return bool(self._tablets)
+
     def table_has_tablets(self, keyspace, table) -> bool:
         return bool(self._tablets.get((keyspace, table), []))
 
@@ -86,6 +89,7 @@ def drop_tablets_by_host_id(self, host_id: Optional[UUID]):
         if host_id is None:
             return
         with self._lock:
+            empty_keys = []
             for key, tablets in self._tablets.items():
                 to_be_deleted = []
                 for tablet_id, tablet in enumerate(tablets):
@@ -95,6 +99,18 @@ def drop_tablets_by_host_id(self, host_id: Optional[UUID]):
                 for tablet_id in reversed(to_be_deleted):
                     tablets.pop(tablet_id)
 
+                if not tablets:
+                    empty_keys.append(key)
+
+            # Delete keys whose tablet list is now empty (can't do this
+            # while iterating self._tablets.items() above -- that would
+            # mutate the dict's size mid-iteration). Leaving an empty list
+            # behind would keep bool(self._tablets) truthy forever for a
+            # table that no longer has any tablets, defeating the
+            # table_has_tablets()/bool(tablets) no-tablet fast path.
+            for key in empty_keys:
+                del self._tablets[key]
+
     def add_tablet(self, keyspace, table, tablet):
         with self._lock:
             tablets_for_table = self._tablets.setdefault((keyspace, table), [])
diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py
index dd1d7d19fd..4208ea9ef1 100644
--- a/tests/unit/test_policies.py
+++ b/tests/unit/test_policies.py
@@ -859,6 +859,7 @@ def test_wrap_round_robin(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
@@ -895,6 +896,7 @@ def test_wrap_dc_aware(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
@@ -947,6 +949,7 @@ def test_wrap_rack_aware(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(8)]
@@ -1153,6 +1156,7 @@ def test_statement_keyspace(self):
         replicas = hosts[2:]
         cluster.metadata.get_replicas.return_value = replicas
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
@@ -1252,6 +1256,7 @@ def _prepare_cluster_with_vnodes(self):
         cluster.metadata.all_hosts.return_value = hosts
         cluster.metadata.get_replicas.return_value = hosts[2:]
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
@@ -1267,6 +1272,7 @@ def _prepare_cluster_with_tablets(self):
         cluster.metadata.all_hosts.return_value = hosts
         cluster.metadata.get_replicas.return_value = hosts[2:]
         cluster.metadata._tablets.get_tablet_for_key.return_value = Tablet(replicas=[(h.host_id, 0) for h in hosts[2:]])
+        cluster.metadata._tablets.table_has_tablets.return_value = True
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
@@ -1332,6 +1338,7 @@ def _make_cache_cluster(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.return_value = hosts[2:]
@@ -1482,6 +1489,7 @@ def test_tablet_path_not_cached(self):
         cluster.metadata = Mock(spec=Metadata)
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = Tablet(replicas=[(h.host_id, 0) for h in hosts[2:]])
+        cluster.metadata._tablets.table_has_tablets.return_value = True
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.return_value = hosts[2:]
@@ -1528,6 +1536,144 @@ def test_cache_invalidation_on_keyspace_replication_change(self):
         # Should have re-fetched replicas because the ks map id changed.
         assert cluster.metadata.token_map.get_replicas.call_count == 2
 
+    def test_tablet_learned_after_vnode_cache_populated_is_used(self):
+        """
+        Regression test for tablet cache poisoning (PR #651 review threads):
+        a cache entry populated with vnode-derived replicas before a tablet
+        was known for that (keyspace, table, routing_key) must NOT
+        permanently shadow a tablet that is later learned for that exact
+        key. Reproduces:
+          (a) query a key before its tablet is known -> gets cached as
+              vnode-derived replicas (the only thing that could have
+              happened, since no tablet existed yet);
+          (b) a tablet covering that key is learned, e.g. via
+              cluster.metadata._tablets.add_tablet(...) as triggered from
+              cassandra/cluster.py on a server response;
+          (c) the same key is queried again -> the NEWLY learned
+              tablet-derived replicas must be used, not the stale cached
+              vnode ones.
+        """
+        hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)]
+        for host in hosts:
+            host.set_up()
+
+        cluster = Mock(spec=Cluster)
+        cluster.metadata = Mock(spec=Metadata)
+        # A real (not mocked) Tablets instance, starting empty -- exactly
+        # the state at cluster startup before any tablets have been
+        # learned from server responses.
+        real_tablets = Tablets({})
+        cluster.metadata._tablets = real_tablets
+
+        vnode_replicas = hosts[:2]
+        tablet_replicas = hosts[2:]
+
+        class _FakeToken:
+            """Stand-in for a real Token: get_tablet_for_key only needs `.value`."""
+            def __init__(self, value):
+                self.value = value
+
+        cluster.metadata.token_map = Mock()
+        cluster.metadata.token_map.token_class.from_key.return_value = _FakeToken(500)
+        cluster.metadata.token_map.get_replicas.return_value = vnode_replicas
+        cluster.metadata.token_map.tokens_to_hosts_by_ks = {'ks': {}}
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'routing-key-1', keyspace='ks', table='t')
+
+        # (a) No tablet known yet for ('ks', 't') -> falls back to vnode
+        # replicas, which get cached.
+        plan1 = list(policy.make_query_plan(None, query))
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+        assert plan1[:2] == vnode_replicas
+        assert len(policy._replica_cache) == 1
+
+        # (b) Learn a tablet for that same key, exactly as cluster.py does
+        # via cluster.metadata._tablets.add_tablet(keyspace, table, tablet)
+        # when handling a server response (see ResponseFuture in
+        # cassandra/cluster.py). The tablet's (first_token, last_token]
+        # range covers the token (500) our fake token_class resolves to.
+        tablet = Tablet(first_token=0, last_token=1000,
+                         replicas=[(h.host_id, 0) for h in tablet_replicas])
+        real_tablets.add_tablet('ks', 't', tablet)
+
+        # (c) Querying the same key again must now use the newly learned
+        # tablet replicas -- NOT the stale cached vnode replicas.
+        plan2 = list(policy.make_query_plan(None, query))
+        assert plan2[:2] == tablet_replicas
+        assert set(plan2[:2]) != set(vnode_replicas)
+        # The vnode fallback (and therefore the cache) must not have been
+        # consulted again -- the tablet path was used instead.
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+        # Once a table has known tablets we never write to the vnode
+        # cache for it again, so the cache should not have grown.
+        assert len(policy._replica_cache) == 1
+
+    def test_put_cached_replicas_captures_ks_map_identity_before_get_replicas(self):
+        """
+        Regression test for a TOCTOU race in _put_cached_replicas: the
+        keyspace replica-map identity used to later validate a cache entry
+        must be captured BEFORE get_replicas() resolves the replicas being
+        cached, not after. Otherwise a concurrent ALTER KEYSPACE rebuild
+        happening while get_replicas() is executing would tag stale
+        replicas with the identity of the NEW map, and they would
+        incorrectly validate as current forever afterward.
+
+        We simulate the race deterministically (no real thread timing) by
+        having get_replicas() itself perform the "concurrent" rebuild of
+        tokens_to_hosts_by_ks['ks'] as a side effect, i.e. exactly mid-way
+        through the window whose ordering item 2 is concerned with.
+        """
+        cluster, hosts = self._make_cache_cluster()
+
+        old_map = {}
+        cluster.metadata.token_map.tokens_to_hosts_by_ks['ks'] = old_map
+        new_map = {'rebuilt': True}
+
+        def get_replicas_side_effect(keyspace, token):
+            # Simulate a concurrent ALTER KEYSPACE rebuilding the
+            # per-keyspace map WHILE this call is resolving replicas for
+            # the current query.
+            cluster.metadata.token_map.tokens_to_hosts_by_ks['ks'] = new_map
+            return hosts[2:]
+
+        cluster.metadata.token_map.get_replicas.side_effect = get_replicas_side_effect
+
+        child_policy = Mock()
+        child_policy.make_query_plan.return_value = hosts
+        child_policy.make_query_plan_with_exclusion.side_effect = lambda k, q, e: [h for h in hosts if h not in e]
+        child_policy.distance.return_value = HostDistance.LOCAL
+
+        policy = TokenAwarePolicy(child_policy, shuffle_replicas=False)
+        policy.populate(cluster, hosts)
+
+        query = Statement(routing_key=b'key1', keyspace='ks')
+        list(policy.make_query_plan(None, query))
+
+        # The identity stored alongside the cached entry must be the map
+        # that existed BEFORE get_replicas() ran (old_map), since that is
+        # the map the returned replicas were actually resolved against --
+        # not new_map, which only came into existence afterward.
+        cache_key = ('ks', None, b'key1')
+        _, stored_ks_map_ref = policy._replica_cache[cache_key]
+        assert stored_ks_map_ref is old_map
+        assert stored_ks_map_ref is not new_map
+
+        # Because the map was rebuilt during resolution, the entry must be
+        # correctly treated as stale (a cache miss) on the very next
+        # lookup -- not incorrectly validated as still current.
+        cluster.metadata.token_map.get_replicas.reset_mock(side_effect=True)
+        cluster.metadata.token_map.get_replicas.return_value = hosts[:2]
+        list(policy.make_query_plan(None, query))
+        assert cluster.metadata.token_map.get_replicas.call_count == 1
+
     # --- LWT determinism tests ---
 
     def _make_lwt_query(self, routing_key, keyspace='ks'):
@@ -2367,6 +2513,7 @@ def get_replicas(keyspace, packed_key):
         cluster.metadata.get_replicas.side_effect = get_replicas
         cluster.metadata._tablets = Mock(spec=Tablets)
         cluster.metadata._tablets.get_tablet_for_key.return_value = None
+        cluster.metadata._tablets.table_has_tablets.return_value = False
         cluster.metadata.token_map = Mock()
         cluster.metadata.token_map.token_class.from_key.side_effect = lambda key: key
         cluster.metadata.token_map.get_replicas.side_effect = cluster.metadata.get_replicas
diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py
index 7a40e7de4d..a2c6b38049 100644
--- a/tests/unit/test_tablets.py
+++ b/tests/unit/test_tablets.py
@@ -1,4 +1,5 @@
 import unittest
+import uuid
 
 from cassandra.tablets import Tablets, Tablet
 
@@ -124,3 +125,64 @@ def __init__(self, v):
         # Token value 50 is not > first_token (100) of the tablet whose
         # last_token (200) is >= 50, so no match.
         self.assertIsNone(tablets.get_tablet_for_key("ks", "tb", Token(50)))
+
+
+class DropTabletsByHostIdTest(unittest.TestCase):
+    """
+    Regression tests: drop_tablets_by_host_id must delete a table's key
+    from _tablets entirely once its tablet list becomes empty, so that
+    table_has_tablets() (and any other truthiness/membership check on
+    _tablets, e.g. bool(tablets)) correctly reflects that no tablets are
+    left -- rather than leaving a stale empty list behind.
+    """
+
+    def test_drop_last_tablet_removes_table_key(self):
+        host_id = uuid.uuid4()
+        t1 = Tablet(0, 100, [(host_id, 0)])
+        tablets = Tablets({("ks", "tb"): [t1]})
+
+        assert tablets.table_has_tablets("ks", "tb") is True
+
+        tablets.drop_tablets_by_host_id(host_id)
+
+        # The no-tablet fast path must now be correctly signalled: the key
+        # must be gone from the dict (not just left as an empty list), and
+        # table_has_tablets must report False.
+        assert ("ks", "tb") not in tablets._tablets
+        assert tablets.table_has_tablets("ks", "tb") is False
+        assert bool(tablets) is False
+
+    def test_drop_some_tablets_keeps_remaining(self):
+        removed_host_id = uuid.uuid4()
+        remaining_host_id = uuid.uuid4()
+        t1 = Tablet(0, 100, [(removed_host_id, 0)])
+        t2 = Tablet(100, 200, [(remaining_host_id, 0)])
+        tablets = Tablets({("ks", "tb"): [t1, t2]})
+
+        tablets.drop_tablets_by_host_id(removed_host_id)
+
+        assert ("ks", "tb") in tablets._tablets
+        assert tablets._tablets[("ks", "tb")] == [t2]
+        assert tablets.table_has_tablets("ks", "tb") is True
+        assert bool(tablets) is True
+
+    def test_drop_last_tablet_for_one_table_keeps_other_tables(self):
+        host_id = uuid.uuid4()
+        other_host_id = uuid.uuid4()
+        t1 = Tablet(0, 100, [(host_id, 0)])
+        t2 = Tablet(0, 100, [(other_host_id, 0)])
+        tablets = Tablets({("ks", "tb1"): [t1], ("ks", "tb2"): [t2]})
+
+        tablets.drop_tablets_by_host_id(host_id)
+
+        assert ("ks", "tb1") not in tablets._tablets
+        assert tablets.table_has_tablets("ks", "tb1") is False
+        assert ("ks", "tb2") in tablets._tablets
+        assert tablets.table_has_tablets("ks", "tb2") is True
+        # Overall dict is still non-empty because tb2 still has tablets.
+        assert bool(tablets) is True
+
+    def test_drop_by_host_id_none_is_noop(self):
+        tablets = Tablets({("ks", "tb"): []})
+        tablets.drop_tablets_by_host_id(None)
+        assert ("ks", "tb") in tablets._tablets