From 748290c48919a521dae87b8915a28bcec258d8b0 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Fri, 3 Apr 2026 20:03:21 +0300 Subject: [PATCH 01/11] 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 f1bfefb41d..ab6aede369 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 546d38af19cde947c7e6b5f90ea6f5ab204f4a34 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Fri, 3 Apr 2026 20:04:08 +0300 Subject: [PATCH 02/11] 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 ab6aede369..296bafe37d 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 886b723b69af7c5f52d9a1b52d74c6c735cdcacb Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Fri, 3 Apr 2026 20:04:51 +0300 Subject: [PATCH 03/11] 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 296bafe37d..2e0cd63958 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 170924e743d24c42763a419f10eb0ea6efe738fc Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Fri, 3 Apr 2026 20:05:15 +0300 Subject: [PATCH 04/11] 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 2e0cd63958..85431a65b2 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 768e7bc0b5ef7f7a85fa9745d453fd21b91391e5 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Fri, 3 Apr 2026 20:06:28 +0300 Subject: [PATCH 05/11] 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/policies.py | 279 +++++++++++++++++++++++++++--------- tests/unit/test_policies.py | 58 ++++++++ 2 files changed, 268 insertions(+), 69 deletions(-) diff --git a/cassandra/policies.py b/cassandra/policies.py index 85431a65b2..3d6b4456ca 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -623,6 +623,14 @@ class TokenAwarePolicy(LoadBalancingPolicy): would contact it at all -- so the request reaches it directly instead of being forwarded there by another coordinator. The private ``_prefer_tablet_leader`` option turns that off. + + 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. It is only consulted for + non-tablet keyspaces; tablet replicas are always resolved fresh since + they carry leader information that can change independently of the + token map. """ _child_policy = None @@ -663,10 +671,25 @@ class TokenAwarePolicy(LoadBalancingPolicy): and the default may change before it is part of the public API. """ - def __init__(self, child_policy, shuffle_replicas=True, _prefer_tablet_leader=True): + def __init__(self, child_policy, shuffle_replicas=True, _prefer_tablet_leader=True, cache_replicas_size=1024): + super().__init__() self._child_policy = child_policy self.shuffle_replicas = shuffle_replicas self._prefer_tablet_leader = _prefer_tablet_leader + 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 @@ -684,13 +707,49 @@ 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 # Deferred: cassandra.metadata reaches this module through @@ -699,80 +758,162 @@ def make_query_plan(self, working_keyspace=None, query=None): # its functions. from cassandra.metadata import _ConsistencyMode + cluster_metadata = self._cluster_metadata + token_map = cluster_metadata.token_map replicas = [] leader_host = None - token = self._cluster_metadata.token_map.token_class.from_key(query.routing_key) - tablet = self._cluster_metadata._tablets.get_tablet_for_key(keyspace, query.table, token) - - 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] - - # The leader concept only exists for strongly-consistent keyspaces, - # which today means exactly the keyspaces whose consistency mode is - # GLOBAL: it is the only mode ScyllaDB implements so far, so LOCAL - # (reserved, unimplemented) and EVENTUAL both have no leader. This - # comparison has to widen once 'local' consistency exists. - # TABLETS_ROUTING_V2 assigns a tablet_version to *every* tablet table - # (eventually- and strongly-consistent alike), so the version alone - # must not be used to infer a leader. Conversely, replicas[0] is only - # leader-ordered for a tablet that came from a V2 payload, so a - # versionless tablet (V1-sourced, or stale across a consistency flip) - # must not be treated as a leader hint either. Require both a - # strongly-consistent keyspace and a versioned tablet; otherwise keep - # normal token-aware/shuffled ordering. - ks_meta = self._cluster_metadata.keyspaces.get(keyspace) - if (self._prefer_tablet_leader - and ks_meta is not None and ks_meta._consistency_mode == _ConsistencyMode.GLOBAL - and tablet.tablet_version is not None): - # Even for a leader-eligible tablet, a request at consistency - # level ONE or LOCAL_ONE is satisfied by any single replica, so - # preferring the leader would only concentrate load into a - # hotspot without buying any consistency; spread those instead. - # TODO: This reads the level off the statement, so a request that - # inherits its consistency level from an execution profile - # looks unset here and is routed to the leader anyway. See - # https://github.com/scylladb/python-driver/issues/953 - effective_cl = query.consistency_level - prefer_leader = effective_cl not in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE) - if prefer_leader: - leader_host_id = tablet.leader - # A tablet with no replicas reports no leader; guard against - # matching a host whose own host_id is still unknown. - if leader_host_id is not None: - for host in replicas: - if host.host_id == leader_host_id: - leader_host = host - break + if token_map: + 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] + + # The leader concept only exists for strongly-consistent keyspaces, + # which today means exactly the keyspaces whose consistency mode is + # GLOBAL: it is the only mode ScyllaDB implements so far, so LOCAL + # (reserved, unimplemented) and EVENTUAL both have no leader. This + # comparison has to widen once 'local' consistency exists. + # TABLETS_ROUTING_V2 assigns a tablet_version to *every* tablet table + # (eventually- and strongly-consistent alike), so the version alone + # must not be used to infer a leader. Conversely, replicas[0] is only + # leader-ordered for a tablet that came from a V2 payload, so a + # versionless tablet (V1-sourced, or stale across a consistency flip) + # must not be treated as a leader hint either. Require both a + # strongly-consistent keyspace and a versioned tablet; otherwise keep + # normal token-aware/shuffled ordering. + ks_meta = cluster_metadata.keyspaces.get(keyspace) + if (self._prefer_tablet_leader + and ks_meta is not None and ks_meta._consistency_mode == _ConsistencyMode.GLOBAL + and tablet.tablet_version is not None): + # Even for a leader-eligible tablet, a request at consistency + # level ONE or LOCAL_ONE is satisfied by any single replica, so + # preferring the leader would only concentrate load into a + # hotspot without buying any consistency; spread those instead. + # TODO: This reads the level off the statement, so a request that + # inherits its consistency level from an execution profile + # looks unset here and is routed to the leader anyway. See + # https://github.com/scylladb/python-driver/issues/953 + effective_cl = query.consistency_level + prefer_leader = effective_cl not in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE) + if prefer_leader: + leader_host_id = tablet.leader + # A tablet with no replicas reports no leader; guard against + # matching a host whose own host_id is still unknown. + if leader_host_id is not None: + for host in replicas: + if host.host_id == leader_host_id: + leader_host = host + break + else: + cached = self._get_cached_replicas(keyspace, query.routing_key, token_map) + if cached is not None: + _, 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 + ) else: - replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key) + replicas = 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): + replicas = list(replicas) 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 - - # If we have a leader hint, yield it first -- but respect the child - # policy's own filter: never front-run a host the child policy would - # exclude (e.g. one a custom policy reports as IGNORED). - if (leader_host is not None and leader_host.is_up - and child.distance(leader_host) != HostDistance.IGNORED): - yield leader_host - - # yield replicas: local_rack, local, remote (skipping leader already yielded) - for host in yield_in_order(replicas): - if host is not leader_host: - yield host + local_rack = [] + local = [] + remote = [] + + child_distance = child.distance - # yield rest of the cluster: local_rack, local, remote - # Note: The leader is always a replica, so we don't need to filter it out here. - yield from yield_in_order([host for host in child.make_query_plan(keyspace, query) if host not in replicas]) + # Leader is yielded separately (below), so it's excluded here even + # if the child policy would otherwise contact it -- matching the + # pre-existing behavior of never yielding the leader twice. + for replica in replicas: + if replica is leader_host: + continue + 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) + + leader_yieldable = (leader_host is not None and leader_host.is_up + and child_distance(leader_host) != HostDistance.IGNORED) + + if local_rack or local or remote or leader_yieldable: + yielded = set() + + # If we have a leader hint, yield it first -- but respect the child + # policy's own filter: never front-run a host the child policy + # would exclude (e.g. one a custom policy reports as IGNORED). + if leader_yieldable: + yielded.add(leader_host) + yield leader_host + elif leader_host is not None: + # Leader exists but isn't yieldable (down or ignored); it was + # already excluded from replicas above and stays excluded. + yielded.add(leader_host) + + 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: + 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/test_policies.py b/tests/unit/test_policies.py index 35c1a96f87..fdd019c74b 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 12e80d99dd739bdd61f1a719005b9db1dfc3e004 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Fri, 3 Apr 2026 20:07:03 +0300 Subject: [PATCH 06/11] 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 3d6b4456ca..89a3852e10 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -1089,6 +1089,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() @@ -1740,6 +1750,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 462f0baeb99cd99266c074e72697131169af3844 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Fri, 3 Apr 2026 20:36:09 +0300 Subject: [PATCH 07/11] 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 | 416 +++++++++++++++++++++++++++++++++++- 1 file changed, 408 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py index fdd019c74b..b1724a9924 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,9 +1108,20 @@ 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 @@ -1422,6 +1547,276 @@ def test_leader_routing_can_be_disabled(self): # The leader is still a replica of the tablet, so it stays in the plan -- # it just no longer jumps the queue. self.assertEqual(qplan.count(leader), 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): @@ -1440,6 +1835,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) @@ -2167,8 +2564,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 0c9f42a4a02b822e3d43e844fed0f2448a45058b Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Mon, 20 Apr 2026 00:00:10 +0300 Subject: [PATCH 08/11] 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 | 179 +++++++++++++++++++++--------------- tests/unit/test_policies.py | 28 ++++++ 2 files changed, 133 insertions(+), 74 deletions(-) diff --git a/cassandra/policies.py b/cassandra/policies.py index 89a3852e10..785314a9fe 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))): @@ -707,11 +707,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 @@ -720,17 +725,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 @@ -738,8 +753,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) @@ -763,66 +779,80 @@ def make_query_plan(self, working_keyspace=None, query=None): replicas = [] leader_host = None if token_map: - 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] - - # The leader concept only exists for strongly-consistent keyspaces, - # which today means exactly the keyspaces whose consistency mode is - # GLOBAL: it is the only mode ScyllaDB implements so far, so LOCAL - # (reserved, unimplemented) and EVENTUAL both have no leader. This - # comparison has to widen once 'local' consistency exists. - # TABLETS_ROUTING_V2 assigns a tablet_version to *every* tablet table - # (eventually- and strongly-consistent alike), so the version alone - # must not be used to infer a leader. Conversely, replicas[0] is only - # leader-ordered for a tablet that came from a V2 payload, so a - # versionless tablet (V1-sourced, or stale across a consistency flip) - # must not be treated as a leader hint either. Require both a - # strongly-consistent keyspace and a versioned tablet; otherwise keep - # normal token-aware/shuffled ordering. - ks_meta = cluster_metadata.keyspaces.get(keyspace) - if (self._prefer_tablet_leader - and ks_meta is not None and ks_meta._consistency_mode == _ConsistencyMode.GLOBAL - and tablet.tablet_version is not None): - # Even for a leader-eligible tablet, a request at consistency - # level ONE or LOCAL_ONE is satisfied by any single replica, so - # preferring the leader would only concentrate load into a - # hotspot without buying any consistency; spread those instead. - # TODO: This reads the level off the statement, so a request that - # inherits its consistency level from an execution profile - # looks unset here and is routed to the leader anyway. See - # https://github.com/scylladb/python-driver/issues/953 - effective_cl = query.consistency_level - prefer_leader = effective_cl not in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE) - if prefer_leader: - leader_host_id = tablet.leader - # A tablet with no replicas reports no leader; guard against - # matching a host whose own host_id is still unknown. - if leader_host_id is not None: - for host in replicas: - if host.host_id == leader_host_id: - leader_host = host - break - else: - cached = self._get_cached_replicas(keyspace, query.routing_key, token_map) + try: + # Check the LRU cache first -- avoids the hash (from_key) + # and token-map lookup on repeated routing keys. A cache hit + # only ever comes from the non-tablet path below (tablet + # replicas are never cached), so there is no leader to + # compute in that case. + cached = self._get_cached_replicas(keyspace, query.table, query.routing_key, token_map) if cached is not None: - _, replicas = cached + 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 + 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] + + # The leader concept only exists for strongly-consistent keyspaces, + # which today means exactly the keyspaces whose consistency mode is + # GLOBAL: it is the only mode ScyllaDB implements so far, so LOCAL + # (reserved, unimplemented) and EVENTUAL both have no leader. This + # comparison has to widen once 'local' consistency exists. + # TABLETS_ROUTING_V2 assigns a tablet_version to *every* tablet table + # (eventually- and strongly-consistent alike), so the version alone + # must not be used to infer a leader. Conversely, replicas[0] is only + # leader-ordered for a tablet that came from a V2 payload, so a + # versionless tablet (V1-sourced, or stale across a consistency flip) + # must not be treated as a leader hint either. Require both a + # strongly-consistent keyspace and a versioned tablet; otherwise keep + # normal token-aware/shuffled ordering. + ks_meta = cluster_metadata.keyspaces.get(keyspace) + if (self._prefer_tablet_leader + and ks_meta is not None and ks_meta._consistency_mode == _ConsistencyMode.GLOBAL + and tablet.tablet_version is not None): + # Even for a leader-eligible tablet, a request at consistency + # level ONE or LOCAL_ONE is satisfied by any single replica, so + # preferring the leader would only concentrate load into a + # hotspot without buying any consistency; spread those instead. + # TODO: This reads the level off the statement, so a request that + # inherits its consistency level from an execution profile + # looks unset here and is routed to the leader anyway. See + # https://github.com/scylladb/python-driver/issues/953 + effective_cl = query.consistency_level + prefer_leader = effective_cl not in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE) + if prefer_leader: + leader_host_id = tablet.leader + # A tablet with no replicas reports no leader; guard against + # matching a host whose own host_id is still unknown. + if leader_host_id is not None: + for host in replicas: + if host.host_id == leader_host_id: + leader_host = host + break + 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.table, query.routing_key, replicas, token_map + ) + except Exception: + log.debug( + "Failed to resolve token or tablet for query plan, " + "falling back to child policy", + exc_info=True, + ) else: replicas = cluster_metadata.get_replicas(keyspace, query.routing_key) @@ -1765,9 +1795,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 b1724a9924..137bb87724 100644 --- a/tests/unit/test_policies.py +++ b/tests/unit/test_policies.py @@ -1561,6 +1561,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): @@ -1639,6 +1641,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)) @@ -1708,6 +1711,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 @@ -1726,6 +1730,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 6cac29e02e10e93e091557675c048abfc7ca6a1b Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Mon, 20 Apr 2026 09:07:49 +0300 Subject: [PATCH 09/11] 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 785314a9fe..2ff0e37de6 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 137bb87724..782645a70a 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 cb94c53da305046439f2b6b270350472805cfb92 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Mon, 20 Apr 2026 11:03:58 +0300 Subject: [PATCH 10/11] 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 | 175 ++++++++++++++++++++++------------ cassandra/tablets.py | 16 ++++ tests/unit/test_policies.py | 183 ++++++++++++++++++++++++++++++++++++ tests/unit/test_tablets.py | 62 ++++++++++++ 4 files changed, 376 insertions(+), 60 deletions(-) diff --git a/cassandra/policies.py b/cassandra/policies.py index 2ff0e37de6..e7699f7426 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -701,10 +701,12 @@ 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. It is only consulted for - non-tablet keyspaces; tablet replicas are always resolved fresh since - they carry leader information that can change independently of the - token map. + 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, and + the leader information a tablet carries is never masked by a stale + cache entry either. """ _child_policy = None @@ -791,6 +793,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 @@ -814,12 +824,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 @@ -828,7 +848,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: @@ -854,62 +873,92 @@ def make_query_plan(self, working_keyspace=None, query=None): leader_host = None if token_map: try: - # Check the LRU cache first -- avoids the hash (from_key) - # and token-map lookup on repeated routing keys. A cache hit - # only ever comes from the non-tablet path below (tablet - # replicas are never cached), so there is no leader to - # compute in that case. - 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] - - # The leader concept only exists for strongly-consistent keyspaces, - # which today means exactly the keyspaces whose consistency mode is - # GLOBAL: it is the only mode ScyllaDB implements so far, so LOCAL - # (reserved, unimplemented) and EVENTUAL both have no leader. This - # comparison has to widen once 'local' consistency exists. - # TABLETS_ROUTING_V2 assigns a tablet_version to *every* tablet table - # (eventually- and strongly-consistent alike), so the version alone - # must not be used to infer a leader. Conversely, replicas[0] is only - # leader-ordered for a tablet that came from a V2 payload, so a - # versionless tablet (V1-sourced, or stale across a consistency flip) - # must not be treated as a leader hint either. Require both a - # strongly-consistent keyspace and a versioned tablet; otherwise keep - # normal token-aware/shuffled ordering. - ks_meta = cluster_metadata.keyspaces.get(keyspace) - if (self._prefer_tablet_leader - and ks_meta is not None and ks_meta._consistency_mode == _ConsistencyMode.GLOBAL - and tablet.tablet_version is not None): - # Even for a leader-eligible tablet, a request at consistency - # level ONE or LOCAL_ONE is satisfied by any single replica, so - # preferring the leader would only concentrate load into a - # hotspot without buying any consistency; spread those instead. - # TODO: This reads the level off the statement, so a request that - # inherits its consistency level from an execution profile - # looks unset here and is routed to the leader anyway. See - # https://github.com/scylladb/python-driver/issues/953 - effective_cl = query.consistency_level - prefer_leader = effective_cl not in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE) - if prefer_leader: - leader_host_id = tablet.leader - # A tablet with no replicas reports no leader; guard against - # matching a host whose own host_id is still unknown. - if leader_host_id is not None: - for host in replicas: - if host.host_id == leader_host_id: - leader_host = host - break - else: + tablet_found = False + if table_has_tablets: + tablet = tablets.get_tablet_for_key(keyspace, query.table, token) + if tablet is not None: + tablet_found = True + 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] + + # The leader concept only exists for strongly-consistent keyspaces, + # which today means exactly the keyspaces whose consistency mode is + # GLOBAL: it is the only mode ScyllaDB implements so far, so LOCAL + # (reserved, unimplemented) and EVENTUAL both have no leader. This + # comparison has to widen once 'local' consistency exists. + # TABLETS_ROUTING_V2 assigns a tablet_version to *every* tablet table + # (eventually- and strongly-consistent alike), so the version alone + # must not be used to infer a leader. Conversely, replicas[0] is only + # leader-ordered for a tablet that came from a V2 payload, so a + # versionless tablet (V1-sourced, or stale across a consistency flip) + # must not be treated as a leader hint either. Require both a + # strongly-consistent keyspace and a versioned tablet; otherwise keep + # normal token-aware/shuffled ordering. + ks_meta = cluster_metadata.keyspaces.get(keyspace) + if (self._prefer_tablet_leader + and ks_meta is not None and ks_meta._consistency_mode == _ConsistencyMode.GLOBAL + and tablet.tablet_version is not None): + # Even for a leader-eligible tablet, a request at consistency + # level ONE or LOCAL_ONE is satisfied by any single replica, so + # preferring the leader would only concentrate load into a + # hotspot without buying any consistency; spread those instead. + # TODO: This reads the level off the statement, so a request that + # inherits its consistency level from an execution profile + # looks unset here and is routed to the leader anyway. See + # https://github.com/scylladb/python-driver/issues/953 + effective_cl = query.consistency_level + prefer_leader = effective_cl not in (ConsistencyLevel.ONE, ConsistencyLevel.LOCAL_ONE) + if prefer_leader: + leader_host_id = tablet.leader + # A tablet with no replicas reports no leader; guard against + # matching a host whose own host_id is still unknown. + if leader_host_id is not None: + for host in replicas: + if host.host_id == leader_host_id: + leader_host = host + break + + if not replicas and not tablet_found: + # 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: @@ -918,9 +967,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 b386d1a372..21edd4cb5a 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -118,6 +118,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), [])) @@ -149,6 +152,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): @@ -158,6 +162,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 782645a70a..749ac10105 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 @@ -1753,6 +1759,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:] @@ -1903,6 +1910,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:] @@ -1925,6 +1933,42 @@ def test_tablet_path_not_cached(self): # Cache should remain empty (tablet results are not cached) assert len(policy._replica_cache) == 0 + def test_tablet_found_but_no_matching_replicas_does_not_fall_back_to_vnode(self): + """A tablet found with no replicas present in the child plan must not + fall back to vnode replicas -- the tablet route was resolved, just to + an empty set (e.g. none of its replica host_ids are in the live plan).""" + 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) + # Tablet's replicas are host_ids not present in the child's plan. + cluster.metadata._tablets.get_tablet_for_key.return_value = Tablet(replicas=[(uuid.uuid4(), 0)]) + 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 + 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'key1', keyspace='ks') + plan = list(policy.make_query_plan(None, query)) + + # The tablet was found, so the vnode fallback must not have kicked in. + assert cluster.metadata.token_map.get_replicas.call_count == 0 + # No replicas matched the tablet, so the plan falls back to the child + # policy's un-prioritized hosts, not the vnode replica set. + assert set(plan) == set(hosts) + 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() @@ -1949,6 +1993,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'): @@ -2787,6 +2969,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 656ae42da7..50f4e83d73 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -1,6 +1,7 @@ import unittest from io import BytesIO from uuid import uuid4 +import uuid from cassandra import ConsistencyLevel, ProtocolVersion from cassandra.protocol import ExecuteMessage @@ -279,3 +280,64 @@ def test_same_message_encodes_consistently_across_connections(self): first_again = self._encode_body(message, ProtocolFeatures(tablets_routing_v2=True)) self.assertEqual(first, first_again) self.assertEqual(first, second_plain + bytes([0x3C])) + + +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 From 0d9bf9b752945153da4788cc1aef68362dac348d Mon Sep 17 00:00:00 2001 From: Yaniv Kaul Date: Sun, 23 Aug 2026 10:25:24 +0300 Subject: [PATCH 11/11] fix: mock make_query_plan_with_exclusion in leader-routing tests The rebase onto master merged the leader-first tablet routing feature with the query-plan-exclusion optimization: the non-DCAware/RackAware fallback path now calls child.make_query_plan_with_exclusion instead of child.make_query_plan. The leader-routing tests used a bare Mock() child policy without stubbing that method, so it returned a non-iterable Mock and every test in that group failed after rebase. Co-Authored-By: Claude Sonnet 5 --- tests/unit/test_policies.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py index 749ac10105..0f36695ea6 100644 --- a/tests/unit/test_policies.py +++ b/tests/unit/test_policies.py @@ -1359,6 +1359,8 @@ def test_leader_aware_routing_with_tablet_version(self): # by distance (LOCAL vs LOCAL_RACK). Without leader-first routing, # other_replica would be yielded before the leader. child_policy.make_query_plan.return_value = [hosts[0], hosts[1], other_replica, leader] + child_policy.make_query_plan_with_exclusion.side_effect = \ + lambda k, q, e: [h for h in child_policy.make_query_plan.return_value if h not in e] distances = { leader: HostDistance.LOCAL, other_replica: HostDistance.LOCAL_RACK, @@ -1414,6 +1416,8 @@ def test_leader_fallback_when_leader_is_down(self): 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) @@ -1463,6 +1467,8 @@ def test_no_leader_routing_without_tablet_version(self): # Order the child plan so the second replica comes before replicas[0]; if # leader-first wrongly triggered, first_replica would be forced to front. child_policy.make_query_plan.return_value = [second_replica, first_replica, hosts[0], hosts[1]] + child_policy.make_query_plan_with_exclusion.side_effect = \ + lambda k, q, e: [h for h in child_policy.make_query_plan.return_value if h not in e] child_policy.distance.return_value = HostDistance.LOCAL # shuffle_replicas=False keeps replica ordering deterministic so we can @@ -1511,6 +1517,8 @@ def test_no_leader_routing_for_eventually_consistent_keyspace(self): # leader-first logic wrongly triggered, first_replica would be forced to # the front instead. child_policy.make_query_plan.return_value = [second_replica, first_replica, hosts[0], hosts[1]] + child_policy.make_query_plan_with_exclusion.side_effect = \ + lambda k, q, e: [h for h in child_policy.make_query_plan.return_value if h not in e] child_policy.distance.return_value = HostDistance.LOCAL # shuffle_replicas=False keeps replica ordering deterministic so we can @@ -1554,6 +1562,8 @@ def test_no_leader_routing_when_keyspace_metadata_missing(self): child_policy = Mock() child_policy.make_query_plan.return_value = [second_replica, first_replica, hosts[0], hosts[1]] + child_policy.make_query_plan_with_exclusion.side_effect = \ + lambda k, q, e: [h for h in child_policy.make_query_plan.return_value if h not in e] child_policy.distance.return_value = HostDistance.LOCAL # shuffle_replicas=False keeps replica ordering deterministic so we can @@ -1598,6 +1608,8 @@ def test_leader_skipped_when_child_policy_ignores_it(self): # The child policy yields the leader but reports it as IGNORED, i.e. it # would never actually route to it. child_policy.make_query_plan.return_value = [leader, other_replica, hosts[0], hosts[1]] + child_policy.make_query_plan_with_exclusion.side_effect = \ + lambda k, q, e: [h for h in child_policy.make_query_plan.return_value if h not in e] distances = { leader: HostDistance.IGNORED, other_replica: HostDistance.LOCAL, @@ -1660,6 +1672,8 @@ def _make_leader_routing_setup(self, *, consistency_mode=_ConsistencyMode.GLOBAL # Leader is last in the child plan and the farther replica (LOCAL vs # LOCAL_RACK); without leader-first, other_replica is yielded first. child_policy.make_query_plan.return_value = [hosts[0], hosts[1], other_replica, leader] + child_policy.make_query_plan_with_exclusion.side_effect = \ + lambda k, q, e: [h for h in child_policy.make_query_plan.return_value if h not in e] distances = { leader: HostDistance.LOCAL, other_replica: HostDistance.LOCAL_RACK,