Make NodeSorters iterable

Author: nicktindall

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java

@@ -384,7 +384,7 @@ private void balance() {
             }
 
             // Balance each partition
-            for (NodeSorter nodeSorter : nodeSorters.allNodeSorters()) {
+            for (NodeSorter nodeSorter : nodeSorters) {
                 balanceByWeights(nodeSorter);
             }
         }

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/GlobalBalancingWeightsFactory.java

@@ -11,9 +11,9 @@
 
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.common.collect.Iterators;
 
-import java.util.Collection;
-import java.util.List;
+import java.util.Iterator;
 
 public class GlobalBalancingWeightsFactory implements BalancingWeightsFactory {
 
@@ -59,13 +59,13 @@ public NodeSorters createNodeSorters(BalancedShardsAllocator.ModelNode[] modelNo
         private record GlobalNodeSorters(BalancedShardsAllocator.NodeSorter nodeSorter) implements NodeSorters {
 
             @Override
-            public Collection<BalancedShardsAllocator.NodeSorter> allNodeSorters() {
-                return List.of(nodeSorter);
+            public BalancedShardsAllocator.NodeSorter sorterForShard(ShardRouting shard) {
+                return nodeSorter;
             }
 
             @Override
-            public BalancedShardsAllocator.NodeSorter sorterForShard(ShardRouting shard) {
-                return nodeSorter;
+            public Iterator<BalancedShardsAllocator.NodeSorter> iterator() {
+                return Iterators.single(nodeSorter);
             }
         }
     }

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/NodeSorters.java

@@ -11,19 +11,14 @@
 
 import org.elasticsearch.cluster.routing.ShardRouting;
 
-import java.util.Collection;
-
 /**
  * NodeSorters is just a cache of
  * {@link org.elasticsearch.cluster.routing.allocation.allocator.BalancedShardsAllocator.NodeSorter}
  * instances for each cluster partition
+ * <p>
+ * The returned iterator will return a node sorter for each partition in the cluster.
  */
-public interface NodeSorters {
-
-    /**
-     * Get the {@link BalancedShardsAllocator.NodeSorter}s for all partitions
-     */
-    Collection<BalancedShardsAllocator.NodeSorter> allNodeSorters();
+public interface NodeSorters extends Iterable<BalancedShardsAllocator.NodeSorter> {
 
     /**
      * Get the {@link BalancedShardsAllocator.NodeSorter} for the specified shard

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/BalancedShardsAllocatorTests.java

@@ -51,9 +51,9 @@
 import org.elasticsearch.test.gateway.TestGatewayAllocator;
 
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -844,9 +844,10 @@ public NodeSorters createNodeSorters(
                     );
                 }
                 return new NodeSorters() {
+
                     @Override
-                    public Collection<BalancedShardsAllocator.NodeSorter> allNodeSorters() {
-                        return prefixNodeSorters.values();
+                    public Iterator<BalancedShardsAllocator.NodeSorter> iterator() {
+                        return prefixNodeSorters.values().iterator();
                     }
 
                     @Override