Further conversions to ChunkedXContentBuilder (elastic#114237)

server/src/main/java/org/elasticsearch/action/admin/cluster/reroute/ClusterRerouteResponse.java

@@ -13,20 +13,18 @@
 import org.elasticsearch.action.support.master.IsAcknowledgedSupplier;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.routing.allocation.RoutingExplanations;
-import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.logging.DeprecationCategory;
 import org.elasticsearch.common.logging.DeprecationLogger;
-import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
+import org.elasticsearch.common.xcontent.ChunkedToXContent;
 import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
 import org.elasticsearch.core.RestApiVersion;
 import org.elasticsearch.core.UpdateForV10;
 import org.elasticsearch.rest.action.search.RestSearchAction;
 import org.elasticsearch.xcontent.ToXContent;
 
 import java.io.IOException;
-import java.util.Collections;
 import java.util.Iterator;
 import java.util.Objects;
 
@@ -98,20 +96,15 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerP
     }
 
     @Override
-    public Iterator<? extends ToXContent> toXContentChunkedV7(ToXContent.Params outerParams) {
-        return Iterators.concat(
-            Iterators.single((builder, params) -> builder.startObject().field(ACKNOWLEDGED_KEY, isAcknowledged())),
-            emitState(outerParams)
-                ? ChunkedToXContentHelper.wrapWithObject("state", state.toXContentChunked(outerParams))
-                : Collections.emptyIterator(),
-            Iterators.single((builder, params) -> {
-                if (params.paramAsBoolean("explain", false)) {
-                    explanations.toXContent(builder, params);
-                }
-
-                builder.endObject();
-                return builder;
-            })
-        );
+    public Iterator<? extends ToXContent> toXContentChunkedV7(ToXContent.Params params) {
+        return ChunkedToXContent.builder(params).object(b -> {
+            b.field(ACKNOWLEDGED_KEY, isAcknowledged());
+            if (emitState(params)) {
+                b.xContentObject("state", state);
+            }
+            if (params.paramAsBoolean("explain", false)) {
+                b.append(explanations);
+            }
+        });
     }
 }

server/src/main/java/org/elasticsearch/cluster/ClusterState.java

@@ -759,10 +759,8 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerP
 
             // customs
             metrics.contains(Metric.CUSTOMS)
-                ? Iterators.flatMap(
-                    customs.entrySet().iterator(),
-                    cursor -> ChunkedToXContentHelper.wrapWithObject(cursor.getKey(), cursor.getValue().toXContentChunked(outerParams))
-                )
+                ? ChunkedToXContent.builder(outerParams)
+                    .forEach(customs.entrySet().iterator(), (b, e) -> b.xContentObject(e.getKey(), e.getValue()))
                 : Collections.emptyIterator()
         );
     }

server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentBuilder.java

@@ -38,6 +38,10 @@ private void addChunk(ToXContent content) {
         builder.add(Objects.requireNonNull(content));
     }
 
+    public ToXContent.Params params() {
+        return params;
+    }
+
     private void startObject() {
         addChunk((b, p) -> b.startObject());
     }
@@ -259,6 +263,16 @@ public <T> ChunkedToXContentBuilder array(Iterator<T> items, BiConsumer<ChunkedT
         return this;
     }
 
+    /**
+     * Creates an array with the contents set by appending together the contents of {@code items}
+     */
+    public ChunkedToXContentBuilder array(Iterator<? extends ToXContent> items) {
+        startArray();
+        items.forEachRemaining(this::append);
+        endArray();
+        return this;
+    }
+
     /**
      * Creates an array, with the contents set by appending together
      * the return values of {@code create} called on each item returned by {@code items}
@@ -351,6 +365,12 @@ public ChunkedToXContentBuilder field(String name, ToXContent value) {
         return this;
     }
 
+    public ChunkedToXContentBuilder field(String name, ChunkedToXContent value) {
+        addChunk((b, p) -> b.field(name));
+        append(value);
+        return this;
+    }
+
     public ChunkedToXContentBuilder field(String name, Object value) {
         addChunk((b, p) -> b.field(name, value));
         return this;

server/src/main/java/org/elasticsearch/common/xcontent/ChunkedToXContentHelper.java

@@ -53,26 +53,10 @@ public static Iterator<ToXContent> field(String name, String value) {
         return Iterators.single(((builder, params) -> builder.field(name, value)));
     }
 
-    /**
-     * Creates an Iterator to serialize a named field where the value is represented by a {@link ChunkedToXContentObject}.
-     * Chunked equivalent for {@code XContentBuilder field(String name, ToXContent value)}
-     * @param name name of the field
-     * @param value value for this field
-     * @param params params to propagate for XContent serialization
-     * @return Iterator composing field name and value serialization
-     */
-    public static Iterator<ToXContent> field(String name, ChunkedToXContentObject value, ToXContent.Params params) {
-        return Iterators.concat(Iterators.single((builder, innerParam) -> builder.field(name)), value.toXContentChunked(params));
-    }
-
     public static Iterator<ToXContent> array(String name, Iterator<? extends ToXContent> contents) {
         return Iterators.concat(ChunkedToXContentHelper.startArray(name), contents, ChunkedToXContentHelper.endArray());
     }
 
-    public static <T extends ToXContent> Iterator<ToXContent> wrapWithObject(String name, Iterator<T> iterator) {
-        return Iterators.concat(startObject(name), iterator, endObject());
-    }
-
     /**
      * Creates an Iterator of a single ToXContent object that serializes the given object as a single chunk. Just wraps {@link
      * Iterators#single}, but still useful because it avoids any type ambiguity.

server/src/main/java/org/elasticsearch/script/ScriptCacheStats.java

@@ -21,15 +21,8 @@
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
 
-import static org.elasticsearch.common.collect.Iterators.concat;
-import static org.elasticsearch.common.collect.Iterators.flatMap;
-import static org.elasticsearch.common.collect.Iterators.single;
-import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.endArray;
-import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.endObject;
-import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.field;
-import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.startArray;
-import static org.elasticsearch.common.xcontent.ChunkedToXContentHelper.startObject;
 import static org.elasticsearch.script.ScriptCacheStats.Fields.SCRIPT_CACHE_STATS;
 
 // This class is deprecated in favor of ScriptStats and ScriptContextStats
@@ -76,35 +69,25 @@ public void writeTo(StreamOutput out) throws IOException {
 
     @Override
     public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
-        return concat(
-            startObject(SCRIPT_CACHE_STATS),
-            startObject(Fields.SUM),
-            general != null
-                ? concat(
-                    field(ScriptStats.Fields.COMPILATIONS, general.getCompilations()),
-                    field(ScriptStats.Fields.CACHE_EVICTIONS, general.getCacheEvictions()),
-                    field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, general.getCompilationLimitTriggered()),
-                    endObject(),
-                    endObject()
-                )
-                : concat(single((builder, params) -> {
-                    var sum = sum();
-                    return builder.field(ScriptStats.Fields.COMPILATIONS, sum.getCompilations())
-                        .field(ScriptStats.Fields.CACHE_EVICTIONS, sum.getCacheEvictions())
-                        .field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, sum.getCompilationLimitTriggered())
-                        .endObject();
-                }), startArray(Fields.CONTEXTS), flatMap(context.keySet().stream().sorted().iterator(), ctx -> {
-                    var stats = context.get(ctx);
-                    return concat(
-                        startObject(),
-                        field(Fields.CONTEXT, ctx),
-                        field(ScriptStats.Fields.COMPILATIONS, stats.getCompilations()),
-                        field(ScriptStats.Fields.CACHE_EVICTIONS, stats.getCacheEvictions()),
-                        field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, stats.getCompilationLimitTriggered()),
-                        endObject()
-                    );
-                }), endArray(), endObject())
-        );
+        Function<ScriptStats, ToXContent> statsFields = s -> (b, p) -> b.field(ScriptStats.Fields.COMPILATIONS, s.getCompilations())
+            .field(ScriptStats.Fields.CACHE_EVICTIONS, s.getCacheEvictions())
+            .field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, s.getCompilationLimitTriggered());
+
+        return ChunkedToXContent.builder(outerParams).object(SCRIPT_CACHE_STATS, sb -> {
+            if (general != null) {
+                sb.xContentObject(Fields.SUM, statsFields.apply(general));
+            } else {
+                sb.xContentObject(Fields.SUM, statsFields.apply(sum()));
+                sb.array(
+                    Fields.CONTEXTS,
+                    context.entrySet().stream().sorted(Map.Entry.comparingByKey()).iterator(),
+                    (eb, e) -> eb.object(ebo -> {
+                        ebo.field(Fields.CONTEXT, e.getKey());
+                        ebo.append(statsFields.apply(e.getValue()));
+                    })
+                );
+            }
+        });
     }
 
     /**

server/src/test/java/org/elasticsearch/action/admin/cluster/reroute/ClusterRerouteResponseTests.java

@@ -313,11 +313,15 @@ private void assertXContent(
             fail(e);
         }
 
-        final var expectedChunks = Objects.equals(params.param("metric"), "none")
-            ? 2
-            : 4 + ClusterStateTests.expectedChunkCount(params, response.getState());
+        int[] expectedChunks = new int[] { 3 };
+        if (Objects.equals(params.param("metric"), "none") == false) {
+            expectedChunks[0] += 2 + ClusterStateTests.expectedChunkCount(params, response.getState());
+        }
+        if (params.paramAsBoolean("explain", false)) {
+            expectedChunks[0]++;
+        }
 
-        AbstractChunkedSerializingTestCase.assertChunkCount(response, params, ignored -> expectedChunks);
+        AbstractChunkedSerializingTestCase.assertChunkCount(response, params, o -> expectedChunks[0]);
         assertCriticalWarnings(criticalDeprecationWarnings);
 
         // check the v7 API too
@@ -331,7 +335,7 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerP
             public boolean isFragment() {
                 return response.isFragment();
             }
-        }, params, ignored -> expectedChunks);
+        }, params, o -> expectedChunks[0]++);
         // the v7 API should not emit any deprecation warnings
         assertCriticalWarnings();
     }

server/src/test/java/org/elasticsearch/common/xcontent/ChunkedToXContentBuilderTests.java

@@ -0,0 +1,88 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common.xcontent;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.collect.Iterators;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.ToXContent;
+
+import java.util.function.IntFunction;
+import java.util.stream.IntStream;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class ChunkedToXContentBuilderTests extends ESTestCase {
+
+    public void testFieldWithInnerChunkedObject() {
+
+        ToXContent innerXContent = (b, p) -> {
+            b.startObject();
+            b.field("field1", 10);
+            b.field("field2", "aaa");
+            b.endObject();
+            return b;
+        };
+
+        ToXContent outerXContent = (b, p) -> b.field("field3", 10).field("field4", innerXContent);
+
+        String expectedContent = Strings.toString(outerXContent);
+
+        ChunkedToXContentObject innerChunkedContent = params -> new ChunkedToXContentBuilder(params).object(
+            o -> o.field("field1", 10).field("field2", "aaa")
+        );
+
+        ChunkedToXContent outerChunkedContent = params -> new ChunkedToXContentBuilder(params).field("field3", 10)
+            .field("field4", innerChunkedContent);
+
+        assertThat(Strings.toString(outerChunkedContent), equalTo(expectedContent));
+    }
+
+    public void testFieldWithInnerChunkedArray() {
+
+        ToXContent innerXContent = (b, p) -> {
+            b.startArray();
+            b.value(10);
+            b.value(20);
+            b.endArray();
+            return b;
+        };
+
+        ToXContent outerXContent = (b, p) -> b.field("field3", 10).field("field4", innerXContent);
+
+        String expectedContent = Strings.toString(outerXContent);
+
+        IntFunction<ToXContent> value = v -> (b, p) -> b.value(v);
+
+        ChunkedToXContentObject innerChunkedContent = params -> new ChunkedToXContentBuilder(params).array(
+            IntStream.of(10, 20).mapToObj(value).iterator()
+        );
+
+        ChunkedToXContent outerChunkedContent = params -> new ChunkedToXContentBuilder(params).field("field3", 10)
+            .field("field4", innerChunkedContent);
+
+        assertThat(Strings.toString(outerChunkedContent), equalTo(expectedContent));
+    }
+
+    public void testFieldWithInnerChunkedField() {
+
+        ToXContent innerXContent = (b, p) -> b.value(10);
+        ToXContent outerXContent = (b, p) -> b.field("field3", 10).field("field4", innerXContent);
+
+        String expectedContent = Strings.toString(outerXContent);
+
+        ChunkedToXContentObject innerChunkedContent = params -> Iterators.single((b, p) -> b.value(10));
+
+        ChunkedToXContent outerChunkedContent = params -> new ChunkedToXContentBuilder(params).field("field3", 10)
+            .field("field4", innerChunkedContent);
+
+        assertThat(Strings.toString(outerChunkedContent), equalTo(expectedContent));
+    }
+}