Add documentation around desired balance (#119902)

Relates ES-10341

Author: DiannaHohensee

server/src/main/java/org/elasticsearch/cluster/routing/allocation/AllocationService.java

@@ -158,9 +158,7 @@ public ClusterState applyStartedShards(ClusterState clusterState, List<ShardRout
     }
 
     private static ClusterState buildResultAndLogHealthChange(ClusterState oldState, RoutingAllocation allocation, String reason) {
-        final RoutingTable oldRoutingTable = oldState.routingTable();
-        final RoutingNodes newRoutingNodes = allocation.routingNodes();
-        final RoutingTable newRoutingTable = RoutingTable.of(newRoutingNodes);
+        final RoutingTable newRoutingTable = RoutingTable.of(allocation.routingNodes());
         final Metadata newMetadata = allocation.updateMetadataWithRoutingChanges(newRoutingTable);
         assert newRoutingTable.validate(newMetadata); // validates the routing table is coherent with the cluster state metadata
 
@@ -271,8 +269,7 @@ public ClusterState applyFailedShards(
     }
 
     /**
-     * unassigned an shards that are associated with nodes that are no longer part of the cluster, potentially promoting replicas
-     * if needed.
+     * Unassign any shards that are associated with nodes that are no longer part of the cluster, potentially promoting replicas if needed.
      */
     public ClusterState disassociateDeadNodes(ClusterState clusterState, boolean reroute, String reason) {
         RoutingAllocation allocation = createRoutingAllocation(clusterState, currentNanoTime());
@@ -284,7 +281,7 @@ public ClusterState disassociateDeadNodes(ClusterState clusterState, boolean rer
             clusterState = buildResultAndLogHealthChange(clusterState, allocation, reason);
         }
         if (reroute) {
-            return reroute(clusterState, reason, rerouteCompletionIsNotRequired());// this is not triggered by a user request
+            return reroute(clusterState, reason, rerouteCompletionIsNotRequired() /* this is not triggered by a user request */);
         } else {
             return clusterState;
         }

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

@@ -19,6 +19,10 @@
 /**
  * The desired balance of the cluster, indicating which nodes should hold a copy of each shard.
  *
+ * @param lastConvergedIndex Identifies what input data the balancer computation round used to produce this {@link DesiredBalance}. See
+ *                           {@link DesiredBalanceInput#index()} for details. Each reroute request in the same master term is assigned a
+ *                           strictly increasing sequence number. A new master term restarts the index values from zero. The balancer,
+ *                           which runs async to reroute, uses the latest request's data to compute the desired balance.
  * @param assignments a set of the (persistent) node IDs to which each {@link ShardId} should be allocated
  * @param weightsPerNode The node weights calculated based on
  * {@link org.elasticsearch.cluster.routing.allocation.allocator.WeightFunction#calculateNodeWeight}
@@ -31,8 +35,11 @@ public record DesiredBalance(
 ) {
 
     enum ComputationFinishReason {
+        /** Computation ran to completion */
         CONVERGED,
+        /** Computation exited and published early because a new cluster event occurred that affects computation */
         YIELD_TO_NEW_INPUT,
+        /** Computation stopped and published early to avoid delaying new shard assignment */
         STOP_EARLY
     }
 
@@ -44,6 +51,7 @@ public DesiredBalance(long lastConvergedIndex, Map<ShardId, ShardAssignment> ass
      * The placeholder value for {@link DesiredBalance} when the node stands down as master.
      */
     public static final DesiredBalance NOT_MASTER = new DesiredBalance(-2, Map.of());
+
     /**
      * The starting value for {@link DesiredBalance} when the node becomes the master.
      */
@@ -57,6 +65,10 @@ public static boolean hasChanges(DesiredBalance a, DesiredBalance b) {
         return Objects.equals(a.assignments, b.assignments) == false;
     }
 
+    /**
+     * Returns the sum of shard movements needed to reach the new desired balance. Doesn't count new shard copies as a move, nor removal or
+     * unassignment of a shard copy.
+     */
     public static int shardMovements(DesiredBalance old, DesiredBalance updated) {
         var intersection = Sets.intersection(old.assignments().keySet(), updated.assignments().keySet());
         int movements = 0;
@@ -70,8 +82,15 @@ public static int shardMovements(DesiredBalance old, DesiredBalance updated) {
         return movements;
     }
 
+    /**
+     * Returns the number of shard movements needed to reach the new shard assignment. Doesn't count new shard copies as a move, nor removal
+     * or unassignment of a shard copy.
+     */
     private static int shardMovements(ShardAssignment old, ShardAssignment updated) {
-        var movements = Math.min(0, old.assigned() - updated.assigned());// compensate newly started shards
+        // A shard move should retain the same number of assigned nodes, just swap out one node for another. We will compensate for newly
+        // started shards -- adding a shard copy is not a move -- by initializing the count with a negative value so that incrementing later
+        // for a new node zeros out.
+        var movements = Math.min(0, old.assigned() - updated.assigned());
         for (String nodeId : updated.nodeIds()) {
             if (old.nodeIds().contains(nodeId) == false) {
                 movements++;

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

@@ -415,11 +415,14 @@ boolean hasEnoughIterations(int currentIteration) {
     }
 
     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()));
+        final var allAssignedShards = routingNodes.getAssignedShards().entrySet();
+        assert allAssignedShards.stream().flatMap(t -> t.getValue().stream()).allMatch(ShardRouting::started) : routingNodes;
+        final Map<ShardId, ShardAssignment> res = Maps.newHashMapWithExpectedSize(allAssignedShards.size());
+        for (var shardIdAndShardRoutings : allAssignedShards) {
+            res.put(
+                shardIdAndShardRoutings.getKey(),
+                ShardAssignment.createFromAssignedShardRoutingsList(shardIdAndShardRoutings.getValue())
+            );
         }
         return res;
     }

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

@@ -54,6 +54,10 @@ public class DesiredBalanceReconciler {
 
     private static final Logger logger = LogManager.getLogger(DesiredBalanceReconciler.class);
 
+    /**
+     * The minimum interval that log messages will be written if the number of undesired shard allocations reaches the percentage of total
+     * shards set by {@link #UNDESIRED_ALLOCATIONS_LOG_THRESHOLD_SETTING}.
+     */
     public static final Setting<TimeValue> UNDESIRED_ALLOCATIONS_LOG_INTERVAL_SETTING = Setting.timeSetting(
         "cluster.routing.allocation.desired_balance.undesired_allocations.log_interval",
         TimeValue.timeValueHours(1),
@@ -62,6 +66,10 @@ public class DesiredBalanceReconciler {
         Setting.Property.NodeScope
     );
 
+    /**
+     * Warning log messages may be periodically written if the number of shards that are on undesired nodes reaches this percentage setting.
+     * Works together with {@link #UNDESIRED_ALLOCATIONS_LOG_INTERVAL_SETTING} to log on a periodic basis.
+     */
     public static final Setting<Double> UNDESIRED_ALLOCATIONS_LOG_THRESHOLD_SETTING = Setting.doubleSetting(
         "cluster.routing.allocation.desired_balance.undesired_allocations.threshold",
         0.1,
@@ -96,6 +104,13 @@ public DesiredBalanceReconciler(
         this.nodeAllocationStatsAndWeightsCalculator = nodeAllocationStatsAndWeightsCalculator;
     }
 
+    /**
+     * Applies a desired shard allocation to the routing table by initializing and relocating shards in the cluster state.
+     *
+     * @param desiredBalance The new desired cluster shard allocation
+     * @param allocation Cluster state information with which to make decisions, contains routing table metadata that will be modified to
+     *                   reach the given desired balance.
+     */
     public void reconcile(DesiredBalance desiredBalance, RoutingAllocation allocation) {
         var nodeIds = allocation.routingNodes().getAllNodeIds();
         allocationOrdering.retainNodes(nodeIds);

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.cluster.ClusterStateTaskListener;
 import org.elasticsearch.cluster.metadata.SingleNodeShutdownMetadata;
 import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.allocation.AllocationService;
 import org.elasticsearch.cluster.routing.allocation.AllocationService.RerouteStrategy;
 import org.elasticsearch.cluster.routing.allocation.NodeAllocationStatsAndWeightsCalculator;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
@@ -56,11 +57,30 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
 
     private final ShardsAllocator delegateAllocator;
     private final ThreadPool threadPool;
+    /**
+     * This is a callback to run {@link AllocationService#executeWithRoutingAllocation(ClusterState, String, RerouteStrategy)}, which
+     * produces a new ClusterState with the changes made by {@link DesiredBalanceReconciler#reconcile}. The {@link RerouteStrategy} provided
+     * to the callback calls into {@link #desiredBalanceReconciler} for the changes. The {@link #masterServiceTaskQueue} will publish the
+     * new cluster state after the cluster state is constructed by the {@link ReconcileDesiredBalanceExecutor}.
+     */
     private final DesiredBalanceReconcilerAction reconciler;
     private final DesiredBalanceComputer desiredBalanceComputer;
+    /**
+     * Reconciliation ({@link DesiredBalanceReconciler#reconcile(DesiredBalance, RoutingAllocation)}) takes the {@link DesiredBalance}
+     * output of {@link DesiredBalanceComputer#compute} and identifies how shards need to be added, moved or removed to go from the current
+     * cluster shard allocation to the new desired allocation.
+     */
     private final DesiredBalanceReconciler desiredBalanceReconciler;
     private final ContinuousComputation<DesiredBalanceInput> desiredBalanceComputation;
-    private final PendingListenersQueue queue;
+    /**
+     * Saves and runs listeners after DesiredBalance computations complete.
+     */
+    private final PendingListenersQueue pendingListenersQueue;
+    /**
+     * Each reroute request gets assigned a monotonically increasing sequence number. Many reroute requests may arrive before the balancer
+     * asynchronously runs a computation. The balancer will use the latest request and save this sequence number to track back to the
+     * request.
+     */
     private final AtomicLong indexGenerator = new AtomicLong(-1);
     private final ConcurrentLinkedQueue<List<MoveAllocationCommand>> pendingDesiredBalanceMoves = new ConcurrentLinkedQueue<>();
     private final MasterServiceTaskQueue<ReconcileDesiredBalanceTask> masterServiceTaskQueue;
@@ -199,7 +219,7 @@ public String toString() {
                 return "DesiredBalanceShardsAllocator#allocate";
             }
         };
-        this.queue = new PendingListenersQueue();
+        this.pendingListenersQueue = new PendingListenersQueue();
         this.masterServiceTaskQueue = clusterService.createTaskQueue(
             "reconcile-desired-balance",
             Priority.URGENT,
@@ -235,7 +255,7 @@ public void allocate(RoutingAllocation allocation, ActionListener<Void> listener
 
         var index = indexGenerator.incrementAndGet();
         logger.debug("Executing allocate for [{}]", index);
-        queue.add(index, listener);
+        pendingListenersQueue.add(index, listener);
         // This can only run on master, so unset not-master if exists
         if (currentDesiredBalanceRef.compareAndSet(DesiredBalance.NOT_MASTER, DesiredBalance.BECOME_MASTER_INITIAL)) {
             logger.debug("initialized desired balance for becoming master");
@@ -378,7 +398,7 @@ public DesiredBalanceStats getStats() {
     private void onNoLongerMaster() {
         if (indexGenerator.getAndSet(-1) != -1) {
             currentDesiredBalanceRef.set(DesiredBalance.NOT_MASTER);
-            queue.completeAllAsNotMaster();
+            pendingListenersQueue.completeAllAsNotMaster();
             pendingDesiredBalanceMoves.clear();
             desiredBalanceReconciler.clear();
             desiredBalanceMetrics.zeroAllMetrics();
@@ -428,7 +448,7 @@ private ClusterState applyBalance(
                     batchExecutionContext.initialState(),
                     createReconcileAllocationAction(latest.getTask().desiredBalance)
                 );
-                latest.success(() -> queue.complete(latest.getTask().desiredBalance.lastConvergedIndex()));
+                latest.success(() -> pendingListenersQueue.complete(latest.getTask().desiredBalance.lastConvergedIndex()));
                 return newState;
             }
         }
@@ -447,7 +467,7 @@ private static void discardSupersededTasks(
 
     // only for tests - in production, this happens after reconciliation
     protected final void completeToLastConvergedIndex() {
-        queue.complete(currentDesiredBalanceRef.get().lastConvergedIndex());
+        pendingListenersQueue.complete(currentDesiredBalanceRef.get().lastConvergedIndex());
     }
 
     private void recordTime(CounterMetric metric, Runnable action) {

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

@@ -24,6 +24,10 @@
 import static org.elasticsearch.cluster.service.ClusterApplierService.CLUSTER_UPDATE_THREAD_NAME;
 import static org.elasticsearch.cluster.service.MasterService.MASTER_UPDATE_THREAD_NAME;
 
+/**
+ * Registers listeners with an `index` number ({@link #add(long, ActionListener)}) and then completes them whenever the latest index number
+ * is greater or equal to a listener's index value ({@link #complete(long)}).
+ */
 public class PendingListenersQueue {
 
     private static final Logger logger = LogManager.getLogger(PendingListenersQueue.class);

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

@@ -17,6 +17,14 @@
 
 import static java.util.Collections.unmodifiableSet;
 
+/**
+ * Simple shard assignment summary of shard copies for a particular index shard.
+ *
+ * @param nodeIds The node IDs of nodes holding a shard copy.
+ * @param total The total number of shard copies.
+ * @param unassigned The number of unassigned shard copies.
+ * @param ignored The number of ignored shard copies.
+ */
 public record ShardAssignment(Set<String> nodeIds, int total, int unassigned, int ignored) {
 
     public ShardAssignment {
@@ -28,9 +36,13 @@ public int assigned() {
         return nodeIds.size();
     }
 
-    public static ShardAssignment ofAssignedShards(List<ShardRouting> routings) {
+    /**
+     * Helper method to instantiate a new ShardAssignment from a given list of ShardRouting instances. Assumes all shards are assigned.
+     */
+    public static ShardAssignment createFromAssignedShardRoutingsList(List<ShardRouting> routings) {
         var nodeIds = new LinkedHashSet<String>();
         for (ShardRouting routing : routings) {
+            assert routing.unassignedInfo() == null : "Expected assigned shard copies only, unassigned info: " + routing.unassignedInfo();
             nodeIds.add(routing.currentNodeId());
         }
         return new ShardAssignment(unmodifiableSet(nodeIds), routings.size(), 0, 0);

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

@@ -46,19 +46,19 @@ public void testShardMovements() {
         );
 
         assertThat(
-            "1 shard movements when existing shard is moved and new shard copy is unassigned",
+            "1 shard movements when an existing shard copy is moved and new shard copy is unassigned",
             shardMovements(new ShardAssignment(Set.of("a", "b"), 2, 0, 0), new ShardAssignment(Set.of("a", "c"), 3, 1, 0)),
             equalTo(1)
         );
 
         assertThat(
-            "1 shard movement",
+            "1 shard movement when an existing shard copy is moved",
             shardMovements(new ShardAssignment(Set.of("a", "b"), 2, 0, 0), new ShardAssignment(Set.of("a", "c"), 2, 0, 0)),
             equalTo(1)
         );
 
         assertThat(
-            "2 shard movement",
+            "2 shard movements when both shard copies are move to new nodes",
             shardMovements(new ShardAssignment(Set.of("a", "b"), 2, 0, 0), new ShardAssignment(Set.of("c", "d"), 2, 0, 0)),
             equalTo(2)
         );
@@ -77,10 +77,10 @@ public void testShardMovements() {
     }
 
     private static int shardMovements(ShardAssignment old, ShardAssignment updated) {
-        return DesiredBalance.shardMovements(of(old), of(updated));
+        return DesiredBalance.shardMovements(createDesiredBalanceWith(old), createDesiredBalanceWith(updated));
     }
 
-    private static DesiredBalance of(ShardAssignment assignment) {
+    private static DesiredBalance createDesiredBalanceWith(ShardAssignment assignment) {
         return new DesiredBalance(1, Map.of(new ShardId("index", "_na_", 0), assignment));
     }
 }