Use chunked encoding for RestGetHealthAction (#91515)

The response of the health API grows when the diagnosis is performed due
to the diagnosis currently reporting all affected resources. In
particular, the payload can grow by a few MB/diagnosis when 50k indices
are reported.

docs/changelog/91515.yaml

@@ -0,0 +1,6 @@
+pr: 91515
+summary: Use chunked encoding for `RestGetHealthAction`
+area: Health
+type: feature
+issues:
+ - 90223

server/src/internalClusterTest/java/org/elasticsearch/discovery/StableMasterDisruptionIT.java

@@ -29,6 +29,7 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.set.Sets;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.health.GetHealthAction;
@@ -43,7 +44,6 @@
 import org.elasticsearch.test.transport.MockTransportService;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xcontent.ToXContent;
-import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.json.JsonXContent;
 import org.hamcrest.Matcher;
@@ -142,9 +142,16 @@ private void assertMasterStability(Client client, HealthStatus expectedStatus, M
         });
     }
 
-    private String xContentToString(ToXContentObject xContent) throws IOException {
+    private String xContentToString(ChunkedToXContent xContent) throws IOException {
         XContentBuilder builder = JsonXContent.contentBuilder();
-        xContent.toXContent(builder, ToXContent.EMPTY_PARAMS);
+        xContent.toXContentChunked().forEachRemaining(xcontent -> {
+            try {
+                xcontent.toXContent(builder, ToXContent.EMPTY_PARAMS);
+            } catch (IOException e) {
+                logger.error(e.getMessage(), e);
+                fail(e.getMessage());
+            }
+        });
         return BytesReference.bytes(builder).utf8ToString();
     }
 

server/src/main/java/org/elasticsearch/health/Diagnosis.java

@@ -9,15 +9,19 @@
 package org.elasticsearch.health;
 
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.common.collect.Iterators;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.core.Nullable;
-import org.elasticsearch.xcontent.ToXContentFragment;
-import org.elasticsearch.xcontent.ToXContentObject;
-import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.ToXContent;
 
-import java.io.IOException;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.stream.StreamSupport;
 
 import static org.elasticsearch.health.HealthService.HEALTH_API_ID_PREFIX;
 
@@ -27,13 +31,13 @@
  * @param definition The definition of the diagnosis (e.g. message, helpURL)
  * @param affectedResources Optional list of "things" that are affected by this condition (e.g. shards, indices, or policies).
  */
-public record Diagnosis(Definition definition, @Nullable List<Resource> affectedResources) implements ToXContentObject {
+public record Diagnosis(Definition definition, @Nullable List<Resource> affectedResources) implements ChunkedToXContent {
 
     /**
      * Represents a type of affected resource, together with the resources/abstractions that
      * are affected.
      */
-    public static class Resource implements ToXContentFragment {
+    public static class Resource implements ChunkedToXContent {
 
         public static final String ID_FIELD = "id";
         public static final String NAME_FIELD = "name";
@@ -73,23 +77,26 @@ public Resource(Collection<DiscoveryNode> nodes) {
         }
 
         @Override
-        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        public Iterator<? extends ToXContent> toXContentChunked() {
+            Iterator<? extends ToXContent> valuesIterator;
             if (nodes != null) {
-                // we report both node ids and names so we need a bit of structure
-                builder.startArray(type.displayValue);
-                for (DiscoveryNode node : nodes) {
+                valuesIterator = nodes.stream().map(node -> (ToXContent) (builder, params) -> {
                     builder.startObject();
                     builder.field(ID_FIELD, node.getId());
                     if (node.getName() != null) {
                         builder.field(NAME_FIELD, node.getName());
                     }
                     builder.endObject();
-                }
-                builder.endArray();
+                    return builder;
+                }).iterator();
             } else {
-                builder.field(type.displayValue, values);
+                valuesIterator = values.stream().map(value -> (ToXContent) (builder, params) -> builder.value(value)).iterator();
             }
-            return builder;
+            return Iterators.concat(
+                Iterators.single((ToXContent) (builder, params) -> builder.startArray(type.displayValue)),
+                valuesIterator,
+                Iterators.single((builder, params) -> builder.endArray())
+            );
         }
 
         @Override
@@ -136,21 +143,30 @@ public Collection<DiscoveryNode> getNodes() {
     public record Definition(String indicatorName, String id, String cause, String action, String helpURL) {}
 
     @Override
-    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject();
-        builder.field("id", HEALTH_API_ID_PREFIX + definition.indicatorName + ":diagnosis:" + definition.id);
-        builder.field("cause", definition.cause);
-        builder.field("action", definition.action);
-
+    public Iterator<? extends ToXContent> toXContentChunked() {
+        Iterator<? extends ToXContent> resourcesIterator = Collections.emptyIterator();
         if (affectedResources != null && affectedResources.size() > 0) {
-            builder.startObject("affected_resources");
-            for (Resource affectedResource : affectedResources) {
-                affectedResource.toXContent(builder, params);
+            resourcesIterator = affectedResources.stream()
+                .flatMap(s -> StreamSupport.stream(Spliterators.spliteratorUnknownSize(s.toXContentChunked(), Spliterator.ORDERED), false))
+                .iterator();
+        }
+        return Iterators.concat(Iterators.single((ToXContent) (builder, params) -> {
+            builder.startObject();
+            builder.field("id", HEALTH_API_ID_PREFIX + definition.indicatorName + ":diagnosis:" + definition.id);
+            builder.field("cause", definition.cause);
+            builder.field("action", definition.action);
+            builder.field("help_url", definition.helpURL);
+
+            if (affectedResources != null && affectedResources.size() > 0) {
+                builder.startObject("affected_resources");
+            }
+            return builder;
+        }), resourcesIterator, Iterators.single((builder, params) -> {
+            if (affectedResources != null && affectedResources.size() > 0) {
+                builder.endObject();
             }
             builder.endObject();
-        }
-
-        builder.field("help_url", definition.helpURL);
-        return builder.endObject();
+            return builder;
+        }));
     }
 }

server/src/main/java/org/elasticsearch/health/GetHealthAction.java

@@ -17,17 +17,18 @@
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.transport.TransportService;
 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.NoSuchElementException;
 import java.util.Objects;
@@ -41,7 +42,7 @@ private GetHealthAction() {
         super(NAME, GetHealthAction.Response::new);
     }
 
-    public static class Response extends ActionResponse implements ToXContentObject {
+    public static class Response extends ActionResponse implements ChunkedToXContent {
 
         private final ClusterName clusterName;
         @Nullable
@@ -83,18 +84,32 @@ public void writeTo(StreamOutput out) throws IOException {
         }
 
         @Override
-        public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
-            builder.startObject();
-            if (status != null) {
-                builder.field("status", status.xContentValue());
-            }
-            builder.field("cluster_name", clusterName.value());
-            builder.startObject("indicators");
-            for (HealthIndicatorResult result : indicators) {
-                builder.field(result.name(), result, params);
-            }
-            builder.endObject();
-            return builder.endObject();
+        @SuppressWarnings("unchecked")
+        public Iterator<? extends ToXContent> toXContentChunked() {
+            return Iterators.concat(Iterators.single((ToXContent) (builder, params) -> {
+                builder.startObject();
+                if (status != null) {
+                    builder.field("status", status.xContentValue());
+                }
+                builder.field("cluster_name", clusterName.value());
+                builder.startObject("indicators");
+                return builder;
+            }),
+                Iterators.concat(
+                    indicators.stream()
+                        .map(
+                            indicator -> Iterators.concat(
+                                // having the indicator name printed here prevents us from flat mapping all
+                                // indicators however the affected resources which are the O(indices) fields are
+                                // flat mapped over all diagnoses within the indicator
+                                Iterators.single((ToXContent) (builder, params) -> builder.field(indicator.name())),
+                                indicator.toXContentChunked()
+                            )
+                        )
+                        .toArray(Iterator[]::new)
+                ),
+                Iterators.single((b, p) -> b.endObject().endObject())
+            );
         }
 
         @Override

server/src/main/java/org/elasticsearch/health/HealthIndicatorResult.java

@@ -8,11 +8,16 @@
 
 package org.elasticsearch.health;
 
-import org.elasticsearch.xcontent.ToXContentObject;
-import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.common.collect.Iterators;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.xcontent.ToXContent;
 
-import java.io.IOException;
+import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Spliterator;
+import java.util.Spliterators;
+import java.util.stream.StreamSupport;
 
 public record HealthIndicatorResult(
     String name,
@@ -21,22 +26,35 @@ public record HealthIndicatorResult(
     HealthIndicatorDetails details,
     List<HealthIndicatorImpact> impacts,
     List<Diagnosis> diagnosisList
-) implements ToXContentObject {
-
+) implements ChunkedToXContent {
     @Override
-    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-        builder.startObject();
-        builder.field("status", status.xContentValue());
-        builder.field("symptom", symptom);
-        if (details != null && HealthIndicatorDetails.EMPTY.equals(details) == false) {
-            builder.field("details", details, params);
-        }
-        if (impacts != null && impacts.isEmpty() == false) {
-            builder.field("impacts", impacts);
-        }
+    public Iterator<? extends ToXContent> toXContentChunked() {
+        Iterator<? extends ToXContent> diagnosisIterator = Collections.emptyIterator();
         if (diagnosisList != null && diagnosisList.isEmpty() == false) {
-            builder.field("diagnosis", diagnosisList);
+            diagnosisIterator = diagnosisList.stream()
+                .flatMap(s -> StreamSupport.stream(Spliterators.spliteratorUnknownSize(s.toXContentChunked(), Spliterator.ORDERED), false))
+                .iterator();
         }
-        return builder.endObject();
+        return Iterators.concat(Iterators.single((ToXContent) (builder, params) -> {
+            builder.startObject();
+            builder.field("status", status.xContentValue());
+            builder.field("symptom", symptom);
+            if (details != null && HealthIndicatorDetails.EMPTY.equals(details) == false) {
+                builder.field("details", details, params);
+            }
+            if (impacts != null && impacts.isEmpty() == false) {
+                builder.field("impacts", impacts);
+            }
+            if (diagnosisList != null && diagnosisList.isEmpty() == false) {
+                builder.startArray("diagnosis");
+            }
+            return builder;
+        }), diagnosisIterator, Iterators.single((builder, params) -> {
+            if (diagnosisList != null && diagnosisList.isEmpty() == false) {
+                builder.endArray();
+            }
+            builder.endObject();
+            return builder;
+        }));
     }
 }

server/src/main/java/org/elasticsearch/health/RestGetHealthAction.java

@@ -11,7 +11,7 @@
 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;
@@ -38,6 +38,6 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
         String indicatorName = request.param("indicator");
         boolean verbose = request.paramAsBoolean(VERBOSE_PARAM, true);
         GetHealthAction.Request getHealthRequest = new GetHealthAction.Request(indicatorName, verbose);
-        return channel -> client.execute(GetHealthAction.INSTANCE, getHealthRequest, new RestToXContentListener<>(channel));
+        return channel -> client.execute(GetHealthAction.INSTANCE, getHealthRequest, new RestChunkedToXContentListener<>(channel));
     }
 }