Add cluster_info to the GET /_internal/desired-balance endpoint output (#94272)

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

@@ -170,6 +170,42 @@ The API returns the following result:
         }
       }
     }
+  },
+  "cluster_info" : {
+      "nodes" : {
+        "UPYt8VwWTt-IADAEbqpLxA" : {
+          "node_name" : "node-1",
+          "least_available" : {
+            "path" : "/data",
+            "total_bytes" : 1440713945,
+            "used_bytes" : 1222486407,
+            "free_bytes" : 218227538,
+            "free_disk_percent" : 15.1,
+            "used_disk_percent" : 84.9
+          },
+          "most_available" : {
+            "path" : "/data",
+            "total_bytes" : 1440713945,
+            "used_bytes" : 1222486407,
+            "free_bytes" : 218227538,
+            "free_disk_percent" : 15.1,
+            "used_disk_percent" : 84.9
+          }
+        }
+      },
+      "shard_sizes" : {
+        "[test][0][p]_bytes" : 1720826288,
+        "[test][1][p]_bytes" : 1720826288
+      },
+      "shard_data_set_sizes" : {
+        "[test][0][p]_bytes" : 1720826288,
+        "[test][1][p]_bytes" : 1720826288
+      },
+      "shard_paths" : {
+        "NodeAndShard[nodeId=UPYt8VwWTt-IADAEbqpLxA, shardId=[test][0]]" : "/data",
+        "NodeAndShard[nodeId=bgC66tboTIeFQ0VgRGI4Gg, shardId=[test][0]]" : "/data"
+      },
+      "reserved_sizes" : []
   }
 }
 --------------------------------------------------

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

@@ -111,3 +111,15 @@ setup:
   - gte: { 'cluster_balance_stats.nodes.$node_name.forecast_write_load': 0.0 }
   - gte: { 'cluster_balance_stats.nodes.$node_name.forecast_disk_usage_bytes' : 0 }
   - gte: { 'cluster_balance_stats.nodes.$node_name.actual_disk_usage_bytes' : 0 }
+
+---
+"Test cluster_info":
+
+  - skip:
+      version: " - 8.7.99"
+      reason: "Field added in in 8.8.0"
+
+  - do:
+      _internal.get_desired_balance: { }
+
+  - is_true: 'cluster_info'

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

@@ -67,6 +67,18 @@ setup:
   - gte: { 'cluster_balance_stats.nodes.$node_name.forecast_disk_usage_bytes' : 0 }
   - gte: { 'cluster_balance_stats.nodes.$node_name.actual_disk_usage_bytes' : 0 }
 
+---
+"Test cluster_info":
+
+  - skip:
+      version: " - 8.7.99"
+      reason: "Field added in in 8.8.0"
+
+  - do:
+      _internal.get_desired_balance: { }
+
+  - is_true: 'cluster_info'
+
 ---
 "Test cluster_balance_stats contains node ID":
 

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

@@ -9,6 +9,7 @@ package org.elasticsearch.action.admin.cluster.allocation;
 
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.cluster.ClusterInfo;
 import org.elasticsearch.cluster.routing.AllocationId;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.allocation.allocator.ClusterBalanceStats;
@@ -35,19 +36,23 @@ import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.singleCh
 public class DesiredBalanceResponse extends ActionResponse implements ChunkedToXContentObject {
 
     private static final TransportVersion CLUSTER_BALANCE_STATS_VERSION = TransportVersion.V_8_7_0;
+    private static final TransportVersion CLUSTER_INFO_VERSION = TransportVersion.V_8_8_0;
 
     private final DesiredBalanceStats stats;
     private final ClusterBalanceStats clusterBalanceStats;
     private final Map<String, Map<Integer, DesiredShards>> routingTable;
+    private final ClusterInfo clusterInfo;
 
     public DesiredBalanceResponse(
         DesiredBalanceStats stats,
         ClusterBalanceStats clusterBalanceStats,
-        Map<String, Map<Integer, DesiredShards>> routingTable
+        Map<String, Map<Integer, DesiredShards>> routingTable,
+        ClusterInfo clusterInfo
     ) {
         this.stats = stats;
         this.clusterBalanceStats = clusterBalanceStats;
         this.routingTable = routingTable;
+        this.clusterInfo = clusterInfo;
     }
 
     public static DesiredBalanceResponse from(StreamInput in) throws IOException {
@@ -56,7 +61,8 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
             in.getTransportVersion().onOrAfter(CLUSTER_BALANCE_STATS_VERSION)
                 ? ClusterBalanceStats.readFrom(in)
                 : ClusterBalanceStats.EMPTY,
-            in.readImmutableMap(StreamInput::readString, v -> v.readImmutableMap(StreamInput::readVInt, DesiredShards::from))
+            in.readImmutableMap(StreamInput::readString, v -> v.readImmutableMap(StreamInput::readVInt, DesiredShards::from)),
+            in.getTransportVersion().onOrAfter(CLUSTER_INFO_VERSION) ? new ClusterInfo(in) : ClusterInfo.EMPTY
         );
     }
 
@@ -64,7 +70,7 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
     public void writeTo(StreamOutput out) throws IOException {
         stats.writeTo(out);
         if (out.getTransportVersion().onOrAfter(CLUSTER_BALANCE_STATS_VERSION)) {
-            clusterBalanceStats.writeTo(out);
+            out.writeWriteable(clusterBalanceStats);
         }
         out.writeMap(
             routingTable,
@@ -75,6 +81,9 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
                 (desiredShardsOut, desiredShards) -> desiredShards.writeTo(desiredShardsOut)
             )
         );
+        if (out.getTransportVersion().onOrAfter(CLUSTER_INFO_VERSION)) {
+            out.writeWriteable(clusterInfo);
+        }
     }
 
     @Override
@@ -87,7 +96,11 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
                 (builder, p) -> builder.startObject("routing_table")
             ),
             routingTableToXContentChunked(),
-            singleChunk((builder, p) -> builder.endObject(), (builder, p) -> builder.endObject())
+            singleChunk(
+                (builder, p) -> builder.endObject(),
+                (builder, p) -> builder.startObject("cluster_info").value(clusterInfo).endObject(),
+                (builder, p) -> builder.endObject()
+            )
         );
     }
 
@@ -114,18 +127,23 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
         return routingTable;
     }
 
+    public ClusterInfo getClusterInfo() {
+        return clusterInfo;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;
         return o instanceof DesiredBalanceResponse that
             && Objects.equals(stats, that.stats)
             && Objects.equals(clusterBalanceStats, that.clusterBalanceStats)
-            && Objects.equals(routingTable, that.routingTable);
+            && Objects.equals(routingTable, that.routingTable)
+            && Objects.equals(clusterInfo, that.clusterInfo);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(stats, clusterBalanceStats, routingTable);
+        return Objects.hash(stats, clusterBalanceStats, routingTable, clusterInfo);
     }
 
     @Override
@@ -136,6 +154,8 @@ public class DesiredBalanceResponse extends ActionResponse implements ChunkedToX
             + clusterBalanceStats
             + ", routingTable="
             + routingTable
+            + ", clusterInfo="
+            + clusterInfo
             + "}";
     }
 

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

@@ -91,11 +91,13 @@ public class TransportGetDesiredBalanceAction extends TransportMasterNodeReadAct
             listener.onFailure(new ResourceNotFoundException("Desired balance is not computed yet"));
             return;
         }
+        var clusterInfo = clusterInfoService.getClusterInfo();
         listener.onResponse(
             new DesiredBalanceResponse(
                 desiredBalanceShardsAllocator.getStats(),
-                ClusterBalanceStats.createFrom(state, clusterInfoService.getClusterInfo(), writeLoadForecaster),
-                createRoutingTable(state, latestDesiredBalance)
+                ClusterBalanceStats.createFrom(state, clusterInfo, writeLoadForecaster),
+                createRoutingTable(state, latestDesiredBalance),
+                clusterInfo
             )
         );
     }

server/src/main/java/org/elasticsearch/cluster/ClusterInfo.java

@@ -87,21 +87,15 @@ public class ClusterInfo implements ToXContentFragment, Writeable {
         this.leastAvailableSpaceUsage = in.readImmutableMap(StreamInput::readString, DiskUsage::new);
         this.mostAvailableSpaceUsage = in.readImmutableMap(StreamInput::readString, DiskUsage::new);
         this.shardSizes = in.readImmutableMap(StreamInput::readString, StreamInput::readLong);
-        if (in.getTransportVersion().onOrAfter(DATA_SET_SIZE_SIZE_VERSION)) {
-            this.shardDataSetSizes = in.readImmutableMap(ShardId::new, StreamInput::readLong);
-        } else {
-            this.shardDataSetSizes = Map.of();
-        }
-        if (in.getTransportVersion().onOrAfter(DATA_PATH_NEW_KEY_VERSION)) {
-            this.dataPath = in.readImmutableMap(NodeAndShard::new, StreamInput::readString);
-        } else {
-            this.dataPath = in.readImmutableMap(nested -> NodeAndShard.from(new ShardRouting(nested)), StreamInput::readString);
-        }
-        if (in.getTransportVersion().onOrAfter(StoreStats.RESERVED_BYTES_VERSION)) {
-            this.reservedSpace = in.readImmutableMap(NodeAndPath::new, ReservedSpace::new);
-        } else {
-            this.reservedSpace = Map.of();
-        }
+        this.shardDataSetSizes = in.getTransportVersion().onOrAfter(DATA_SET_SIZE_SIZE_VERSION)
+            ? in.readImmutableMap(ShardId::new, StreamInput::readLong)
+            : Map.of();
+        this.dataPath = in.getTransportVersion().onOrAfter(DATA_PATH_NEW_KEY_VERSION)
+            ? in.readImmutableMap(NodeAndShard::new, StreamInput::readString)
+            : in.readImmutableMap(nested -> NodeAndShard.from(new ShardRouting(nested)), StreamInput::readString);
+        this.reservedSpace = in.getTransportVersion().onOrAfter(StoreStats.RESERVED_BYTES_VERSION)
+            ? in.readImmutableMap(NodeAndPath::new, ReservedSpace::new)
+            : Map.of();
     }
 
     @Override

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

@@ -7,6 +7,8 @@
  */
 package org.elasticsearch.action.admin.cluster.allocation;
 
+import org.elasticsearch.cluster.ClusterInfo;
+import org.elasticsearch.cluster.ClusterInfoTests;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.cluster.routing.ShardRoutingState;
 import org.elasticsearch.cluster.routing.allocation.allocator.ClusterBalanceStats;
@@ -39,7 +41,12 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
 
     @Override
     protected DesiredBalanceResponse createTestInstance() {
-        return new DesiredBalanceResponse(randomDesiredBalanceStats(), randomClusterBalanceStats(), randomRoutingTable());
+        return new DesiredBalanceResponse(
+            randomDesiredBalanceStats(),
+            randomClusterBalanceStats(),
+            randomRoutingTable(),
+            randomClusterInfo()
+        );
     }
 
     private DesiredBalanceStats randomDesiredBalanceStats() {
@@ -90,6 +97,10 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
         );
     }
 
+    private ClusterInfo randomClusterInfo() {
+        return ClusterInfoTests.randomClusterInfo();
+    }
+
     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++) {
@@ -132,21 +143,30 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
 
     @Override
     protected DesiredBalanceResponse mutateInstance(DesiredBalanceResponse instance) {
-        return switch (randomInt(3)) {
+        return switch (randomInt(4)) {
             case 0 -> new DesiredBalanceResponse(
                 randomValueOtherThan(instance.getStats(), this::randomDesiredBalanceStats),
                 instance.getClusterBalanceStats(),
-                instance.getRoutingTable()
+                instance.getRoutingTable(),
+                instance.getClusterInfo()
             );
             case 1 -> new DesiredBalanceResponse(
                 instance.getStats(),
                 randomValueOtherThan(instance.getClusterBalanceStats(), this::randomClusterBalanceStats),
-                instance.getRoutingTable()
+                instance.getRoutingTable(),
+                instance.getClusterInfo()
             );
             case 2 -> new DesiredBalanceResponse(
                 instance.getStats(),
                 instance.getClusterBalanceStats(),
-                randomValueOtherThan(instance.getRoutingTable(), this::randomRoutingTable)
+                randomValueOtherThan(instance.getRoutingTable(), this::randomRoutingTable),
+                instance.getClusterInfo()
+            );
+            case 3 -> new DesiredBalanceResponse(
+                instance.getStats(),
+                instance.getClusterBalanceStats(),
+                instance.getRoutingTable(),
+                randomValueOtherThan(instance.getClusterInfo(), this::randomClusterInfo)
             );
             default -> randomValueOtherThan(instance, this::createTestInstance);
         };
@@ -157,13 +177,14 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
         DesiredBalanceResponse response = new DesiredBalanceResponse(
             randomDesiredBalanceStats(),
             randomClusterBalanceStats(),
-            randomRoutingTable()
+            randomRoutingTable(),
+            randomClusterInfo()
         );
 
         Map<String, Object> json = createParser(
             ChunkedToXContent.wrapAsToXContent(response).toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)
         ).map();
-        assertThat(json.keySet(), containsInAnyOrder("stats", "cluster_balance_stats", "routing_table"));
+        assertThat(json.keySet(), containsInAnyOrder("stats", "cluster_balance_stats", "routing_table", "cluster_info"));
 
         // stats
         Map<String, Object> stats = (Map<String, Object>) json.get("stats");
@@ -279,7 +300,7 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
 
     public void testChunking() {
         AbstractChunkedSerializingTestCase.assertChunkCount(
-            new DesiredBalanceResponse(randomDesiredBalanceStats(), randomClusterBalanceStats(), randomRoutingTable()),
+            new DesiredBalanceResponse(randomDesiredBalanceStats(), randomClusterBalanceStats(), randomRoutingTable(), randomClusterInfo()),
             response -> response.getRoutingTable().size() + 2
         );
     }

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

@@ -54,6 +54,7 @@ import static org.elasticsearch.cluster.ClusterModule.BALANCED_ALLOCATOR;
 import static org.elasticsearch.cluster.ClusterModule.DESIRED_BALANCE_ALLOCATOR;
 import static org.elasticsearch.cluster.ClusterModule.SHARDS_ALLOCATOR_TYPE_SETTING;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.notNullValue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -207,7 +208,8 @@ public class TransportGetDesiredBalanceActionTests extends ESAllocationTestCase
             randomInt(Integer.MAX_VALUE)
         );
         when(desiredBalanceShardsAllocator.getStats()).thenReturn(desiredBalanceStats);
-        when(clusterInfoService.getClusterInfo()).thenReturn(ClusterInfo.EMPTY);
+        ClusterInfo clusterInfo = ClusterInfo.EMPTY;
+        when(clusterInfoService.getClusterInfo()).thenReturn(clusterInfo);
 
         var clusterState = ClusterState.builder(ClusterName.DEFAULT)
             .metadata(metadataBuilder.build())
@@ -220,8 +222,9 @@ public class TransportGetDesiredBalanceActionTests extends ESAllocationTestCase
         ArgumentCaptor<DesiredBalanceResponse> desiredBalanceResponseCaptor = ArgumentCaptor.forClass(DesiredBalanceResponse.class);
         verify(listener).onResponse(desiredBalanceResponseCaptor.capture());
         DesiredBalanceResponse desiredBalanceResponse = desiredBalanceResponseCaptor.getValue();
-        assertEquals(desiredBalanceStats, desiredBalanceResponse.getStats());
-        assertNotNull(desiredBalanceResponse.getClusterBalanceStats());
+        assertThat(desiredBalanceResponse.getStats(), equalTo(desiredBalanceStats));
+        assertThat(desiredBalanceResponse.getClusterBalanceStats(), notNullValue());
+        assertThat(desiredBalanceResponse.getClusterInfo(), equalTo(clusterInfo));
         assertEquals(indexShards.keySet(), desiredBalanceResponse.getRoutingTable().keySet());
         for (var e : desiredBalanceResponse.getRoutingTable().entrySet()) {
             String index = e.getKey();

server/src/test/java/org/elasticsearch/cluster/ClusterInfoTests.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.cluster;
 
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 
@@ -23,6 +24,15 @@ public class ClusterInfoTests extends AbstractWireSerializingTestCase<ClusterInf
 
     @Override
     protected ClusterInfo createTestInstance() {
+        return randomClusterInfo();
+    }
+
+    @Override
+    protected ClusterInfo mutateInstance(ClusterInfo instance) {
+        return randomClusterInfo();
+    }
+
+    public static ClusterInfo randomClusterInfo() {
         return new ClusterInfo(
             randomDiskUsage(),
             randomDiskUsage(),
@@ -33,11 +43,6 @@ public class ClusterInfoTests extends AbstractWireSerializingTestCase<ClusterInf
         );
     }
 
-    @Override
-    protected ClusterInfo mutateInstance(ClusterInfo instance) {
-        return createTestInstance();
-    }
-
     private static Map<String, DiskUsage> randomDiskUsage() {
         int numEntries = randomIntBetween(0, 128);
         Map<String, DiskUsage> builder = new HashMap<>(numEntries);
@@ -58,26 +63,26 @@ public class ClusterInfoTests extends AbstractWireSerializingTestCase<ClusterInf
 
     private static Map<String, Long> randomShardSizes() {
         int numEntries = randomIntBetween(0, 128);
-        Map<String, Long> builder = new HashMap<>(numEntries);
+        var builder = Maps.<String, Long>newMapWithExpectedSize(numEntries);
         for (int i = 0; i < numEntries; i++) {
-            String key = randomAlphaOfLength(32);
-            long shardSize = randomIntBetween(0, Integer.MAX_VALUE);
-            builder.put(key, shardSize);
+            builder.put(ClusterInfo.shardIdentifierFromRouting(randomShardId(), randomBoolean()), randomLongBetween(0, Integer.MAX_VALUE));
         }
         return builder;
     }
 
     private static Map<ShardId, Long> randomDataSetSizes() {
         int numEntries = randomIntBetween(0, 128);
-        Map<ShardId, Long> builder = new HashMap<>(numEntries);
+        var builder = Maps.<ShardId, Long>newMapWithExpectedSize(numEntries);
         for (int i = 0; i < numEntries; i++) {
-            ShardId key = new ShardId(randomAlphaOfLength(10), randomAlphaOfLength(10), between(0, Integer.MAX_VALUE));
-            long shardSize = randomIntBetween(0, Integer.MAX_VALUE);
-            builder.put(key, shardSize);
+            builder.put(randomShardId(), randomLongBetween(0, Integer.MAX_VALUE));
         }
         return builder;
     }
 
+    private static ShardId randomShardId() {
+        return new ShardId(randomAlphaOfLength(10), randomAlphaOfLength(10), between(0, Integer.MAX_VALUE));
+    }
+
     private static Map<ClusterInfo.NodeAndShard, String> randomRoutingToDataPath() {
         int numEntries = randomIntBetween(0, 128);
         Map<ClusterInfo.NodeAndShard, String> builder = new HashMap<>(numEntries);
@@ -101,5 +106,4 @@ public class ClusterInfoTests extends AbstractWireSerializingTestCase<ClusterInf
         }
         return builder;
     }
-
 }