Make DesiredBalanceResponse JSON representation chunked (#91766)

DesiredBalanceResponse dumps the cluster's routing table which can be quite large. Let's allow us to stream it to the client in chunks.

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

@@ -11,20 +11,24 @@ 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.collect.Iterators;
 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.xcontent.ChunkedToXContent;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 
-public class DesiredBalanceResponse extends ActionResponse implements ToXContentObject {
+public class DesiredBalanceResponse extends ActionResponse implements ChunkedToXContent {
 
     private final DesiredBalanceStats stats;
     private final Map<String, Map<Integer, DesiredShards>> routingTable;
@@ -56,26 +60,21 @@ public class DesiredBalanceResponse extends ActionResponse implements ToXContent
     }
 
     @Override
-    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject();
-        {
+    public Iterator<? extends ToXContent> toXContentChunked() {
+        return Iterators.concat(Iterators.single((builder, params) -> {
+            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();
+            return builder.startObject("routing_table");
+        }), routingTable.entrySet().stream().map(indexEntry -> (ToXContent) (builder, params) -> {
+            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();
-        }
-        return builder.endObject();
+            return builder.endObject();
+        }).iterator(), Iterators.single((builder, params) -> builder.endObject().endObject()));
     }
 
     public DesiredBalanceStats getStats() {

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

@@ -13,7 +13,7 @@ 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 org.elasticsearch.rest.action.RestChunkedToXContentListener;
 
 import java.io.IOException;
 import java.util.List;
@@ -35,7 +35,7 @@ public class RestGetDesiredBalanceAction extends BaseRestHandler {
         return restChannel -> client.execute(
             GetDesiredBalanceAction.INSTANCE,
             new DesiredBalanceRequest(),
-            new RestToXContentListener<>(restChannel)
+            new RestChunkedToXContentListener<>(restChannel)
         );
     }
 }

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

@@ -11,6 +11,7 @@ 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.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentFactory;
@@ -104,7 +105,9 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
     public void testToXContent() throws IOException {
         DesiredBalanceResponse response = new DesiredBalanceResponse(randomStats(), randomRoutingTable());
 
-        Map<String, Object> json = createParser(response.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)).map();
+        Map<String, Object> json = createParser(
+            ChunkedToXContent.wrapAsXContentObject(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");
@@ -152,4 +155,15 @@ public class DesiredBalanceResponseTests extends AbstractWireSerializingTestCase
             }
         }
     }
+
+    public void testToChunkedXContent() {
+        DesiredBalanceResponse response = new DesiredBalanceResponse(randomStats(), randomRoutingTable());
+        var toXContentChunked = response.toXContentChunked();
+        int chunks = 0;
+        while (toXContentChunked.hasNext()) {
+            toXContentChunked.next();
+            chunks++;
+        }
+        assertEquals(response.getRoutingTable().size() + 2, chunks);
+    }
 }