(improvement)Optimize DCAware/RackAware/TokenAware/HostFilter policies with host distance caching and overall perf. improvements#651
Open
mykaul wants to merge 7 commits intoscylladb:masterfrom
Open
Conversation
Author
|
This is interesting, my change has exposed this - Need to understand this better :-/ |
Author
if host is None:
host = self._cluster.metadata.get_host_by_host_id(host_id)
if host and host.endpoint != endpoint:
log.debug("[control connection] Updating host ip from %s to %s for (%s)", host.endpoint, endpoint, host_id)
old_endpoint = host.endpoint
host.endpoint = endpoint
self._cluster.metadata.update_host(host, old_endpoint)
reconnector = host.get_and_set_reconnection_handler(None)
if reconnector:
reconnector.cancel()
self._cluster.on_down(host, is_host_addition=False, expect_host_to_be_down=True)So first we update the host with the new endpoint, then mark it as down? |
Author
|
This fixes it for me: which also makes sense to me. |
mykaul
force-pushed
the
query_plan_opt_2
branch
from
76ee195 to
edab823
Compare
Author
|
I think CI failure is unrelated and is #359 |
mykaul
force-pushed
the
query_plan_opt_2
branch
from
edab823 to
1884f59
Compare
Author
|
By using the (not amazing) benchmark from #653 , I got the following results: For master branch as a baseline: This branch (with just DC aware improvements): ** 433 -> 781 Kops/sec improvement ** With improvement to rack aware (on top of master), I got: ** 277 -> 324 Kops/sec improvement ** And on top of this branch: ** 277 -> 344 Kops/sec improvement ** And finally, for #650 : which kinda makes me suspect that branch is no good :-/ |
Author
Sent separate PR - #654 |
mykaul
changed the title
Author
|
With rack aware added (3rd commit), these are the current numbers: |
mykaul
force-pushed
the
query_plan_opt_2
branch
from
cc0204d to
6282e6f
Compare
Author
Now that I also cache non-local hosts, not just remote (duh!), perf. is better: |
Author
|
Added for TokenAware as well some optimization (need to improve commit message). So reasonable improvement, at least in this micro-benchmark. |
mykaul
force-pushed
the
query_plan_opt_2
branch
from
8f96d39 to
87c6a01
Compare
8 tasks
mykaul
changed the title
Author
|
Last push, I think I'm done: |
mykaul
changed the title
mykaul
force-pushed
the
query_plan_opt_2
branch
2 times, most recently
from
5f283d1 to
bd6a9c5
Compare
Author
|
Latest numbers: |
There was a problem hiding this comment.
Pull request overview
This PR optimizes load balancing policies (DCAwareRoundRobinPolicy, RackAwareRoundRobinPolicy, TokenAwarePolicy, HostFilterPolicy) with host distance caching and general performance improvements. The key insight is caching computed host distance data (remote hosts, non-local-rack hosts) and replica lookups to avoid repeated computation in the hot query-planning path.
Changes:
- Introduce
_remote_hosts(COW dict) onDCAwareRoundRobinPolicyandRackAwareRoundRobinPolicyfor O(1) distance lookups, plus_non_local_rack_hostsfor rack-aware iteration; both refreshed on topology changes. - Add an LRU replica cache to
TokenAwarePolicy(keyed by(keyspace, routing_key), invalidated on token map changes) and restructuremake_query_planto use direct distance bucketing instead of repeatedyield_in_orderscans. - Add
make_query_plan_with_exclusionAPI toLoadBalancingPolicyand all subclasses, enablingTokenAwarePolicyto skip already-yielded replicas when querying the child policy for remaining hosts.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
cassandra/policies.py |
Core optimization: distance caching in DC/RackAware policies, LRU replica cache in TokenAwarePolicy, new make_query_plan_with_exclusion method across policies, formatting cleanup |
tests/unit/test_policies.py |
New tests for make_query_plan_with_exclusion, replica cache (hit/miss/eviction/invalidation/disabled), LWT determinism, tablet bypass; test mocks updated for new token_map-based replica resolution; formatting cleanup |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
This PR optimizes the load balancing policies in the ScyllaDB Python driver by introducing distance caching via Copy-On-Write (COW) strategy, an LRU cache for token-to-replica lookups in TokenAwarePolicy, and a new make_query_plan_with_exclusion API to avoid redundant iteration.
Changes:
- Introduces
_remote_hostsand_non_local_rack_hostscached dictionaries/lists inDCAwareRoundRobinPolicyandRackAwareRoundRobinPolicyfor O(1) distance lookups, rebuilt on topology changes. - Adds an LRU cache (
OrderedDict) inTokenAwarePolicyfor token-to-replica lookups, invalidated by token_map object identity change; includes LWT deterministic ordering (no shuffle). - Adds
make_query_plan_with_exclusion()toLoadBalancingPolicyand its subclasses to efficiently skip already-yielded replicas inTokenAwarePolicy.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
cassandra/policies.py |
Core performance optimizations: COW distance caching, LRU replica cache, make_query_plan_with_exclusion API, LWT shuffle skip, code formatting |
tests/unit/test_policies.py |
New tests for exclusion-based query planning, LRU cache behavior, LWT determinism, cache invalidation, and test infrastructure updates for new mock requirements |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
mykaul
added a commit
to mykaul/python-driver
that referenced
this pull request
Mar 25, 2026
Add micro-benchmarks measuring query plan generation throughput for DCAwareRoundRobinPolicy, RackAwareRoundRobinPolicy, TokenAwarePolicy, DefaultLoadBalancingPolicy, and HostFilterPolicy. Uses pytest-benchmark for accurate timing and statistical reporting with a simulated 45-node cluster topology (5 DCs x 3 racks x 3 nodes) and 100,000 deterministic queries. Also rename tests/integration/standard/column_encryption/test_policies.py to test_encrypted_policies.py to avoid module name conflicts when running the full test suite. Run with: pytest -m benchmark tests/performance/ Benchmark results comparing master vs PR scylladb#651 optimizations (Python 3.14.3, pytest-benchmark 5.2.3, GC disabled, median): Policy | master (Kops/s) | PR#651 (Kops/s) | Speedup --------------------------|-----------------|-----------------|-------- DCAware | 833 | 1898 | 2.3x RackAware | 542 | 1589 | 2.9x TokenAware(DCAware) | 135 | 572 | 4.2x TokenAware(RackAware) | 123 | 539 | 4.4x Default(DCAware) | 674 | 1257 | 1.9x HostFilter(DCAware) | 394 | 579 | 1.5x
This was referenced Apr 1, 2026
Introduce _remote_hosts dict to cache REMOTE hosts, enabling O(1) distance lookups instead of scanning per-DC host lists. Replace islice(cycle(...)) with index arithmetic in make_query_plan. Call _refresh_remote_hosts() on topology changes.
…dRobin, and DCAware Add a new make_query_plan_with_exclusion() method that skips hosts in an exclusion set. The base class provides a default filtering implementation; RoundRobin and DCAware override for efficiency.
Cache remote hosts and non-local-rack hosts to enable O(1) distance lookups. Replace islice(cycle(...)) with index arithmetic. Reorder on_up/on_down to update DC-level hosts before rack-level for correct cache invalidation.
Optimized exclusion-aware query plan that avoids re-computing non-local-rack and remote host lists.
- Add LRU cache (default 1024 entries) for token-to-replicas lookups, auto-invalidated on topology changes (token_map identity check). - Sort replicas by distance (LOCAL_RACK > LOCAL > REMOTE) in a single pass instead of iterating three times. - Skip distance re-sorting for DCAware/RackAware child policies since they already yield in distance order; fallback re-sort for others. - LWT queries skip replica shuffling for deterministic plans. - Use make_query_plan_with_exclusion to avoid re-yielding replicas.
…ultLoadBalancingPolicy Both delegate to their child policy's exclusion-aware query plan while preserving their specific filtering/targeting behavior.
…erminism Add tests for make_query_plan_with_exclusion in RoundRobin, DCAware, and RackAware policies. Add cache tests (hit, miss, eviction, topology invalidation, disabled) and LWT determinism tests for TokenAwarePolicy. Update existing tests to set up token_map mocks and shuffle_replicas=False to match the new TokenAwarePolicy implementation.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refactor
DCAwareRoundRobinPolicyto use a Copy-On-Write (COW) strategy for managing host distances.Results (5 DCs × 3 racks × 3 nodes = 45 nodes, 100K queries, median of 5 iterations):
Key changes:
_remote_hoststo cacheREMOTEhosts, enabling O(1) distance lookups during query planning for distance.IGNOREDhosts do not need to be stored in the cache.For 'LOCAL' we do a simple comparison.
_refresh_remote_hoststo handle node changes.TokenAwarePolicy(default 1024 entries, auto-invalidated on topology change).TokenAwarePolicyskips distance re-sorting for DCAware/RackAware child policies (they already yield in distance order), with a fallback re-sort for other child policies.TokenAwarePolicyno longer uses__slots__to avoid breaking downstream subclasses.This is a different attempt from #650 to add caching to host distance to make query planning faster.
Pre-review checklist
./docs/source/.Fixes:annotations to PR description.