feat: add make_query_plan_with_exclusion to RackAwareRoundRobinPolicy

Optimized exclusion-aware query plan that avoids re-computing
non-local-rack and remote host lists.

Author: mykaul

cassandra/policies.py

@@ -503,6 +503,56 @@ def make_query_plan(self, working_keyspace=None, query=None):
         for host in remote_hosts:
             yield host
 
+    def make_query_plan_with_exclusion(self, working_keyspace=None, query=None, excluded=()):
+        pos = self._position
+        self._position += 1
+
+        local_rack_live = self._live_hosts.get((self.local_dc, self.local_rack), ())
+        length = len(local_rack_live)
+        remote_hosts = self._remote_hosts
+        if not excluded:
+            if length:
+                p = pos % length
+                for i in range(length):
+                    yield local_rack_live[(p + i) % length]
+
+            local_non_rack = self._non_local_rack_hosts
+            length = len(local_non_rack)
+            if length:
+                p = pos % length
+                for i in range(length):
+                    yield local_non_rack[(p + i) % length]
+
+            for host in remote_hosts:
+                yield host
+            return
+
+        if not isinstance(excluded, set):
+            excluded = set(excluded)
+
+        if length:
+            p = pos % length
+            for i in range(length):
+                host = local_rack_live[(p + i) % length]
+                if host in excluded:
+                    continue
+                yield host
+
+        local_non_rack = self._non_local_rack_hosts
+        length = len(local_non_rack)
+        if length:
+            p = pos % length
+            for i in range(length):
+                host = local_non_rack[(p + i) % length]
+                if host in excluded:
+                    continue
+                yield host
+
+        for host in remote_hosts:
+            if host in excluded:
+                continue
+            yield host
+
     def on_up(self, host):
         dc = self._dc(host)
         rack = self._rack(host)