Expose per node counts (#93439)

docs/reference/cluster/get-desired-balance.asciidoc

@@ -33,53 +33,70 @@ The API returns the following result:
     "reconciliation_time_in_millis": 0
   },
   "cluster_balance_stats" : {
-    {
+    "tiers": {
       "data_hot" : {
-        "total_shard_size" : {
-          "total" : 36.0,
-          "min" : 10.0,
-          "max" : 16.0,
-          "average" : 12.0,
-          "std_dev" : 2.8284271247461903
+        "shard_count" : {
+          "total" : 7.0,
+          "min" : 2.0,
+          "max" : 3.0,
+          "average" : 2.3333333333333335,
+          "std_dev" : 0.4714045207910317
         },
-        "total_write_load" : {
+        "forecast_write_load" : {
           "total" : 21.0,
           "min" : 6.0,
           "max" : 8.5,
           "average" : 7.0,
           "std_dev" : 1.0801234497346435
         },
-        "shard_count" : {
-          "total" : 7.0,
-          "min" : 2.0,
-          "max" : 3.0,
-          "average" : 2.3333333333333335,
-          "std_dev" : 0.4714045207910317
+        "forecast_disk_usage" : {
+          "total" : 36.0,
+          "min" : 10.0,
+          "max" : 16.0,
+          "average" : 12.0,
+          "std_dev" : 2.8284271247461903
         }
       },
       "data_warm" : {
-        "total_shard_size" : {
-          "total" : 42.0,
-          "min" : 12.0,
-          "max" : 18.0,
-          "average" : 14.0,
-          "std_dev" : 2.8284271247461903
+        "shard_count" : {
+          "total" : 3.0,
+          "min" : 1.0,
+          "max" : 1.0,
+          "average" : 1.0,
+          "std_dev" : 0.0
         },
-        "total_write_load" : {
+        "forecast_write_load" : {
           "total" : 0.0,
           "min" : 0.0,
           "max" : 0.0,
           "average" : 0.0,
           "std_dev" : 0.0
         },
-        "shard_count" : {
-          "total" : 3.0,
-          "min" : 1.0,
-          "max" : 1.0,
-          "average" : 1.0,
-          "std_dev" : 0.0
+        "forecast_disk_usage" : {
+          "total" : 42.0,
+          "min" : 12.0,
+          "max" : 18.0,
+          "average" : 14.0,
+          "std_dev" : 2.8284271247461903
         }
       }
+    },
+    "nodes": {
+      "node-1": {
+        "shard_count": 10,
+        "forecast_write_load": 8.5,
+        "forecast_disk_usage_bytes": 498435
+      },
+      "node-2": {
+        "shard_count": 15,
+        "forecast_write_load": 3.25,
+        "forecast_disk_usage_bytes": 384935
+      },
+      "node-3": {
+        "shard_count": 12,
+        "forecast_write_load": 6.0,
+        "forecast_disk_usage_bytes": 648766
+      }
     }
   },
   "routing_table": {
@@ -95,8 +112,8 @@ The API returns the following result:
             "relocating_node_is_desired": false,
             "shard_id": 0,
             "index": "test",
-            "forecasted_write_load": 8.0,
-            "forecasted_shard_size_in_bytes": 1024
+            "forecast_write_load": 8.0,
+            "forecast_shard_size_in_bytes": 1024
           }
         ],
         "desired": {
@@ -119,8 +136,8 @@ The API returns the following result:
             "relocating_node_is_desired": false,
             "shard_id": 1,
             "index": "test",
-            "forecasted_write_load": null,
-            "forecasted_shard_size_in_bytes": null
+            "forecast_write_load": null,
+            "forecast_shard_size_in_bytes": null
           }
         ],
         "desired": {

qa/smoke-test-multinode/src/yamlRestTest/resources/rest-api-spec/test/smoke_test_multinode/30_desired_balance.yml

@@ -75,21 +75,27 @@ setup:
       _internal.get_desired_balance: { }
 
   - is_true: 'cluster_balance_stats'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.total'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.min'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.max'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.average'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.std_dev'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.total'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.min'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.max'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.average'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.std_dev'
-  - is_true: 'cluster_balance_stats.data_content.shard_count'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.total'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.min'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.max'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.average'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.std_dev'
+  - is_true: 'cluster_balance_stats.tiers'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.total'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.min'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.max'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.average'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.std_dev'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.total'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.min'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.max'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.average'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.std_dev'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.total'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.min'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.max'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.average'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.std_dev'
+  - is_true: 'cluster_balance_stats.nodes'
+  - is_true: 'cluster_balance_stats.nodes.test-cluster-0'
+  - gte: { 'cluster_balance_stats.nodes.test-cluster-0.shard_count' : 0 }
+  - gte: { 'cluster_balance_stats.nodes.test-cluster-0.forecast_write_load': 0.0 }
+  - gte: { 'cluster_balance_stats.nodes.test-cluster-0.forecast_disk_usage_bytes' : 0 }

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

@@ -30,24 +30,30 @@ setup:
       _internal.get_desired_balance: { }
 
   - is_true: 'cluster_balance_stats'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.total'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.min'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.max'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.average'
-  - is_true: 'cluster_balance_stats.data_content.total_shard_size.std_dev'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.total'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.min'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.max'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.average'
-  - is_true: 'cluster_balance_stats.data_content.total_write_load.std_dev'
-  - is_true: 'cluster_balance_stats.data_content.shard_count'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.total'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.min'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.max'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.average'
-  - is_true: 'cluster_balance_stats.data_content.shard_count.std_dev'
+  - is_true: 'cluster_balance_stats.tiers'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.total'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.min'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.max'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.average'
+  - is_true: 'cluster_balance_stats.tiers.data_content.shard_count.std_dev'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.total'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.min'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.max'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.average'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_write_load.std_dev'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.total'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.min'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.max'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.average'
+  - is_true: 'cluster_balance_stats.tiers.data_content.forecast_disk_usage.std_dev'
+  - is_true: 'cluster_balance_stats.nodes'
+  - is_true: 'cluster_balance_stats.nodes.test-cluster-0'
+  - gte: { 'cluster_balance_stats.nodes.test-cluster-0.shard_count' : 0 }
+  - gte: { 'cluster_balance_stats.nodes.test-cluster-0.forecast_write_load': 0.0 }
+  - gte: { 'cluster_balance_stats.nodes.test-cluster-0.forecast_disk_usage_bytes' : 0 }
 
 ---
 "Test get desired balance for single shard":
@@ -81,8 +87,8 @@ setup:
   - is_true: 'routing_table.test.0.current.0.node_is_desired'
   - is_false: 'routing_table.test.0.current.0.relocating_node'
   - is_false: 'routing_table.test.0.current.0.relocating_node_is_desired'
-  - is_false: 'routing_table.test.0.current.0.forecasted_write_load'
-  - is_false: 'routing_table.test.0.current.0.forecasted_shard_size_in_bytes'
+  - is_false: 'routing_table.test.0.current.0.forecast_write_load'
+  - is_false: 'routing_table.test.0.current.0.forecast_shard_size_in_bytes'
   - match: { routing_table.test.0.desired.total: 1 }
   - gte: { routing_table.test.0.desired.unassigned: 0 }
   - gte: { routing_table.test.0.desired.ignored: 0 }

server/src/main/java/org/elasticsearch/action/admin/cluster/allocation/DesiredBalanceResponse.java

@@ -174,8 +174,8 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
         boolean relocatingNodeIsDesired,
         int shardId,
         String index,
-        @Nullable Double forecastedWriteLoad,
-        @Nullable Long forecastedShardSizeInBytes
+        @Nullable Double forecastWriteLoad,
+        @Nullable Long forecastShardSizeInBytes
     ) implements Writeable, ToXContentObject {
 
         private static final TransportVersion ADD_FORECASTS_VERSION = TransportVersion.V_8_7_0;
@@ -223,8 +223,8 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
             out.writeVInt(shardId);
             out.writeString(index);
             if (out.getTransportVersion().onOrAfter(ADD_FORECASTS_VERSION)) {
-                out.writeOptionalDouble(forecastedWriteLoad);
-                out.writeOptionalLong(forecastedShardSizeInBytes);
+                out.writeOptionalDouble(forecastWriteLoad);
+                out.writeOptionalLong(forecastShardSizeInBytes);
             } else {
                 out.writeMissingWriteable(AllocationId.class);
             }
@@ -241,8 +241,8 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
                 .field("relocating_node_is_desired", relocatingNodeIsDesired)
                 .field("shard_id", shardId)
                 .field("index", index)
-                .field("forecasted_write_load", forecastedWriteLoad)
-                .field("forecasted_shard_size_in_bytes", forecastedShardSizeInBytes)
+                .field("forecast_write_load", forecastWriteLoad)
+                .field("forecast_shard_size_in_bytes", forecastShardSizeInBytes)
                 .endObject();
         }
     }

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

@@ -17,6 +17,7 @@ import org.elasticsearch.cluster.routing.allocation.WriteLoadForecaster;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -28,49 +29,58 @@ import java.util.List;
 import java.util.Map;
 import java.util.function.ToDoubleFunction;
 
-public record ClusterBalanceStats(Map<String, TierBalanceStats> tiers) implements Writeable, ToXContentObject {
+public record ClusterBalanceStats(Map<String, TierBalanceStats> tiers, Map<String, NodeBalanceStats> nodes)
+    implements
+        Writeable,
+        ToXContentObject {
 
-    public static ClusterBalanceStats EMPTY = new ClusterBalanceStats(Map.of());
+    public static ClusterBalanceStats EMPTY = new ClusterBalanceStats(Map.of(), Map.of());
 
     public static ClusterBalanceStats createFrom(ClusterState clusterState, WriteLoadForecaster writeLoadForecaster) {
-        var tierToNodeStats = new HashMap<String, List<NodeStats>>();
+        var tierToNodeStats = new HashMap<String, List<NodeBalanceStats>>();
+        var nodes = new HashMap<String, NodeBalanceStats>();
         for (RoutingNode routingNode : clusterState.getRoutingNodes()) {
             var dataRoles = routingNode.node().getRoles().stream().filter(DiscoveryNodeRole::canContainData).toList();
             if (dataRoles.isEmpty()) {
                 continue;
             }
-            var nodeStats = NodeStats.createFrom(routingNode, clusterState.metadata(), writeLoadForecaster);
+            var nodeStats = NodeBalanceStats.createFrom(routingNode, clusterState.metadata(), writeLoadForecaster);
+            nodes.put(routingNode.node().getName(), nodeStats);
             for (DiscoveryNodeRole role : dataRoles) {
                 tierToNodeStats.computeIfAbsent(role.roleName(), ignored -> new ArrayList<>()).add(nodeStats);
             }
         }
-        return new ClusterBalanceStats(Maps.transformValues(tierToNodeStats, TierBalanceStats::createFrom));
+        return new ClusterBalanceStats(Maps.transformValues(tierToNodeStats, TierBalanceStats::createFrom), nodes);
     }
 
     public static ClusterBalanceStats readFrom(StreamInput in) throws IOException {
-        return new ClusterBalanceStats(in.readImmutableMap(StreamInput::readString, TierBalanceStats::readFrom));
+        return new ClusterBalanceStats(
+            in.readImmutableMap(StreamInput::readString, TierBalanceStats::readFrom),
+            in.readImmutableMap(StreamInput::readString, NodeBalanceStats::readFrom)
+        );
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         out.writeMap(tiers, StreamOutput::writeString, StreamOutput::writeWriteable);
+        out.writeMap(nodes, StreamOutput::writeString, StreamOutput::writeWriteable);
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        return builder.map(tiers);
+        return builder.startObject().field("tiers").map(tiers).field("nodes").map(nodes).endObject();
     }
 
-    public record TierBalanceStats(MetricStats shardCount, MetricStats totalWriteLoad, MetricStats totalShardSize)
+    public record TierBalanceStats(MetricStats shardCount, MetricStats forecastWriteLoad, MetricStats forecastShardSize)
         implements
             Writeable,
             ToXContentObject {
 
-        private static TierBalanceStats createFrom(List<NodeStats> nodes) {
+        private static TierBalanceStats createFrom(List<NodeBalanceStats> nodes) {
             return new TierBalanceStats(
                 MetricStats.createFrom(nodes, it -> it.shards),
-                MetricStats.createFrom(nodes, it -> it.totalWriteLoad),
-                MetricStats.createFrom(nodes, it -> it.totalShardSize)
+                MetricStats.createFrom(nodes, it -> it.forecastWriteLoad),
+                MetricStats.createFrom(nodes, it -> it.forecastShardSize)
             );
         }
 
@@ -81,30 +91,30 @@ public record ClusterBalanceStats(Map<String, TierBalanceStats> tiers) implement
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             shardCount.writeTo(out);
-            totalWriteLoad.writeTo(out);
-            totalShardSize.writeTo(out);
+            forecastWriteLoad.writeTo(out);
+            forecastShardSize.writeTo(out);
         }
 
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             return builder.startObject()
                 .field("shard_count", shardCount)
-                .field("total_write_load", totalWriteLoad)
-                .field("total_shard_size", totalShardSize)
+                .field("forecast_write_load", forecastWriteLoad)
+                .field("forecast_disk_usage", forecastShardSize)
                 .endObject();
         }
     }
 
     public record MetricStats(double total, double min, double max, double average, double stdDev) implements Writeable, ToXContentObject {
 
-        private static MetricStats createFrom(List<NodeStats> nodes, ToDoubleFunction<NodeStats> metricExtractor) {
+        private static MetricStats createFrom(List<NodeBalanceStats> nodes, ToDoubleFunction<NodeBalanceStats> metricExtractor) {
             assert nodes.isEmpty() == false : "Stats must be created from non empty nodes";
             double total = 0.0;
             double total2 = 0.0;
             double min = Double.POSITIVE_INFINITY;
             double max = Double.NEGATIVE_INFINITY;
             int count = 0;
-            for (NodeStats node : nodes) {
+            for (NodeBalanceStats node : nodes) {
                 var metric = metricExtractor.applyAsDouble(node);
                 if (Double.isNaN(metric)) {
                     continue;
@@ -145,9 +155,9 @@ public record ClusterBalanceStats(Map<String, TierBalanceStats> tiers) implement
         }
     }
 
-    private record NodeStats(int shards, double totalWriteLoad, long totalShardSize) {
+    public record NodeBalanceStats(int shards, double forecastWriteLoad, long forecastShardSize) implements Writeable, ToXContentObject {
 
-        private static NodeStats createFrom(RoutingNode routingNode, Metadata metadata, WriteLoadForecaster writeLoadForecaster) {
+        private static NodeBalanceStats createFrom(RoutingNode routingNode, Metadata metadata, WriteLoadForecaster writeLoadForecaster) {
             double totalWriteLoad = 0.0;
             long totalShardSize = 0L;
 
@@ -158,7 +168,27 @@ public record ClusterBalanceStats(Map<String, TierBalanceStats> tiers) implement
                 totalShardSize += indexMetadata.getForecastedShardSizeInBytes().orElse(0);
             }
 
-            return new NodeStats(routingNode.size(), totalWriteLoad, totalShardSize);
+            return new NodeBalanceStats(routingNode.size(), totalWriteLoad, totalShardSize);
+        }
+
+        public static NodeBalanceStats readFrom(StreamInput in) throws IOException {
+            return new NodeBalanceStats(in.readInt(), in.readDouble(), in.readLong());
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeInt(shards);
+            out.writeDouble(forecastWriteLoad);
+            out.writeLong(forecastShardSize);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            return builder.startObject()
+                .field("shard_count", shards)
+                .field("forecast_write_load", forecastWriteLoad)
+                .humanReadableField("forecast_disk_usage_bytes", "forecast_disk_usage", ByteSizeValue.ofBytes(forecastShardSize))
+                .endObject();
         }
     }
 }

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

@@ -66,7 +66,8 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
                         DiscoveryNodeRole.DATA_COLD_NODE_ROLE,
                         DiscoveryNodeRole.DATA_FROZEN_NODE_ROLE
                     )
-                ).stream().map(DiscoveryNodeRole::roleName).collect(toMap(identity(), ignore -> randomTierBalanceStats()))
+                ).stream().map(DiscoveryNodeRole::roleName).collect(toMap(identity(), ignore -> randomTierBalanceStats())),
+            randomList(10, () -> randomAlphaOfLength(10)).stream().collect(toMap(identity(), ignore -> randomNodeBalanceStats()))
         );
     }
 
@@ -78,6 +79,14 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
         );
     }
 
+    private ClusterBalanceStats.NodeBalanceStats randomNodeBalanceStats() {
+        return new ClusterBalanceStats.NodeBalanceStats(
+            randomIntBetween(0, Integer.MAX_VALUE),
+            randomDouble(),
+            randomLongBetween(0, Long.MAX_VALUE)
+        );
+    }
+
     private Map<String, Map<Integer, DesiredBalanceResponse.DesiredShards>> randomRoutingTable() {
         Map<String, Map<Integer, DesiredBalanceResponse.DesiredShards>> routingTable = new HashMap<>();
         for (int i = 0; i < randomInt(8); i++) {
@@ -166,10 +175,14 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
 
         // cluster balance stats
         Map<String, Object> clusterBalanceStats = (Map<String, Object>) json.get("cluster_balance_stats");
-        assertEquals(clusterBalanceStats.keySet(), response.getClusterBalanceStats().tiers().keySet());
+        assertEquals(Set.of("tiers", "nodes"), clusterBalanceStats.keySet());
+
+        // tier balance stats
+        Map<String, Object> tiers = (Map<String, Object>) clusterBalanceStats.get("tiers");
+        assertEquals(tiers.keySet(), response.getClusterBalanceStats().tiers().keySet());
         for (var entry : response.getClusterBalanceStats().tiers().entrySet()) {
-            Map<String, Object> tierStats = (Map<String, Object>) clusterBalanceStats.get(entry.getKey());
-            assertEquals(Set.of("shard_count", "total_write_load", "total_shard_size"), tierStats.keySet());
+            Map<String, Object> tierStats = (Map<String, Object>) tiers.get(entry.getKey());
+            assertEquals(Set.of("shard_count", "forecast_write_load", "forecast_disk_usage"), tierStats.keySet());
 
             Map<String, Object> shardCountStats = (Map<String, Object>) tierStats.get("shard_count");
             assertEquals(Set.of("total", "average", "min", "max", "std_dev"), shardCountStats.keySet());
@@ -179,21 +192,32 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
             assertEquals(shardCountStats.get("max"), entry.getValue().shardCount().max());
             assertEquals(shardCountStats.get("std_dev"), entry.getValue().shardCount().stdDev());
 
-            Map<String, Object> totalWriteLoadStats = (Map<String, Object>) tierStats.get("total_write_load");
+            Map<String, Object> totalWriteLoadStats = (Map<String, Object>) tierStats.get("forecast_write_load");
             assertEquals(Set.of("total", "average", "min", "max", "std_dev"), totalWriteLoadStats.keySet());
-            assertEquals(totalWriteLoadStats.get("total"), entry.getValue().totalWriteLoad().total());
-            assertEquals(totalWriteLoadStats.get("average"), entry.getValue().totalWriteLoad().average());
-            assertEquals(totalWriteLoadStats.get("min"), entry.getValue().totalWriteLoad().min());
-            assertEquals(totalWriteLoadStats.get("max"), entry.getValue().totalWriteLoad().max());
-            assertEquals(totalWriteLoadStats.get("std_dev"), entry.getValue().totalWriteLoad().stdDev());
+            assertEquals(totalWriteLoadStats.get("total"), entry.getValue().forecastWriteLoad().total());
+            assertEquals(totalWriteLoadStats.get("average"), entry.getValue().forecastWriteLoad().average());
+            assertEquals(totalWriteLoadStats.get("min"), entry.getValue().forecastWriteLoad().min());
+            assertEquals(totalWriteLoadStats.get("max"), entry.getValue().forecastWriteLoad().max());
+            assertEquals(totalWriteLoadStats.get("std_dev"), entry.getValue().forecastWriteLoad().stdDev());
 
-            Map<String, Object> totalShardStats = (Map<String, Object>) tierStats.get("total_shard_size");
+            Map<String, Object> totalShardStats = (Map<String, Object>) tierStats.get("forecast_disk_usage");
             assertEquals(Set.of("total", "average", "min", "max", "std_dev"), totalShardStats.keySet());
-            assertEquals(totalShardStats.get("total"), entry.getValue().totalShardSize().total());
-            assertEquals(totalShardStats.get("average"), entry.getValue().totalShardSize().average());
-            assertEquals(totalShardStats.get("min"), entry.getValue().totalShardSize().min());
-            assertEquals(totalShardStats.get("max"), entry.getValue().totalShardSize().max());
-            assertEquals(totalShardStats.get("std_dev"), entry.getValue().totalShardSize().stdDev());
+            assertEquals(totalShardStats.get("total"), entry.getValue().forecastShardSize().total());
+            assertEquals(totalShardStats.get("average"), entry.getValue().forecastShardSize().average());
+            assertEquals(totalShardStats.get("min"), entry.getValue().forecastShardSize().min());
+            assertEquals(totalShardStats.get("max"), entry.getValue().forecastShardSize().max());
+            assertEquals(totalShardStats.get("std_dev"), entry.getValue().forecastShardSize().stdDev());
+        }
+        // node balance stats
+        Map<String, Object> nodes = (Map<String, Object>) clusterBalanceStats.get("nodes");
+        assertEquals(nodes.keySet(), response.getClusterBalanceStats().nodes().keySet());
+        for (var entry : response.getClusterBalanceStats().nodes().entrySet()) {
+            Map<String, Object> nodesStats = (Map<String, Object>) nodes.get(entry.getKey());
+            assertEquals(Set.of("shard_count", "forecast_write_load", "forecast_disk_usage_bytes"), nodesStats.keySet());
+
+            assertEquals(nodesStats.get("shard_count"), entry.getValue().shards());
+            assertEquals(nodesStats.get("forecast_write_load"), entry.getValue().forecastWriteLoad());
+            assertEquals(nodesStats.get("forecast_disk_usage_bytes"), entry.getValue().forecastShardSize());
         }
 
         // routing table
@@ -221,8 +245,8 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
                     assertEquals(jsonShard.get("relocating_node_is_desired"), shardView.relocatingNodeIsDesired());
                     assertEquals(jsonShard.get("shard_id"), shardView.shardId());
                     assertEquals(jsonShard.get("index"), shardView.index());
-                    assertEquals(jsonShard.get("forecasted_write_load"), shardView.forecastedWriteLoad());
-                    assertEquals(jsonShard.get("forecasted_shard_size_in_bytes"), shardView.forecastedShardSizeInBytes());
+                    assertEquals(jsonShard.get("forecast_write_load"), shardView.forecastWriteLoad());
+                    assertEquals(jsonShard.get("forecast_shard_size_in_bytes"), shardView.forecastShardSizeInBytes());
                 }
 
                 Map<String, Object> jsonDesired = (Map<String, Object>) jsonDesiredShard.get("desired");

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

@@ -236,14 +236,11 @@ public class TransportGetDesiredBalanceActionTests extends ESAllocationTestCase
                     assertEquals(shard.index().getName(), shardView.index());
                     assertEquals(shard.shardId().id(), shardView.shardId());
                     var forecastedWriteLoad = TEST_WRITE_LOAD_FORECASTER.getForecastedWriteLoad(indexMetadata);
-                    assertEquals(
-                        forecastedWriteLoad.isPresent() ? forecastedWriteLoad.getAsDouble() : null,
-                        shardView.forecastedWriteLoad()
-                    );
+                    assertEquals(forecastedWriteLoad.isPresent() ? forecastedWriteLoad.getAsDouble() : null, shardView.forecastWriteLoad());
                     var forecastedShardSizeInBytes = indexMetadata.getForecastedShardSizeInBytes();
                     assertEquals(
                         forecastedShardSizeInBytes.isPresent() ? forecastedShardSizeInBytes.getAsLong() : null,
-                        shardView.forecastedShardSizeInBytes()
+                        shardView.forecastShardSizeInBytes()
                     );
                     Set<String> desiredNodeIds = Optional.ofNullable(shardAssignments.get(shard.shardId()))
                         .map(ShardAssignment::nodeIds)

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

@@ -39,9 +39,9 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
 
         var clusterState = createClusterState(
             List.of(
-                newNode("node-1", Set.of(DATA_CONTENT_NODE_ROLE)),
-                newNode("node-2", Set.of(DATA_CONTENT_NODE_ROLE)),
-                newNode("node-3", Set.of(DATA_CONTENT_NODE_ROLE))
+                newNode("node-1", "node-1", Set.of(DATA_CONTENT_NODE_ROLE)),
+                newNode("node-2", "node-2", Set.of(DATA_CONTENT_NODE_ROLE)),
+                newNode("node-3", "node-3", Set.of(DATA_CONTENT_NODE_ROLE))
             ),
             List.of(
                 startedIndex("index-1", null, null, "node-1", "node-2"),
@@ -63,6 +63,14 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
                             new ClusterBalanceStats.MetricStats(0.0, 0.0, 0.0, 0.0, 0.0),
                             new ClusterBalanceStats.MetricStats(0.0, 0.0, 0.0, 0.0, 0.0)
                         )
+                    ),
+                    Map.of(
+                        "node-1",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 0.0, 0L),
+                        "node-2",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 0.0, 0L),
+                        "node-3",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 0.0, 0L)
                     )
                 )
             )
@@ -73,9 +81,9 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
 
         var clusterState = createClusterState(
             List.of(
-                newNode("node-1", Set.of(DATA_CONTENT_NODE_ROLE)),
-                newNode("node-2", Set.of(DATA_CONTENT_NODE_ROLE)),
-                newNode("node-3", Set.of(DATA_CONTENT_NODE_ROLE))
+                newNode("node-1", "node-1", Set.of(DATA_CONTENT_NODE_ROLE)),
+                newNode("node-2", "node-2", Set.of(DATA_CONTENT_NODE_ROLE)),
+                newNode("node-3", "node-3", Set.of(DATA_CONTENT_NODE_ROLE))
             ),
             List.of(
                 startedIndex("index-1", 1.5, 8L, "node-1", "node-2"),
@@ -97,6 +105,14 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
                             new ClusterBalanceStats.MetricStats(12.0, 3.5, 4.5, 4.0, stdDev(3.5, 4.0, 4.5)),
                             new ClusterBalanceStats.MetricStats(36.0, 10.0, 14.0, 12.0, stdDev(10.0, 12.0, 14.0))
                         )
+                    ),
+                    Map.of(
+                        "node-1",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 3.5, 14L),
+                        "node-2",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 4.0, 12L),
+                        "node-3",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 4.5, 10L)
                     )
                 )
             )
@@ -107,12 +123,12 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
 
         var clusterState = createClusterState(
             List.of(
-                newNode("node-hot-1", Set.of(DATA_CONTENT_NODE_ROLE, DATA_HOT_NODE_ROLE)),
-                newNode("node-hot-2", Set.of(DATA_CONTENT_NODE_ROLE, DATA_HOT_NODE_ROLE)),
-                newNode("node-hot-3", Set.of(DATA_CONTENT_NODE_ROLE, DATA_HOT_NODE_ROLE)),
-                newNode("node-warm-1", Set.of(DATA_WARM_NODE_ROLE)),
-                newNode("node-warm-2", Set.of(DATA_WARM_NODE_ROLE)),
-                newNode("node-warm-3", Set.of(DATA_WARM_NODE_ROLE))
+                newNode("node-hot-1", "node-hot-1", Set.of(DATA_CONTENT_NODE_ROLE, DATA_HOT_NODE_ROLE)),
+                newNode("node-hot-2", "node-hot-2", Set.of(DATA_CONTENT_NODE_ROLE, DATA_HOT_NODE_ROLE)),
+                newNode("node-hot-3", "node-hot-3", Set.of(DATA_CONTENT_NODE_ROLE, DATA_HOT_NODE_ROLE)),
+                newNode("node-warm-1", "node-warm-1", Set.of(DATA_WARM_NODE_ROLE)),
+                newNode("node-warm-2", "node-warm-2", Set.of(DATA_WARM_NODE_ROLE)),
+                newNode("node-warm-3", "node-warm-3", Set.of(DATA_WARM_NODE_ROLE))
             ),
             List.of(
                 startedIndex("index-hot-1", 4.0, 4L, "node-hot-1", "node-hot-2", "node-hot-3"),
@@ -148,6 +164,20 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
                             new ClusterBalanceStats.MetricStats(0.0, 0.0, 0.0, 0.0, 0.0),
                             new ClusterBalanceStats.MetricStats(42.0, 12.0, 18.0, 14.0, stdDev(12.0, 12.0, 18.0))
                         )
+                    ),
+                    Map.of(
+                        "node-hot-1",
+                        new ClusterBalanceStats.NodeBalanceStats(3, 8.5, 16L),
+                        "node-hot-2",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 6.0, 10L),
+                        "node-hot-3",
+                        new ClusterBalanceStats.NodeBalanceStats(2, 6.5, 10L),
+                        "node-warm-1",
+                        new ClusterBalanceStats.NodeBalanceStats(1, 0.0, 12L),
+                        "node-warm-2",
+                        new ClusterBalanceStats.NodeBalanceStats(1, 0.0, 12L),
+                        "node-warm-3",
+                        new ClusterBalanceStats.NodeBalanceStats(1, 0.0, 18L)
                     )
                 )
             )
@@ -158,9 +188,9 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
 
         var clusterState = createClusterState(
             List.of(
-                newNode("node-1", Set.of(DATA_CONTENT_NODE_ROLE)),
-                newNode("node-2", Set.of(DATA_CONTENT_NODE_ROLE)),
-                newNode("node-3", Set.of(DATA_CONTENT_NODE_ROLE))
+                newNode("node-1", "node-1", Set.of(DATA_CONTENT_NODE_ROLE)),
+                newNode("node-2", "node-2", Set.of(DATA_CONTENT_NODE_ROLE)),
+                newNode("node-3", "node-3", Set.of(DATA_CONTENT_NODE_ROLE))
             ),
             List.of()
         );
@@ -178,6 +208,14 @@ public class ClusterBalanceStatsTests extends ESAllocationTestCase {
                             new ClusterBalanceStats.MetricStats(0.0, 0.0, 0.0, 0.0, 0.0),
                             new ClusterBalanceStats.MetricStats(0.0, 0.0, 0.0, 0.0, 0.0)
                         )
+                    ),
+                    Map.of(
+                        "node-1",
+                        new ClusterBalanceStats.NodeBalanceStats(0, 0.0, 0L),
+                        "node-2",
+                        new ClusterBalanceStats.NodeBalanceStats(0, 0.0, 0L),
+                        "node-3",
+                        new ClusterBalanceStats.NodeBalanceStats(0, 0.0, 0L)
                     )
                 )
             )