DesiredBalance: expose it via _internal/desired_balance (#91038)

Add an internal endpoint for exposed the desired balance and computation stats at a master node as GET _internal/desired_balance and returns
```
{
  "stats": {
    "computation_active": false,
    "computation_submitted": 5,
    "computation_executed": 5,
    "computation_converged": 5,
    "computation_iterations": 4,
    "computation_converged_index": 4,
    "computation_time_in_millis": 0,
    "reconciliation_time_in_millis": 0
  },
  "routing_table": {
    "test": {
      "0": {
        "current": [
          {
            "state": "STARTED",
            "primary": true,
            "node": "UPYt8VwWTt-IADAEbqpLxA",
            "node_is_desired": true,
            "relocating_node": null,
            "relocating_node_is_desired": false,
            "shard_id": 0,
            "index": "test"
          }
        ],
        "desired": {
          "node_ids": [
            "UPYt8VwWTt-IADAEbqpLxA"
          ],
          "total": 1,
          "unassigned": 0,
          "ignored": 0
        }
      },
      "1": {
        "current": [
          {
            "state": "STARTED",
            "primary": true,
            "node": "2x1VTuSOQdeguXPdN73yRw",
            "node_is_desired": true,
            "relocating_node": null,
            "relocating_node_is_desired": false,
            "shard_id": 1,
            "index": "test"
          }
        ],
        "desired": {
          "node_ids": [
            "2x1VTuSOQdeguXPdN73yRw"
          ],
          "total": 1,
          "unassigned": 0,
          "ignored": 0
        }
      }
    }
  }
}
```

Fixes #90583

docs/changelog/91038.yaml

@@ -0,0 +1,6 @@
+pr: 91038
+summary: '`DesiredBalance:` expose it via _internal/desired_balance'
+area: Allocation
+type: enhancement
+issues:
+ - 90583

docs/reference/cluster.asciidoc

@@ -115,3 +115,5 @@ include::cluster/update-desired-nodes.asciidoc[]
 include::cluster/get-desired-nodes.asciidoc[]
 
 include::cluster/delete-desired-nodes.asciidoc[]
+
+include::cluster/get-desired-balance.asciidoc[]

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

@@ -0,0 +1,85 @@
+[[get-desired-balance]]
+=== Get desired balance API
+++++
+<titleabbrev>Get desired balance</titleabbrev>
+++++
+
+NOTE: {cloud-only}
+
+Exposes the desired balance and basic metrics.
+
+[[get-desired-balance-request]]
+==== {api-request-title}
+
+[source,console]
+--------------------------------------------------
+GET /_internal/desired_balance
+--------------------------------------------------
+// TEST[skip:Can't reliably test desired balance]
+
+The API returns the following result:
+
+[source,console-result]
+--------------------------------------------------
+{
+  "stats": {
+    "computation_active": false,
+    "computation_submitted": 5,
+    "computation_executed": 5,
+    "computation_converged": 5,
+    "computation_iterations": 4,
+    "computation_converged_index": 4,
+    "computation_time_in_millis": 0,
+    "reconciliation_time_in_millis": 0
+  },
+  "routing_table": {
+    "test": {
+      "0": {
+        "current": [
+          {
+            "state": "STARTED",
+            "primary": true,
+            "node": "UPYt8VwWTt-IADAEbqpLxA",
+            "node_is_desired": true,
+            "relocating_node": null,
+            "relocating_node_is_desired": false,
+            "shard_id": 0,
+            "index": "test"
+          }
+        ],
+        "desired": {
+          "node_ids": [
+            "UPYt8VwWTt-IADAEbqpLxA"
+          ],
+          "total": 1,
+          "unassigned": 0,
+          "ignored": 0
+        }
+      },
+      "1": {
+        "current": [
+          {
+            "state": "STARTED",
+            "primary": true,
+            "node": "2x1VTuSOQdeguXPdN73yRw",
+            "node_is_desired": true,
+            "relocating_node": null,
+            "relocating_node_is_desired": false,
+            "shard_id": 1,
+            "index": "test"
+          }
+        ],
+        "desired": {
+          "node_ids": [
+            "2x1VTuSOQdeguXPdN73yRw"
+          ],
+          "total": 1,
+          "unassigned": 0,
+          "ignored": 0
+        }
+      }
+    }
+  }
+}
+--------------------------------------------------
+// TEST[skip:Can't reliably test desired balance]

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

@@ -0,0 +1,66 @@
+---
+setup:
+  - skip:
+      version: " - 8.5.99"
+      reason: "API added in in 8.6.0"
+
+---
+"Test get desired balance for multiple shards and replicas":
+  - do:
+      indices.create:
+        index: test
+        body:
+          settings:
+            number_of_shards: 2
+            number_of_replicas: 1
+
+  - do:
+      cluster.health:
+        index: test
+        wait_for_status: green
+
+  - do:
+      _internal.get_desired_balance: { }
+
+  - gte: { stats.computation_submitted: 0 }
+  - gte: { stats.computation_executed: 0 }
+  - gte: { stats.computation_converged: 0 }
+  - gte: { stats.computation_iterations: 0 }
+  - gte: { stats.computation_converged_index: 0 }
+  - gte: { stats.computation_time_in_millis: 0 }
+  - gte: { stats.reconciliation_time_in_millis: 0 }
+
+  # Can't match node ids and desired node status for replicas since we don't know
+  # whether the test is run on a single or multi-node cluster. Also can't match the status
+  # because we can't control the ordering of replicas and primaries in the output
+
+  - match: { routing_table.test.0.current.0.state: STARTED }
+  - match: { routing_table.test.0.current.0.shard_id: 0 }
+  - match: { routing_table.test.0.current.0.index: test }
+  - is_false: 'routing_table.test.0.current.0.relocating_node'
+  - is_false: 'routing_table.test.0.current.0.relocating_node_is_desired'
+  - match: { routing_table.test.0.current.1.state: STARTED }
+  - match: { routing_table.test.0.current.1.shard_id: 0 }
+  - match: { routing_table.test.0.current.1.index: test }
+  - is_false: 'routing_table.test.0.current.1.relocating_node'
+  - is_false: 'routing_table.test.0.current.1.relocating_node_is_desired'
+  - match: { routing_table.test.0.desired.total: 2 }
+  - gte: { routing_table.test.0.desired.unassigned: 0 }
+  - gte: { routing_table.test.0.desired.ignored: 0 }
+  - is_true: 'routing_table.test.0.desired.node_ids'
+
+  - match: { routing_table.test.1.current.0.state: STARTED }
+  - match: { routing_table.test.1.current.0.shard_id: 1 }
+  - match: { routing_table.test.1.current.0.index: test }
+  - is_false: 'routing_table.test.1.current.0.relocating_node'
+  - is_false: 'routing_table.test.1.current.0.relocating_node_is_desired'
+  - match: { routing_table.test.1.current.1.state: STARTED }
+  - match: { routing_table.test.1.current.1.shard_id: 1 }
+  - match: { routing_table.test.1.current.1.index: test }
+  - is_false: 'routing_table.test.1.current.1.relocating_node'
+  - is_false: 'routing_table.test.1.current.1.relocating_node_is_desired'
+  - match: { routing_table.test.1.desired.total: 2 }
+  - gte: { routing_table.test.1.desired.unassigned: 0 }
+  - gte: { routing_table.test.1.desired.ignored: 0 }
+  - is_true: 'routing_table.test.1.desired.node_ids'
+

rest-api-spec/src/main/resources/rest-api-spec/api/_internal.get_desired_balance.json

@@ -0,0 +1,23 @@
+{
+  "_internal.get_desired_balance":{
+    "documentation":{
+      "url":"https://www.elastic.co/guide/en/elasticsearch/reference/current/get-desired-balance.html",
+      "description": "This API is a diagnostics API and the output should not be relied upon for building applications."
+    },
+    "stability":"experimental",
+    "visibility":"private",
+    "headers":{
+      "accept": [ "application/json"]
+    },
+    "url":{
+      "paths":[
+        {
+          "path":"/_internal/desired_balance",
+          "methods":[
+            "GET"
+          ]
+        }
+      ]
+    }
+  }
+}

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

@@ -0,0 +1,58 @@
+---
+setup:
+  - skip:
+      version: " - 8.5.99"
+      reason: "API added in in 8.6.0"
+
+---
+"Test empty desired balance":
+
+  - do:
+      _internal.get_desired_balance: { }
+
+  - gte: { stats.computation_submitted: 0 }
+  - gte: { stats.computation_executed: 0 }
+  - gte: { stats.computation_converged: 0 }
+  - gte: { stats.computation_iterations: 0 }
+  - gte: { stats.computation_converged_index: 0 }
+  - gte: { stats.computation_time_in_millis: 0 }
+  - gte: { stats.reconciliation_time_in_millis: 0 }
+  - match: { routing_table: {} }
+
+---
+"Test get desired balance for single shard":
+  - 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:
+      _internal.get_desired_balance: { }
+
+  - gte: { stats.computation_submitted: 0 }
+  - gte: { stats.computation_executed: 0 }
+  - gte: { stats.computation_converged: 0 }
+  - gte: { stats.computation_iterations: 0 }
+  - gte: { stats.computation_converged_index: 0 }
+  - gte: { stats.computation_time_in_millis: 0 }
+  - gte: { stats.reconciliation_time_in_millis: 0 }
+
+  - match: { routing_table.test.0.current.0.state: 'STARTED' }
+  - match: { routing_table.test.0.current.0.shard_id: 0 }
+  - match: { routing_table.test.0.current.0.index: test }
+  - 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'
+  - match: { routing_table.test.0.desired.total: 1 }
+  - gte: { routing_table.test.0.desired.unassigned: 0 }
+  - gte: { routing_table.test.0.desired.ignored: 0 }
+  - is_true: 'routing_table.test.0.desired.node_ids'
+

server/src/main/java/org/elasticsearch/action/ActionModule.java

@@ -11,7 +11,9 @@ package org.elasticsearch.action;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.admin.cluster.allocation.ClusterAllocationExplainAction;
+import org.elasticsearch.action.admin.cluster.allocation.GetDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.allocation.TransportClusterAllocationExplainAction;
+import org.elasticsearch.action.admin.cluster.allocation.TransportGetDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.configuration.AddVotingConfigExclusionsAction;
 import org.elasticsearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsAction;
 import org.elasticsearch.action.admin.cluster.configuration.TransportAddVotingConfigExclusionsAction;
@@ -302,6 +304,7 @@ import org.elasticsearch.rest.action.admin.cluster.RestDeleteDesiredNodesAction;
 import org.elasticsearch.rest.action.admin.cluster.RestDeleteRepositoryAction;
 import org.elasticsearch.rest.action.admin.cluster.RestDeleteSnapshotAction;
 import org.elasticsearch.rest.action.admin.cluster.RestDeleteStoredScriptAction;
+import org.elasticsearch.rest.action.admin.cluster.RestGetDesiredBalanceAction;
 import org.elasticsearch.rest.action.admin.cluster.RestGetDesiredNodesAction;
 import org.elasticsearch.rest.action.admin.cluster.RestGetFeatureUpgradeStatusAction;
 import org.elasticsearch.rest.action.admin.cluster.RestGetRepositoriesAction;
@@ -563,6 +566,7 @@ public class ActionModule extends AbstractModule {
         actions.register(AddVotingConfigExclusionsAction.INSTANCE, TransportAddVotingConfigExclusionsAction.class);
         actions.register(ClearVotingConfigExclusionsAction.INSTANCE, TransportClearVotingConfigExclusionsAction.class);
         actions.register(ClusterAllocationExplainAction.INSTANCE, TransportClusterAllocationExplainAction.class);
+        actions.register(GetDesiredBalanceAction.INSTANCE, TransportGetDesiredBalanceAction.class);
         actions.register(ClusterStatsAction.INSTANCE, TransportClusterStatsAction.class);
         actions.register(ClusterStateAction.INSTANCE, TransportClusterStateAction.class);
         actions.register(ClusterHealthAction.INSTANCE, TransportClusterHealthAction.class);
@@ -731,6 +735,7 @@ public class ActionModule extends AbstractModule {
         registerHandler.accept(new RestNodesUsageAction());
         registerHandler.accept(new RestNodesHotThreadsAction());
         registerHandler.accept(new RestClusterAllocationExplainAction());
+        registerHandler.accept(new RestGetDesiredBalanceAction());
         registerHandler.accept(new RestClusterStatsAction());
         registerHandler.accept(new RestClusterStateAction(settingsFilter, threadPool));
         registerHandler.accept(new RestClusterHealthAction());

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.action.support.master.MasterNodeReadRequest;
+import org.elasticsearch.common.io.stream.StreamInput;
+
+import java.io.IOException;
+
+public class DesiredBalanceRequest extends MasterNodeReadRequest<DesiredBalanceRequest> {
+    public DesiredBalanceRequest() {}
+
+    public DesiredBalanceRequest(StreamInput in) throws IOException {
+        super(in);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        return null;
+    }
+}

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

@@ -0,0 +1,212 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.cluster.routing.AllocationId;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceStats;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+
+public class DesiredBalanceResponse extends ActionResponse implements ToXContentObject {
+
+    private final DesiredBalanceStats stats;
+    private final Map<String, Map<Integer, DesiredShards>> routingTable;
+
+    public DesiredBalanceResponse(DesiredBalanceStats stats, Map<String, Map<Integer, DesiredShards>> routingTable) {
+        this.stats = stats;
+        this.routingTable = routingTable;
+    }
+
+    public static DesiredBalanceResponse from(StreamInput in) throws IOException {
+        return new DesiredBalanceResponse(
+            DesiredBalanceStats.readFrom(in),
+            in.readImmutableMap(StreamInput::readString, v -> v.readImmutableMap(StreamInput::readVInt, DesiredShards::from))
+        );
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        stats.writeTo(out);
+        out.writeMap(
+            routingTable,
+            StreamOutput::writeString,
+            (shardsOut, shards) -> shardsOut.writeMap(
+                shards,
+                StreamOutput::writeVInt,
+                (desiredShardsOut, desiredShards) -> desiredShards.writeTo(desiredShardsOut)
+            )
+        );
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        {
+            builder.startObject("stats");
+            stats.toXContent(builder, params);
+            builder.endObject();
+        }
+        {
+            builder.startObject("routing_table");
+            for (Map.Entry<String, Map<Integer, DesiredShards>> indexEntry : routingTable.entrySet()) {
+                builder.startObject(indexEntry.getKey());
+                for (Map.Entry<Integer, DesiredShards> shardEntry : indexEntry.getValue().entrySet()) {
+                    builder.field(String.valueOf(shardEntry.getKey()));
+                    shardEntry.getValue().toXContent(builder, params);
+                }
+                builder.endObject();
+            }
+            builder.endObject();
+        }
+        return builder.endObject();
+    }
+
+    public DesiredBalanceStats getStats() {
+        return stats;
+    }
+
+    public Map<String, Map<Integer, DesiredShards>> getRoutingTable() {
+        return routingTable;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        return o instanceof DesiredBalanceResponse that
+            && Objects.equals(stats, that.stats)
+            && Objects.equals(routingTable, that.routingTable);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(stats, routingTable);
+    }
+
+    @Override
+    public String toString() {
+        return "DesiredBalanceResponse{stats=" + stats + ", routingTable=" + routingTable + "}";
+    }
+
+    public record DesiredShards(List<ShardView> current, ShardAssignmentView desired) implements Writeable, ToXContentObject {
+
+        public static DesiredShards from(StreamInput in) throws IOException {
+            return new DesiredShards(in.readList(ShardView::from), ShardAssignmentView.from(in));
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeList(current);
+            desired.writeTo(out);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            builder.startObject();
+            builder.startArray("current");
+            for (ShardView shardView : current) {
+                shardView.toXContent(builder, params);
+            }
+            builder.endArray();
+            desired.toXContent(builder.field("desired"), params);
+            return builder.endObject();
+        }
+
+    }
+
+    public record ShardView(
+        ShardRoutingState state,
+        boolean primary,
+        String node,
+        boolean nodeIsDesired,
+        @Nullable String relocatingNode,
+        boolean relocatingNodeIsDesired,
+        int shardId,
+        String index,
+        AllocationId allocationId
+    ) implements Writeable, ToXContentObject {
+
+        public static ShardView from(StreamInput in) throws IOException {
+            return new ShardView(
+                ShardRoutingState.fromValue(in.readByte()),
+                in.readBoolean(),
+                in.readOptionalString(),
+                in.readBoolean(),
+                in.readOptionalString(),
+                in.readBoolean(),
+                in.readVInt(),
+                in.readString(),
+                in.readOptionalWriteable(AllocationId::new)
+            );
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeByte(state.value());
+            out.writeBoolean(primary);
+            out.writeOptionalString(node);
+            out.writeBoolean(nodeIsDesired);
+            out.writeOptionalString(relocatingNode);
+            out.writeBoolean(relocatingNodeIsDesired);
+            out.writeVInt(shardId);
+            out.writeString(index);
+            out.writeOptionalWriteable(allocationId);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            return builder.startObject()
+                .field("state", state.toString())
+                .field("primary", primary)
+                .field("node", node)
+                .field("node_is_desired", nodeIsDesired)
+                .field("relocating_node", relocatingNode)
+                .field("relocating_node_is_desired", relocatingNodeIsDesired)
+                .field("shard_id", shardId)
+                .field("index", index)
+                .endObject();
+        }
+    }
+
+    public record ShardAssignmentView(Set<String> nodeIds, int total, int unassigned, int ignored) implements Writeable, ToXContentObject {
+
+        public static ShardAssignmentView from(StreamInput in) throws IOException {
+            return new ShardAssignmentView(in.readSet(StreamInput::readString), in.readVInt(), in.readVInt(), in.readVInt());
+        }
+
+        @Override
+        public void writeTo(StreamOutput out) throws IOException {
+            out.writeCollection(nodeIds, StreamOutput::writeString);
+            out.writeVInt(total);
+            out.writeVInt(unassigned);
+            out.writeVInt(ignored);
+        }
+
+        @Override
+        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+            return builder.startObject()
+                .array("node_ids", nodeIds.toArray(String[]::new))
+                .field("total", total)
+                .field("unassigned", unassigned)
+                .field("ignored", ignored)
+                .endObject();
+        }
+    }
+
+}

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

@@ -0,0 +1,20 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.action.ActionType;
+
+public class GetDesiredBalanceAction extends ActionType<DesiredBalanceResponse> {
+    public static final GetDesiredBalanceAction INSTANCE = new GetDesiredBalanceAction();
+    public static final String NAME = "cluster:admin/desired_balance/get";
+
+    GetDesiredBalanceAction() {
+        super(NAME, DesiredBalanceResponse::from);
+    }
+
+}

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

@@ -0,0 +1,141 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.ResourceNotFoundException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.master.TransportMasterNodeReadAction;
+import org.elasticsearch.cluster.ClusterModule;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.block.ClusterBlockException;
+import org.elasticsearch.cluster.block.ClusterBlockLevel;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.routing.IndexRoutingTable;
+import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalance;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator;
+import org.elasticsearch.cluster.routing.allocation.allocator.ShardAssignment;
+import org.elasticsearch.cluster.routing.allocation.allocator.ShardsAllocator;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.inject.Inject;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.TransportService;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class TransportGetDesiredBalanceAction extends TransportMasterNodeReadAction<DesiredBalanceRequest, DesiredBalanceResponse> {
+
+    @Nullable
+    private final DesiredBalanceShardsAllocator desiredBalanceShardsAllocator;
+
+    @Inject
+    public TransportGetDesiredBalanceAction(
+        TransportService transportService,
+        ClusterService clusterService,
+        ThreadPool threadPool,
+        ActionFilters actionFilters,
+        IndexNameExpressionResolver indexNameExpressionResolver,
+        ShardsAllocator shardsAllocator
+    ) {
+        super(
+            GetDesiredBalanceAction.NAME,
+            transportService,
+            clusterService,
+            threadPool,
+            actionFilters,
+            DesiredBalanceRequest::new,
+            indexNameExpressionResolver,
+            DesiredBalanceResponse::from,
+            ThreadPool.Names.MANAGEMENT
+        );
+        this.desiredBalanceShardsAllocator = shardsAllocator instanceof DesiredBalanceShardsAllocator
+            ? (DesiredBalanceShardsAllocator) shardsAllocator
+            : null;
+    }
+
+    @Override
+    protected void masterOperation(
+        Task task,
+        DesiredBalanceRequest request,
+        ClusterState state,
+        ActionListener<DesiredBalanceResponse> listener
+    ) throws Exception {
+        String allocatorName = ClusterModule.SHARDS_ALLOCATOR_TYPE_SETTING.get(state.metadata().settings());
+        if (allocatorName.equals(ClusterModule.DESIRED_BALANCE_ALLOCATOR) == false || desiredBalanceShardsAllocator == null) {
+            listener.onFailure(
+                new ResourceNotFoundException(
+                    "Expected the shard balance allocator to be `desired_balance`, but got `" + allocatorName + "`"
+                )
+            );
+            return;
+        }
+
+        DesiredBalance latestDesiredBalance = desiredBalanceShardsAllocator.getDesiredBalance();
+        if (latestDesiredBalance == null) {
+            listener.onFailure(new ResourceNotFoundException("Desired balance is not computed yet"));
+            return;
+        }
+        Map<String, Map<Integer, DesiredBalanceResponse.DesiredShards>> routingTable = new HashMap<>();
+        for (IndexRoutingTable indexRoutingTable : state.routingTable()) {
+            Map<Integer, DesiredBalanceResponse.DesiredShards> indexDesiredShards = new HashMap<>();
+            for (int shardId = 0; shardId < indexRoutingTable.size(); shardId++) {
+                IndexShardRoutingTable shardRoutingTable = indexRoutingTable.shard(shardId);
+                ShardAssignment shardAssignment = latestDesiredBalance.assignments().get(shardRoutingTable.shardId());
+                List<DesiredBalanceResponse.ShardView> shardViews = new ArrayList<>();
+                for (int idx = 0; idx < shardRoutingTable.size(); idx++) {
+                    ShardRouting shard = shardRoutingTable.shard(idx);
+                    shardViews.add(
+                        new DesiredBalanceResponse.ShardView(
+                            shard.state(),
+                            shard.primary(),
+                            shard.currentNodeId(),
+                            shard.currentNodeId() != null
+                                && shardAssignment != null
+                                && shardAssignment.nodeIds().contains(shard.currentNodeId()),
+                            shard.relocatingNodeId(),
+                            shard.relocatingNodeId() != null
+                                && shardAssignment != null
+                                && shardAssignment.nodeIds().contains(shard.relocatingNodeId()),
+                            shard.shardId().id(),
+                            shard.getIndexName(),
+                            shard.allocationId()
+                        )
+                    );
+                }
+                indexDesiredShards.put(
+                    shardId,
+                    new DesiredBalanceResponse.DesiredShards(
+                        shardViews,
+                        shardAssignment != null
+                            ? new DesiredBalanceResponse.ShardAssignmentView(
+                                shardAssignment.nodeIds(),
+                                shardAssignment.total(),
+                                shardAssignment.unassigned(),
+                                shardAssignment.ignored()
+                            )
+                            : null
+                    )
+                );
+            }
+            routingTable.put(indexRoutingTable.getIndex().getName(), indexDesiredShards);
+        }
+        listener.onResponse(new DesiredBalanceResponse(desiredBalanceShardsAllocator.getStats(), routingTable));
+    }
+
+    @Override
+    protected ClusterBlockException checkBlock(DesiredBalanceRequest request, ClusterState state) {
+        return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_READ);
+    }
+}

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

@@ -63,7 +63,7 @@ public class AllocationId implements ToXContentObject, Writeable {
     @Nullable
     private final String relocationId;
 
-    AllocationId(StreamInput in) throws IOException {
+    public AllocationId(StreamInput in) throws IOException {
         this.id = in.readString();
         this.relocationId = in.readOptionalString();
     }

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.rest.action.admin.cluster;
+
+import org.elasticsearch.action.admin.cluster.allocation.DesiredBalanceRequest;
+import org.elasticsearch.action.admin.cluster.allocation.GetDesiredBalanceAction;
+import org.elasticsearch.client.internal.node.NodeClient;
+import org.elasticsearch.rest.BaseRestHandler;
+import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.rest.action.RestToXContentListener;
+
+import java.io.IOException;
+import java.util.List;
+
+public class RestGetDesiredBalanceAction extends BaseRestHandler {
+
+    @Override
+    public String getName() {
+        return "get_desired_balance";
+    }
+
+    @Override
+    public List<Route> routes() {
+        return List.of(new Route(RestRequest.Method.GET, "_internal/desired_balance"));
+    }
+
+    @Override
+    protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        return restChannel -> client.execute(
+            GetDesiredBalanceAction.INSTANCE,
+            new DesiredBalanceRequest(),
+            new RestToXContentListener<>(restChannel)
+        );
+    }
+}

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

@@ -0,0 +1,155 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.cluster.routing.AllocationId;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceStats;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.test.AbstractWireSerializingTestCase;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xcontent.XContentFactory;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase<DesiredBalanceResponse> {
+
+    @Override
+    protected Writeable.Reader<DesiredBalanceResponse> instanceReader() {
+        return DesiredBalanceResponse::from;
+    }
+
+    @Override
+    protected DesiredBalanceResponse createTestInstance() {
+        return new DesiredBalanceResponse(randomStats(), randomRoutingTable());
+    }
+
+    private DesiredBalanceStats randomStats() {
+        return new DesiredBalanceStats(
+            randomNonNegativeLong(),
+            randomBoolean(),
+            randomNonNegativeLong(),
+            randomNonNegativeLong(),
+            randomNonNegativeLong(),
+            randomNonNegativeLong(),
+            randomNonNegativeLong(),
+            randomNonNegativeLong()
+        );
+    }
+
+    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++) {
+            String indexName = randomAlphaOfLength(8);
+            Map<Integer, DesiredBalanceResponse.DesiredShards> desiredShards = new HashMap<>();
+            for (int j = 0; j < randomInt(8); j++) {
+                int shardId = randomInt(1024);
+                desiredShards.put(
+                    shardId,
+                    new DesiredBalanceResponse.DesiredShards(
+                        IntStream.range(0, randomIntBetween(1, 4))
+                            .mapToObj(
+                                k -> new DesiredBalanceResponse.ShardView(
+                                    randomFrom(ShardRoutingState.STARTED, ShardRoutingState.UNASSIGNED, ShardRoutingState.INITIALIZING),
+                                    randomBoolean(),
+                                    randomAlphaOfLength(8),
+                                    randomBoolean(),
+                                    randomAlphaOfLength(8),
+                                    randomBoolean(),
+                                    shardId,
+                                    indexName,
+                                    AllocationId.newInitializing()
+                                )
+                            )
+                            .toList(),
+                        new DesiredBalanceResponse.ShardAssignmentView(
+                            randomUnique(() -> randomAlphaOfLength(8), randomIntBetween(1, 8)),
+                            randomInt(8),
+                            randomInt(8),
+                            randomInt(8)
+                        )
+                    )
+                );
+            }
+            routingTable.put(indexName, Collections.unmodifiableMap(desiredShards));
+        }
+        return Collections.unmodifiableMap(routingTable);
+    }
+
+    @Override
+    protected DesiredBalanceResponse mutateInstance(DesiredBalanceResponse instance) throws IOException {
+        return switch (randomInt(2)) {
+            case 0 -> new DesiredBalanceResponse(
+                instance.getStats(),
+                randomValueOtherThan(instance.getRoutingTable(), this::randomRoutingTable)
+            );
+            case 1 -> new DesiredBalanceResponse(randomStats(), instance.getRoutingTable());
+            default -> randomValueOtherThan(instance, this::createTestInstance);
+        };
+    }
+
+    @SuppressWarnings("unchecked")
+    public void testToXContent() throws IOException {
+        DesiredBalanceResponse response = new DesiredBalanceResponse(randomStats(), randomRoutingTable());
+
+        Map<String, Object> json = createParser(response.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)).map();
+        assertEquals(Set.of("stats", "routing_table"), json.keySet());
+
+        Map<String, Object> stats = (Map<String, Object>) json.get("stats");
+        assertEquals(stats.get("computation_active"), response.getStats().computationActive());
+        assertEquals(stats.get("computation_submitted"), response.getStats().computationSubmitted());
+        assertEquals(stats.get("computation_executed"), response.getStats().computationExecuted());
+        assertEquals(stats.get("computation_converged"), response.getStats().computationConverged());
+        assertEquals(stats.get("computation_iterations"), response.getStats().computationIterations());
+        assertEquals(stats.get("computation_converged_index"), response.getStats().lastConvergedIndex());
+        assertEquals(stats.get("computation_time_in_millis"), response.getStats().cumulativeComputationTime());
+        assertEquals(stats.get("reconciliation_time_in_millis"), response.getStats().cumulativeReconciliationTime());
+
+        Map<String, Object> jsonRoutingTable = (Map<String, Object>) json.get("routing_table");
+        assertEquals(jsonRoutingTable.keySet(), response.getRoutingTable().keySet());
+        for (var indexEntry : response.getRoutingTable().entrySet()) {
+            Map<String, Object> jsonIndexShards = (Map<String, Object>) jsonRoutingTable.get(indexEntry.getKey());
+            assertEquals(
+                jsonIndexShards.keySet(),
+                indexEntry.getValue().keySet().stream().map(String::valueOf).collect(Collectors.toSet())
+            );
+            for (var shardEntry : indexEntry.getValue().entrySet()) {
+                DesiredBalanceResponse.DesiredShards desiredShards = shardEntry.getValue();
+                Map<String, Object> jsonDesiredShard = (Map<String, Object>) jsonIndexShards.get(String.valueOf(shardEntry.getKey()));
+                assertEquals(Set.of("current", "desired"), jsonDesiredShard.keySet());
+                List<Map<String, Object>> jsonCurrent = (List<Map<String, Object>>) jsonDesiredShard.get("current");
+                for (int i = 0; i < jsonCurrent.size(); i++) {
+                    Map<String, Object> jsonShard = jsonCurrent.get(i);
+                    DesiredBalanceResponse.ShardView shardView = desiredShards.current().get(i);
+                    assertEquals(jsonShard.get("state"), shardView.state().toString());
+                    assertEquals(jsonShard.get("primary"), shardView.primary());
+                    assertEquals(jsonShard.get("node"), shardView.node());
+                    assertEquals(jsonShard.get("node_is_desired"), shardView.nodeIsDesired());
+                    assertEquals(jsonShard.get("relocating_node"), shardView.relocatingNode());
+                    assertEquals(jsonShard.get("relocating_node_is_desired"), shardView.relocatingNodeIsDesired());
+                    assertEquals(jsonShard.get("shard_id"), shardView.shardId());
+                    assertEquals(jsonShard.get("index"), shardView.index());
+                }
+
+                Map<String, Object> jsonDesired = (Map<String, Object>) jsonDesiredShard.get("desired");
+                List<String> nodeIds = (List<String>) jsonDesired.get("node_ids");
+                assertEquals(nodeIds, List.copyOf(desiredShards.desired().nodeIds()));
+                assertEquals(jsonDesired.get("total"), desiredShards.desired().total());
+                assertEquals(jsonDesired.get("unassigned"), desiredShards.desired().unassigned());
+                assertEquals(jsonDesired.get("ignored"), desiredShards.desired().ignored());
+            }
+        }
+    }
+}

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

@@ -0,0 +1,152 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest;
+import org.elasticsearch.action.bulk.BulkRequestBuilder;
+import org.elasticsearch.action.index.IndexRequest;
+import org.elasticsearch.cluster.health.ClusterHealthStatus;
+import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+@ESIntegTestCase.ClusterScope(scope = ESIntegTestCase.Scope.TEST, numDataNodes = 0)
+public class TransportGetDesiredBalanceActionIntegrationTests extends ESIntegTestCase {
+
+    public void testDesiredBalanceOnMultiNodeCluster() throws Exception {
+        internalCluster().startMasterOnlyNode();
+        internalCluster().startDataOnlyNodes(randomIntBetween(2, 5));
+
+        String index = "test";
+        int numberOfShards = 2;
+        int numberOfReplicas = 1;
+        createIndex(index, Settings.builder().put("index.number_of_shards", 2).put("index.number_of_replicas", 1).build());
+
+        indexData(index);
+
+        var clusterHealthResponse = client().admin()
+            .cluster()
+            .health(new ClusterHealthRequest().waitForStatus(ClusterHealthStatus.GREEN))
+            .get();
+        assertEquals(RestStatus.OK, clusterHealthResponse.status());
+
+        DesiredBalanceResponse desiredBalanceResponse = client().execute(GetDesiredBalanceAction.INSTANCE, new DesiredBalanceRequest())
+            .get();
+
+        assertEquals(1, desiredBalanceResponse.getRoutingTable().size());
+        Map<Integer, DesiredBalanceResponse.DesiredShards> shardsMap = desiredBalanceResponse.getRoutingTable().get(index);
+        assertEquals(numberOfShards, shardsMap.size());
+        for (var entry : shardsMap.entrySet()) {
+            Integer shardId = entry.getKey();
+            DesiredBalanceResponse.DesiredShards desiredShards = entry.getValue();
+            IndexShardRoutingTable shardRoutingTable = client().admin()
+                .cluster()
+                .prepareState()
+                .get()
+                .getState()
+                .routingTable()
+                .shardRoutingTable(index, shardId);
+            for (int i = 0; i < shardRoutingTable.size(); i++) {
+                assertShard(shardRoutingTable.shard(i), desiredShards.current().get(i));
+            }
+            assertEquals(
+                new DesiredBalanceResponse.ShardAssignmentView(getShardNodeIds(shardRoutingTable), numberOfReplicas + 1, 0, 0),
+                desiredShards.desired()
+            );
+        }
+    }
+
+    public void testDesiredBalanceWithUnassignedShards() throws Exception {
+        internalCluster().startNode();
+
+        String index = "test";
+        int numberOfShards = 2;
+        int numberOfReplicas = 1;
+        createIndex(
+            index,
+            Settings.builder().put("index.number_of_shards", numberOfShards).put("index.number_of_replicas", numberOfReplicas).build()
+        );
+        indexData(index);
+        var clusterHealthResponse = client().admin()
+            .cluster()
+            .health(new ClusterHealthRequest(index).waitForStatus(ClusterHealthStatus.YELLOW))
+            .get();
+        assertEquals(RestStatus.OK, clusterHealthResponse.status());
+
+        DesiredBalanceResponse desiredBalanceResponse = client().execute(GetDesiredBalanceAction.INSTANCE, new DesiredBalanceRequest())
+            .get();
+
+        assertEquals(1, desiredBalanceResponse.getRoutingTable().size());
+        Map<Integer, DesiredBalanceResponse.DesiredShards> shardsMap = desiredBalanceResponse.getRoutingTable().get(index);
+        assertEquals(numberOfShards, shardsMap.size());
+        for (var entry : shardsMap.entrySet()) {
+            Integer shardId = entry.getKey();
+            DesiredBalanceResponse.DesiredShards desiredShards = entry.getValue();
+            IndexShardRoutingTable shardRoutingTable = client().admin()
+                .cluster()
+                .prepareState()
+                .get()
+                .getState()
+                .routingTable()
+                .shardRoutingTable(index, shardId);
+            for (int i = 0; i < shardRoutingTable.size(); i++) {
+                assertShard(shardRoutingTable.shard(i), desiredShards.current().get(i));
+            }
+            assertEquals(
+                new DesiredBalanceResponse.ShardAssignmentView(
+                    getShardNodeIds(shardRoutingTable),
+                    numberOfReplicas + 1,
+                    numberOfReplicas,
+                    numberOfReplicas
+                ),
+                desiredShards.desired()
+            );
+        }
+    }
+
+    private void assertShard(ShardRouting shard, DesiredBalanceResponse.ShardView shardView) {
+        assertEquals(shard.state(), shardView.state());
+        assertEquals(shard.primary(), shardView.primary());
+        assertEquals(shard.shardId().id(), shardView.shardId());
+        assertEquals(shard.shardId().getIndexName(), shardView.index());
+        assertEquals(shard.currentNodeId(), shardView.node());
+        if (shardView.state() == ShardRoutingState.STARTED) {
+            assertTrue(shardView.nodeIsDesired());
+        } else if (shardView.state() == ShardRoutingState.UNASSIGNED) {
+            assertFalse(shardView.nodeIsDesired());
+        }
+        assertEquals(shard.relocatingNodeId(), shardView.relocatingNode());
+        assertFalse(shardView.relocatingNodeIsDesired());
+    }
+
+    private static Set<String> getShardNodeIds(IndexShardRoutingTable shardRoutingTable) {
+        return IntStream.range(0, shardRoutingTable.size())
+            .mapToObj(shardRoutingTable::shard)
+            .map(ShardRouting::currentNodeId)
+            .filter(Objects::nonNull)
+            .collect(Collectors.toSet());
+    }
+
+    private static void indexData(String index) {
+        BulkRequestBuilder bulkRequestBuilder = client().prepareBulk();
+        for (int i = 0; i < randomIntBetween(5, 32); i++) {
+            bulkRequestBuilder.add(new IndexRequest(index).id(String.valueOf(i)).source("field", "foo " + i));
+        }
+        var bulkResponse = bulkRequestBuilder.get();
+        assertFalse(bulkResponse.hasFailures());
+    }
+}

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

@@ -0,0 +1,227 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.ResourceNotFoundException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.cluster.routing.IndexRoutingTable;
+import org.elasticsearch.cluster.routing.IndexShardRoutingTable;
+import org.elasticsearch.cluster.routing.RoutingTable;
+import org.elasticsearch.cluster.routing.ShardRouting;
+import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.cluster.routing.TestShardRouting;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalance;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceStats;
+import org.elasticsearch.cluster.routing.allocation.allocator.ShardAssignment;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.index.Index;
+import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.TransportService;
+import org.junit.Before;
+import org.mockito.ArgumentCaptor;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class TransportGetDesiredBalanceActionTests extends ESTestCase {
+
+    private final DesiredBalanceShardsAllocator desiredBalanceShardsAllocator = mock(DesiredBalanceShardsAllocator.class);
+    private final ClusterState clusterState = mock(ClusterState.class);
+    private final Metadata metadata = mock(Metadata.class);
+    private final TransportGetDesiredBalanceAction transportGetDesiredBalanceAction = new TransportGetDesiredBalanceAction(
+        mock(TransportService.class),
+        mock(ClusterService.class),
+        mock(ThreadPool.class),
+        mock(ActionFilters.class),
+        mock(IndexNameExpressionResolver.class),
+        desiredBalanceShardsAllocator
+    );
+    @SuppressWarnings("unchecked")
+    private final ActionListener<DesiredBalanceResponse> listener = mock(ActionListener.class);
+
+    @Before
+    public void setUpMocks() throws Exception {
+        when(clusterState.metadata()).thenReturn(metadata);
+    }
+
+    public void testReturnsErrorIfAllocatorIsNotDesiredBalanced() throws Exception {
+        when(metadata.settings()).thenReturn(Settings.builder().put("cluster.routing.allocation.type", "balanced").build());
+
+        transportGetDesiredBalanceAction.masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
+
+        ArgumentCaptor<ResourceNotFoundException> exceptionArgumentCaptor = ArgumentCaptor.forClass(ResourceNotFoundException.class);
+        verify(listener).onFailure(exceptionArgumentCaptor.capture());
+
+        assertEquals(
+            "Expected the shard balance allocator to be `desired_balance`, but got `balanced`",
+            exceptionArgumentCaptor.getValue().getMessage()
+        );
+    }
+
+    public void testReturnsErrorIfDesiredBalanceIsNotAvailable() throws Exception {
+        when(metadata.settings()).thenReturn(Settings.builder().put("cluster.routing.allocation.type", "desired_balance").build());
+
+        transportGetDesiredBalanceAction.masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
+
+        ArgumentCaptor<ResourceNotFoundException> exceptionArgumentCaptor = ArgumentCaptor.forClass(ResourceNotFoundException.class);
+        verify(listener).onFailure(exceptionArgumentCaptor.capture());
+
+        assertEquals("Desired balance is not computed yet", exceptionArgumentCaptor.getValue().getMessage());
+    }
+
+    public void testGetDesiredBalance() throws Exception {
+        Set<String> nodeIds = randomUnique(() -> randomAlphaOfLength(8), randomIntBetween(1, 32));
+        RoutingTable.Builder routingTableBuilder = RoutingTable.builder();
+        for (int i = 0; i < randomInt(8); i++) {
+            Index index = new Index(randomAlphaOfLength(8), UUIDs.randomBase64UUID());
+            IndexRoutingTable.Builder indexRoutingTableBuilder = IndexRoutingTable.builder(index);
+            for (int j = 0; j < randomIntBetween(1, 16); j++) {
+                String nodeId = randomFrom(nodeIds);
+                switch (randomInt(3)) {
+                    case 0 -> indexRoutingTableBuilder.addShard(
+                        TestShardRouting.newShardRouting(new ShardId(index, j), nodeId, true, ShardRoutingState.STARTED)
+                    );
+                    case 1 -> {
+                        indexRoutingTableBuilder.addShard(
+                            TestShardRouting.newShardRouting(new ShardId(index, j), nodeId, true, ShardRoutingState.STARTED)
+                        );
+                        if (nodeIds.size() > 1) {
+                            indexRoutingTableBuilder.addShard(
+                                TestShardRouting.newShardRouting(
+                                    new ShardId(index, j),
+                                    randomValueOtherThan(nodeId, () -> randomFrom(nodeIds)),
+                                    false,
+                                    ShardRoutingState.STARTED
+                                )
+                            );
+                        }
+                        break;
+                    }
+                    case 2 -> indexRoutingTableBuilder.addShard(
+                        TestShardRouting.newShardRouting(new ShardId(index, j), null, false, ShardRoutingState.UNASSIGNED)
+                    );
+                    case 3 -> {
+                        ShardRouting shard = TestShardRouting.newShardRouting(
+                            new ShardId(index, j),
+                            nodeId,
+                            true,
+                            ShardRoutingState.STARTED
+                        );
+                        if (nodeIds.size() > 1) {
+                            shard = TestShardRouting.relocate(
+                                shard,
+                                randomValueOtherThan(nodeId, () -> randomFrom(nodeIds)),
+                                ShardRouting.UNAVAILABLE_EXPECTED_SHARD_SIZE
+                            );
+                        }
+                        indexRoutingTableBuilder.addShard(shard);
+                    }
+                }
+            }
+            routingTableBuilder.add(indexRoutingTableBuilder.build());
+        }
+        RoutingTable routingTable = routingTableBuilder.build();
+        when(clusterState.routingTable()).thenReturn(routingTable);
+
+        List<ShardId> shardIds = routingTable.allShards().stream().map(ShardRouting::shardId).toList();
+        Map<String, Set<ShardId>> indexShards = shardIds.stream()
+            .collect(Collectors.groupingBy(e -> e.getIndex().getName(), Collectors.toSet()));
+        Map<ShardId, ShardAssignment> shardAssignments = new HashMap<>();
+        if (shardIds.size() > 0) {
+            for (int i = 0; i < randomInt(8); i++) {
+                int total = randomIntBetween(1, 1024);
+                Set<String> shardNodeIds = randomUnique(() -> randomFrom(nodeIds), randomInt(8));
+                shardAssignments.put(
+                    randomFrom(shardIds),
+                    new ShardAssignment(shardNodeIds, total, total - shardNodeIds.size(), randomInt(1024))
+                );
+            }
+        }
+
+        when(desiredBalanceShardsAllocator.getDesiredBalance()).thenReturn(new DesiredBalance(randomInt(1024), shardAssignments));
+        DesiredBalanceStats desiredBalanceStats = new DesiredBalanceStats(
+            randomInt(Integer.MAX_VALUE),
+            randomBoolean(),
+            randomInt(Integer.MAX_VALUE),
+            randomInt(Integer.MAX_VALUE),
+            randomInt(Integer.MAX_VALUE),
+            randomInt(Integer.MAX_VALUE),
+            randomInt(Integer.MAX_VALUE),
+            randomInt(Integer.MAX_VALUE)
+        );
+        when(desiredBalanceShardsAllocator.getStats()).thenReturn(desiredBalanceStats);
+        when(metadata.settings()).thenReturn(Settings.builder().put("cluster.routing.allocation.type", "desired_balance").build());
+
+        transportGetDesiredBalanceAction.masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
+
+        ArgumentCaptor<DesiredBalanceResponse> desiredBalanceResponseCaptor = ArgumentCaptor.forClass(DesiredBalanceResponse.class);
+        verify(listener).onResponse(desiredBalanceResponseCaptor.capture());
+        DesiredBalanceResponse desiredBalanceResponse = desiredBalanceResponseCaptor.getValue();
+        assertEquals(desiredBalanceStats, desiredBalanceResponse.getStats());
+        assertEquals(indexShards.keySet(), desiredBalanceResponse.getRoutingTable().keySet());
+        for (var e : desiredBalanceResponse.getRoutingTable().entrySet()) {
+            String index = e.getKey();
+            Map<Integer, DesiredBalanceResponse.DesiredShards> shardsMap = e.getValue();
+            assertEquals(indexShards.get(index).stream().map(ShardId::id).collect(Collectors.toSet()), shardsMap.keySet());
+            for (var shardDesiredBalance : shardsMap.entrySet()) {
+                DesiredBalanceResponse.DesiredShards desiredShard = shardDesiredBalance.getValue();
+                int shardId = shardDesiredBalance.getKey();
+                IndexShardRoutingTable indexShardRoutingTable = routingTable.shardRoutingTable(index, shardId);
+                for (int idx = 0; idx < indexShardRoutingTable.size(); idx++) {
+                    ShardRouting shard = indexShardRoutingTable.shard(idx);
+                    DesiredBalanceResponse.ShardView shardView = desiredShard.current().get(idx);
+                    assertEquals(shard.state(), shardView.state());
+                    assertEquals(shard.primary(), shardView.primary());
+                    assertEquals(shard.currentNodeId(), shardView.node());
+                    assertEquals(shard.relocatingNodeId(), shardView.relocatingNode());
+                    assertEquals(shard.index().getName(), shardView.index());
+                    assertEquals(shard.shardId().id(), shardView.shardId());
+                    assertEquals(shard.allocationId(), shardView.allocationId());
+                    Set<String> desiredNodeIds = Optional.ofNullable(shardAssignments.get(shard.shardId()))
+                        .map(ShardAssignment::nodeIds)
+                        .orElse(Set.of());
+                    assertEquals(
+                        shard.currentNodeId() != null && desiredNodeIds.contains(shard.currentNodeId()),
+                        shardView.nodeIsDesired()
+                    );
+                    assertEquals(
+                        shard.relocatingNodeId() != null && desiredNodeIds.contains(shard.relocatingNodeId()),
+                        shardView.relocatingNodeIsDesired()
+                    );
+                }
+                Optional<ShardAssignment> shardAssignment = Optional.ofNullable(shardAssignments.get(indexShardRoutingTable.shardId()));
+                if (shardAssignment.isPresent()) {
+                    assertEquals(shardAssignment.get().nodeIds(), desiredShard.desired().nodeIds());
+                    assertEquals(shardAssignment.get().total(), desiredShard.desired().total());
+                    assertEquals(shardAssignment.get().unassigned(), desiredShard.desired().unassigned());
+                    assertEquals(shardAssignment.get().ignored(), desiredShard.desired().ignored());
+                } else {
+                    assertNull(desiredShard.desired());
+                }
+            }
+        }
+    }
+}

x-pack/docs/en/security/operator-privileges/operator-only-functionality.asciidoc

@@ -25,6 +25,7 @@ given {es} version.
 * <<update-desired-nodes>>
 * <<get-desired-nodes>>
 * <<delete-desired-nodes>>
+* <<get-desired-balance>>
 
 [[operator-only-dynamic-cluster-settings]]
 ==== Operator-only dynamic cluster settings

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/operator/OperatorOnlyRegistry.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.security.operator;
 
+import org.elasticsearch.action.admin.cluster.allocation.GetDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.configuration.AddVotingConfigExclusionsAction;
 import org.elasticsearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsAction;
 import org.elasticsearch.action.admin.cluster.desirednodes.DeleteDesiredNodesAction;
@@ -46,7 +47,8 @@ public class OperatorOnlyRegistry {
         // Desired Nodes API
         DeleteDesiredNodesAction.NAME,
         GetDesiredNodesAction.NAME,
-        UpdateDesiredNodesAction.NAME
+        UpdateDesiredNodesAction.NAME,
+        GetDesiredBalanceAction.NAME
     );
 
     private final ClusterSettings clusterSettings;