(improvement)Optimize DCAwareRoundRobinPolicy with host distance caching

Refactor `DCAwareRoundRobinPolicy` to use a Copy-On-Write (COW) strategy for managing host distances.

Key changes:
- Introduce `_hosts_by_distance` to cache `LOCAL` and `REMOTE` hosts, enabling O(1) distance lookups and thread-safe iteration during query planning.
`IGNORED` hosts do not need to be stored in the cache.
- Optimize control plane operations (`on_up`, `on_down`):
  - Add `_add_local_distance` and `_remove_local_distance` to handle local DC node changes via atomic dictionary swaps, avoiding full cache rebuilds.
  - Update `_refresh_distances` to support partial refreshes, reusing existing `LOCAL` or `REMOTE` structures when only one DC changes.

Signed-off-by: Yaniv Kaul <yaniv.kaul@scylladb.com>

Author: mykaul

cassandra/policies.py

@@ -244,34 +244,81 @@ def __init__(self, local_dc='', used_hosts_per_remote_dc=0):
         self.local_dc = local_dc
         self.used_hosts_per_remote_dc = used_hosts_per_remote_dc
         self._dc_live_hosts = {}
+        self._hosts_by_distance = {
+            HostDistance.LOCAL: {},
+            HostDistance.REMOTE: {}
+        }
         self._position = 0
         LoadBalancingPolicy.__init__(self)
 
     def _dc(self, host):
         return host.datacenter or self.local_dc
 
+    def _refresh_distances(self, dc=None):
+        # Always called under self._hosts_lock, so it's safe to iterate over self._dc_live_hosts
+
+        if dc is None:
+            # Full rebuild
+            local_hosts = dict.fromkeys(self._dc_live_hosts.get(self.local_dc, ()))
+            remote_hosts = {}
+            if self.used_hosts_per_remote_dc > 0:
+                for datacenter, hosts in self._dc_live_hosts.items():
+                    if datacenter != self.local_dc:
+                        remote_hosts.update(dict.fromkeys(hosts[:self.used_hosts_per_remote_dc]))
+        elif dc == self.local_dc:
+            # Rebuild local, reuse remote
+            local_hosts = dict.fromkeys(self._dc_live_hosts.get(self.local_dc, ()))
+            remote_hosts = self._hosts_by_distance[HostDistance.REMOTE]
+        else:
+            # Rebuild remote, reuse local
+            local_hosts = self._hosts_by_distance[HostDistance.LOCAL]
+            remote_hosts = {}
+            if self.used_hosts_per_remote_dc > 0:
+                for datacenter, hosts in self._dc_live_hosts.items():
+                    if datacenter != self.local_dc:
+                        remote_hosts.update(dict.fromkeys(hosts[:self.used_hosts_per_remote_dc]))
+
+        self._hosts_by_distance = {
+            HostDistance.LOCAL: local_hosts,
+            HostDistance.REMOTE: remote_hosts
+        }
+
+    def _add_local_distance(self, host):
+        distances = self._hosts_by_distance
+        if host not in distances[HostDistance.LOCAL]:
+            new_local = distances[HostDistance.LOCAL].copy()
+            new_local[host] = None
+            self._hosts_by_distance = {
+                HostDistance.LOCAL: new_local,
+                HostDistance.REMOTE: distances[HostDistance.REMOTE]
+            }
+
+    def _remove_local_distance(self, host):
+        distances = self._hosts_by_distance
+        if host in distances[HostDistance.LOCAL]:
+            new_local = distances[HostDistance.LOCAL].copy()
+            del new_local[host]
+            self._hosts_by_distance = {
+                HostDistance.LOCAL: new_local,
+                HostDistance.REMOTE: distances[HostDistance.REMOTE]
+            }
+
     def populate(self, cluster, hosts):
         for dc, dc_hosts in groupby(hosts, lambda h: self._dc(h)):
             self._dc_live_hosts[dc] = tuple({*dc_hosts, *self._dc_live_hosts.get(dc, [])})
 
         self._position = randint(0, len(hosts) - 1) if hosts else 0
+        self._refresh_distances()
 
     def distance(self, host):
-        dc = self._dc(host)
-        if dc == self.local_dc:
+        # Capture the reference to the dict to ensure we check a consistent snapshot
+        # (avoiding the case where we check LOCAL, fail, context switch, update happens, check REMOTE, fail)
+        distances = self._hosts_by_distance
+        if host in distances[HostDistance.LOCAL]:
             return HostDistance.LOCAL
-
-        if not self.used_hosts_per_remote_dc:
-            return HostDistance.IGNORED
-        else:
-            dc_hosts = self._dc_live_hosts.get(dc)
-            if not dc_hosts:
-                return HostDistance.IGNORED
-
-            if host in list(dc_hosts)[:self.used_hosts_per_remote_dc]:
-                return HostDistance.REMOTE
-            else:
-                return HostDistance.IGNORED
+        elif host in distances[HostDistance.REMOTE]:
+            return HostDistance.REMOTE
+        return HostDistance.IGNORED
 
     def make_query_plan(self, working_keyspace=None, query=None):
         # not thread-safe, but we don't care much about lost increments
@@ -284,29 +331,34 @@ def make_query_plan(self, working_keyspace=None, query=None):
         for host in islice(cycle(local_live), pos, pos + len(local_live)):
             yield host
 
-        # the dict can change, so get candidate DCs iterating over keys of a copy
-        other_dcs = [dc for dc in self._dc_live_hosts.copy().keys() if dc != self.local_dc]
-        for dc in other_dcs:
-            remote_live = self._dc_live_hosts.get(dc, ())
-            for host in remote_live[:self.used_hosts_per_remote_dc]:
-                yield host
+        for host in self._hosts_by_distance[HostDistance.REMOTE]:
+            yield host
 
     def on_up(self, host):
         # not worrying about threads because this will happen during
         # control connection startup/refresh
+        refresh_all = False
         if not self.local_dc and host.datacenter:
             self.local_dc = host.datacenter
             log.info("Using datacenter '%s' for DCAwareRoundRobinPolicy (via host '%s'); "
                         "if incorrect, please specify a local_dc to the constructor, "
                         "or limit contact points to local cluster nodes" %
                         (self.local_dc, host.endpoint))
+            refresh_all = True
 
         dc = self._dc(host)
         with self._hosts_lock:
             current_hosts = self._dc_live_hosts.get(dc, ())
             if host not in current_hosts:
                 self._dc_live_hosts[dc] = current_hosts + (host, )
 
+            if refresh_all:
+                self._refresh_distances()
+            elif dc == self.local_dc:
+                self._add_local_distance(host)
+            else:
+                self._refresh_distances(dc)
+
     def on_down(self, host):
         dc = self._dc(host)
         with self._hosts_lock:
@@ -318,6 +370,11 @@ def on_down(self, host):
                 else:
                     del self._dc_live_hosts[dc]
 
+                if dc == self.local_dc:
+                     self._remove_local_distance(host)
+                else:
+                    self._refresh_distances(dc)
+
     def on_add(self, host):
         self.on_up(host)