More efficient sort in tryRelocateShard (#128063) (#128159)

No need to do this via an allocation-heavy `Stream`, we can just put the
objects straight into an array, sort them in-place, and keep hold of the
array to avoid having to allocate anything on the next iteration.

Also slims down `BY_DESCENDING_SHARD_ID`: it's always sorting the same
index so we don't need to look at `ShardId#index` in the comparison, nor
do we really need multiple layers of vtable lookups, we can just compare
the shard IDs directly.

Relates #128021

Author: DaveCTurner

docs/changelog/128063.yaml

@@ -0,0 +1,5 @@
+pr: 128063
+summary: More efficient sort in `tryRelocateShard`
+area: Allocation
+type: enhancement
+issues: []

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

@@ -54,7 +54,6 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.function.BiFunction;
-import java.util.stream.StreamSupport;
 
 import static org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata.Type.REPLACE;
 import static org.elasticsearch.cluster.routing.ExpectedShardSizeEstimator.getExpectedShardSize;
@@ -1206,7 +1205,13 @@ private AllocateUnassignedDecision decideAllocateUnassigned(final ShardRouting s
             return AllocateUnassignedDecision.fromDecision(decision, minNode != null ? minNode.routingNode.node() : null, nodeDecisions);
         }
 
-        private static final Comparator<ShardRouting> BY_DESCENDING_SHARD_ID = Comparator.comparing(ShardRouting::shardId).reversed();
+        private static final Comparator<ShardRouting> BY_DESCENDING_SHARD_ID = (s1, s2) -> Integer.compare(s2.id(), s1.id());
+
+        /**
+         * Scratch space for accumulating/sorting the {@link ShardRouting} instances when contemplating moving the shards away from a node
+         * in {@link #tryRelocateShard} - re-used to avoid extraneous allocations etc.
+         */
+        private ShardRouting[] shardRoutingsOnMaxWeightNode;
 
         /**
          * Tries to find a relocation from the max node to the minimal node for an arbitrary shard of the given index on the
@@ -1217,13 +1222,24 @@ private boolean tryRelocateShard(ModelNode minNode, ModelNode maxNode, String id
             final ModelIndex index = maxNode.getIndex(idx);
             if (index != null) {
                 logger.trace("Try relocating shard of [{}] from [{}] to [{}]", idx, maxNode.getNodeId(), minNode.getNodeId());
-                final Iterable<ShardRouting> shardRoutings = StreamSupport.stream(index.spliterator(), false)
-                    .filter(ShardRouting::started) // cannot rebalance unassigned, initializing or relocating shards anyway
-                    .sorted(BY_DESCENDING_SHARD_ID) // check in descending order of shard id so that the decision is deterministic
-                ::iterator;
+                if (shardRoutingsOnMaxWeightNode == null || shardRoutingsOnMaxWeightNode.length < index.numShards()) {
+                    shardRoutingsOnMaxWeightNode = new ShardRouting[index.numShards() * 2]; // oversized so reuse is more likely
+                }
+
+                int startedShards = 0;
+                for (final var shardRouting : index) {
+                    if (shardRouting.started()) { // cannot rebalance unassigned, initializing or relocating shards anyway
+                        shardRoutingsOnMaxWeightNode[startedShards] = shardRouting;
+                        startedShards += 1;
+                    }
+                }
+                // check in descending order of shard id so that the decision is deterministic
+                ArrayUtil.timSort(shardRoutingsOnMaxWeightNode, 0, startedShards, BY_DESCENDING_SHARD_ID);
 
                 final AllocationDeciders deciders = allocation.deciders();
-                for (ShardRouting shard : shardRoutings) {
+                for (int shardIndex = 0; shardIndex < startedShards; shardIndex++) {
+                    final ShardRouting shard = shardRoutingsOnMaxWeightNode[shardIndex];
+
                     final Decision rebalanceDecision = deciders.canRebalance(shard, allocation);
                     if (rebalanceDecision.type() == Type.NO) {
                         continue;