Tidy up DesiredBalanceMetrics

Author: nicktindall

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

@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.ToLongFunction;
 
 /**
  * Maintains balancer metrics and makes them accessible to the {@link MeterRegistry} and APM reporting. Metrics are updated
@@ -53,6 +54,19 @@ public long undesiredAllocationsExcludingShuttingDownNodes() {
                 .mapToLong(RoleAllocationStats::undesiredAllocationsExcludingShuttingDownNodes)
                 .sum();
         }
+
+        /**
+         * Return the ratio of undesired allocations to the total number of allocations.
+         *
+         * @return a value in [0.0, 1.0]
+         */
+        public double undesiredAllocationsRatio() {
+            long totalAllocations = totalAllocations();
+            if (totalAllocations == 0) {
+                return 0;
+            }
+            return undesiredAllocationsExcludingShuttingDownNodes() / (double) totalAllocations;
+        }
     }
 
     /**
@@ -86,7 +100,7 @@ public record NodeWeightStats(long shardCount, double diskUsageInBytes, double w
     public static final String UNASSIGNED_SHARDS_METRIC_NAME = "es.allocator.desired_balance.shards.unassigned.current";
     /** See {@link #totalAllocations} */
     public static final String TOTAL_SHARDS_METRIC_NAME = "es.allocator.desired_balance.shards.current";
-    /** See {@link #undesiredAllocationsExcludingShuttingDownNodes} */
+    /** See {@link #undesiredAllocations} */
     public static final String UNDESIRED_ALLOCATION_COUNT_METRIC_NAME = "es.allocator.desired_balance.allocations.undesired.current";
     /** {@link #UNDESIRED_ALLOCATION_COUNT_METRIC_NAME} / {@link #TOTAL_SHARDS_METRIC_NAME} */
     public static final String UNDESIRED_ALLOCATION_RATIO_METRIC_NAME = "es.allocator.desired_balance.allocations.undesired.ratio";
@@ -110,31 +124,14 @@ public record NodeWeightStats(long shardCount, double diskUsageInBytes, double w
     public static final String CURRENT_NODE_FORECASTED_DISK_USAGE_METRIC_NAME =
         "es.allocator.allocations.node.forecasted_disk_usage_bytes.current";
 
-    public static final AllocationStats EMPTY_ALLOCATION_STATS = new AllocationStats(-1, Map.of());
+    public static final AllocationStats EMPTY_ALLOCATION_STATS = new AllocationStats(0, Map.of());
 
     private volatile boolean nodeIsMaster = false;
 
     /**
-     * Number of unassigned shards during last reconciliation
-     */
-    private volatile long unassignedShards;
-
-    /**
-     * Total number of assigned shards during last reconciliation
-     */
-    private volatile long totalAllocations;
-
-    /**
-     * Number of assigned shards during last reconciliation that are not allocated on a desired node and need to be moved.
-     * This excludes shards that must be reassigned due to a shutting down node.
-     */
-    private volatile long undesiredAllocationsExcludingShuttingDownNodes;
-
-    /**
-     * A breakdown of shards assigned and the undesired allocations from the last reconciliation
-     * broken down by {@link ShardRouting.Role}.
+     * The stats from the most recent reconciliation
      */
-    private volatile Map<ShardRouting.Role, RoleAllocationStats> allocationStatsByRole;
+    private volatile AllocationStats lastReconciliationAllocationStats = EMPTY_ALLOCATION_STATS;
 
     private final AtomicReference<Map<DiscoveryNode, NodeWeightStats>> weightStatsPerNodeRef = new AtomicReference<>(Map.of());
     private final AtomicReference<Map<DiscoveryNode, NodeAllocationStatsAndWeight>> allocationStatsPerNodeRef = new AtomicReference<>(
@@ -149,10 +146,7 @@ public void updateMetrics(
         assert allocationStats != null : "allocation stats cannot be null";
         assert weightStatsPerNode != null : "node balance weight stats cannot be null";
         if (allocationStats != EMPTY_ALLOCATION_STATS) {
-            this.unassignedShards = allocationStats.unassignedShards;
-            this.totalAllocations = allocationStats.totalAllocations();
-            this.undesiredAllocationsExcludingShuttingDownNodes = allocationStats.undesiredAllocationsExcludingShuttingDownNodes();
-            this.allocationStatsByRole = allocationStats.allocationStatsByRole();
+            this.lastReconciliationAllocationStats = allocationStats;
         }
         weightStatsPerNodeRef.set(weightStatsPerNode);
         allocationStatsPerNodeRef.set(nodeAllocationStats);
@@ -251,23 +245,23 @@ public void setNodeIsMaster(boolean nodeIsMaster) {
     }
 
     public long unassignedShards() {
-        return unassignedShards;
+        return lastReconciliationAllocationStats.unassignedShards();
     }
 
     public long totalAllocations() {
-        return totalAllocations;
+        return lastReconciliationAllocationStats.totalAllocations();
     }
 
     public long undesiredAllocations() {
-        return undesiredAllocationsExcludingShuttingDownNodes;
+        return lastReconciliationAllocationStats.undesiredAllocationsExcludingShuttingDownNodes();
     }
 
     public Map<ShardRouting.Role, RoleAllocationStats> allocationStatsByRole() {
-        return allocationStatsByRole;
+        return lastReconciliationAllocationStats.allocationStatsByRole();
     }
 
     private List<LongWithAttributes> getUnassignedShardsMetrics() {
-        return getIfPublishing(unassignedShards);
+        return getIfPublishing(AllocationStats::unassignedShards);
     }
 
     private List<DoubleWithAttributes> getDesiredBalanceNodeWeightMetrics() {
@@ -396,25 +390,25 @@ private Map<String, Object> getNodeAttributes(DiscoveryNode node) {
     }
 
     private List<LongWithAttributes> getTotalAllocationsMetrics() {
-        return getIfPublishing(totalAllocations);
+        return getIfPublishing(AllocationStats::totalAllocations);
     }
 
     private List<LongWithAttributes> getUndesiredAllocationsExcludingShuttingDownNodesMetrics() {
-        return getIfPublishing(undesiredAllocationsExcludingShuttingDownNodes);
+        return getIfPublishing(AllocationStats::undesiredAllocationsExcludingShuttingDownNodes);
     }
 
-    private List<LongWithAttributes> getIfPublishing(long value) {
-        if (nodeIsMaster) {
-            return List.of(new LongWithAttributes(value));
+    private List<LongWithAttributes> getIfPublishing(ToLongFunction<AllocationStats> value) {
+        var currentStats = lastReconciliationAllocationStats;
+        if (currentStats != EMPTY_ALLOCATION_STATS) {
+            return List.of(new LongWithAttributes(value.applyAsLong(currentStats)));
         }
         return List.of();
     }
 
     private List<DoubleWithAttributes> getUndesiredAllocationsRatioMetrics() {
-        if (nodeIsMaster) {
-            var total = totalAllocations;
-            var undesired = undesiredAllocationsExcludingShuttingDownNodes;
-            return List.of(new DoubleWithAttributes(total != 0 ? (double) undesired / total : 0.0));
+        var currentStats = lastReconciliationAllocationStats;
+        if (currentStats != EMPTY_ALLOCATION_STATS) {
+            return List.of(new DoubleWithAttributes(currentStats.undesiredAllocationsRatio()));
         }
         return List.of();
     }
@@ -424,9 +418,7 @@ private List<DoubleWithAttributes> getUndesiredAllocationsRatioMetrics() {
      * This is best-effort because it is possible for {@link #updateMetrics} to race with this method.
      */
     public void zeroAllMetrics() {
-        unassignedShards = 0;
-        totalAllocations = 0;
-        undesiredAllocationsExcludingShuttingDownNodes = 0;
+        lastReconciliationAllocationStats = EMPTY_ALLOCATION_STATS;
         weightStatsPerNodeRef.set(Map.of());
         allocationStatsPerNodeRef.set(Map.of());
     }