Add health xcontent objects

thecoop · Oct 8, 2024 · 264a63d · 264a63d
1 parent 074e141
commit 264a63d
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 59 deletions.
diff --git a/server/src/main/java/org/elasticsearch/health/Diagnosis.java b/server/src/main/java/org/elasticsearch/health/Diagnosis.java
@@ -10,13 +10,12 @@
 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.common.xcontent.ChunkedToXContentBuilder;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.util.Collection;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Objects;
@@ -77,8 +76,8 @@ public Resource(Collection<DiscoveryNode> nodes) {
         }
 
         @Override
-        public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
-            var builder = ChunkedToXContent.builder(outerParams);
+        public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
+            var builder = ChunkedToXContent.builder(params);
             if (nodes != null) {
                 return builder.array(type.displayValue, nodes.iterator(), node -> (b, p) -> {
                     b.startObject();
@@ -141,30 +140,18 @@ public String getUniqueId() {
     }
 
     @Override
-    public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
-        final Iterator<? extends ToXContent> resourcesIterator;
-        if (affectedResources == null) {
-            resourcesIterator = Collections.emptyIterator();
-        } else {
-            resourcesIterator = Iterators.flatMap(affectedResources.iterator(), s -> s.toXContentChunked(outerParams));
-        }
-        return Iterators.concat(Iterators.single((ToXContent) (builder, params) -> {
-            builder.startObject();
-            builder.field("id", definition.getUniqueId());
-            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();
+    public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
+        return ChunkedToXContent.builder(params).object(ob -> {
+            ob.append((b, p) -> {
+                b.field("id", definition.getUniqueId());
+                b.field("cause", definition.cause);
+                b.field("action", definition.action);
+                b.field("help_url", definition.helpURL);
+                return b;
+            });
+            if (affectedResources != null && affectedResources.isEmpty() == false) {
+                ob.object("affected_resources", affectedResources.iterator(), ChunkedToXContentBuilder::append);
             }
-            builder.endObject();
-            return builder;
-        }));
+        });
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/health/HealthIndicatorResult.java b/server/src/main/java/org/elasticsearch/health/HealthIndicatorResult.java
@@ -9,11 +9,11 @@
 
 package org.elasticsearch.health;
 
-import org.elasticsearch.common.collect.Iterators;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
+import org.elasticsearch.common.xcontent.ChunkedToXContentBuilder;
 import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.xcontent.ToXContent;
 
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 
@@ -26,33 +26,22 @@ public record HealthIndicatorResult(
     List<Diagnosis> diagnosisList
 ) implements ChunkedToXContentObject {
     @Override
-    public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
-        final Iterator<? extends ToXContent> diagnosisIterator;
-        if (diagnosisList == null) {
-            diagnosisIterator = Collections.emptyIterator();
-        } else {
-            diagnosisIterator = Iterators.flatMap(diagnosisList.iterator(), s -> s.toXContentChunked(outerParams));
-        }
-        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) -> {
+    public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
+        return ChunkedToXContent.builder(params).object(ob -> {
+            ob.append((b, p) -> {
+                b.field("status", status.xContentValue());
+                b.field("symptom", symptom);
+                if (details != null && HealthIndicatorDetails.EMPTY.equals(details) == false) {
+                    b.field("details", details, p);
+                }
+                if (impacts != null && impacts.isEmpty() == false) {
+                    b.field("impacts", impacts);
+                }
+                return b;
+            });
             if (diagnosisList != null && diagnosisList.isEmpty() == false) {
-                builder.endArray();
+                ob.array("diagnosis", diagnosisList.iterator(), ChunkedToXContentBuilder::append);
             }
-            builder.endObject();
-            return builder;
-        }));
+        });
     }
 }
diff --git a/server/src/test/java/org/elasticsearch/health/HealthIndicatorResultTests.java b/server/src/test/java/org/elasticsearch/health/HealthIndicatorResultTests.java
@@ -198,8 +198,8 @@ public void testChunkCount() {
         diagnosisList.add(diagnosis2);
         HealthIndicatorResult result = new HealthIndicatorResult(name, status, symptom, details, impacts, diagnosisList);
 
-        // -> each Diagnosis yields 5 chunks => 10 chunks from both diagnosis
-        // -> HealthIndicatorResult surrounds the diagnosis list by 2 chunks
-        AbstractChunkedSerializingTestCase.assertChunkCount(result, ignored -> 12);
+        // -> each Diagnosis yields 6 chunks => 12 chunks from both diagnosis
+        // -> HealthIndicatorResult surrounds the diagnosis list by 5 chunks
+        AbstractChunkedSerializingTestCase.assertChunkCount(result, ignored -> (6 * 2) + 5);
     }
 }