diff --git a/cassandra/datastax/insights/serializers.py b/cassandra/datastax/insights/serializers.py index 289c165e8a..7a0abedefe 100644 --- a/cassandra/datastax/insights/serializers.py +++ b/cassandra/datastax/insights/serializers.py @@ -70,7 +70,8 @@ def token_aware_policy_insights_serializer(policy): 'namespace': namespace(policy.__class__), 'options': {'child_policy': insights_registry.serialize(policy._child_policy, policy=True), - 'shuffle_replicas': policy.shuffle_replicas} + 'shuffle_replicas': policy.shuffle_replicas, + 'cache_replicas_size': policy.cache_replicas_size} } @insights_registry.register_serializer_for(WhiteListRoundRobinPolicy) diff --git a/cassandra/policies.py b/cassandra/policies.py index 89702e8c89..4323e67fef 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 @@ -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, )) @@ -222,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): """ @@ -241,37 +266,72 @@ 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 + 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,21 +340,62 @@ 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)): + length = len(local_live) + if length: + pos %= length + for i in range(length): + yield local_live[(pos + i) % length] + + # Read _remote_hosts late so topology changes during local + # iteration are visible. + for host in self._remote_hosts: yield host - # 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]: + 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) + if not excluded: + if length: + pos %= length + for i in range(length): + yield local_live[(pos + i) % length] + # Read _remote_hosts late so topology changes during local + # iteration are visible. + for host in self._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 self._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 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, " @@ -306,6 +407,8 @@ def on_up(self, host): 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: + self._refresh_remote_hosts() def on_down(self, host): dc = self._dc(host) @@ -318,6 +421,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) @@ -331,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): """ @@ -348,21 +453,82 @@ 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 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 = 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))): self._live_hosts[(dc, rack)] = tuple({*rack_hosts, *self._live_hosts.get((dc, rack), [])}) @@ -370,71 +536,118 @@ 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)): + 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] + + # Read _remote_hosts late so topology changes during local + # iteration are visible. + for host in self._remote_hosts: yield host - # 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 + def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excluded=()): + pos = self._position + self._position += 1 + + local_rack_live = self._live_hosts.get((self.local_dc, self.local_rack), ()) + length = len(local_rack_live) + if 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] + + # Read _remote_hosts late so topology changes during local + # iteration are visible. + for host in self._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 + + # 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 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) @@ -443,6 +656,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) @@ -464,6 +690,14 @@ class TokenAwarePolicy(LoadBalancingPolicy): If no :attr:`~.Statement.routing_key` is set on the query, the child policy's query plan will be used as is. + + An LRU cache of size :attr:`cache_replicas_size` (default 1024) avoids + repeated token-to-replica lookups for the same (keyspace, routing_key) + pair. Set to 0 to disable caching. The cache is automatically + invalidated when the cluster topology changes. It only ever holds + vnode-derived replica sets; tables with known tablets always resolve + replicas via the tablet metadata directly, so tablet-aware routing is + never masked by a cache entry created before a tablet was learned. """ _child_policy = None @@ -473,9 +707,24 @@ class TokenAwarePolicy(LoadBalancingPolicy): Yield local replicas in a random order. """ - def __init__(self, child_policy, shuffle_replicas=True): + def __init__(self, child_policy, shuffle_replicas=True, cache_replicas_size=1024): + super().__init__() self._child_policy = child_policy self.shuffle_replicas = shuffle_replicas + self._cluster_metadata = None + self._cache_replicas_size = max(0, cache_replicas_size) + self._replica_cache = OrderedDict() + self._replica_cache_token_map_ref = None + self._cache_lock = Lock() + + @property + def cache_replicas_size(self): + """ + The configured size of the replica LRU cache, as passed to the + `cache_replicas_size` constructor argument (default 1024). A value + of 0 means caching is disabled. Read-only. + """ + return self._cache_replicas_size def populate(self, cluster, hosts): self._cluster_metadata = cluster.metadata @@ -493,40 +742,230 @@ def check_supported(self): def distance(self, *args, **kwargs): return self._child_policy.distance(*args, **kwargs) + def _get_cached_replicas(self, keyspace, table, routing_key_bytes, token_map): + """ + 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. + + 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 + 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, 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 replicas + return None + + 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 + 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, table, routing_key_bytes) + 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) + def make_query_plan(self, working_keyspace=None, query=None): keyspace = query.keyspace if query and query.keyspace else working_keyspace child = self._child_policy if query is None or query.routing_key is None or keyspace is None: - for host in child.make_query_plan(keyspace, query): - yield host + yield from child.make_query_plan(keyspace, query) return + cluster_metadata = self._cluster_metadata + token_map = cluster_metadata.token_map replicas = [] - tablet = self._cluster_metadata._tablets.get_tablet_for_key( - keyspace, query.table, self._cluster_metadata.token_map.token_class.from_key(query.routing_key)) - - if tablet is not None: - replicas_mapped = set(map(lambda r: r[0], tablet.replicas)) - child_plan = child.make_query_plan(keyspace, query) - - replicas = [host for host in child_plan if host.host_id in replicas_mapped] + if token_map: + try: + 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) + + if table_has_tablets: + tablet = tablets.get_tablet_for_key(keyspace, query.table, token) + if tablet is not None: + replicas_mapped = {r[0] for r in tablet.replicas} + child_plan = child.make_query_plan(keyspace, query) + replicas = [host for host in child_plan if host.host_id in replicas_mapped] + + if not replicas: + # Snapshot the keyspace replica-map identity BEFORE + # resolving replicas (not after) so a concurrent + # ALTER KEYSPACE rebuild in this window is detected + # as invalidating the entry we're about to cache + # rather than accidentally validating it later. See + # _put_cached_replicas. + ks_map_ref = token_map.tokens_to_hosts_by_ks.get(keyspace) + try: + replicas = token_map.get_replicas(keyspace, token) + except Exception: + log.debug( + "Failed to get replicas from token_map, " + "falling back to cluster metadata" + ) + replicas = cluster_metadata.get_replicas(keyspace, query.routing_key) + # 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, " + "falling back to child policy", + exc_info=True, + ) + + if self.shuffle_replicas: + if not query.is_lwt() and not ConsistencyLevel.is_serial(query.consistency_level): + replicas = list(replicas) + shuffle(replicas) + + local_rack = [] + local = [] + remote = [] + + child_distance = child.distance + + for replica in replicas: + if replica.is_up: + d = child_distance(replica) + if d == HostDistance.LOCAL_RACK: + local_rack.append(replica) + elif d == HostDistance.LOCAL: + local.append(replica) + elif d == HostDistance.REMOTE: + remote.append(replica) + + if local_rack or local or remote: + yielded = set() + + for replica in local_rack: + yielded.add(replica) + yield replica + + for replica in local: + yielded.add(replica) + yield replica + + for replica in remote: + yielded.add(replica) + yield replica + + # Yield the rest of the cluster (non-replica hosts). + # DCAware and RackAware already yield in distance order + # (local_rack -> local -> remote), so we can stream directly. + # For other child policies we must re-sort by distance. + if isinstance(child, (DCAwareRoundRobinPolicy, RackAwareRoundRobinPolicy)): + yield from child.make_query_plan_with_exclusion(keyspace, query, yielded) + else: + remaining_local_rack = [] + remaining_local = [] + remaining_remote = [] + # Hosts whose distance is neither LOCAL_RACK/LOCAL/REMOTE + # (e.g. IGNORED, or a distance that flipped mid-plan due to + # a concurrent topology refresh) are collected here rather + # than dropped, and yielded last -- preserving their + # relative order from the child plan, matching the + # pre-existing (non-optimized) behavior of never shrinking + # the plan compared to what the child policy yielded. + remaining_other = [] + for host in child.make_query_plan_with_exclusion(keyspace, query, yielded): + d = child_distance(host) + if d == HostDistance.LOCAL_RACK: + remaining_local_rack.append(host) + elif d == HostDistance.LOCAL: + remaining_local.append(host) + elif d == HostDistance.REMOTE: + remaining_remote.append(host) + else: + remaining_other.append(host) + yield from remaining_local_rack + yield from remaining_local + yield from remaining_remote + yield from remaining_other else: - replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key) - - if self.shuffle_replicas and not query.is_lwt() and not ConsistencyLevel.is_serial(query.consistency_level): - shuffle(replicas) - - def yield_in_order(hosts): - for distance in [HostDistance.LOCAL_RACK, HostDistance.LOCAL, HostDistance.REMOTE]: - for replica in hosts: - if replica.is_up and child.distance(replica) == distance: - yield replica - - # yield replicas: local_rack, local, remote - yield from yield_in_order(replicas) - # yield rest of the cluster: local_rack, local, remote - yield from yield_in_order([host for host in child.make_query_plan(keyspace, query) if host not in replicas]) + yield from child.make_query_plan(keyspace, query) def on_up(self, *args, **kwargs): return self._child_policy.on_up(*args, **kwargs) @@ -702,6 +1141,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() @@ -1353,6 +1802,28 @@ 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 + # 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) + # TODO for backward compatibility, remove in next major class DSELoadBalancingPolicy(DefaultLoadBalancingPolicy): diff --git a/cassandra/tablets.py b/cassandra/tablets.py index 96e61a50c2..616f64cd15 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -55,6 +55,9 @@ def __init__(self, tablets): self._tablets = tablets self._lock = Lock() + def __bool__(self): + return bool(self._tablets) + def table_has_tablets(self, keyspace, table) -> bool: return bool(self._tablets.get((keyspace, table), [])) @@ -86,6 +89,7 @@ def drop_tablets_by_host_id(self, host_id: Optional[UUID]): if host_id is None: return with self._lock: + empty_keys = [] for key, tablets in self._tablets.items(): to_be_deleted = [] for tablet_id, tablet in enumerate(tablets): @@ -95,6 +99,18 @@ def drop_tablets_by_host_id(self, host_id: Optional[UUID]): for tablet_id in reversed(to_be_deleted): tablets.pop(tablet_id) + if not tablets: + empty_keys.append(key) + + # Delete keys whose tablet list is now empty (can't do this + # while iterating self._tablets.items() above -- that would + # mutate the dict's size mid-iteration). Leaving an empty list + # behind would keep bool(self._tablets) truthy forever for a + # table that no longer has any tablets, defeating the + # table_has_tablets()/bool(tablets) no-tablet fast path. + for key in empty_keys: + del self._tablets[key] + def add_tablet(self, keyspace, table, tablet): with self._lock: tablets_for_table = self._tablets.setdefault((keyspace, table), []) diff --git a/tests/unit/advanced/test_insights.py b/tests/unit/advanced/test_insights.py index 2050439804..0773ff965a 100644 --- a/tests/unit/advanced/test_insights.py +++ b/tests/unit/advanced/test_insights.py @@ -120,7 +120,8 @@ def test_execution_profile(self): 'options': {'local_dc': '', 'used_hosts_per_remote_dc': 0}, 'type': 'DCAwareRoundRobinPolicy'}, - 'shuffle_replicas': True}, + 'shuffle_replicas': True, + 'cache_replicas_size': 1024}, 'type': 'TokenAwarePolicy'}, 'readTimeout': 10.0, 'retry': {'namespace': 'cassandra.policies', 'options': {}, 'type': 'RetryPolicy'}, @@ -139,7 +140,8 @@ def test_graph_execution_profile(self): 'options': {'local_dc': '', 'used_hosts_per_remote_dc': 0}, 'type': 'DCAwareRoundRobinPolicy'}, - 'shuffle_replicas': True}, + 'shuffle_replicas': True, + 'cache_replicas_size': 1024}, 'type': 'TokenAwarePolicy'}, 'readTimeout': 30.0, 'retry': {'namespace': 'cassandra.policies', 'options': {}, 'type': 'NeverRetryPolicy'}, @@ -161,7 +163,8 @@ def test_graph_analytics_execution_profile(self): 'options': {'local_dc': '', 'used_hosts_per_remote_dc': 0}, 'type': 'DCAwareRoundRobinPolicy'}, - 'shuffle_replicas': True}, + 'shuffle_replicas': True, + 'cache_replicas_size': 1024}, 'type': 'TokenAwarePolicy'}}, 'type': 'DefaultLoadBalancingPolicy'}, 'readTimeout': 604800.0, @@ -189,7 +192,8 @@ def test_token_aware_policy(self): 'options': {'child_policy': {'namespace': 'cassandra.policies', 'options': {}, 'type': 'LoadBalancingPolicy'}, - 'shuffle_replicas': True}, + 'shuffle_replicas': True, + 'cache_replicas_size': 1024}, 'type': 'TokenAwarePolicy'} def test_whitelist_round_robin_policy(self): diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py index 63a3c3d12d..4208ea9ef1 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]: @@ -537,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): @@ -570,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): @@ -577,6 +859,9 @@ 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)] for host in hosts: host.set_up() @@ -586,8 +871,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 +896,9 @@ 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)] for host in hosts: host.set_up() @@ -627,8 +916,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 +949,9 @@ 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)] for host in hosts: host.set_up() @@ -680,8 +973,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): @@ -710,6 +1004,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) @@ -804,12 +1156,17 @@ 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 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 @@ -848,7 +1205,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): """ @@ -897,6 +1256,10 @@ 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 return cluster def _prepare_cluster_with_tablets(self): @@ -909,14 +1272,23 @@ 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 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) @@ -926,6 +1298,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: @@ -936,13 +1309,463 @@ def _assert_shuffle(self, patched_shuffle, cluster, keyspace, routing_key): else: assert set(replicas) == set(qplan[:2]) assert hosts[:2] == qplan[2:] + if is_tablets: + # Tablet path: make_query_plan called once for replica ordering, + # make_query_plan_with_exclusion called once for remaining hosts child_policy.make_query_plan.assert_called_with(keyspace, query) - assert child_policy.make_query_plan.call_count == 2 + child_policy.make_query_plan_with_exclusion.assert_called() + elif child_policy.make_query_plan_with_exclusion.called: + # Non-tablet path with replicas: exclusion set should contain + # the replicas that were already yielded + exc_call = child_policy.make_query_plan_with_exclusion.call_args + excluded_hosts = exc_call[0][2] if len(exc_call[0]) > 2 else exc_call[1].get('excluded', set()) + assert set(replicas).issubset(excluded_hosts), \ + 'Exclusion set should contain the yielded replicas, ' \ + 'got %s, expected superset of %s' % (excluded_hosts, set(replicas)) else: child_policy.make_query_plan.assert_called_once_with(keyspace, query) assert patched_shuffle.call_count == 1 + # --- Replica cache tests --- + + def _make_cache_cluster(self): + """Create a mock cluster suitable for cache tests.""" + hosts = [Host(DefaultEndPoint(str(i)), SimpleConvictionPolicy, host_id=uuid.uuid4()) for i in range(4)] + for host in hosts: + host.set_up() + cluster = Mock(spec=Cluster) + cluster.metadata = Mock(spec=Metadata) + cluster.metadata._tablets = Mock(spec=Tablets) + cluster.metadata._tablets.get_tablet_for_key.return_value = None + cluster.metadata._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:] + # 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): + """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:] + new_token_map.tokens_to_hosts_by_ks = {'ks': {}} + 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._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:] + 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') + 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 + + 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 + + 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'): + """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): @@ -961,6 +1784,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) @@ -1688,8 +2513,12 @@ 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 - child_policy = TokenAwarePolicy(RoundRobinPolicy()) + child_policy = TokenAwarePolicy(RoundRobinPolicy(), shuffle_replicas=False) hfp = HostFilterPolicy( child_policy=child_policy, diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index 7a40e7de4d..a2c6b38049 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -1,4 +1,5 @@ import unittest +import uuid from cassandra.tablets import Tablets, Tablet @@ -124,3 +125,64 @@ def __init__(self, v): # Token value 50 is not > first_token (100) of the tablet whose # last_token (200) is >= 50, so no match. self.assertIsNone(tablets.get_tablet_for_key("ks", "tb", Token(50))) + + +class DropTabletsByHostIdTest(unittest.TestCase): + """ + Regression tests: drop_tablets_by_host_id must delete a table's key + from _tablets entirely once its tablet list becomes empty, so that + table_has_tablets() (and any other truthiness/membership check on + _tablets, e.g. bool(tablets)) correctly reflects that no tablets are + left -- rather than leaving a stale empty list behind. + """ + + def test_drop_last_tablet_removes_table_key(self): + host_id = uuid.uuid4() + t1 = Tablet(0, 100, [(host_id, 0)]) + tablets = Tablets({("ks", "tb"): [t1]}) + + assert tablets.table_has_tablets("ks", "tb") is True + + tablets.drop_tablets_by_host_id(host_id) + + # The no-tablet fast path must now be correctly signalled: the key + # must be gone from the dict (not just left as an empty list), and + # table_has_tablets must report False. + assert ("ks", "tb") not in tablets._tablets + assert tablets.table_has_tablets("ks", "tb") is False + assert bool(tablets) is False + + def test_drop_some_tablets_keeps_remaining(self): + removed_host_id = uuid.uuid4() + remaining_host_id = uuid.uuid4() + t1 = Tablet(0, 100, [(removed_host_id, 0)]) + t2 = Tablet(100, 200, [(remaining_host_id, 0)]) + tablets = Tablets({("ks", "tb"): [t1, t2]}) + + tablets.drop_tablets_by_host_id(removed_host_id) + + assert ("ks", "tb") in tablets._tablets + assert tablets._tablets[("ks", "tb")] == [t2] + assert tablets.table_has_tablets("ks", "tb") is True + assert bool(tablets) is True + + def test_drop_last_tablet_for_one_table_keeps_other_tables(self): + host_id = uuid.uuid4() + other_host_id = uuid.uuid4() + t1 = Tablet(0, 100, [(host_id, 0)]) + t2 = Tablet(0, 100, [(other_host_id, 0)]) + tablets = Tablets({("ks", "tb1"): [t1], ("ks", "tb2"): [t2]}) + + tablets.drop_tablets_by_host_id(host_id) + + assert ("ks", "tb1") not in tablets._tablets + assert tablets.table_has_tablets("ks", "tb1") is False + assert ("ks", "tb2") in tablets._tablets + assert tablets.table_has_tablets("ks", "tb2") is True + # Overall dict is still non-empty because tb2 still has tablets. + assert bool(tablets) is True + + def test_drop_by_host_id_none_is_noop(self): + tablets = Tablets({("ks", "tb"): []}) + tablets.drop_tablets_by_host_id(None) + assert ("ks", "tb") in tablets._tablets