Add node weight changes to balance round summaries

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

@@ -18,6 +18,8 @@
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -83,6 +85,65 @@ public AllocationBalancingRoundSummaryService(ThreadPool threadPool, ClusterSett
         });
     }
 
+    /**
+     * Summarizes the work required to move from an old to new desired balance shard allocation.
+     */
+    public static BalancingRoundSummary createBalancerRoundSummary(DesiredBalance oldDesiredBalance, DesiredBalance newDesiredBalance) {
+        return new BalancingRoundSummary(
+            createWeightsSummary(oldDesiredBalance, newDesiredBalance),
+            DesiredBalance.shardMovements(oldDesiredBalance, newDesiredBalance)
+        );
+    }
+
+    /**
+     * Creates a summary of the node weight changes from {@code oldDesiredBalance} to {@code newDesiredBalance}.
+     * See {@link BalancingRoundSummary.NodesWeightsChanges} for content details.
+     */
+    private static Map<String, BalancingRoundSummary.NodesWeightsChanges> createWeightsSummary(
+        DesiredBalance oldDesiredBalance,
+        DesiredBalance newDesiredBalance
+    ) {
+        var oldWeightsPerNode = oldDesiredBalance.weightsPerNode();
+        var newWeightsPerNode = newDesiredBalance.weightsPerNode();
+
+        Map<String, BalancingRoundSummary.NodesWeightsChanges> nodeNameToWeightInfo = new HashMap<>(oldWeightsPerNode.size());
+        for (var nodeAndWeights : oldWeightsPerNode.entrySet()) {
+            var discoveryNode = nodeAndWeights.getKey();
+            var oldNodeWeightStats = nodeAndWeights.getValue();
+
+            // The node may no longer exists in the new DesiredBalance. If so, the new weights for that node are effectively zero. New
+            // weights of zero will result in correctly negative weight diffs for the removed node.
+            var newNodeWeightStats = newWeightsPerNode.getOrDefault(discoveryNode, DesiredBalanceMetrics.NodeWeightStats.ZERO);
+
+            nodeNameToWeightInfo.put(
+                discoveryNode.getName(),
+                new BalancingRoundSummary.NodesWeightsChanges(
+                    oldNodeWeightStats,
+                    BalancingRoundSummary.NodeWeightsDiff.create(oldNodeWeightStats, newNodeWeightStats)
+                )
+            );
+        }
+
+        // There may be a new node in the new DesiredBalance that was not in the old DesiredBalance. So we'll need to iterate the nodes in
+        // the new DesiredBalance to check.
+        for (var nodeAndWeights : newWeightsPerNode.entrySet()) {
+            var discoveryNode = nodeAndWeights.getKey();
+            if (nodeNameToWeightInfo.containsKey(discoveryNode.getName()) == false) {
+                // This node is new in the new DesiredBalance, there was no entry added during iteration of the nodes in the old
+                // DesiredBalance. So we'll make a new entry with a base of zero value weights and a weights diff of the new node's weights.
+                nodeNameToWeightInfo.put(
+                    discoveryNode.getName(),
+                    new BalancingRoundSummary.NodesWeightsChanges(
+                        DesiredBalanceMetrics.NodeWeightStats.ZERO,
+                        BalancingRoundSummary.NodeWeightsDiff.create(DesiredBalanceMetrics.NodeWeightStats.ZERO, nodeAndWeights.getValue())
+                    )
+                );
+            }
+        }
+
+        return nodeNameToWeightInfo;
+    }
+
     /**
      * Adds the summary of a balancing round. If summaries are enabled, this will eventually be reported (logging, etc.). If balancer round
      * summaries are not enabled in the cluster, then the summary is immediately discarded (so as not to fill up a data structure that will
@@ -110,7 +171,7 @@ private void reportSummariesAndThenReschedule() {
      */
     private void drainAndReportSummaries() {
         var combinedSummaries = drainSummaries();
-        if (combinedSummaries == CombinedBalancingRoundSummary.EMPTY_RESULTS) {
+        if (combinedSummaries == BalancingRoundSummary.CombinedBalancingRoundSummary.EMPTY_RESULTS) {
             return;
         }
 
@@ -120,14 +181,15 @@ private void drainAndReportSummaries() {
     /**
      * Returns a combined summary of all unreported allocation round summaries: may summarize a single balancer round, multiple, or none.
      *
-     * @return {@link CombinedBalancingRoundSummary#EMPTY_RESULTS} if there are no balancing round summaries waiting to be reported.
+     * @return {@link BalancingRoundSummary.CombinedBalancingRoundSummary#EMPTY_RESULTS} if there are no balancing round summaries waiting
+     * to be reported.
      */
-    private CombinedBalancingRoundSummary drainSummaries() {
+    private BalancingRoundSummary.CombinedBalancingRoundSummary drainSummaries() {
         ArrayList<BalancingRoundSummary> batchOfSummaries = new ArrayList<>();
         while (summaries.isEmpty() == false) {
             batchOfSummaries.add(summaries.poll());
         }
-        return CombinedBalancingRoundSummary.combine(batchOfSummaries);
+        return BalancingRoundSummary.CombinedBalancingRoundSummary.combine(batchOfSummaries);
     }
 
     /**
@@ -186,7 +248,9 @@ private void rescheduleReporting() {
         }
     }
 
-    // @VisibleForTesting
+    /**
+     * Checks that the number of entries in {@link #summaries} matches the given {@code numberOfSummaries}.
+     */
     protected void verifyNumberOfSummaries(int numberOfSummaries) {
         assert numberOfSummaries == summaries.size();
     }

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

@@ -162,6 +162,7 @@ public void allocate(RoutingAllocation allocation) {
         balancer.moveShards();
         balancer.balance();
 
+        // Node weights are calculated after each internal balancing round and saved to the RoutingNodes copy.
         collectAndRecordNodeWeightStats(balancer, weightFunction, allocation);
     }
 

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

@@ -9,16 +9,149 @@
 
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
 /**
  * Summarizes the impact to the cluster as a result of a rebalancing round.
  *
- * @param numberOfShardsToMove The number of shard moves required to move from the previous desired balance to the new one.
+ * @param nodeNameToWeightChanges The shard balance weight changes for each node (by name), comparing a previous DesiredBalance shard
+ *                                allocation to a new DesiredBalance allocation.
+ * @param numberOfShardsToMove The number of shard moves required to move from the previous desired balance to the new one. Does not include
+ *                             new (index creation) or removed (index deletion) shard assignements.
  */
-public record BalancingRoundSummary(long numberOfShardsToMove) {
+public record BalancingRoundSummary(Map<String, NodesWeightsChanges> nodeNameToWeightChanges, long numberOfShardsToMove) {
+
+    /**
+     * Represents the change in weights for a node going from an old DesiredBalance to a new DesiredBalance
+     * Saves the node weights of an old DesiredBalance, along with a diff against a newer DesiredBalance.
+     *
+     * @param baseWeights The starting {@link DesiredBalanceMetrics.NodeWeightStats} of a previous DesiredBalance.
+     * @param weightsDiff The difference between the {@code baseWeights} and a new DesiredBalance.
+     */
+    record NodesWeightsChanges(DesiredBalanceMetrics.NodeWeightStats baseWeights, NodeWeightsDiff weightsDiff) {}
+
+    /**
+     * Represents the change of shard balance weights for a node, comparing an older DesiredBalance with the latest DesiredBalance.
+     *
+     * @param shardCountDiff How many more, or less, shards are assigned to the node in the latest DesiredBalance.
+     * @param diskUsageInBytesDiff How much more, or less, disk is used by shards assigned to the node in the latest DesiredBalance.
+     * @param writeLoadDiff How much more, or less, write load is estimated for shards assigned to the node in the latest DesiredBalance.
+     * @param totalWeightDiff How much more, or less, the total weight is of shards assigned to the node in the latest DesiredBalance.
+     */
+    record NodeWeightsDiff(long shardCountDiff, double diskUsageInBytesDiff, double writeLoadDiff, double totalWeightDiff) {
+
+        /**
+         * Creates a diff where the {@code base} weights will be subtracted from the {@code next} weights, to show the changes made to reach
+         * the {@code next} weights.
+         *
+         * @param base has the original weights
+         * @param next has the new weights
+         * @return The diff of ({@code next} - {@code base})
+         */
+        public static NodeWeightsDiff create(DesiredBalanceMetrics.NodeWeightStats base, DesiredBalanceMetrics.NodeWeightStats next) {
+            return new NodeWeightsDiff(
+                next.shardCount() - base.shardCount(),
+                next.diskUsageInBytes() - base.diskUsageInBytes(),
+                next.writeLoad() - base.writeLoad(),
+                next.nodeWeight() - base.nodeWeight()
+            );
+        }
+
+        /**
+         * Creates a new {@link NodeWeightsDiff} summing this instance's values with {@code otherDiff}'s values.
+         */
+        public NodeWeightsDiff combine(NodeWeightsDiff otherDiff) {
+            return new NodeWeightsDiff(
+                this.shardCountDiff + otherDiff.shardCountDiff,
+                this.diskUsageInBytesDiff + otherDiff.diskUsageInBytesDiff,
+                this.writeLoadDiff + otherDiff.writeLoadDiff,
+                this.totalWeightDiff + otherDiff.totalWeightDiff
+            );
+        }
+    }
 
     @Override
     public String toString() {
-        return "BalancingRoundSummary{" + "numberOfShardsToMove=" + numberOfShardsToMove + '}';
+        return "BalancingRoundSummary{"
+            + "nodeNameToWeightChanges"
+            + nodeNameToWeightChanges
+            + ", numberOfShardsToMove="
+            + numberOfShardsToMove
+            + '}';
+    }
+
+    /**
+     * Holds combined {@link BalancingRoundSummary} results. Essentially holds a list of the balancing events and the summed up changes
+     * across all those events: what allocation work was done across some period of time.
+     * TODO: WIP ES-10341
+     *
+     * Note that each balancing round summary is the difference between, at the time, latest desired balance and the previous desired
+     * balance. Each summary represents a step towards the next desired balance, which is based on presuming the previous desired balance is
+     * reached. So combining them is roughly the difference between the first summary's previous desired balance and the last summary's
+     * latest desired balance.
+     *
+     * @param numberOfBalancingRounds How many balancing round summaries are combined in this report.
+     * @param nodeNameToWeightChanges
+     * @param numberOfShardMoves The sum of shard moves for each balancing round being combined into a single summary.
+     */
+    public record CombinedBalancingRoundSummary(
+        int numberOfBalancingRounds,
+        Map<String, NodesWeightsChanges> nodeNameToWeightChanges,
+        long numberOfShardMoves
+    ) {
+
+        public static final CombinedBalancingRoundSummary EMPTY_RESULTS = new CombinedBalancingRoundSummary(0, new HashMap<>(), 0);
+
+        /**
+         * Merges multiple {@link BalancingRoundSummary} summaries into a single {@link CombinedBalancingRoundSummary}.
+         */
+        public static CombinedBalancingRoundSummary combine(List<BalancingRoundSummary> summaries) {
+            if (summaries.isEmpty()) {
+                return EMPTY_RESULTS;
+            }
+
+            // We will loop through the summaries and sum the weight diffs for each node entry.
+            Map<String, NodesWeightsChanges> combinedNodeNameToWeightChanges = new HashMap<>();
+
+            // Number of shards moves are simply summed across summaries. Each new balancing round is built upon the last one, so it is
+            // possible that a shard is reassigned back to a node before it even moves away, and that will still be counted as 2 moves here.
+            long numberOfShardMoves = 0;
+
+            // Total number of summaries that are being combined.
+            int numSummaries = 0;
+
+            var iterator = summaries.iterator();
+            while (iterator.hasNext()) {
+                var summary = iterator.next();
+
+                // We'll build the weight changes by keeping the node weight base from the first summary in which a node appears and then
+                // summing the weight diffs in each summary to get total weight diffs across summaries.
+                for (var nodeNameAndWeights : summary.nodeNameToWeightChanges.entrySet()) {
+                    var combined = combinedNodeNameToWeightChanges.get(nodeNameAndWeights.getKey());
+                    if (combined == null) {
+                        // Either this is the first summary, and combinedNodeNameToWeightChanges hasn't been initialized yet for this node;
+                        // or a later balancing round had a new node. Either way, initialize the node entry with the weight changes from the
+                        // first summary in which it appears.
+                        combinedNodeNameToWeightChanges.put(nodeNameAndWeights.getKey(), nodeNameAndWeights.getValue());
+                    } else {
+                        // We have at least two summaries containing this node, so let's combine them.
+                        var newCombinedChanges = new NodesWeightsChanges(
+                            combined.baseWeights,
+                            combined.weightsDiff.combine(nodeNameAndWeights.getValue().weightsDiff())
+                        );
+                        combinedNodeNameToWeightChanges.put(nodeNameAndWeights.getKey(), newCombinedChanges);
+                    }
+                }
+
+                ++numSummaries;
+                numberOfShardMoves += summary.numberOfShardsToMove();
+            }
+
+            return new CombinedBalancingRoundSummary(numSummaries, combinedNodeNameToWeightChanges, numberOfShardMoves);
+        }
+
     }
 
 }

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

@@ -38,7 +38,9 @@ public class DesiredBalanceMetrics {
      */
     public record AllocationStats(long unassignedShards, long totalAllocations, long undesiredAllocationsExcludingShuttingDownNodes) {}
 
-    public record NodeWeightStats(long shardCount, double diskUsageInBytes, double writeLoad, double nodeWeight) {}
+    public record NodeWeightStats(long shardCount, double diskUsageInBytes, double writeLoad, double nodeWeight) {
+        public static final NodeWeightStats ZERO = new NodeWeightStats(0, 0, 0, 0);
+    }
 
     // Reconciliation metrics.
     /** See {@link #unassignedShards} */

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

@@ -324,7 +324,9 @@ private void setCurrentDesiredBalance(DesiredBalance newDesiredBalance) {
             }
 
             if (currentDesiredBalanceRef.compareAndSet(oldDesiredBalance, newDesiredBalance)) {
-                balancerRoundSummaryService.addBalancerRoundSummary(calculateBalancingRoundSummary(oldDesiredBalance, newDesiredBalance));
+                balancerRoundSummaryService.addBalancerRoundSummary(
+                    AllocationBalancingRoundSummaryService.createBalancerRoundSummary(oldDesiredBalance, newDesiredBalance)
+                );
                 if (logger.isTraceEnabled()) {
                     var diff = DesiredBalance.hasChanges(oldDesiredBalance, newDesiredBalance)
                         ? "Diff: " + DesiredBalance.humanReadableDiff(oldDesiredBalance, newDesiredBalance)
@@ -339,13 +341,6 @@ private void setCurrentDesiredBalance(DesiredBalance newDesiredBalance) {
         }
     }
 
-    /**
-     * Summarizes the work required to move from an old to new desired balance shard allocation.
-     */
-    private BalancingRoundSummary calculateBalancingRoundSummary(DesiredBalance oldDesiredBalance, DesiredBalance newDesiredBalance) {
-        return new BalancingRoundSummary(DesiredBalance.shardMovements(oldDesiredBalance, newDesiredBalance));
-    }
-
     /**
      * Submits the desired balance to be reconciled (applies the desired changes to the routing table) and creates and publishes a new
      * cluster state. The data nodes will receive and apply the new cluster state to start/move/remove shards.