From f74be25d1603de2b6d160a7ad06d26951f4eefe9 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%) --- cassandra/tablets.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index 96e61a50c2..4250acff4b 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -15,9 +15,7 @@ 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 + __slots__ = ('first_token', 'last_token', 'replicas') def __init__(self, first_token=0, last_token=0, replicas=None): self.first_token = first_token From e183cc92f5c2c87e023d8540315400158163e5d8 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 --- cassandra/tablets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index 4250acff4b..2c90cf02b3 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -20,7 +20,7 @@ class Tablet(object): def __init__(self, first_token=0, last_token=0, replicas=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 def __str__(self): return "" \ From 079236aa44e28b52c6169adca8bc2ecb2a4244e3 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. --- cassandra/policies.py | 4 +-- cassandra/pool.py | 5 +-- cassandra/tablets.py | 18 +++++++---- tests/unit/test_tablets.py | 66 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 81 insertions(+), 12 deletions(-) diff --git a/cassandra/policies.py b/cassandra/policies.py index 89702e8c89..7e77ea5df2 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -507,10 +507,10 @@ def make_query_plan(self, working_keyspace=None, query=None): 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)) + 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] else: replicas = self._cluster_metadata.get_replicas(keyspace, query.routing_key) diff --git a/cassandra/pool.py b/cassandra/pool.py index 176751f60a..68ce80e4e8 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -460,10 +460,7 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t) 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: shard_id = self.host.sharding_info.shard_id_from_token(t.value) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index 2c90cf02b3..c3114c3b15 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -15,12 +15,18 @@ class Tablet(object): It stores information about each replica, its host and shard, and the token interval in the format (first_token, last_token]. """ - __slots__ = ('first_token', 'last_token', 'replicas') + __slots__ = ('first_token', 'last_token', 'replicas', '_replica_dict') def __init__(self, first_token=0, last_token=0, replicas=None): self.first_token = first_token self.last_token = last_token - self.replicas = tuple(replicas) if replicas is not None else None + if replicas is not None: + replicas_tuple = tuple(replicas) + self.replicas = replicas_tuple + self._replica_dict = {r[0]: r[1] for r in replicas_tuple} + else: + self.replicas = None + self._replica_dict = {} def __str__(self): return "" \ @@ -39,10 +45,10 @@ def from_row(first_token, last_token, replicas): return None 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 7a40e7de4d..fa47a5a1a3 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -1,4 +1,5 @@ import unittest +from uuid import UUID from cassandra.tablets import Tablets, Tablet @@ -124,3 +125,68 @@ 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 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 f292effc0498db595df67cedf542543bb8ac2e30 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). --- cassandra/tablets.py | 16 +++++++------- tests/unit/test_tablets.py | 44 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 8 deletions(-) diff --git a/cassandra/tablets.py b/cassandra/tablets.py index c3114c3b15..04c1341353 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -33,16 +33,16 @@ def __str__(self): % (self.first_token, self.last_token, self.replicas) __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): - if Tablet._is_valid_tablet(replicas): - tablet = Tablet(first_token, last_token, replicas) - 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 + return Tablet(first_token, last_token, replicas_tuple) def replica_contains_host_id(self, uuid: UUID) -> bool: return uuid in self._replica_dict diff --git a/tests/unit/test_tablets.py b/tests/unit/test_tablets.py index fa47a5a1a3..57ee0fb3bd 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -127,6 +127,50 @@ def __init__(self, v): self.assertIsNone(tablets.get_tablet_for_key("ks", "tb", Token(50))) +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 cfb96157a560f34a049ee42f471b3d1407f9aa9f 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. --- 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 04c1341353..0cdad3d220 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -1,13 +1,8 @@ from bisect import bisect_left -from operator import attrgetter 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") - class Tablet(object): """ @@ -52,30 +47,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 = [] @@ -85,6 +96,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: @@ -98,23 +111,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 57ee0fb3bd..4cbd60459b 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -89,6 +89,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 01b77e21e75e8f5fbda716d7892157e98afb2e05 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. --- 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 0cdad3d220..eb7d2025b5 100644 --- a/cassandra/tablets.py +++ b/cassandra/tablets.py @@ -104,15 +104,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 4cbd60459b..27496361ef 100644 --- a/tests/unit/test_tablets.py +++ b/tests/unit/test_tablets.py @@ -259,3 +259,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 0a34aa5d400e54dc7827cf5d37a7080d6e155a6c 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. --- cassandra/cluster.py | 5 +++- cassandra/policies.py | 15 ++++++++++ cassandra/pool.py | 23 +++++++++------ tests/unit/test_policies.py | 58 +++++++++++++++++++++++++++++++++++++ 4 files changed, 91 insertions(+), 10 deletions(-) diff --git a/cassandra/cluster.py b/cassandra/cluster.py index 88c8d2707a..e959ef9cdb 100644 --- a/cassandra/cluster.py +++ b/cassandra/cluster.py @@ -5033,7 +5033,10 @@ def _query(self, host, message=None, cb=None): try: # TODO get connectTimeout from cluster settings if self.query: - connection, request_id = pool.borrow_connection(timeout=2.0, routing_key=self.query.routing_key, keyspace=self.query.keyspace, table=self.query.table) + connection, request_id = pool.borrow_connection( + timeout=2.0, routing_key=self.query.routing_key, + keyspace=self.query.keyspace, table=self.query.table, + 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 7e77ea5df2..bb32ea1b82 100644 --- a/cassandra/policies.py +++ b/cassandra/policies.py @@ -498,6 +498,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 @@ -511,8 +516,18 @@ 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 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 68ce80e4e8..87f709712a 100644 --- a/cassandra/pool.py +++ b/cassandra/pool.py @@ -440,7 +440,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): + def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table=None, tablet=None): if self.is_shutdown: raise ConnectionException( "Pool for %s is shutdown" % (self.host,), self.host) @@ -454,13 +454,18 @@ def _get_connection_for_routing_key(self, routing_key=None, keyspace=None, table shard_id = None if self.tablets_routing_v1 and table is not None: - if keyspace is None: - keyspace = self._keyspace - - tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t) - + # Reuse tablet from query planning if available, avoiding + # a redundant bisect lookup in the tablet map. if tablet is not None: shard_id = tablet._replica_dict.get(self.host.host_id) + else: + if keyspace is None: + keyspace = self._keyspace + + tablet = self._session.cluster.metadata._tablets.get_tablet_for_key(keyspace, table, t) + + if tablet is not None: + shard_id = tablet._replica_dict.get(self.host.host_id) if shard_id is None: shard_id = self.host.sharding_info.shard_id_from_token(t.value) @@ -503,15 +508,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): - conn = self._get_connection_for_routing_key(routing_key, keyspace, table) + def borrow_connection(self, timeout, routing_key=None, keyspace=None, table=None, tablet=None): + conn = self._get_connection_for_routing_key(routing_key, keyspace, table, 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) + conn = self._get_connection_for_routing_key(routing_key, keyspace, table, 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 63a3c3d12d..a3db9c890b 100644 --- a/tests/unit/test_policies.py +++ b/tests/unit/test_policies.py @@ -972,6 +972,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 49fa1462701e4ac4797e346b87eb6bb51940d8bf 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. --- tests/unit/test_response_future.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_response_future.py b/tests/unit/test_response_future.py index cf1194a91f..c08d7a8b11 100644 --- a/tests/unit/test_response_future.py +++ b/tests/unit/test_response_future.py @@ -94,7 +94,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) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=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=[]) @@ -285,7 +285,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) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=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={}) @@ -304,7 +304,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) + pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=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): @@ -319,7 +319,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) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=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 @@ -338,7 +338,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) + pool.borrow_connection.assert_called_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=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 @@ -1055,7 +1055,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) + pool.borrow_connection.assert_called_once_with(timeout=ANY, routing_key=ANY, keyspace=ANY, table=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)