Export desired balance node weight and its components as metrics (#115854)

Note that this includes only the three node-level weight components (out of the four), as we were not sure how to aggregate and expose the index-specific component and how useful it will be at all. 

- some of the weight components are also recalculated or exposed as stats (not APM metrics) else where (e.g. `AllocationStatsService`) but since they are available right where we calculate the weight (which we also want), I have just exported all of them together.
- How to pass the weight from the BalancedAllocator which is used as a delegated allocator in the desired balance allocator, and from there to the reconciler where we publish, could probably also be done differently, but using `RoutingNodes` and `DesiredBalance` seemed to make more sense to me. Not sure if it is blasphemy for those more familiar with the allocation code!
- I liked the `DesiredBalanceMetrics` and how its used so I tried to clean up its existing usage a bit and colocate the new metrics.

Relates ES-9866

Author: pxsalehi

server/src/internalClusterTest/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceReconcilerMetricsIT.java

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
+import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.PluginsService;
@@ -17,10 +18,16 @@
 import org.hamcrest.Matcher;
 
 import java.util.Collection;
+import java.util.stream.Collectors;
 
 import static org.hamcrest.Matchers.empty;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.in;
+import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 
+@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST, numDataNodes = 0)
 public class DesiredBalanceReconcilerMetricsIT extends ESIntegTestCase {
 
     @Override
@@ -31,6 +38,7 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
     public void testDesiredBalanceGaugeMetricsAreOnlyPublishedByCurrentMaster() throws Exception {
         internalCluster().ensureAtLeastNumDataNodes(2);
         prepareCreate("test").setSettings(indexSettings(2, 1)).get();
+        indexRandom(randomBoolean(), "test", between(50, 100));
         ensureGreen();
 
         assertOnlyMasterIsPublishingMetrics();
@@ -45,6 +53,59 @@ public void testDesiredBalanceGaugeMetricsAreOnlyPublishedByCurrentMaster() thro
         }
     }
 
+    public void testDesiredBalanceNodeWeightMetrics() {
+        internalCluster().startNodes(2);
+        prepareCreate("test").setSettings(indexSettings(2, 1)).get();
+        indexRandom(randomBoolean(), "test", between(50, 100));
+        ensureGreen();
+        final var telemetryPlugin = getTelemetryPlugin(internalCluster().getMasterName());
+        telemetryPlugin.collect();
+        assertThat(telemetryPlugin.getLongGaugeMeasurement(DesiredBalanceMetrics.UNASSIGNED_SHARDS_METRIC_NAME), not(empty()));
+        assertThat(telemetryPlugin.getLongGaugeMeasurement(DesiredBalanceMetrics.TOTAL_SHARDS_METRIC_NAME), not(empty()));
+        assertThat(telemetryPlugin.getLongGaugeMeasurement(DesiredBalanceMetrics.UNDESIRED_ALLOCATION_COUNT_METRIC_NAME), not(empty()));
+        assertThat(telemetryPlugin.getDoubleGaugeMeasurement(DesiredBalanceMetrics.UNDESIRED_ALLOCATION_RATIO_METRIC_NAME), not(empty()));
+
+        var nodeIds = internalCluster().clusterService().state().nodes().stream().map(DiscoveryNode::getId).collect(Collectors.toSet());
+        var nodeNames = internalCluster().clusterService().state().nodes().stream().map(DiscoveryNode::getName).collect(Collectors.toSet());
+
+        final var nodeWeightsMetrics = telemetryPlugin.getDoubleGaugeMeasurement(
+            DesiredBalanceMetrics.DESIRED_BALANCE_NODE_WEIGHT_METRIC_NAME
+        );
+        assertThat(nodeWeightsMetrics.size(), equalTo(2));
+        for (var nodeStat : nodeWeightsMetrics) {
+            assertThat(nodeStat.value().doubleValue(), greaterThanOrEqualTo(0.0));
+            assertThat((String) nodeStat.attributes().get("node_id"), is(in(nodeIds)));
+            assertThat((String) nodeStat.attributes().get("node_name"), is(in(nodeNames)));
+        }
+        final var nodeShardCountMetrics = telemetryPlugin.getLongGaugeMeasurement(
+            DesiredBalanceMetrics.DESIRED_BALANCE_NODE_SHARD_COUNT_METRIC_NAME
+        );
+        assertThat(nodeShardCountMetrics.size(), equalTo(2));
+        for (var nodeStat : nodeShardCountMetrics) {
+            assertThat(nodeStat.value().longValue(), equalTo(2L));
+            assertThat((String) nodeStat.attributes().get("node_id"), is(in(nodeIds)));
+            assertThat((String) nodeStat.attributes().get("node_name"), is(in(nodeNames)));
+        }
+        final var nodeWriteLoadMetrics = telemetryPlugin.getDoubleGaugeMeasurement(
+            DesiredBalanceMetrics.DESIRED_BALANCE_NODE_WRITE_LOAD_METRIC_NAME
+        );
+        assertThat(nodeWriteLoadMetrics.size(), equalTo(2));
+        for (var nodeStat : nodeWriteLoadMetrics) {
+            assertThat(nodeStat.value().doubleValue(), greaterThanOrEqualTo(0.0));
+            assertThat((String) nodeStat.attributes().get("node_id"), is(in(nodeIds)));
+            assertThat((String) nodeStat.attributes().get("node_name"), is(in(nodeNames)));
+        }
+        final var nodeDiskUsageMetrics = telemetryPlugin.getDoubleGaugeMeasurement(
+            DesiredBalanceMetrics.DESIRED_BALANCE_NODE_DISK_USAGE_METRIC_NAME
+        );
+        assertThat(nodeDiskUsageMetrics.size(), equalTo(2));
+        for (var nodeStat : nodeDiskUsageMetrics) {
+            assertThat(nodeStat.value().doubleValue(), greaterThanOrEqualTo(0.0));
+            assertThat((String) nodeStat.attributes().get("node_id"), is(in(nodeIds)));
+            assertThat((String) nodeStat.attributes().get("node_name"), is(in(nodeNames)));
+        }
+    }
+
     private static void assertOnlyMasterIsPublishingMetrics() {
         String masterNodeName = internalCluster().getMasterName();
         String[] nodeNames = internalCluster().getNodeNames();
@@ -54,16 +115,33 @@ private static void assertOnlyMasterIsPublishingMetrics() {
     }
 
     private static void assertMetricsAreBeingPublished(String nodeName, boolean shouldBePublishing) {
-        final TestTelemetryPlugin testTelemetryPlugin = internalCluster().getInstance(PluginsService.class, nodeName)
-            .filterPlugins(TestTelemetryPlugin.class)
-            .findFirst()
-            .orElseThrow();
+        final TestTelemetryPlugin testTelemetryPlugin = getTelemetryPlugin(nodeName);
         testTelemetryPlugin.resetMeter();
         testTelemetryPlugin.collect();
         Matcher<Collection<?>> matcher = shouldBePublishing ? not(empty()) : empty();
         assertThat(testTelemetryPlugin.getLongGaugeMeasurement(DesiredBalanceMetrics.UNASSIGNED_SHARDS_METRIC_NAME), matcher);
         assertThat(testTelemetryPlugin.getLongGaugeMeasurement(DesiredBalanceMetrics.TOTAL_SHARDS_METRIC_NAME), matcher);
         assertThat(testTelemetryPlugin.getLongGaugeMeasurement(DesiredBalanceMetrics.UNDESIRED_ALLOCATION_COUNT_METRIC_NAME), matcher);
         assertThat(testTelemetryPlugin.getDoubleGaugeMeasurement(DesiredBalanceMetrics.UNDESIRED_ALLOCATION_RATIO_METRIC_NAME), matcher);
+        assertThat(testTelemetryPlugin.getDoubleGaugeMeasurement(DesiredBalanceMetrics.DESIRED_BALANCE_NODE_WEIGHT_METRIC_NAME), matcher);
+        assertThat(
+            testTelemetryPlugin.getDoubleGaugeMeasurement(DesiredBalanceMetrics.DESIRED_BALANCE_NODE_WRITE_LOAD_METRIC_NAME),
+            matcher
+        );
+        assertThat(
+            testTelemetryPlugin.getDoubleGaugeMeasurement(DesiredBalanceMetrics.DESIRED_BALANCE_NODE_DISK_USAGE_METRIC_NAME),
+            matcher
+        );
+        assertThat(
+            testTelemetryPlugin.getLongGaugeMeasurement(DesiredBalanceMetrics.DESIRED_BALANCE_NODE_SHARD_COUNT_METRIC_NAME),
+            matcher
+        );
+    }
+
+    private static TestTelemetryPlugin getTelemetryPlugin(String nodeName) {
+        return internalCluster().getInstance(PluginsService.class, nodeName)
+            .filterPlugins(TestTelemetryPlugin.class)
+            .findFirst()
+            .orElseThrow();
     }
 }

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.routing.UnassignedInfo.AllocationStatus;
 import org.elasticsearch.cluster.routing.allocation.ExistingShardsAllocator;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceMetrics;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.core.Assertions;
@@ -76,6 +77,8 @@ public class RoutingNodes implements Iterable<RoutingNode> {
     private final Map<String, Set<String>> attributeValuesByAttribute;
     private final Map<String, Recoveries> recoveriesPerNode;
 
+    private Map<DiscoveryNode, DesiredBalanceMetrics.NodeWeightStats> balanceWeightStatsPerNode;
+
     /**
      * Creates an immutable instance from the {@link RoutingTable} and {@link DiscoveryNodes} found in a cluster state. Used to initialize
      * the routing nodes in {@link ClusterState#getRoutingNodes()}. This method should not be used directly, use
@@ -89,6 +92,14 @@ public static RoutingNodes mutable(RoutingTable routingTable, DiscoveryNodes dis
         return new RoutingNodes(routingTable, discoveryNodes, false);
     }
 
+    public void setBalanceWeightStatsPerNode(Map<DiscoveryNode, DesiredBalanceMetrics.NodeWeightStats> weightStatsPerNode) {
+        this.balanceWeightStatsPerNode = weightStatsPerNode;
+    }
+
+    public Map<DiscoveryNode, DesiredBalanceMetrics.NodeWeightStats> getBalanceWeightStatsPerNode() {
+        return balanceWeightStatsPerNode;
+    }
+
     private RoutingNodes(RoutingTable routingTable, DiscoveryNodes discoveryNodes, boolean readOnly) {
         this.readOnly = readOnly;
         this.recoveriesPerNode = new HashMap<>();
@@ -97,6 +108,7 @@ private RoutingNodes(RoutingTable routingTable, DiscoveryNodes discoveryNodes, b
         this.unassignedShards = new UnassignedShards(this);
         this.attributeValuesByAttribute = Collections.synchronizedMap(new HashMap<>());
 
+        balanceWeightStatsPerNode = Maps.newMapWithExpectedSize(discoveryNodes.getDataNodes().size());
         nodesToShards = Maps.newMapWithExpectedSize(discoveryNodes.getDataNodes().size());
         // fill in the nodeToShards with the "live" nodes
         var dataNodes = discoveryNodes.getDataNodes().keySet();

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.cluster.ClusterInfo;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.routing.RoutingNode;
 import org.elasticsearch.cluster.routing.RoutingNodes;
 import org.elasticsearch.cluster.routing.ShardRouting;
@@ -159,6 +160,25 @@ public void allocate(RoutingAllocation allocation) {
         balancer.allocateUnassigned();
         balancer.moveShards();
         balancer.balance();
+
+        collectAndRecordNodeWeightStats(balancer, weightFunction, allocation);
+    }
+
+    private void collectAndRecordNodeWeightStats(Balancer balancer, WeightFunction weightFunction, RoutingAllocation allocation) {
+        Map<DiscoveryNode, DesiredBalanceMetrics.NodeWeightStats> nodeLevelWeights = new HashMap<>();
+        for (var entry : balancer.nodes.entrySet()) {
+            var node = entry.getValue();
+            nodeLevelWeights.put(
+                node.routingNode.node(),
+                new DesiredBalanceMetrics.NodeWeightStats(
+                    node.numShards(),
+                    node.diskUsageInBytes(),
+                    node.writeLoad(),
+                    weightFunction.nodeWeight(balancer, node)
+                )
+            );
+        }
+        allocation.routingNodes().setBalanceWeightStatsPerNode(nodeLevelWeights);
     }
 
     @Override
@@ -275,11 +295,15 @@ private static class WeightFunction {
         }
 
         float weight(Balancer balancer, ModelNode node, String index) {
-            final float weightShard = node.numShards() - balancer.avgShardsPerNode();
             final float weightIndex = node.numShards(index) - balancer.avgShardsPerNode(index);
+            return nodeWeight(balancer, node) + theta1 * weightIndex;
+        }
+
+        float nodeWeight(Balancer balancer, ModelNode node) {
+            final float weightShard = node.numShards() - balancer.avgShardsPerNode();
             final float ingestLoad = (float) (node.writeLoad() - balancer.avgWriteLoadPerNode());
             final float diskUsage = (float) (node.diskUsageInBytes() - balancer.avgDiskUsageInBytesPerNode());
-            return theta0 * weightShard + theta1 * weightIndex + theta2 * ingestLoad + theta3 * diskUsage;
+            return theta0 * weightShard + theta2 * ingestLoad + theta3 * diskUsage;
         }
 
         float minWeightDelta(Balancer balancer, String index) {

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

@@ -9,6 +9,7 @@
 
 package org.elasticsearch.cluster.routing.allocation.allocator;
 
+import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.index.shard.ShardId;
 
@@ -19,8 +20,18 @@
  * The desired balance of the cluster, indicating which nodes should hold a copy of each shard.
  *
  * @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.BalancedShardsAllocator.WeightFunction#nodeWeight}
  */
-public record DesiredBalance(long lastConvergedIndex, Map<ShardId, ShardAssignment> assignments) {
+public record DesiredBalance(
+    long lastConvergedIndex,
+    Map<ShardId, ShardAssignment> assignments,
+    Map<DiscoveryNode, DesiredBalanceMetrics.NodeWeightStats> weightsPerNode
+) {
+
+    public DesiredBalance(long lastConvergedIndex, Map<ShardId, ShardAssignment> assignments) {
+        this(lastConvergedIndex, assignments, Map.of());
+    }
 
     public static final DesiredBalance INITIAL = new DesiredBalance(-1, Map.of());
 

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

@@ -368,7 +368,7 @@ public DesiredBalance compute(
         }
 
         long lastConvergedIndex = hasChanges ? previousDesiredBalance.lastConvergedIndex() : desiredBalanceInput.index();
-        return new DesiredBalance(lastConvergedIndex, assignments);
+        return new DesiredBalance(lastConvergedIndex, assignments, routingNodes.getBalanceWeightStatsPerNode());
     }
 
     private static Map<ShardId, ShardAssignment> collectShardAssignments(RoutingNodes routingNodes) {