Chunk profiling stacktrace response (#96340)

With this commit we implement chunking for the get stacktraces API in
the profiling plugin as typical response sizes are dozens of MB large.

docs/changelog/96340.yaml

@@ -0,0 +1,5 @@
+pr: 96340
+summary: Chunk profiling stacktrace response
+area: Application
+type: enhancement
+issues: []

x-pack/plugin/profiler/src/internalClusterTest/java/org/elasticsearch/xpack/profiler/GetProfilingActionIT.java

@@ -7,8 +7,6 @@
 
 package org.elasticsearch.xpack.profiler;
 
-import org.elasticsearch.rest.RestStatus;
-
 import java.util.List;
 
 public class GetProfilingActionIT extends ProfilingTestCase {
@@ -20,7 +18,6 @@ public class GetProfilingActionIT extends ProfilingTestCase {
     public void testGetProfilingDataUnfiltered() throws Exception {
         GetProfilingRequest request = new GetProfilingRequest(1, null);
         GetProfilingResponse response = client().execute(GetProfilingAction.INSTANCE, request).get();
-        assertEquals(RestStatus.OK, response.status());
         assertEquals(1, response.getTotalFrames());
         assertNotNull(response.getStackTraces());
         StackTrace stackTrace = response.getStackTraces().get("QjoLteG7HX3VUUXr-J4kHQ");

x-pack/plugin/profiler/src/main/java/org/elasticsearch/xpack/profiler/GetProfilingResponse.java

@@ -7,22 +7,23 @@
 package org.elasticsearch.xpack.profiler;
 
 import org.elasticsearch.ElasticsearchException;
-import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.common.xcontent.StatusToXContentObject;
+import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
+import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.core.Nullable;
-import org.elasticsearch.rest.RestStatus;
-import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.BiFunction;
 
-import static org.elasticsearch.rest.RestStatus.OK;
-
-public class GetProfilingResponse extends ActionResponse implements StatusToXContentObject {
+public class GetProfilingResponse extends ActionResponse implements ChunkedToXContentObject {
     @Nullable
     private final Map<String, StackTrace> stackTraces;
     @Nullable
@@ -139,11 +140,6 @@ public class GetProfilingResponse extends ActionResponse implements StatusToXCon
         }
     }
 
-    @Override
-    public RestStatus status() {
-        return error != null ? ExceptionsHelper.status(ExceptionsHelper.unwrapCause(error)) : OK;
-    }
-
     public Map<String, StackTrace> getStackTraces() {
         return stackTraces;
     }
@@ -169,36 +165,32 @@ public class GetProfilingResponse extends ActionResponse implements StatusToXCon
     }
 
     @Override
-    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject();
-        if (stackTraces != null) {
-            builder.startObject("stack_traces");
-            builder.mapContents(stackTraces);
-            builder.endObject();
-        }
-        if (stackFrames != null) {
-            builder.startObject("stack_frames");
-            builder.mapContents(stackFrames);
-            builder.endObject();
-        }
-        if (executables != null) {
-            builder.startObject("executables");
-            builder.mapContents(executables);
-            builder.endObject();
-        }
-        if (stackTraceEvents != null) {
-            builder.startObject("stack_trace_events");
-            builder.mapContents(stackTraceEvents);
-            builder.endObject();
-        }
-        builder.field("total_frames", totalFrames);
+    public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
         if (error != null) {
-            builder.startObject("error");
-            ElasticsearchException.generateThrowableXContent(builder, params, error);
-            builder.endObject();
+            return Iterators.concat(
+                ChunkedToXContentHelper.startObject(),
+                Iterators.single((b, p) -> ElasticsearchException.generateFailureXContent(b, params, error, true)),
+                ChunkedToXContentHelper.endObject()
+            );
+        } else {
+            return Iterators.concat(
+                ChunkedToXContentHelper.startObject(),
+                optional("stack_traces", stackTraces, ChunkedToXContentHelper::xContentValuesMap),
+                optional("stack_frames", stackFrames, ChunkedToXContentHelper::xContentValuesMap),
+                optional("executables", executables, ChunkedToXContentHelper::map),
+                optional("stack_trace_events", stackTraceEvents, ChunkedToXContentHelper::map),
+                Iterators.single((b, p) -> b.field("total_frames", totalFrames)),
+                ChunkedToXContentHelper.endObject()
+            );
         }
-        builder.endObject();
-        return builder;
+    }
+
+    private <T> Iterator<? extends ToXContent> optional(
+        String name,
+        Map<String, T> values,
+        BiFunction<String, Map<String, T>, Iterator<? extends ToXContent>> supplier
+    ) {
+        return (values != null) ? supplier.apply(name, values) : Collections.emptyIterator();
     }
 
     @Override

x-pack/plugin/profiler/src/main/java/org/elasticsearch/xpack/profiler/RestGetProfilingAction.java

@@ -9,8 +9,9 @@ package org.elasticsearch.xpack.profiler;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.rest.BaseRestHandler;
 import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.rest.action.RestActionListener;
 import org.elasticsearch.rest.action.RestCancellableNodeClient;
-import org.elasticsearch.rest.action.RestStatusToXContentListener;
+import org.elasticsearch.rest.action.RestChunkedToXContentListener;
 
 import java.io.IOException;
 import java.util.List;
@@ -30,7 +31,7 @@ public class RestGetProfilingAction extends BaseRestHandler {
         request.applyContentParser(getProfilingRequest::parseXContent);
 
         return channel -> {
-            RestStatusToXContentListener<GetProfilingResponse> listener = new RestStatusToXContentListener<>(channel);
+            RestActionListener<GetProfilingResponse> listener = new RestChunkedToXContentListener<>(channel);
             RestCancellableNodeClient cancelClient = new RestCancellableNodeClient(client, request.getHttpChannel());
             cancelClient.execute(GetProfilingAction.INSTANCE, getProfilingRequest, listener);
         };

x-pack/plugin/profiler/src/test/java/org/elasticsearch/xpack/profiler/GetProfilingResponseTests.java

@@ -8,17 +8,13 @@
 package org.elasticsearch.xpack.profiler;
 
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.test.AbstractWireSerializingTestCase;
 
 import java.util.List;
 import java.util.Map;
-import java.util.function.Supplier;
 
 public class GetProfilingResponseTests extends AbstractWireSerializingTestCase<GetProfilingResponse> {
-    private <T> T randomNullable(Supplier<T> v) {
-        return randomBoolean() ? v.get() : null;
-    }
-
     private <T> T randomNullable(T v) {
         return randomBoolean() ? v : null;
     }
@@ -60,4 +56,21 @@ public class GetProfilingResponseTests extends AbstractWireSerializingTestCase<G
     protected Writeable.Reader<GetProfilingResponse> instanceReader() {
         return GetProfilingResponse::new;
     }
+
+    public void testChunking() {
+        AbstractChunkedSerializingTestCase.assertChunkCount(createTestInstance(), instance -> {
+            // start, end, total_frames
+            int chunks = 3;
+            chunks += size(instance.getExecutables());
+            chunks += size(instance.getStackFrames());
+            chunks += size(instance.getStackTraces());
+            chunks += size(instance.getStackTraceEvents());
+            return chunks;
+        });
+    }
+
+    private int size(Map<?, ?> m) {
+        // if there is a map, we also need to take into account start and end object
+        return m != null ? 2 + m.size() : 0;
+    }
 }