Add Node Weight to GetDesiredBalance (#131025)

Extends the `_internal/desired_balance` API to return the node weights.
These are found within the `DesiredBalanceResponse.ClusterBalanceStats
.NodeBalanceStats` object. If no node weights have been calculated, then
 this value defaults to null.

Issue: #126579

Author: joshua-adams-1

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/cluster.desired_balance/10_basic.yml

@@ -2,11 +2,10 @@
 setup:
   - requires:
       cluster_features: ["gte_v8.6.0"]
-      reason: "API added in in 8.6.0"
+      reason: "API added in 8.6.0"
 
 ---
 "Test get empty desired balance":
-
   - do:
       _internal.get_desired_balance: { }
 
@@ -63,7 +62,7 @@ setup:
 
   - requires:
       cluster_features: ["gte_v8.7.0"]
-      reason: "cluster_balance_stats added in in 8.7.0"
+      reason: "cluster_balance_stats added in 8.7.0"
 
   - do:
       cluster.state: {}
@@ -111,7 +110,7 @@ setup:
 
   - requires:
       cluster_features: [gte_v8.8.0]
-      reason: "cluster_info added in in 8.8.0"
+      reason: "cluster_info added in 8.8.0"
 
   - do:
       _internal.get_desired_balance: { }
@@ -123,7 +122,7 @@ setup:
 
   - requires:
       cluster_features: ["gte_v8.8.0"]
-      reason: "node_id and roles added in in 8.8.0"
+      reason: "node_id and roles added in 8.8.0"
 
   - do:
       cluster.state: {}
@@ -142,7 +141,7 @@ setup:
 
   - requires:
       cluster_features: ["gte_v8.8.0"]
-      reason: "tier_preference added in in 8.8.0"
+      reason: "tier_preference added in 8.8.0"
 
   - do:
       indices.create:
@@ -167,7 +166,7 @@ setup:
 
   - requires:
       cluster_features: ["gte_v8.8.0"]
-      reason: "computed_shard_movements added in in 8.8.0"
+      reason: "computed_shard_movements added in 8.8.0"
 
   - do:
       _internal.get_desired_balance: { }
@@ -179,7 +178,7 @@ setup:
 
   - requires:
       cluster_features: ["gte_v8.8.0"]
-      reason: "reset API added in in 8.8.0"
+      reason: "reset API added in 8.8.0"
       test_runner_features: [ capabilities ]
       capabilities:
         - method: DELETE
@@ -194,7 +193,7 @@ setup:
 
   - requires:
       cluster_features: ["gte_v8.12.0"]
-      reason: "undesired_shard_allocation_count added in in 8.12.0"
+      reason: "undesired_shard_allocation_count added in 8.12.0"
 
   - do:
       indices.create:
@@ -232,7 +231,7 @@ setup:
 
   - requires:
       cluster_features: [ "gte_v8.12.0" ]
-      reason: "undesired_shard_allocation_count added in in 8.12.0"
+      reason: "undesired_shard_allocation_count added in 8.12.0"
 
   - do:
       _internal.get_desired_balance: { }
@@ -241,3 +240,37 @@ setup:
   - gte: { 'stats.total_allocations' : 0 }
   - gte: { 'stats.undesired_allocations' : 0 }
   - gte: { 'stats.undesired_allocations_ratio' : 0.0 }
+
+---
+"Test node_weights":
+
+  - requires:
+      capabilities:
+        - method: GET
+          path: /_internal/desired_balance
+          capabilities: [ desired_balance_node_weights_in_response ]
+      test_runner_features: [ capabilities ]
+      reason: "Node weights were added in version 9.2.0"
+
+  - do:
+      indices.create:
+        index: test
+        body:
+          settings:
+            number_of_shards: 1
+            number_of_replicas: 0
+
+  - do:
+      cluster.health:
+        index: test
+        wait_for_status: green
+
+  - do:
+      cluster.state: { }
+  - set: { nodes._arbitrary_key_: node_id }
+  - set: { nodes.$node_id.name: node_name }
+
+  - do:
+      _internal.get_desired_balance: { }
+
+  - exists: cluster_balance_stats.nodes.$node_name.node_weight

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -351,6 +351,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SHARD_WRITE_LOAD_IN_CLUSTER_INFO = def(9_126_0_00);
     public static final TransportVersion ESQL_SAMPLE_OPERATOR_STATUS = def(9_127_0_00);
     public static final TransportVersion ESQL_TOPN_TIMINGS = def(9_128_0_00);
+    public static final TransportVersion NODE_WEIGHTS_ADDED_TO_NODE_BALANCE_STATS = def(9_129_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -30,6 +30,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.ToDoubleFunction;
 
 public record ClusterBalanceStats(
@@ -213,7 +214,8 @@ public record NodeBalanceStats(
         int undesiredShardAllocations,
         double forecastWriteLoad,
         long forecastShardSize,
-        long actualShardSize
+        long actualShardSize,
+        Double nodeWeight
     ) implements Writeable, ToXContentObject {
 
         private static final String UNKNOWN_NODE_ID = "UNKNOWN";
@@ -241,14 +243,20 @@ private static NodeBalanceStats createFrom(
                 }
             }
 
+            assert desiredBalance != null;
+            Double nodeWeight = Optional.ofNullable(desiredBalance.weightsPerNode().get(routingNode.node()))
+                .map(DesiredBalanceMetrics.NodeWeightStats::nodeWeight)
+                .orElse(null);
+
             return new NodeBalanceStats(
                 routingNode.nodeId(),
                 routingNode.node().getRoles().stream().map(DiscoveryNodeRole::roleName).toList(),
                 routingNode.size(),
                 undesired,
                 forecastWriteLoad,
                 forecastShardSize,
-                actualShardSize
+                actualShardSize,
+                nodeWeight
             );
         }
 
@@ -269,7 +277,10 @@ public static NodeBalanceStats readFrom(StreamInput in) throws IOException {
                 in.getTransportVersion().onOrAfter(TransportVersions.V_8_12_0) ? in.readVInt() : -1,
                 in.readDouble(),
                 in.readLong(),
-                in.readLong()
+                in.readLong(),
+                in.getTransportVersion().onOrAfter(TransportVersions.NODE_WEIGHTS_ADDED_TO_NODE_BALANCE_STATS)
+                    ? in.readOptionalDouble()
+                    : null
             );
         }
 
@@ -286,6 +297,9 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeDouble(forecastWriteLoad);
             out.writeLong(forecastShardSize);
             out.writeLong(actualShardSize);
+            if (out.getTransportVersion().onOrAfter(TransportVersions.NODE_WEIGHTS_ADDED_TO_NODE_BALANCE_STATS)) {
+                out.writeOptionalDouble(nodeWeight);
+            }
         }
 
         @Override
@@ -294,13 +308,16 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             if (UNKNOWN_NODE_ID.equals(nodeId) == false) {
                 builder.field("node_id", nodeId);
             }
-            return builder.field("roles", roles)
+            builder.field("roles", roles)
                 .field("shard_count", shards)
                 .field("undesired_shard_allocation_count", undesiredShardAllocations)
                 .field("forecast_write_load", forecastWriteLoad)
                 .humanReadableField("forecast_disk_usage_bytes", "forecast_disk_usage", ByteSizeValue.ofBytes(forecastShardSize))
-                .humanReadableField("actual_disk_usage_bytes", "actual_disk_usage", ByteSizeValue.ofBytes(actualShardSize))
-                .endObject();
+                .humanReadableField("actual_disk_usage_bytes", "actual_disk_usage", ByteSizeValue.ofBytes(actualShardSize));
+            if (nodeWeight != null) {
+                builder.field("node_weight", nodeWeight);
+            }
+            return builder.endObject();
         }
     }
 }

server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestGetDesiredBalanceAction.java

@@ -21,6 +21,7 @@
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Set;
 
 @ServerlessScope(Scope.INTERNAL)
 public class RestGetDesiredBalanceAction extends BaseRestHandler {
@@ -44,4 +45,9 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             new RestRefCountedChunkedToXContentListener<>(restChannel)
         );
     }
+
+    @Override
+    public Set<String> supportedCapabilities() {
+        return Set.of("desired_balance_node_weights_in_response");
+    }
 }

server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponseTests.java

@@ -91,7 +91,8 @@ private ClusterBalanceStats.NodeBalanceStats randomNodeBalanceStats() {
             randomNonNegativeInt(),
             randomDouble(),
             randomNonNegativeLong(),
-            randomNonNegativeLong()
+            randomNonNegativeLong(),
+            randomDouble()
         );
     }
 
@@ -283,7 +284,8 @@ public void testToXContent() throws IOException {
                     "undesired_shard_allocation_count",
                     "forecast_write_load",
                     "forecast_disk_usage_bytes",
-                    "actual_disk_usage_bytes"
+                    "actual_disk_usage_bytes",
+                    "node_weight"
                 )
             );
             assertEquals(nodesStats.get("node_id"), entry.getValue().nodeId());
@@ -293,6 +295,7 @@ public void testToXContent() throws IOException {
             assertEquals(nodesStats.get("forecast_write_load"), entry.getValue().forecastWriteLoad());
             assertEquals(nodesStats.get("forecast_disk_usage_bytes"), entry.getValue().forecastShardSize());
             assertEquals(nodesStats.get("actual_disk_usage_bytes"), entry.getValue().actualShardSize());
+            assertEquals(nodesStats.get("node_weight"), entry.getValue().nodeWeight());
         }
 
         // routing table