Make use of Chunked Encoding in the REST segments API (#90136)

This one returns the largest of all responses in the diagnostics
in most cases, making it chunked by index which should be good enough.

Author: original-brownbear

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

@@ -759,7 +759,7 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster) {
         registerHandler.accept(new RestPostFeatureUpgradeAction());
         registerHandler.accept(new RestGetIndicesAction(threadPool));
         registerHandler.accept(new RestIndicesStatsAction());
-        registerHandler.accept(new RestIndicesSegmentsAction(threadPool));
+        registerHandler.accept(new RestIndicesSegmentsAction());
         registerHandler.accept(new RestIndicesShardStoresAction());
         registerHandler.accept(new RestGetAliasesAction());
         registerHandler.accept(new RestIndexDeleteAliasesAction());

server/src/main/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponse.java

@@ -15,23 +15,27 @@
 import org.apache.lucene.util.Accountable;
 import org.elasticsearch.action.support.DefaultShardOperationFailedException;
 import org.elasticsearch.action.support.broadcast.BroadcastResponse;
+import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.unit.ByteSizeValue;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.core.RestApiVersion;
 import org.elasticsearch.index.engine.Segment;
-import org.elasticsearch.transport.Transports;
+import org.elasticsearch.rest.action.RestActions;
+import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 
-public class IndicesSegmentResponse extends BroadcastResponse {
+public class IndicesSegmentResponse extends BroadcastResponse implements ChunkedToXContent {
 
     private final ShardSegments[] shards;
 
@@ -77,12 +81,12 @@ public void writeTo(StreamOutput out) throws IOException {
     }
 
     @Override
-    protected void addCustomXContentFields(XContentBuilder builder, Params params) throws IOException {
-        assert Transports.assertNotTransportThread("segments are very numerous, too expensive to serialize on a transport thread");
-
-        builder.startObject(Fields.INDICES);
-
-        for (IndexSegments indexSegments : getIndices().values()) {
+    public Iterator<? extends ToXContent> toXContentChunked() {
+        return Iterators.concat(Iterators.single(((builder, params) -> {
+            builder.startObject();
+            RestActions.buildBroadcastShardsHeader(builder, params, this);
+            return builder.startObject(Fields.INDICES);
+        })), getIndices().values().stream().map(indexSegments -> (ToXContent) (builder, params) -> {
             builder.startObject(indexSegments.getIndex());
 
             builder.startObject(Fields.SHARDS);
@@ -141,9 +145,15 @@ protected void addCustomXContentFields(XContentBuilder builder, Params params) t
             builder.endObject();
 
             builder.endObject();
-        }
+            return builder;
+        }).iterator(), Iterators.single((builder, params) -> builder.endObject().endObject()));
+    }
 
-        builder.endObject();
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        // override the BroadcastResponse behavior that is not compatible with the toXContentChunked implementation
+        // TODO: make this not a ToXContent to make this unnecessary
+        return ChunkedToXContent.super.toXContent(builder, params);
     }
 
     private static void toXContent(XContentBuilder builder, Sort sort) throws IOException {

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

@@ -16,9 +16,8 @@
 import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.RestRequest;
-import org.elasticsearch.rest.action.DispatchingRestToXContentListener;
 import org.elasticsearch.rest.action.RestCancellableNodeClient;
-import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.rest.action.RestChunkedToXContentListener;
 
 import java.io.IOException;
 import java.util.List;
@@ -29,11 +28,7 @@ public class RestIndicesSegmentsAction extends BaseRestHandler {
 
     private static final DeprecationLogger DEPRECATION_LOGGER = DeprecationLogger.getLogger(RestIndicesSegmentsAction.class);
 
-    private final ThreadPool threadPool;
-
-    public RestIndicesSegmentsAction(ThreadPool threadPool) {
-        this.threadPool = threadPool;
-    }
+    public RestIndicesSegmentsAction() {}
 
     @Override
     public List<Route> routes() {
@@ -60,9 +55,6 @@ public RestChannelConsumer prepareRequest(final RestRequest request, final NodeC
         indicesSegmentsRequest.indicesOptions(IndicesOptions.fromRequest(request, indicesSegmentsRequest.indicesOptions()));
         return channel -> new RestCancellableNodeClient(client, request.getHttpChannel()).admin()
             .indices()
-            .segments(
-                indicesSegmentsRequest,
-                new DispatchingRestToXContentListener<>(threadPool.executor(ThreadPool.Names.MANAGEMENT), channel, request)
-            );
+            .segments(indicesSegmentsRequest, new RestChunkedToXContentListener<>(channel));
     }
 }

server/src/test/java/org/elasticsearch/action/admin/indices/segments/IndicesSegmentResponseTests.java

@@ -18,7 +18,9 @@
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 
@@ -44,4 +46,30 @@ public void testToXContentSerialiationWithSortedFields() throws Exception {
             response.toXContent(builder, ToXContent.EMPTY_PARAMS);
         }
     }
+
+    public void testSerializesOneChunkPerIndex() {
+        final int indices = randomIntBetween(1, 10);
+        final List<ShardRouting> routings = new ArrayList<>(indices);
+        for (int i = 0; i < indices; i++) {
+            routings.add(TestShardRouting.newShardRouting("index-" + i, 0, "node_id", true, ShardRoutingState.STARTED));
+        }
+        Segment segment = new Segment("my");
+        SortField sortField = new SortField("foo", SortField.Type.STRING);
+        sortField.setMissingValue(SortField.STRING_LAST);
+        segment.segmentSort = new Sort(sortField);
+        IndicesSegmentResponse response = new IndicesSegmentResponse(
+            routings.stream().map(routing -> new ShardSegments(routing, List.of(segment))).toArray(ShardSegments[]::new),
+            indices,
+            indices,
+            0,
+            Collections.emptyList()
+        );
+        int chunks = 0;
+        final var iterator = response.toXContentChunked();
+        while (iterator.hasNext()) {
+            iterator.next();
+            chunks++;
+        }
+        assertEquals(indices + 2, chunks);
+    }
 }