From 31d3472f6b030a80a28209b172d0f89e4fedd9b4 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Thu, 9 Apr 2026 13:43:51 +0300 Subject: [PATCH 1/8] perf: add __slots__ to Tablet to eliminate per-instance __dict__ Add __slots__ to the Tablet class, removing the per-instance __dict__ allocation. Tablets are created frequently (one per token range per table) and are long-lived, so the cumulative memory savings are significant. Before: 416 bytes/tablet (48 instance + 96 __dict__ + 80 replicas + 192 tuples) After: 328 bytes/tablet (56 instance + 0 __dict__ + 80 replicas + 192 tuples) Saving: 88 bytes/tablet (21%) Scale impact (3 replicas/tablet): 12,800 tablets (100 tables x 128): saves 1.1 MB 128,000 tablets (1000 tables x 128): saves 10.7 MB 256,000 tablets (1000 tables x 256): saves 21.5 MB Tablet.from_row construction also improves: Before: 186 ns/call After: 147 ns/call (1.27x faster, -21%) Signed-off-by: Yaniv Kaul --- cassandra/tablets.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index b386d1a372..ba7d84799b 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -42,12 +42,9 @@ class Tablet(object): It stores information about each replica, its host and shard, and the token interval in the format (first_token, last_token]. """ - first_token = 0 - last_token = 0 - replicas = None # uint64 hash; None means unknown -- a cold start, or a tablet learned over # TABLETS_ROUTING_V1, which does not report a version. - tablet_version = None + __slots__ = ('first_token', 'last_token', 'replicas', 'tablet_version') def __init__(self, first_token=0, last_token=0, replicas=None, tablet_version=None): self.first_token = first_token From 2667014055ba57073362ed22e8ab6a9bbdb2f1d1 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Thu, 9 Apr 2026 13:44:48 +0300 Subject: [PATCH 2/8] perf: store Tablet.replicas as tuple instead of list Replicas are never mutated after Tablet construction; convert to tuple in __init__ to save 8 bytes per tablet (list overallocates for future appends that never happen) and communicate immutability. Before: 328 bytes/tablet (replicas container: 80 bytes as list) After: 320 bytes/tablet (replicas container: 72 bytes as tuple) Saving: 8 bytes/tablet (2.4%) Combined with __slots__ (commit 1), total savings so far: 96 bytes/tablet. Scale impact (3 replicas/tablet): 128,000 tablets: saves ~1.0 MB (tuple) + 10.7 MB (slots) = 11.7 MB total 256,000 tablets: saves ~2.0 MB (tuple) + 21.5 MB (slots) = 23.5 MB total Signed-off-by: Yaniv Kaul --- cassandra/tablets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index ba7d84799b..aedbe4bdd2 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -49,7 +49,7 @@ class Tablet(object): def __init__(self, first_token=0, last_token=0, replicas=None, tablet_version=None): self.first_token = first_token self.last_token = last_token - self.replicas = replicas + self.replicas = tuple(replicas) if replicas is not None else None self.tablet_version = tablet_version def __str__(self): From 0f831a3289267d0257da4ccba5f1b46d26c42b65 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Thu, 9 Apr 2026 14:17:19 +0300 Subject: [PATCH 3/8] perf: cache _replica_dict on Tablet for O(1) host/shard lookup Build a {host_id: shard_id} dict once at Tablet construction time so that policies.py and pool.py can replace set(map(lambda ...)) and linear scans with O(1) dict operations. - Add _replica_dict to __slots__ - Build dict from the materialized tuple (not the raw replicas arg) to avoid double-consuming a one-shot iterator - Update DCAwareRoundRobinPolicy to use tablet._replica_dict keys - Update HostConnection to use tablet._replica_dict.get() for shard - Rewrite replica_contains_host_id() to use dict membership - Add 7 unit tests covering dict construction, lookup, host membership, tuple storage, and the iterator edge case Add a public get_replica_shard_id() accessor alongside replica_contains_host_id(), and rewrite TabletReplicaDictTest to assert through that public API instead of poking at the private _replica_dict cache directly (per review feedback: reaching into the private field makes future refactors of the internal representation unnecessarily fragile). One minimal test (test_replica_dict_populated_as_expected) is kept to directly assert on _replica_dict's shape, since the public API alone can't prove the O(1) cache is actually populated as expected. Signed-off-by: Yaniv Kaul --- cassandra/policies.py | 4 +-- cassandra/pool.py | 5 +-- cassandra/tablets.py | 15 +++++---- tests/unit/test_tablets.py | 66 +++++++++++++++++++++++++++++++++++++- 4 files changed, 77 insertions(+), 13 deletions(-) diff --git a/cassandra/policies.py b/cassandra/policies.py index f1bfefb41d..fb0257c438 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -553,10 +553,10 @@ def make_query_plan(self, working_keyspace=None, query=None): 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)) + replica_dict = tablet._replica_dict child_plan = child.make_query_plan(keyspace, query) - replicas = [host for host in child_plan if host.host_id in replicas_mapped] + replicas = [host for host in child_plan if host.host_id in replica_dict] # The leader concept only exists for strongly-consistent keyspaces, # which today means exactly the keyspaces whose consistency mode is diff --git a/cassandra/pool.py b/cassandra/pool.py index 1d90e3233f..f2f2bf4405 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -472,10 +472,7 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table # the shard that this host owns for the tablet. Leader-aware host # selection (V2) happens earlier, in the load balancing policy. if tablet is not None: - for replica in tablet.replicas: - if replica[0] == self.host.host_id: - shard_id = replica[1] - break + shard_id = tablet._replica_dict.get(self.host.host_id) if shard_id is None and t is not None: shard_id = self.host.sharding_info.shard_id_from_token(t.value) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index aedbe4bdd2..a05856742a 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -42,14 +42,17 @@ class Tablet(object): It stores information about each replica, its host and shard, and the token interval in the format (first_token, last_token]. """ - # uint64 hash; None means unknown -- a cold start, or a tablet learned over +# uint64 hash; None means unknown -- a cold start, or a tablet learned over # TABLETS_ROUTING_V1, which does not report a version. - __slots__ = ('first_token', 'last_token', 'replicas', 'tablet_version') + __slots__ = ('first_token', 'last_token', 'replicas', 'tablet_version', '_replica_dict') def __init__(self, first_token=0, last_token=0, replicas=None, tablet_version=None): self.first_token = first_token self.last_token = last_token + # Materialize once: `replicas` may be a one-shot iterator, and both + # the tuple and the lookup dict must come from the same iteration. self.replicas = tuple(replicas) if replicas is not None else None + self._replica_dict = {r[0]: r[1] for r in self.replicas} if self.replicas else {} self.tablet_version = tablet_version def __str__(self): @@ -101,10 +104,10 @@ def leader(self) -> Optional[UUID]: return self.replicas[0][0] def replica_contains_host_id(self, uuid: UUID) -> bool: - for replica in self.replicas: - if replica[0] == uuid: - return True - return False + return uuid in self._replica_dict + + def get_replica_shard_id(self, uuid: UUID) -> Optional[int]: + return self._replica_dict.get(uuid) class Tablets(object): diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index 656ae42da7..a5b3613cc8 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -1,6 +1,6 @@ import unittest from io import BytesIO -from uuid import uuid4 +from uuid import UUID, uuid4 from cassandra import ConsistencyLevel, ProtocolVersion from cassandra.protocol import ExecuteMessage @@ -279,3 +279,67 @@ 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 TabletReplicaDictTest(unittest.TestCase): + """Tests for Tablet's replica/shard lookup behavior, backed internally + by a cached _replica_dict for O(1) host/shard lookup. + + Most of these tests go through the public API (replica_contains_host_id + and get_replica_shard_id) so they keep working across internal + refactors of the cache; see test_replica_dict_populated_as_expected + for the one targeted check of the internal structure itself. + """ + + def test_replica_contains_host_id(self): + u1 = UUID('12345678-1234-5678-1234-567812345678') + u2 = UUID('87654321-4321-8765-4321-876543218765') + u3 = UUID('aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee') + t = Tablet(0, 100, [(u1, 3), (u2, 7)]) + self.assertTrue(t.replica_contains_host_id(u1)) + self.assertTrue(t.replica_contains_host_id(u2)) + self.assertFalse(t.replica_contains_host_id(u3)) + + def test_replica_contains_host_id_false_when_no_replicas(self): + u1 = UUID('12345678-1234-5678-1234-567812345678') + t = Tablet(0, 100, None) + self.assertFalse(t.replica_contains_host_id(u1)) + + def test_get_replica_shard_id(self): + u1 = UUID('12345678-1234-5678-1234-567812345678') + u2 = UUID('87654321-4321-8765-4321-876543218765') + u3 = UUID('aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee') + t = Tablet(0, 100, [(u1, 3), (u2, 7)]) + self.assertEqual(t.get_replica_shard_id(u1), 3) + self.assertEqual(t.get_replica_shard_id(u2), 7) + self.assertIsNone(t.get_replica_shard_id(u3)) + + def test_replicas_stored_as_tuple(self): + t = Tablet(0, 100, [("host1", 0), ("host2", 1)]) + self.assertIsInstance(t.replicas, tuple) + + def test_replica_lookup_from_iterator(self): + """Ensure replica lookups work correctly even when replicas is a + one-shot iterator (generator), not a reusable list.""" + u1 = UUID('12345678-1234-5678-1234-567812345678') + u2 = UUID('87654321-4321-8765-4321-876543218765') + + def gen(): + yield (u1, 3) + yield (u2, 7) + + t = Tablet(0, 100, gen()) + self.assertEqual(t.replicas, ((u1, 3), (u2, 7))) + self.assertTrue(t.replica_contains_host_id(u1)) + self.assertTrue(t.replica_contains_host_id(u2)) + self.assertEqual(t.get_replica_shard_id(u1), 3) + self.assertEqual(t.get_replica_shard_id(u2), 7) + + def test_replica_dict_populated_as_expected(self): + """Minimal targeted regression test for the internal _replica_dict + cache: confirms the O(1)-lookup structure this optimization relies + on is actually populated as {host_id: shard_id}, which the public + API alone does not prove.""" + u1 = UUID('12345678-1234-5678-1234-567812345678') + u2 = UUID('87654321-4321-8765-4321-876543218765') + t = Tablet(0, 100, [(u1, 3), (u2, 7)]) + self.assertEqual(t._replica_dict, {u1: 3, u2: 7}) From 62346f50927a00cc3693ef95dbb5397371bd2039 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Thu, 9 Apr 2026 14:31:39 +0300 Subject: [PATCH 4/8] perf: streamline Tablet.from_row by inlining validation Remove the _is_valid_tablet staticmethod indirection and replace the two-step from_row -> _is_valid_tablet -> Tablet() chain with a single truthiness guard and direct construction. Saves ~54 ns/call (12%) by eliminating a staticmethod descriptor lookup, an extra function call, and redundant 'is not None' check (replicas from CQL deserialization is always a list or None). Fix a real correctness regression introduced by this same change: the inlined `if not replicas:` truthiness check is always False for a one-shot iterator/generator, even an empty one (iterators have no __len__/__bool__ so the default 'always truthy' rule applies). The previous _is_valid_tablet() helper used len(replicas) != 0, which would at least raise a TypeError for a generator rather than silently constructing a Tablet with empty replicas/_replica_dict instead of returning None. Materialize replicas into a tuple once and check that for emptiness, then hand the already-materialized tuple to Tablet() so it isn't consumed twice. Add TabletFromRowTest covering the empty list/generator/None cases (return None) and non-empty list/generator cases (Tablet is built correctly, including from a single-use generator). Signed-off-by: Yaniv Kaul --- cassandra/tablets.py | 26 +++++++++++----------- tests/unit/test_tablets.py | 44 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 13 deletions(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index a05856742a..8190058901 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -60,21 +60,21 @@ def __str__(self): % (self.first_token, self.last_token, self.replicas, self.tablet_version) __repr__ = __str__ - @staticmethod - def _is_valid_tablet(replicas): - return replicas is not None and len(replicas) != 0 - @staticmethod def from_row(first_token, last_token, replicas, tablet_version=None): - if Tablet._is_valid_tablet(replicas): - if tablet_version is not None: - # tablet_version is an unsigned 64-bit value, but it is - # deserialized from the wire as a signed LongType; normalize it - # back to unsigned so it matches the server's representation. - tablet_version &= 0xFFFFFFFFFFFFFFFF - tablet = Tablet(first_token, last_token, replicas, tablet_version) - return tablet - return None + # Materialize once: `replicas` may be a one-shot iterator (e.g. a + # generator), and a plain `if not replicas` truthiness check would + # always be False for such an object even when it yields nothing, + # since iterators have no __len__/__bool__ and are always truthy. + replicas_tuple = tuple(replicas) if replicas is not None else () + if not replicas_tuple: + return None + if tablet_version is not None: + # tablet_version is an unsigned 64-bit value, but it is + # deserialized from the wire as a signed LongType; normalize it + # back to unsigned so it matches the server's representation. + tablet_version &= 0xFFFFFFFFFFFFFFFF + return Tablet(first_token, last_token, replicas_tuple, tablet_version) @property def leader(self) -> Optional[UUID]: diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index a5b3613cc8..8072c4e9f0 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -280,6 +280,50 @@ def test_same_message_encodes_consistently_across_connections(self): self.assertEqual(first, first_again) self.assertEqual(first, second_plain + bytes([0x3C])) +class TabletFromRowTest(unittest.TestCase): + """Tests for Tablet.from_row, in particular that emptiness is detected + correctly regardless of whether `replicas` is a reusable sequence or a + one-shot iterator/generator.""" + + def test_empty_list_returns_none(self): + self.assertIsNone(Tablet.from_row(0, 100, [])) + + def test_empty_generator_returns_none(self): + # A generator is always truthy, even when empty, so a naive + # `if not replicas` check would fail to detect this case. + self.assertIsNone(Tablet.from_row(0, 100, (x for x in []))) + + def test_none_returns_none(self): + self.assertIsNone(Tablet.from_row(0, 100, None)) + + def test_non_empty_list_builds_tablet(self): + u1 = UUID('12345678-1234-5678-1234-567812345678') + u2 = UUID('87654321-4321-8765-4321-876543218765') + tablet = Tablet.from_row(0, 100, [(u1, 3), (u2, 7)]) + self.assertIsNotNone(tablet) + self.assertEqual(tablet.replicas, ((u1, 3), (u2, 7))) + self.assertTrue(tablet.replica_contains_host_id(u1)) + self.assertEqual(tablet.get_replica_shard_id(u2), 7) + + def test_non_empty_generator_builds_tablet(self): + # Generators are single-use: confirm the fix materializes the + # replicas exactly once and doesn't lose data by iterating twice. + u1 = UUID('12345678-1234-5678-1234-567812345678') + u2 = UUID('87654321-4321-8765-4321-876543218765') + + def gen(): + yield (u1, 3) + yield (u2, 7) + + tablet = Tablet.from_row(0, 100, gen()) + self.assertIsNotNone(tablet) + self.assertEqual(tablet.replicas, ((u1, 3), (u2, 7))) + self.assertTrue(tablet.replica_contains_host_id(u1)) + self.assertTrue(tablet.replica_contains_host_id(u2)) + self.assertEqual(tablet.get_replica_shard_id(u1), 3) + self.assertEqual(tablet.get_replica_shard_id(u2), 7) + + class TabletReplicaDictTest(unittest.TestCase): """Tests for Tablet's replica/shard lookup behavior, backed internally by a cached _replica_dict for O(1) host/shard lookup. From 24d54707576282ddaee5e467ee4cc3b9284e6e28 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Thu, 9 Apr 2026 14:52:24 +0300 Subject: [PATCH 5/8] perf: eliminate bisect key= callback via parallel token index lists Maintain parallel _first_tokens and _last_tokens dicts alongside _tablets, each mapping (keyspace, table) to a plain list[int]. This lets bisect_left run entirely in C on native ints instead of calling an attrgetter callback on every comparison during binary search. Follow-up to PR #757 which identified the opportunity: its own benchmarks showed bisect_left without key= is 2.7-5.7x faster than with key=attrgetter. Results (best-of-5, Python 3.14): get_tablet_for_key (hit): Tablets Before After Saved Speedup 10 293ns 216ns 78ns 1.36x 100 351ns 233ns 118ns 1.51x 1,000 448ns 267ns 181ns 1.68x 10,000 537ns 282ns 255ns 1.90x All three dicts are kept in sync by add_tablet, drop_tablets, and drop_tablets_by_host_id. The attrgetter imports are no longer needed and have been removed. Also drop the mutable {} class-level defaults this same change added (_first_tokens, _last_tokens), plus the pre-existing _tablets = {} they were modeled on: leaving mutable dicts at class scope is a latent shared-state hazard (e.g. if a future alternative constructor bypassed __init__), even though __init__ already reassigns them per instance today. All three are now instance-only, initialized solely in __init__. Add TabletsInstanceStateTest to lock this in: one check that the class body itself carries no such attributes, one check that two instances never observe each other's state. Signed-off-by: Yaniv Kaul --- cassandra/tablets.py | 62 ++++++++++++++++++++++++++------------ tests/unit/test_tablets.py | 25 +++++++++++++++ 2 files changed, 67 insertions(+), 20 deletions(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index 8190058901..5a5c4f6af8 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -1,14 +1,9 @@ from bisect import bisect_left -from operator import attrgetter from random import getrandbits from threading import Lock from typing import Optional from uuid import UUID -# C-accelerated attrgetter avoids per-call lambda allocation overhead -_get_first_token = attrgetter("first_token") -_get_last_token = attrgetter("last_token") - def choose_tablet_version_block(tablet_version: int) -> int: """ @@ -111,30 +106,46 @@ def get_replica_shard_id(self, uuid: UUID) -> Optional[int]: class Tablets(object): - _lock = None - _tablets = {} - def __init__(self, tablets): - self._tablets = tablets + # NOTE: these are intentionally instance attributes only (not class + # attributes) to avoid mutable class-level dicts being shared across + # instances, e.g. if a future alternative constructor were to bypass + # __init__. self._lock = Lock() + self._tablets = tablets + # Build parallel token index lists from any pre-populated data + # (keyspace, table) -> list[int] for both _first_tokens/_last_tokens + self._first_tokens = { + key: [t.first_token for t in tlist] + for key, tlist in tablets.items() + } + self._last_tokens = { + key: [t.last_token for t in tlist] + for key, tlist in tablets.items() + } def table_has_tablets(self, keyspace, table) -> bool: return bool(self._tablets.get((keyspace, table), [])) def get_tablet_for_key(self, keyspace, table, t): - tablet = self._tablets.get((keyspace, table), []) - if not tablet: + key = (keyspace, table) + last_tokens = self._last_tokens.get(key) + if not last_tokens: return None - id = bisect_left(tablet, t.value, key=_get_last_token) - if id < len(tablet) and t.value > tablet[id].first_token: - return tablet[id] + token_value = t.value + id = bisect_left(last_tokens, token_value) + if id < len(last_tokens) and token_value > self._first_tokens[key][id]: + return self._tablets[key][id] return None def drop_tablets(self, keyspace: str, table: Optional[str] = None): with self._lock: if table is not None: - self._tablets.pop((keyspace, table), None) + key = (keyspace, table) + self._tablets.pop(key, None) + self._first_tokens.pop(key, None) + self._last_tokens.pop(key, None) return to_be_deleted = [] @@ -144,6 +155,8 @@ def drop_tablets(self, keyspace: str, table: Optional[str] = None): for key in to_be_deleted: del self._tablets[key] + self._first_tokens.pop(key, None) + self._last_tokens.pop(key, None) def drop_tablets_by_host_id(self, host_id: Optional[UUID]): if host_id is None: @@ -157,23 +170,32 @@ def drop_tablets_by_host_id(self, host_id: Optional[UUID]): for tablet_id in reversed(to_be_deleted): tablets.pop(tablet_id) + self._first_tokens[key].pop(tablet_id) + self._last_tokens[key].pop(tablet_id) def add_tablet(self, keyspace, table, tablet): with self._lock: - tablets_for_table = self._tablets.setdefault((keyspace, table), []) + key = (keyspace, table) + tablets_for_table = self._tablets.setdefault(key, []) + first_tokens = self._first_tokens.setdefault(key, []) + last_tokens = self._last_tokens.setdefault(key, []) # find first overlapping range - start = bisect_left(tablets_for_table, tablet.first_token, key=_get_first_token) - if start > 0 and tablets_for_table[start - 1].last_token > tablet.first_token: + start = bisect_left(first_tokens, tablet.first_token) + if start > 0 and last_tokens[start - 1] > tablet.first_token: start = start - 1 # find last overlapping range - end = bisect_left(tablets_for_table, tablet.last_token, key=_get_last_token) - if end < len(tablets_for_table) and tablets_for_table[end].first_token >= tablet.last_token: + end = bisect_left(last_tokens, tablet.last_token) + if end < len(last_tokens) and first_tokens[end] >= tablet.last_token: end = end - 1 if start <= end: del tablets_for_table[start:end + 1] + del first_tokens[start:end + 1] + del last_tokens[start:end + 1] tablets_for_table.insert(start, tablet) + first_tokens.insert(start, tablet.first_token) + last_tokens.insert(start, tablet.last_token) diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index 8072c4e9f0..1f0333ce35 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -93,6 +93,31 @@ def test_add_tablet_intersecting_with_last(self): (-5011686018427387905, -2987529027641081857)]) +class TabletsInstanceStateTest(unittest.TestCase): + """Tests that Tablets' internal dicts are per-instance state, not + shared mutable class attributes (a well-known Python footgun).""" + + def test_internal_dicts_are_not_class_attributes(self): + self.assertNotIn('_tablets', vars(Tablets)) + self.assertNotIn('_first_tokens', vars(Tablets)) + self.assertNotIn('_last_tokens', vars(Tablets)) + + def test_instances_do_not_share_internal_dicts(self): + a = Tablets({}) + b = Tablets({}) + self.assertIsNot(a._tablets, b._tablets) + self.assertIsNot(a._first_tokens, b._first_tokens) + self.assertIsNot(a._last_tokens, b._last_tokens) + + t1 = Tablet(0, 100, [("host1", 0)]) + a.add_tablet("ks", "tb", t1) + # Mutating `a` must not be visible through `b`. + self.assertFalse(b.table_has_tablets("ks", "tb")) + self.assertEqual(b._tablets, {}) + self.assertEqual(b._first_tokens, {}) + self.assertEqual(b._last_tokens, {}) + + class GetTabletForKeyTest(unittest.TestCase): """Tests for Tablets.get_tablet_for_key.""" From a7bd30462cd5aeaeef35265904a70b11c82527c4 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Thu, 9 Apr 2026 20:32:21 +0300 Subject: [PATCH 6/8] perf: batch-filter drop_tablets_by_host_id instead of triple pop Replace the per-tablet reversed pop() loop (O(k*n) for each of three parallel lists) with a single-pass index filter that rebuilds the lists once. This avoids repeated list element shifting and scales better when many tablets are dropped at once. Benchmark (3 replicas/tablet, ~33% dropped): Tablets Old (triple-pop) New (batch-filter) Speedup 100 123 us 128 us ~1.0x 1,000 1,375 us 1,113 us 1.24x 10,000 25,429 us 13,079 us 1.94x Add 3 unit tests for drop_tablets_by_host_id covering matching, None host_id, and nonexistent host_id. Signed-off-by: Yaniv Kaul --- cassandra/tablets.py | 19 ++++++++++--------- tests/unit/test_tablets.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 9 deletions(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index 5a5c4f6af8..b0bafd3d26 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -163,15 +163,16 @@ def drop_tablets_by_host_id(self, host_id: Optional[UUID]): return with self._lock: for key, tablets in self._tablets.items(): - to_be_deleted = [] - for tablet_id, tablet in enumerate(tablets): - if tablet.replica_contains_host_id(host_id): - to_be_deleted.append(tablet_id) - - for tablet_id in reversed(to_be_deleted): - tablets.pop(tablet_id) - self._first_tokens[key].pop(tablet_id) - self._last_tokens[key].pop(tablet_id) + # Filter in one pass instead of popping one-by-one (O(n) vs O(k*n)) + keep = [i for i, t in enumerate(tablets) + if not t.replica_contains_host_id(host_id)] + if len(keep) == len(tablets): + continue # nothing to drop + self._tablets[key] = [tablets[i] for i in keep] + first = self._first_tokens[key] + last = self._last_tokens[key] + self._first_tokens[key] = [first[i] for i in keep] + self._last_tokens[key] = [last[i] for i in keep] def add_tablet(self, keyspace, table, tablet): with self._lock: diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index 1f0333ce35..ee4adf8424 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -412,3 +412,38 @@ def test_replica_dict_populated_as_expected(self): u2 = UUID('87654321-4321-8765-4321-876543218765') t = Tablet(0, 100, [(u1, 3), (u2, 7)]) self.assertEqual(t._replica_dict, {u1: 3, u2: 7}) + + +class DropTabletsByHostIdTest(unittest.TestCase): + """Tests for Tablets.drop_tablets_by_host_id batch-filter path.""" + + def test_drop_removes_matching_tablets(self): + u1 = UUID('12345678-1234-5678-1234-567812345678') + u2 = UUID('87654321-4321-8765-4321-876543218765') + t1 = Tablet(0, 100, [(u1, 0)]) + t2 = Tablet(100, 200, [(u2, 0)]) + t3 = Tablet(200, 300, [(u1, 1), (u2, 1)]) + tablets = Tablets({("ks", "tb"): [t1, t2, t3]}) + + tablets.drop_tablets_by_host_id(u1) + + remaining = tablets._tablets[("ks", "tb")] + self.assertEqual(len(remaining), 1) + self.assertIs(remaining[0], t2) + # Verify token index lists are in sync + self.assertEqual(tablets._first_tokens[("ks", "tb")], [100]) + self.assertEqual(tablets._last_tokens[("ks", "tb")], [200]) + + def test_drop_none_host_id_is_noop(self): + t1 = Tablet(0, 100, [("host1", 0)]) + tablets = Tablets({("ks", "tb"): [t1]}) + tablets.drop_tablets_by_host_id(None) + self.assertEqual(len(tablets._tablets[("ks", "tb")]), 1) + + def test_drop_nonexistent_host_id_is_noop(self): + u1 = UUID('12345678-1234-5678-1234-567812345678') + u_missing = UUID('aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee') + t1 = Tablet(0, 100, [(u1, 0)]) + tablets = Tablets({("ks", "tb"): [t1]}) + tablets.drop_tablets_by_host_id(u_missing) + self.assertEqual(len(tablets._tablets[("ks", "tb")]), 1) From 29119b5f168e75ae65efabffb789ec270cccdc89 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Tue, 21 Apr 2026 11:57:44 +0300 Subject: [PATCH 7/8] perf: avoid redundant tablet lookup in shard-aware connection selection When tablets are in use, get_tablet_for_key() was called twice per request: once in TokenAwarePolicy.make_query_plan() to find the replica, and again in HostConnection._get_connection_for_routing_key() to determine the shard. Stash the tablet found during query planning on the query object (query._tablet) and pass it through to borrow_connection(), which skips the second lookup when a tablet is already available. This eliminates redundant bisect_left calls and associated dict lookups. A Statement (e.g. BoundStatement) can be rebound and re-executed by the caller, so the same query object may be passed to make_query_plan() again with a different routing key. Clear query._tablet whenever the current lookup finds no tablet (or there is no routing key at all), so a tablet stashed for an earlier, unrelated routing key can't leak into shard-aware connection selection for the new one. Signed-off-by: Yaniv Kaul --- cassandra/cluster.py | 7 +++-- cassandra/policies.py | 14 +++++++++ cassandra/pool.py | 17 ++++++----- tests/unit/test_policies.py | 58 +++++++++++++++++++++++++++++++++++++ 4 files changed, 87 insertions(+), 9 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index bcc7852c33..f0e58864c4 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -5114,11 +5114,14 @@ def _query(self, host, message=None, cb=None): # TODO get connectTimeout from cluster settings if self.query: # Pass the ring token computed once for this request so the pool - # can select the shard without re-hashing the routing key. + # can select the shard without re-hashing the routing key, and + # the tablet found during query planning so the pool can skip a + # redundant lookup in the tablet map. connection, request_id = pool.borrow_connection( timeout=2.0, routing_key=self.query.routing_key, keyspace=self.query.keyspace, table=self.query.table, - routing_token=self._routing_token) + routing_token=self._routing_token, + tablet=getattr(self.query, '_tablet', None)) else: connection, request_id = pool.borrow_connection(timeout=2.0) self._connection = connection diff --git a/cassandra/policies.py b/cassandra/policies.py index fb0257c438..525d56c5b6 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -537,6 +537,11 @@ def make_query_plan(self, working_keyspace=None, query=None): child = self._child_policy if query is None or query.routing_key is None or keyspace is None: + if query is not None: + # A Statement (e.g. BoundStatement) can be rebound and + # re-executed by the caller; make sure a tablet stashed by + # an earlier, unrelated execution isn't picked up below. + query._tablet = None for host in child.make_query_plan(keyspace, query): yield host return @@ -557,6 +562,9 @@ def make_query_plan(self, working_keyspace=None, query=None): child_plan = child.make_query_plan(keyspace, query) replicas = [host for host in child_plan if host.host_id in replica_dict] + # Stash the tablet so that downstream shard-aware connection + # selection can reuse it instead of repeating the bisect lookup. + query._tablet = tablet # The leader concept only exists for strongly-consistent keyspaces, # which today means exactly the keyspaces whose consistency mode is @@ -596,6 +604,12 @@ def make_query_plan(self, working_keyspace=None, query=None): break else: replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key) + # Clear any tablet stashed by a previous execution of this same + # query object (statements may be rebound and reused, e.g. via + # BoundStatement.bind()) so a stale tablet -- for a different + # routing key -- isn't reused for shard-aware connection + # selection below. + query._tablet = None if self.shuffle_replicas and not query.is_lwt() and not ConsistencyLevel.is_serial(query.consistency_level): shuffle(replicas) diff --git a/cassandra/pool.py b/cassandra/pool.py index f2f2bf4405..176ce3710e 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -442,7 +442,7 @@ def __init__(self, host, host_distance, session): log.debug("Finished initializing connection for host %s", self.host) - def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None, routing_token=None): + def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None, routing_token=None, tablet=None): if self.is_shutdown: raise ConnectionException( "Pool for %s is shutdown" % (self.host,), self.host) @@ -463,10 +463,13 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table if t is None and metadata.token_map is not None and metadata.can_support_partitioner(): t = metadata.token_map.token_class.from_key(routing_key) if t is not None and self.supports_tablet_routing and table is not None: - if keyspace is None: - keyspace = self._keyspace + # Reuse the tablet found during query planning when available, + # avoiding a redundant bisect lookup in the tablet map. + if tablet is None: + if keyspace is None: + keyspace = self._keyspace - tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t) + tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t) # In both V1 and V2 the request is sent to this host, so we pick # the shard that this host owns for the tablet. Leader-aware host @@ -515,15 +518,15 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table return random.choice(active_connections) return random.choice(list(self._connections.values())) - def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None, routing_token=None): - conn = self._get_connection_for_routing_key(routing_key, keyspace, table, routing_token) + def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None, routing_token=None, tablet=None): + conn = self._get_connection_for_routing_key(routing_key, keyspace, table, routing_token, tablet) start = time.time() remaining = timeout last_retry = False while True: if conn.is_closed: # The connection might have been closed in the meantime - if so, try again - conn = self._get_connection_for_routing_key(routing_key, keyspace, table, routing_token) + conn = self._get_connection_for_routing_key(routing_key, keyspace, table, routing_token, tablet) with conn.lock: if (not conn.is_closed or last_retry) and conn.in_flight < conn.max_request_id: # On last retry we ignore connection status, since it is better to return closed connection than diff --git a/tests/unit/test_policies.py b/tests/unit/test_policies.py index 35c1a96f87..44b832c24a 100644 --- a/tests/unit/test_policies.py +++ b/tests/unit/test_policies.py @@ -1393,6 +1393,64 @@ def test_no_shuffle_for_serial_consistency(self, patched_shuffle): assert patched_shuffle.call_count == 0, \ "shuffle should not be called for consistency level %s" % cl + def test_stale_tablet_not_reused_across_query_plans(self): + """ + A Statement (e.g. a BoundStatement) may be rebound and re-executed by + the caller, so the same query object can be passed to + make_query_plan() multiple times with a different routing key each + time. Verify that a tablet stashed on the query object for shard-aware + connection selection (query._tablet) from one call doesn't leak into + a later call for which no tablet is found -- otherwise downstream + shard selection could pick a shard belonging to an unrelated, + previously-looked-up tablet. + """ + cluster = self._prepare_cluster_with_tablets() + hosts = cluster.metadata.all_hosts() + tablet = cluster.metadata._tablets.get_tablet_for_key.return_value + + child_policy = Mock() + child_policy.make_query_plan.return_value = hosts + child_policy.distance.return_value = HostDistance.LOCAL + + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + + query = Statement(routing_key='routing_key', keyspace='keyspace') + list(policy.make_query_plan('keyspace', query)) + self.assertIs(query._tablet, tablet) + + # Same (reused) query object, but this time no tablet is found for + # the (new) routing key -- e.g. it hasn't been discovered yet, or + # the table isn't tablets-based. + cluster.metadata._tablets.get_tablet_for_key.return_value = None + list(policy.make_query_plan('keyspace', query)) + self.assertIsNone(query._tablet) + + def test_stale_tablet_not_reused_when_no_routing_key(self): + """ + Same as above, but covers the early-return path (no routing key / + no keyspace), which must also clear any previously stashed tablet. + """ + cluster = self._prepare_cluster_with_tablets() + hosts = cluster.metadata.all_hosts() + tablet = cluster.metadata._tablets.get_tablet_for_key.return_value + + child_policy = Mock() + child_policy.make_query_plan.return_value = hosts + child_policy.distance.return_value = HostDistance.LOCAL + + policy = TokenAwarePolicy(child_policy, shuffle_replicas=False) + policy.populate(cluster, hosts) + + query = Statement(routing_key='routing_key', keyspace='keyspace') + list(policy.make_query_plan('keyspace', query)) + self.assertIs(query._tablet, tablet) + + # Reuse the same statement without a routing key this time. + query.routing_key = None + list(policy.make_query_plan('keyspace', query)) + self.assertIsNone(query._tablet) + class ConvictionPolicyTest(unittest.TestCase): def test_not_implemented(self): From 98747ad01dc91fa4c1e70c887e68b7d9a767a228 Mon Sep 17 00:00:00 2001 From: Yaniv Michael Kaul Date: Tue, 30 Jun 2026 17:23:17 +0300 Subject: [PATCH 8/8] fix: update response_future tests for borrow_connection tablet= kwarg The 6th perf commit added a tablet= keyword argument to borrow_connection. Update the 6 mock assertions in test_response_future.py to expect the new parameter. Signed-off-by: Yaniv Kaul --- tests/unit/test_response_future.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index d71943ec04..3cb8a83005 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -98,8 +98,7 @@ def test_result_message(self): rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) - + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY, tablet=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) expected_result = (object(), object()) @@ -292,7 +291,7 @@ def test_retry_policy_says_retry(self): rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY, tablet=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) result = Mock(spec=UnavailableErrorMessage, info={}) @@ -311,7 +310,7 @@ def test_retry_policy_says_retry(self): # it should try again with the same host since this was # an UnavailableException rf.session._pools.get.assert_called_with(host) - pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) + pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY, tablet=ANY) connection.send_msg.assert_called_with(rf.message, 2, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) def test_retry_with_different_host(self): @@ -326,7 +325,7 @@ def test_retry_with_different_host(self): rf.send_request() rf.session._pools.get.assert_called_once_with('ip1') - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY, tablet=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) assert ConsistencyLevel.QUORUM == rf.message.consistency_level @@ -345,7 +344,7 @@ def test_retry_with_different_host(self): # it should try with a different host rf.session._pools.get.assert_called_with('ip2') - pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) + pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY, tablet=ANY) connection.send_msg.assert_called_with(rf.message, 2, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) # the consistency level should be the same @@ -1062,7 +1061,7 @@ def test_single_host_query_plan_exhausted_after_one_retry(self): # Verify initial request was sent rf.session._pools.get.assert_called_once_with(specific_host) - pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=ANY, routing_token=ANY, tablet=ANY) connection.send_msg.assert_called_once_with(rf.message, 1, cb=ANY, encoder=ProtocolHandler.encode_message, decoder=ProtocolHandler.decode_message, result_metadata=[]) # Simulate a ServerError response (which triggers RETRY_NEXT_HOST by default)