Balance priorities during reconciliation (#95454) (#96148)

When reconciling a balance with a lot of shards on undesired nodes there is a
possibility of causing node hot spots due to usage of
nodeInterleavedShardIterator. This iterator orders shards based on nodes they
are located and order nodes based hash map iteration. This means it tends
to pick shards returned first by the iterator. This change uses OrderedShardsIterator
that applies custom shards order based on allocation recency.

(cherry picked from commit 6ecd74d)

Author: idegtiarenko

docs/changelog/95454.yaml

@@ -0,0 +1,5 @@
+pr: 95454
+summary: Balance priorities during reconciliation
+area: Allocation
+type: bug
+issues: []

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

@@ -249,6 +249,10 @@ private ShardRouting findAssignedPrimaryIfPeerRecovery(ShardRouting routing) {
         return primary;
     }
 
+    public Set<String> getAllNodeIds() {
+        return Collections.unmodifiableSet(nodesToShards.keySet());
+    }
+
     @Override
     public Iterator<RoutingNode> iterator() {
         return Collections.unmodifiableCollection(nodesToShards.values()).iterator();

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

@@ -47,16 +47,19 @@ public class DesiredBalanceReconciler {
     private final RoutingAllocation allocation; // name chosen to align with code in BalancedShardsAllocator but TODO rename
     private final RoutingNodes routingNodes;
     private final NodeAllocationOrdering allocationOrdering;
+    private final NodeAllocationOrdering moveOrdering;
 
     DesiredBalanceReconciler(
         DesiredBalance desiredBalance,
         RoutingAllocation routingAllocation,
-        NodeAllocationOrdering allocationOrdering
+        NodeAllocationOrdering allocationOrdering,
+        NodeAllocationOrdering moveOrdering
     ) {
         this.desiredBalance = desiredBalance;
         this.allocation = routingAllocation;
         this.routingNodes = routingAllocation.routingNodes();
         this.allocationOrdering = allocationOrdering;
+        this.moveOrdering = moveOrdering;
     }
 
     void run() {
@@ -220,9 +223,7 @@ private void allocateUnassigned() {
                             final var decision = allocation.deciders().canAllocate(shard, routingNode, allocation);
                             switch (decision.type()) {
                                 case YES -> {
-                                    if (logger.isTraceEnabled()) {
-                                        logger.trace("Assigned shard [{}] to [{}]", shard, desiredNodeId);
-                                    }
+                                    logger.debug("Assigning shard [{}] to [{}]", shard, desiredNodeId);
                                     final long shardSize = DiskThresholdDecider.getExpectedShardSize(
                                         shard,
                                         ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE,
@@ -253,9 +254,7 @@ private void allocateUnassigned() {
                     }
                 }
 
-                if (logger.isTraceEnabled()) {
-                    logger.trace("No eligible node found to assign shard [{}] amongst [{}]", shard, assignment);
-                }
+                logger.debug("No eligible node found to assign shard [{}] amongst [{}]", shard, assignment);
 
                 final UnassignedInfo.AllocationStatus allocationStatus;
                 if (assignment == null || assignment.isIgnored(shard.primary())) {
@@ -284,20 +283,16 @@ private void allocateUnassigned() {
 
     private Iterable<String> getDesiredNodesIds(ShardRouting shard, ShardAssignment assignment) {
         return allocationOrdering.sort(allocation.deciders().getForcedInitialShardAllocationToNodes(shard, allocation).map(forced -> {
-            if (logger.isDebugEnabled()) {
-                logger.debug("Shard [{}] assignment is ignored. Initial allocation forced to {}", shard.shardId(), forced);
-            }
+            logger.debug("Shard [{}] assignment is ignored. Initial allocation forced to {}", shard.shardId(), forced);
             return forced;
         }).orElse(assignment.nodeIds()));
     }
 
     private Iterable<String> getFallbackNodeIds(ShardRouting shard, AtomicBoolean isThrottled) {
         return () -> {
             if (shard.primary() && isThrottled.get() == false) {
-                var fallbackNodeIds = allocation.routingNodes().stream().map(RoutingNode::nodeId).toList();
-                if (logger.isDebugEnabled()) {
-                    logger.trace("Shard [{}] assignment is temporary not possible. Falling back to {}", shard.shardId(), fallbackNodeIds);
-                }
+                var fallbackNodeIds = allocation.routingNodes().getAllNodeIds();
+                logger.debug("Shard [{}] assignment is temporary not possible. Falling back to {}", shard.shardId(), fallbackNodeIds);
                 return allocationOrdering.sort(fallbackNodeIds).iterator();
             } else {
                 return Collections.emptyIterator();
@@ -306,10 +301,9 @@ private Iterable<String> getFallbackNodeIds(ShardRouting shard, AtomicBoolean is
     }
 
     private void moveShards() {
-        // Iterate over the started shards interleaving between nodes, and check if they can remain. In the presence of throttling
-        // shard movements, the goal of this iteration order is to achieve a fairer movement of shards from the nodes that are
-        // offloading the shards.
-        for (final var iterator = routingNodes.nodeInterleavedShardIterator(); iterator.hasNext();) {
+        // Iterate over all started shards and check if they can remain. In the presence of throttling shard movements,
+        // the goal of this iteration order is to achieve a fairer movement of shards from the nodes that are offloading the shards.
+        for (final var iterator = OrderedShardsIterator.create(routingNodes, moveOrdering); iterator.hasNext();) {
             final var shardRouting = iterator.next();
 
             if (shardRouting.started() == false) {
@@ -343,12 +337,15 @@ private void moveShards() {
 
             final var moveTarget = findRelocationTarget(shardRouting, assignment.nodeIds());
             if (moveTarget != null) {
+                logger.debug("Moving shard {} from {} to {}", shardRouting.shardId(), shardRouting.currentNodeId(), moveTarget.getId());
                 routingNodes.relocateOrReinitializeShard(
                     shardRouting,
                     moveTarget.getId(),
                     allocation.clusterInfo().getShardSize(shardRouting, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE),
                     allocation.changes()
                 );
+                iterator.dePrioritizeNode(shardRouting.currentNodeId());
+                moveOrdering.recordAllocation(shardRouting.currentNodeId());
             }
         }
     }
@@ -358,10 +355,9 @@ private void balance() {
             return;
         }
 
-        // Iterate over the started shards interleaving between nodes, and try to move any which are on undesired nodes. In the presence of
-        // throttling shard movements, the goal of this iteration order is to achieve a fairer movement of shards from the nodes that are
-        // offloading the shards.
-        for (final var iterator = routingNodes.nodeInterleavedShardIterator(); iterator.hasNext();) {
+        // Iterate over all started shards and try to move any which are on undesired nodes. In the presence of throttling shard movements,
+        // the goal of this iteration order is to achieve a fairer movement of shards from the nodes that are offloading the shards.
+        for (final var iterator = OrderedShardsIterator.create(routingNodes, moveOrdering); iterator.hasNext();) {
             final var shardRouting = iterator.next();
 
             if (shardRouting.started() == false) {
@@ -392,12 +388,21 @@ private void balance() {
 
             final var rebalanceTarget = findRelocationTarget(shardRouting, assignment.nodeIds(), this::decideCanAllocate);
             if (rebalanceTarget != null) {
+                logger.debug(
+                    "Rebalancing shard {} from {} to {}",
+                    shardRouting.shardId(),
+                    shardRouting.currentNodeId(),
+                    rebalanceTarget.getId()
+                );
+
                 routingNodes.relocateOrReinitializeShard(
                     shardRouting,
                     rebalanceTarget.getId(),
                     allocation.clusterInfo().getShardSize(shardRouting, ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE),
                     allocation.changes()
                 );
+                iterator.dePrioritizeNode(shardRouting.currentNodeId());
+                moveOrdering.recordAllocation(shardRouting.currentNodeId());
             }
         }
     }

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

@@ -14,8 +14,6 @@
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ClusterStateTaskExecutor;
 import org.elasticsearch.cluster.ClusterStateTaskListener;
-import org.elasticsearch.cluster.routing.RoutingNode;
-import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.ShardRouting;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.RoutingExplanations;
@@ -37,13 +35,10 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
 
-import static java.util.stream.Collectors.toSet;
-
 /**
  * A {@link ShardsAllocator} which asynchronously refreshes the desired balance held by the {@link DesiredBalanceComputer} and then takes
  * steps towards the desired balance using the {@link DesiredBalanceReconciler}.
@@ -62,6 +57,7 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
     private final ConcurrentLinkedQueue<List<MoveAllocationCommand>> pendingDesiredBalanceMoves = new ConcurrentLinkedQueue<>();
     private final MasterServiceTaskQueue<ReconcileDesiredBalanceTask> masterServiceTaskQueue;
     private final NodeAllocationOrdering allocationOrdering = new NodeAllocationOrdering();
+    private final NodeAllocationOrdering moveOrdering = new NodeAllocationOrdering();
     private volatile DesiredBalance currentDesiredBalance = DesiredBalance.INITIAL;
     private volatile boolean resetCurrentDesiredBalance = false;
 
@@ -228,8 +224,13 @@ protected void reconcile(DesiredBalance desiredBalance, RoutingAllocation alloca
         } else {
             logger.debug("Reconciling desired balance for [{}]", desiredBalance.lastConvergedIndex());
         }
-        allocationOrdering.retainNodes(getNodeIds(allocation.routingNodes()));
-        recordTime(cumulativeReconciliationTime, new DesiredBalanceReconciler(desiredBalance, allocation, allocationOrdering)::run);
+        var allNodeIds = allocation.routingNodes().getAllNodeIds();
+        allocationOrdering.retainNodes(allNodeIds);
+        moveOrdering.retainNodes(allNodeIds);
+        recordTime(
+            cumulativeReconciliationTime,
+            new DesiredBalanceReconciler(desiredBalance, allocation, allocationOrdering, moveOrdering)::run
+        );
         if (logger.isTraceEnabled()) {
             logger.trace("Reconciled desired balance: {}", desiredBalance);
         } else {
@@ -346,8 +347,4 @@ private void recordTime(CounterMetric metric, Runnable action) {
             metric.inc(finished - started);
         }
     }
-
-    private static Set<String> getNodeIds(RoutingNodes nodes) {
-        return nodes.stream().map(RoutingNode::nodeId).collect(toSet());
-    }
 }

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

@@ -0,0 +1,70 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.cluster.routing.allocation.allocator;
+
+import org.elasticsearch.cluster.routing.RoutingNodes;
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.common.collect.Iterators;
+
+import java.util.ArrayDeque;
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.Objects;
+
+/**
+ * This class iterates all shards from all nodes in order of allocation recency.
+ * Shards from the node that had a new shard allocation would appear in the end of iteration.
+ */
+public class OrderedShardsIterator implements Iterator<ShardRouting> {
+
+    private final ArrayDeque<NodeAndShardIterator> queue;
+
+    public static OrderedShardsIterator create(RoutingNodes routingNodes, NodeAllocationOrdering ordering) {
+        var queue = new ArrayDeque<NodeAndShardIterator>(routingNodes.size());
+        for (var nodeId : ordering.sort(routingNodes.getAllNodeIds())) {
+            var node = routingNodes.node(nodeId);
+            if (node.size() > 0) {
+                queue.add(new NodeAndShardIterator(nodeId, Iterators.forArray(node.copyShards())));
+            }
+        }
+        return new OrderedShardsIterator(queue);
+    }
+
+    private OrderedShardsIterator(ArrayDeque<NodeAndShardIterator> queue) {
+        this.queue = queue;
+    }
+
+    @Override
+    public boolean hasNext() {
+        return queue.isEmpty() == false;
+    }
+
+    @Override
+    public ShardRouting next() {
+        if (queue.isEmpty()) {
+            throw new NoSuchElementException();
+        }
+        var entry = queue.peek();
+        assert entry.iterator.hasNext();
+        final var nextShard = entry.iterator.next();
+        if (entry.iterator.hasNext() == false) {
+            queue.poll();
+        }
+        return nextShard;
+    }
+
+    public void dePrioritizeNode(String nodeId) {
+        var entry = queue.peek();
+        if (entry != null && Objects.equals(nodeId, entry.nodeId)) {
+            queue.offer(queue.poll());
+        }
+    }
+
+    private record NodeAndShardIterator(String nodeId, Iterator<ShardRouting> iterator) {}
+}