feat: Optimize TokenAwarePolicy with thread-safe distance caching #650
Conversation
Implement distance caching in TokenAwarePolicy to eliminate distance calculations on every query execution. Implementation: - Added _hosts_by_distance cache structure grouping hosts by distance - Added _get_child_lock() helper to be thread-safe for all cache operations - Incremental cache updates: in host additions/removals only - Cache population happens once during policy initialization - Copy-on-read snapshots in hot path to avoid holding locks during iteration Performance improvements: - Before: distance() calculation called for every host on every query - After: distance() called only during topology changes and cache population Additional optimizations: - Added __slots__ for reduced memory overhead - Mock compatibility for unit tests with defensive programming The optimization is especially beneficial for datacenter-aware topologies where distance calculations involve string comparisons and dictionary lookups. Signed-off-by: Yaniv Kaul <[email protected]>
| # yield rest of the cluster: local_rack, local, remote | ||
| yield from yield_in_order([host for host in child.make_query_plan(keyspace, query) if host not in replicas]) | ||
|
|
||
| def _populate_distance_cache(self): |
There was a problem hiding this comment.
Cache needs to be implemented on child policy, not on TokenAwarePolicy.
This is crazy way to deal with the problem )
There was a problem hiding this comment.
I was pondering this myself, but couldn't convince myself is the right thing to do - why not do it once, vs. re-implementing per child policy?
There was a problem hiding this comment.
because only child policy knows when to invalidate it and if you do it there whole change will be 32 lines of code with simple lookup in map and delete from it on invalidation.
There was a problem hiding this comment.
OK, I can take a look. From what I've seen there was a completely wasteful work on adding/removing hosts when they were up or down - which is completely useless, as the query plan needs to eventually check if the replica is up anyway.
But I'll give it a shot and see if it makes it easier to change the child policies. I'm quite afraid it's going to be a bigger change, but we'll see.
2 participants
Implement distance caching in TokenAwarePolicy to eliminate distance calculations on every query execution.
Implementation:
Performance improvements:
Additional optimizations:
The optimization is especially beneficial for datacenter-aware topologies where distance calculations involve string comparisons and dictionary lookups.
Pre-review checklist
./docs/source/.Fixes:annotations to PR description.