Speedup computation of ShardAssignment map (#98850)

Random find from running many shards benchmarks: Using a stream in the
nested loop here performs pretty badly, a normal loop inlines fine in
many-shards benchmarks. Also, extracting the loop into a small method
gives better behaviour here as well + makes the profile easier to read.

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

@@ -21,6 +21,7 @@ import org.elasticsearch.cluster.routing.allocation.decider.Decision;
 import org.elasticsearch.common.metrics.MeanMetric;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Setting;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.shard.ShardId;
@@ -329,10 +330,7 @@ public class DesiredBalanceComputer {
         }
         iterations.inc(i);
 
-        final var assignments = new HashMap<ShardId, ShardAssignment>();
-        for (var shardAndAssignments : routingNodes.getAssignedShards().entrySet()) {
-            assignments.put(shardAndAssignments.getKey(), ShardAssignment.ofAssignedShards(shardAndAssignments.getValue()));
-        }
+        final var assignments = collectShardAssignments(routingNodes);
 
         for (var shard : routingNodes.unassigned().ignored()) {
             var info = shard.unassignedInfo();
@@ -362,6 +360,16 @@ public class DesiredBalanceComputer {
         return new DesiredBalance(lastConvergedIndex, assignments);
     }
 
+    private static Map<ShardId, ShardAssignment> collectShardAssignments(RoutingNodes routingNodes) {
+        final var entries = routingNodes.getAssignedShards().entrySet();
+        assert entries.stream().flatMap(t -> t.getValue().stream()).allMatch(ShardRouting::started) : routingNodes;
+        final Map<ShardId, ShardAssignment> res = Maps.newHashMapWithExpectedSize(entries.size());
+        for (var shardAndAssignments : entries) {
+            res.put(shardAndAssignments.getKey(), ShardAssignment.ofAssignedShards(shardAndAssignments.getValue()));
+        }
+        return res;
+    }
+
     private record ShardRoutings(List<ShardRouting> unassigned, List<ShardRouting> assigned) {
 
         private ShardRoutings(ShardId ignored) {

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

@@ -15,7 +15,6 @@ import java.util.List;
 import java.util.Set;
 
 import static java.util.Collections.unmodifiableSet;
-import static java.util.stream.Collectors.toCollection;
 
 public record ShardAssignment(Set<String> nodeIds, int total, int unassigned, int ignored) {
 
@@ -29,8 +28,10 @@ public record ShardAssignment(Set<String> nodeIds, int total, int unassigned, in
     }
 
     public static ShardAssignment ofAssignedShards(List<ShardRouting> routings) {
-        assert routings.stream().allMatch(ShardRouting::started) : routings;
-        var nodeIds = routings.stream().map(ShardRouting::currentNodeId).collect(toCollection(LinkedHashSet::new));
+        var nodeIds = new LinkedHashSet<String>();
+        for (ShardRouting routing : routings) {
+            nodeIds.add(routing.currentNodeId());
+        }
         return new ShardAssignment(unmodifiableSet(nodeIds), routings.size(), 0, 0);
     }
 }