Fix groupby issue in DCAwareRoundRobinPolicy and RackAwareRoundRobinPolicy

- Fixed DCAwareRoundRobinPolicy.populate() to properly collect all hosts per DC
- Fixed RackAwareRoundRobinPolicy.populate() to properly collect all hosts per DC/rack
- Added regression tests for both policies to verify interleaved hosts are not lost
- All existing tests pass (83 tests in test_policies.py)

Co-authored-by: dkropachev <[email protected]>

Author: Copilot

cassandra/policies.py

@@ -253,8 +253,17 @@ def _dc(self, host):
         return host.datacenter or self.local_dc
 
     def populate(self, cluster, hosts):
-        for dc, dc_hosts in groupby(hosts, lambda h: self._dc(h)):
-            self._dc_live_hosts[dc] = tuple(set(dc_hosts))
+        # Group hosts by datacenter without relying on groupby which only groups consecutive items
+        dc_hosts_dict = {}
+        for host in hosts:
+            dc = self._dc(host)
+            if dc not in dc_hosts_dict:
+                dc_hosts_dict[dc] = []
+            dc_hosts_dict[dc].append(host)
+        
+        # Convert lists to tuples with unique hosts
+        for dc, host_list in dc_hosts_dict.items():
+            self._dc_live_hosts[dc] = tuple(set(host_list))
 
         if not self.local_dc:
             self._endpoints = [
@@ -373,10 +382,31 @@ def _dc(self, host):
         return host.datacenter or self.local_dc
 
     def populate(self, cluster, hosts):
-        for (dc, rack), rack_hosts in groupby(hosts, lambda host: (self._dc(host), self._rack(host))):
-            self._live_hosts[(dc, rack)] = tuple(set(rack_hosts))
-        for dc, dc_hosts in groupby(hosts, lambda host: self._dc(host)):
-            self._dc_live_hosts[dc] = tuple(set(dc_hosts))
+        # Group hosts by (dc, rack) and by dc without relying on groupby which only groups consecutive items
+        rack_hosts_dict = {}
+        dc_hosts_dict = {}
+        
+        for host in hosts:
+            dc = self._dc(host)
+            rack = self._rack(host)
+            
+            # Group by (dc, rack)
+            key = (dc, rack)
+            if key not in rack_hosts_dict:
+                rack_hosts_dict[key] = []
+            rack_hosts_dict[key].append(host)
+            
+            # Group by dc
+            if dc not in dc_hosts_dict:
+                dc_hosts_dict[dc] = []
+            dc_hosts_dict[dc].append(host)
+        
+        # Convert lists to tuples with unique hosts
+        for key, host_list in rack_hosts_dict.items():
+            self._live_hosts[key] = tuple(set(host_list))
+        
+        for dc, host_list in dc_hosts_dict.items():
+            self._dc_live_hosts[dc] = tuple(set(host_list))
 
         self._position = randint(0, len(hosts) - 1) if hosts else 0
 

tests/unit/test_policies.py

@@ -531,6 +531,7 @@ def test_wrong_dc(self, policy_specialization, constructor_args):
         qplan = list(policy.make_query_plan())
         assert len(qplan) == 0
 
+
 class DCAwareRoundRobinPolicyTest(unittest.TestCase):
 
     def test_default_dc(self):
@@ -573,6 +574,116 @@ def test_default_dc(self):
         policy.on_add(host_remote)
         assert not policy.local_dc
 
+    def test_populate_with_interleaved_dcs(self):
+        """Test that DCAwareRoundRobinPolicy doesn't lose hosts when DCs are interleaved.
+        
+        This is a regression test for the issue where groupby only groups consecutive items,
+        which caused hosts to be lost when datacenters were not sorted.
+        """
+        # Create hosts with interleaved datacenters (dc1, dc2, dc1, dc2)
+        hosts = []
+        
+        # DC1 host 1
+        h1 = Host(DefaultEndPoint("10.0.0.1"), SimpleConvictionPolicy)
+        h1.set_location_info("dc1", "rack1")
+        hosts.append(h1)
+        
+        # DC2 host 1
+        h2 = Host(DefaultEndPoint("10.0.0.2"), SimpleConvictionPolicy)
+        h2.set_location_info("dc2", "rack1")
+        hosts.append(h2)
+        
+        # DC1 host 2 (interleaved)
+        h3 = Host(DefaultEndPoint("10.0.0.3"), SimpleConvictionPolicy)
+        h3.set_location_info("dc1", "rack1")
+        hosts.append(h3)
+        
+        # DC2 host 2 (interleaved)
+        h4 = Host(DefaultEndPoint("10.0.0.4"), SimpleConvictionPolicy)
+        h4.set_location_info("dc2", "rack1")
+        hosts.append(h4)
+        
+        policy = DCAwareRoundRobinPolicy("dc1")
+        policy.populate(Mock(), hosts)
+        
+        # Check that all hosts are registered
+        dc1_hosts = policy._dc_live_hosts.get("dc1", ())
+        dc2_hosts = policy._dc_live_hosts.get("dc2", ())
+        
+        assert len(dc1_hosts) == 2, "DC1 should have 2 hosts"
+        assert len(dc2_hosts) == 2, "DC2 should have 2 hosts"
+        assert h1 in dc1_hosts
+        assert h3 in dc1_hosts
+        assert h2 in dc2_hosts
+        assert h4 in dc2_hosts
+
+
+class RackAwareRoundRobinPolicyTest(unittest.TestCase):
+
+    def test_populate_with_interleaved_racks(self):
+        """Test that RackAwareRoundRobinPolicy doesn't lose hosts when racks are interleaved.
+        
+        This is a regression test for the issue where groupby only groups consecutive items,
+        which caused hosts to be lost when racks/datacenters were not sorted.
+        """
+        # Create hosts with interleaved racks
+        hosts = []
+        
+        # DC1 Rack1 host 1
+        h1 = Host(DefaultEndPoint("10.0.0.1"), SimpleConvictionPolicy)
+        h1.set_location_info("dc1", "rack1")
+        hosts.append(h1)
+        
+        # DC1 Rack2 host 1
+        h2 = Host(DefaultEndPoint("10.0.0.2"), SimpleConvictionPolicy)
+        h2.set_location_info("dc1", "rack2")
+        hosts.append(h2)
+        
+        # DC1 Rack1 host 2 (interleaved)
+        h3 = Host(DefaultEndPoint("10.0.0.3"), SimpleConvictionPolicy)
+        h3.set_location_info("dc1", "rack1")
+        hosts.append(h3)
+        
+        # DC1 Rack2 host 2 (interleaved)
+        h4 = Host(DefaultEndPoint("10.0.0.4"), SimpleConvictionPolicy)
+        h4.set_location_info("dc1", "rack2")
+        hosts.append(h4)
+        
+        # DC2 Rack1 host 1
+        h5 = Host(DefaultEndPoint("10.0.0.5"), SimpleConvictionPolicy)
+        h5.set_location_info("dc2", "rack1")
+        hosts.append(h5)
+        
+        # DC1 Rack1 host 3 (interleaved again)
+        h6 = Host(DefaultEndPoint("10.0.0.6"), SimpleConvictionPolicy)
+        h6.set_location_info("dc1", "rack1")
+        hosts.append(h6)
+        
+        policy = RackAwareRoundRobinPolicy("dc1", "rack1")
+        policy.populate(Mock(), hosts)
+        
+        # Check that all hosts are registered
+        dc1_rack1_hosts = policy._live_hosts.get(("dc1", "rack1"), ())
+        dc1_rack2_hosts = policy._live_hosts.get(("dc1", "rack2"), ())
+        dc2_rack1_hosts = policy._live_hosts.get(("dc2", "rack1"), ())
+        
+        dc1_hosts = policy._dc_live_hosts.get("dc1", ())
+        dc2_hosts = policy._dc_live_hosts.get("dc2", ())
+        
+        assert len(dc1_rack1_hosts) == 3, "DC1 Rack1 should have 3 hosts"
+        assert len(dc1_rack2_hosts) == 2, "DC1 Rack2 should have 2 hosts"
+        assert len(dc2_rack1_hosts) == 1, "DC2 Rack1 should have 1 host"
+        assert len(dc1_hosts) == 5, "DC1 should have 5 hosts total"
+        assert len(dc2_hosts) == 1, "DC2 should have 1 host total"
+        
+        assert h1 in dc1_rack1_hosts
+        assert h3 in dc1_rack1_hosts
+        assert h6 in dc1_rack1_hosts
+        assert h2 in dc1_rack2_hosts
+        assert h4 in dc1_rack2_hosts
+        assert h5 in dc2_rack1_hosts
+
+
 class TokenAwarePolicyTest(unittest.TestCase):
 
     def test_wrap_round_robin(self):