Make DiskThresholdDecider a little faster (#87821)

This is the most expensive allocator left in many-shards benchmarks.
It gets a little faster with this change iterating over the cheaper
array and inlining a little nicer by having specialized methods
for the two cached states.

server/src/main/java/org/elasticsearch/cluster/routing/RoutingNode.java

@@ -213,29 +213,34 @@ public class RoutingNode implements Iterable<ShardRouting> {
 
     /**
      * Determine the shards with a specific state
-     * @param states set of states which should be listed
+     * @param state state which should be listed
      * @return List of shards
      */
-    public List<ShardRouting> shardsWithState(ShardRoutingState... states) {
-        if (states.length == 1) {
-            if (states[0] == ShardRoutingState.INITIALIZING) {
-                return new ArrayList<>(initializingShards);
-            } else if (states[0] == ShardRoutingState.RELOCATING) {
-                return new ArrayList<>(relocatingShards);
-            }
+    public List<ShardRouting> shardsWithState(ShardRoutingState state) {
+        if (state == ShardRoutingState.INITIALIZING) {
+            return new ArrayList<>(initializingShards);
+        } else if (state == ShardRoutingState.RELOCATING) {
+            return new ArrayList<>(relocatingShards);
         }
-
         List<ShardRouting> shards = new ArrayList<>();
         for (ShardRouting shardEntry : this) {
-            for (ShardRoutingState state : states) {
-                if (shardEntry.state() == state) {
-                    shards.add(shardEntry);
-                }
+            if (shardEntry.state() == state) {
+                shards.add(shardEntry);
             }
         }
         return shards;
     }
 
+    private static final ShardRouting[] EMPTY_SHARD_ROUTING_ARRAY = new ShardRouting[0];
+
+    public ShardRouting[] initializing() {
+        return initializingShards.toArray(EMPTY_SHARD_ROUTING_ARRAY);
+    }
+
+    public ShardRouting[] relocating() {
+        return relocatingShards.toArray(EMPTY_SHARD_ROUTING_ARRAY);
+    }
+
     /**
      * Determine the shards of an index with a specific state
      * @param index id of the index

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDecider.java

@@ -21,7 +21,6 @@ import org.elasticsearch.cluster.routing.RecoverySource;
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.RoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
-import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.allocation.DiskThresholdSettings;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.Strings;
@@ -125,7 +124,7 @@ public class DiskThresholdDecider extends AllocationDecider {
         // no longer initializing because their recovery failed or was cancelled.
 
         // Where reserved space is unavailable (e.g. stats are out-of-sync) compute a conservative estimate for initialising shards
-        for (ShardRouting routing : node.shardsWithState(ShardRoutingState.INITIALIZING)) {
+        for (ShardRouting routing : node.initializing()) {
             if (routing.relocatingNodeId() == null) {
                 // in practice the only initializing-but-not-relocating shards with a nonzero expected shard size will be ones created
                 // by a resize (shrink/split/clone) operation which we expect to happen using hard links, so they shouldn't be taking
@@ -145,7 +144,7 @@ public class DiskThresholdDecider extends AllocationDecider {
         }
 
         if (subtractShardsMovingAway) {
-            for (ShardRouting routing : node.shardsWithState(ShardRoutingState.RELOCATING)) {
+            for (ShardRouting routing : node.relocating()) {
                 String actualPath = clusterInfo.getDataPath(routing);
                 if (actualPath == null) {
                     // we might know the path of this shard from before when it was relocating

server/src/main/java/org/elasticsearch/cluster/routing/allocation/decider/ThrottlingAllocationDecider.java

@@ -13,7 +13,6 @@ import org.apache.logging.log4j.Logger;
 import org.elasticsearch.cluster.routing.RecoverySource;
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.ShardRouting;
-import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.UnassignedInfo;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.common.settings.ClusterSettings;
@@ -129,7 +128,7 @@ public class ThrottlingAllocationDecider extends AllocationDecider {
             // count *just the primaries* currently doing recovery on the node and check against primariesInitialRecoveries
 
             int primariesInRecovery = 0;
-            for (ShardRouting shard : node.shardsWithState(ShardRoutingState.INITIALIZING)) {
+            for (ShardRouting shard : node.initializing()) {
                 // when a primary shard is INITIALIZING, it can be because of *initial recovery* or *relocation from another node*
                 // we only count initial recoveries here, so we need to make sure that relocating node is null
                 if (shard.primary() && shard.relocatingNodeId() == null) {

server/src/test/java/org/elasticsearch/cluster/routing/RoutingNodeTests.java

@@ -100,7 +100,6 @@ public class RoutingNodeTests extends ESTestCase {
     }
 
     public void testShardsWithState() {
-        assertThat(routingNode.shardsWithState(ShardRoutingState.INITIALIZING, ShardRoutingState.STARTED).size(), equalTo(2));
         assertThat(routingNode.shardsWithState(ShardRoutingState.STARTED).size(), equalTo(1));
         assertThat(routingNode.shardsWithState(ShardRoutingState.RELOCATING).size(), equalTo(1));
         assertThat(routingNode.shardsWithState(ShardRoutingState.INITIALIZING).size(), equalTo(1));

test/framework/src/main/java/org/elasticsearch/cluster/routing/RoutingNodesHelper.java

@@ -19,16 +19,13 @@ public final class RoutingNodesHelper {
 
     private RoutingNodesHelper() {}
 
-    public static List<ShardRouting> shardsWithState(RoutingNodes routingNodes, ShardRoutingState... state) {
+    public static List<ShardRouting> shardsWithState(RoutingNodes routingNodes, ShardRoutingState state) {
         List<ShardRouting> shards = new ArrayList<>();
         for (RoutingNode routingNode : routingNodes) {
             shards.addAll(routingNode.shardsWithState(state));
         }
-        for (ShardRoutingState s : state) {
-            if (s == ShardRoutingState.UNASSIGNED) {
-                routingNodes.unassigned().forEach(shards::add);
-                break;
-            }
+        if (state == ShardRoutingState.UNASSIGNED) {
+            routingNodes.unassigned().forEach(shards::add);
         }
         return shards;
     }