Chunked FieldUsageStatsResponse (#91942)

These responses can become extremely large, chunk them.

server/src/main/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponse.java

@@ -9,16 +9,17 @@
 package org.elasticsearch.action.admin.indices.stats;
 
 import org.elasticsearch.action.support.DefaultShardOperationFailedException;
-import org.elasticsearch.action.support.broadcast.BroadcastResponse;
+import org.elasticsearch.action.support.broadcast.ChunkedBroadcastResponse;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 
-public class FieldUsageStatsResponse extends BroadcastResponse {
+public class FieldUsageStatsResponse extends ChunkedBroadcastResponse {
     private final Map<String, List<FieldUsageShardResponse>> stats;
 
     FieldUsageStatsResponse(
@@ -48,19 +49,15 @@ public class FieldUsageStatsResponse extends BroadcastResponse {
     }
 
     @Override
-    protected void addCustomXContentFields(XContentBuilder builder, Params params) throws IOException {
-        final List<Map.Entry<String, List<FieldUsageShardResponse>>> sortedEntries = stats.entrySet()
-            .stream()
-            .sorted(Map.Entry.comparingByKey())
-            .toList();
-        for (Map.Entry<String, List<FieldUsageShardResponse>> entry : sortedEntries) {
+    protected Iterator<ToXContent> customXContentChunks(ToXContent.Params params) {
+        return stats.entrySet().stream().sorted(Map.Entry.comparingByKey()).map(entry -> (ToXContent) (builder, p) -> {
             builder.startObject(entry.getKey());
             builder.startArray("shards");
             for (FieldUsageShardResponse resp : entry.getValue()) {
                 resp.toXContent(builder, params);
             }
             builder.endArray();
-            builder.endObject();
-        }
+            return builder.endObject();
+        }).iterator();
     }
 }

server/src/main/java/org/elasticsearch/rest/action/admin/indices/RestFieldUsageStatsAction.java

@@ -17,7 +17,7 @@ import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.RestHandler;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.action.RestCancellableNodeClient;
-import org.elasticsearch.rest.action.RestToXContentListener;
+import org.elasticsearch.rest.action.RestChunkedToXContentListener;
 
 import java.io.IOException;
 import java.util.List;
@@ -42,7 +42,7 @@ public class RestFieldUsageStatsAction extends BaseRestHandler {
         fusRequest.fields(request.paramAsStringArray("fields", fusRequest.fields()));
         return channel -> {
             final RestCancellableNodeClient cancelClient = new RestCancellableNodeClient(client, request.getHttpChannel());
-            cancelClient.execute(FieldUsageStatsAction.INSTANCE, fusRequest, new RestToXContentListener<>(channel));
+            cancelClient.execute(FieldUsageStatsAction.INSTANCE, fusRequest, new RestChunkedToXContentListener<>(channel));
         };
     }
 }

server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java

@@ -0,0 +1,60 @@
+/*
+ * 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.indices.stats;
+
+import org.elasticsearch.cluster.routing.ShardRoutingState;
+import org.elasticsearch.cluster.routing.TestShardRouting;
+import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.common.util.Maps;
+import org.elasticsearch.index.search.stats.FieldUsageStats;
+import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.ToXContent;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.json.JsonXContent;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+public class FieldUsageStatsResponseTests extends ESTestCase {
+
+    public void testToXContentChunkPerIndex() throws IOException {
+        final int indices = randomIntBetween(0, 100);
+        final Map<String, List<FieldUsageShardResponse>> perIndex = Maps.newMapWithExpectedSize(indices);
+        for (int i = 0; i < indices; i++) {
+            perIndex.put(
+                "index-" + i,
+                List.of(
+                    new FieldUsageShardResponse(
+                        "tracking_id",
+                        TestShardRouting.newShardRouting(
+                            new ShardId("index" + i, UUIDs.randomBase64UUID(random()), 0),
+                            "node_id",
+                            true,
+                            ShardRoutingState.STARTED
+                        ),
+                        0,
+                        new FieldUsageStats()
+                    )
+                )
+            );
+        }
+        final FieldUsageStatsResponse response = new FieldUsageStatsResponse(indices, indices, 0, List.of(), perIndex);
+
+        final XContentBuilder builder = JsonXContent.contentBuilder();
+        final var iterator = response.toXContentChunked(ToXContent.EMPTY_PARAMS);
+        int chunks = 0;
+        while (iterator.hasNext()) {
+            iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS);
+            chunks++;
+        }
+        assertEquals(indices + 2, chunks);
+    }
+}