perf: optimize Tablet memory layout and per-query lookup speed - #812
perf: optimize Tablet memory layout and per-query lookup speed#812mykaul wants to merge 8 commits into
Conversation
2640c18 to
1d25663
Compare
3d5cfbd to
601f065
Compare
601f065 to
6c9a4b0
Compare
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
b5a23f2 to
364ad45
Compare
364ad45 to
f91d7d9
Compare
|
Rebased onto current Replica-ordering check (relevant to the LWT Paxos-leader routing fix in #782 /
Bug found and fixed during self-review: the "avoid redundant tablet lookup" commit stashed Tests: CI: all checks were green on the pre-rebase HEAD ( Force-pushed the rebased + amended branch; still a draft. |
There was a problem hiding this comment.
Pull request overview
Note
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Optimizes the tablet-related hot paths by reducing Tablet memory overhead and speeding up per-query tablet/replica lookups, while plumbing a cached tablet through query planning into shard-aware connection selection.
Changes:
- Add
__slots__, store replicas as tuples, and cache a host→shard lookup dict onTablet - Speed up tablet lookups with parallel
first_token/last_tokenindex lists and streamlinefrom_row - Reuse planned tablet during connection selection and update/extend unit tests accordingly
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| cassandra/tablets.py | Core optimizations: slots, replica dict cache, token index lists, faster bisect-based lookup and updates |
| cassandra/policies.py | Stashes/clears query._tablet to enable reuse downstream and avoids extra replica mapping work |
| cassandra/pool.py | Accepts an optional tablet to skip redundant tablet lookups and do O(1) shard selection |
| cassandra/cluster.py | Plumbs query._tablet into pool.borrow_connection() |
| tests/unit/test_tablets.py | Adds tests for _replica_dict, iterator replicas, tuple storage, and token index sync on drops |
| tests/unit/test_policies.py | Adds regression tests ensuring stale tablets aren’t reused across executions/query plans |
| tests/unit/test_response_future.py | Updates mocks to match the new borrow_connection(..., tablet=...) call shape |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
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%)
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
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.
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).
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 scylladb#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.
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.
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.
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.
f91d7d9 to
49fa146
Compare
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Suppressed comments (1)
cassandra/tablets.py:80
get_tablet_for_key()now combines three independently mutable structures without taking_lock. Writers update them in separate steps: for example,drop_tablets_by_host_id()replaces_tablets[key]at line 112 before replacing the token lists, whiledrop_tablets()can remove the token-list keys after this method has capturedlast_tokens. A concurrent topology update can therefore return a tablet from a different range or raiseKeyError/IndexError. Keep each table's tablets and indexes in one snapshot that writers replace atomically, or synchronize reads with all writer updates.
if id < len(last_tokens) and token_value > self._first_tokens[key][id]:
return self._tablets[key][id]
| # Stash the tablet so that downstream shard-aware | ||
| # connection selection can reuse it instead of | ||
| # repeating the bisect lookup. | ||
| query._tablet = tablet |
Summary
Five incremental optimizations to the tablets hot path, each in a separate commit:
__slots__onTablet— eliminates per-instance__dict__allocationtuplereplicas — replicas are immutable after creation; use tuple instead of list_replica_dict— build{host_id: shard_id}dict once at construction, use it in both per-query hot paths (policies.pyandpool.py). Also fixes a latent iterator-consumption bug where passing a generator toTablet()would silently produce an empty_replica_dict.from_row— inline the_is_valid_tabletcheck to eliminate a staticmethod descriptor lookup, an extra function call, and a redundantis not Noneguard_first_tokensand_last_tokensas plainlist[int]dicts alongside_tablets, sobisect_leftruns entirely in C on native ints instead of calling anattrgettercallback per comparison. Follow-up to PR perf: use stdlib bisect and attrgetter in tablets.py (100's of ns, 1.5-5.6x speedup) #757 which identified this opportunity in its own benchmarks.Benchmarks (Python 3.14, best-of-5 rounds)
Memory per Tablet (deep,
pympler.asizeof, 3 replicas)__slots__,list, no dict)__slots__)tuple)_replica_dict)Commits 1+2 save 472 bytes/tablet. Commit 3 spends 232 bytes back on the
_replica_dictcache. Net: 240 bytes saved per tablet (13%). Commit 5 adds 16 bytes/tablet (two ints in parallel lists) — negligible.Shallow breakdown (sys.getsizeof)
__dict___replica_dict(3 entries)get_tablet_for_key(hit — the primary per-query hot path)Miss path (N=1000): 458 ns -> 229 ns (2.0x).
Other per-query hot paths
policies.py:set(map(lambda r: r[0], tablet.replicas))tablet._replica_dict)pool.py: linear scan for shard_iddict.get)replica_contains_host_idinConstruction (
Tablet.from_row)_replica_dict+ oldfrom_row)from_row)Commit 4 recovers ~54 ns (12%) of the construction regression. The remaining +267 ns vs master is the irreducible cost of building
tuple()+dict()at construction time (~250 ns), which pays for itself on every query.Tests
All 223 unit tests pass (tablets, policies, pool, metadata, cluster, response_future). 7 new tests added for
_replica_dictbehavior, including an iterator edge-case regression test.