Skip to content

Commit 655fec2

Browse files
thecoopthecoop
authored
Simplify some chunked xcontent implementations (#120714)
Remove the use of `ChunkedXContentHelper.field`, as that can easily capture variables that shouldn't be captured
1 parent 436e604 commit 655fec2

File tree

17 files changed

+178
-180
lines changed

17 files changed

+178
-180
lines changed

‎server/src/internalClusterTest/java/org/elasticsearch/rest/StreamingXContentResponseIT.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ private static void handleStreamingXContentRestRequest(
186186
ActionRunnable.run(ActionListener.releaseAfter(refs.acquireListener(), ref), () -> {
187187
Thread.yield();
188188
streamingXContentResponse.writeFragment(
189-
p -> ChunkedToXContentHelper.field(fragment.getKey(), fragment.getValue()),
189+
p -> ChunkedToXContentHelper.chunk((b, xp) -> b.field(fragment.getKey(), fragment.getValue())),
190190
refs.acquire()
191191
);
192192
})

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

Lines changed: 19 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import org.elasticsearch.common.collect.Iterators;
1313
import org.elasticsearch.xcontent.ToXContent;
1414

15-
import java.util.Collections;
1615
import java.util.Iterator;
1716
import java.util.Map;
1817
import java.util.function.Function;
@@ -51,6 +50,13 @@ public static Iterator<ToXContent> object(String name, Map<String, ?> map) {
5150
return object(name, map, e -> (b, p) -> b.field(e.getKey(), e.getValue()));
5251
}
5352

53+
/**
54+
* Defines an object named {@code name}, with the contents set by calling {@code toXContent} on each entry in {@code map}
55+
*/
56+
public static <T> Iterator<ToXContent> object(String name, Map<String, T> map, Function<Map.Entry<String, T>, ToXContent> toXContent) {
57+
return object(name, Iterators.map(map.entrySet().iterator(), toXContent));
58+
}
59+
5460
/**
5561
* Defines an object named {@code name}, with the contents of each field created from each entry in {@code map}
5662
*/
@@ -65,26 +71,6 @@ public static Iterator<ToXContent> xContentObjectFieldObjects(String name, Map<S
6571
return object(name, map, e -> (b, p) -> e.getValue().toXContent(b.startObject(e.getKey()), p).endObject());
6672
}
6773

68-
public static Iterator<ToXContent> field(String name, boolean value) {
69-
return Iterators.single((b, p) -> b.field(name, value));
70-
}
71-
72-
public static Iterator<ToXContent> field(String name, long value) {
73-
return Iterators.single((b, p) -> b.field(name, value));
74-
}
75-
76-
public static Iterator<ToXContent> field(String name, String value) {
77-
return Iterators.single((b, p) -> b.field(name, value));
78-
}
79-
80-
public static Iterator<ToXContent> optionalField(String name, String value) {
81-
if (value == null) {
82-
return Collections.emptyIterator();
83-
} else {
84-
return field(name, value);
85-
}
86-
}
87-
8874
/**
8975
* Creates an Iterator to serialize a named field where the value is represented by a {@link ChunkedToXContentObject}.
9076
* Chunked equivalent for {@code XContentBuilder field(String name, ToXContent value)}
@@ -97,10 +83,22 @@ public static Iterator<ToXContent> field(String name, ChunkedToXContentObject va
9783
return Iterators.concat(Iterators.single((builder, innerParam) -> builder.field(name)), value.toXContentChunked(params));
9884
}
9985

86+
public static Iterator<ToXContent> array(Iterator<? extends ToXContent> contents) {
87+
return Iterators.concat(startArray(), contents, endArray());
88+
}
89+
10090
public static Iterator<ToXContent> array(String name, Iterator<? extends ToXContent> contents) {
10191
return Iterators.concat(startArray(name), contents, endArray());
10292
}
10393

94+
public static <T> Iterator<ToXContent> array(Iterator<T> items, Function<T, ToXContent> toXContent) {
95+
return Iterators.concat(startArray(), Iterators.map(items, toXContent), endArray());
96+
}
97+
98+
public static <T> Iterator<ToXContent> array(String name, Iterator<T> items, Function<T, ToXContent> toXContent) {
99+
return Iterators.concat(startArray(name), Iterators.map(items, toXContent), endArray());
100+
}
101+
104102
/**
105103
* Creates an Iterator to serialize a named field where the value is represented by an iterator of {@link ChunkedToXContentObject}.
106104
* Chunked equivalent for {@code XContentBuilder array(String name, ToXContent value)}
@@ -120,13 +118,6 @@ public static Iterator<ToXContent> object(String name, Iterator<? extends ToXCon
120118
return Iterators.concat(startObject(name), iterator, endObject());
121119
}
122120

123-
/**
124-
* Defines an object named {@code name}, with the contents set by calling {@code toXContent} on each entry in {@code map}
125-
*/
126-
public static <T> Iterator<ToXContent> object(String name, Map<String, T> map, Function<Map.Entry<String, T>, ToXContent> toXContent) {
127-
return object(name, Iterators.map(map.entrySet().iterator(), toXContent));
128-
}
129-
130121
/**
131122
* Creates an Iterator of a single ToXContent object that serializes the given object as a single chunk. Just wraps {@link
132123
* Iterators#single}, but still useful because it avoids any type ambiguity.

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

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -143,31 +143,34 @@ public String getUniqueId() {
143143
}
144144
}
145145

146+
private boolean hasResources() {
147+
return affectedResources != null && affectedResources.isEmpty() == false;
148+
}
149+
146150
@Override
147151
public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
148-
final Iterator<? extends ToXContent> resourcesIterator;
149-
if (affectedResources == null) {
150-
resourcesIterator = Collections.emptyIterator();
151-
} else {
152-
resourcesIterator = Iterators.flatMap(affectedResources.iterator(), s -> s.toXContentChunked(outerParams));
153-
}
154-
return Iterators.concat(Iterators.single((ToXContent) (builder, params) -> {
152+
return Iterators.concat(ChunkedToXContentHelper.chunk((builder, params) -> {
155153
builder.startObject();
156154
builder.field("id", definition.getUniqueId());
157155
builder.field("cause", definition.cause);
158156
builder.field("action", definition.action);
159157
builder.field("help_url", definition.helpURL);
160158

161-
if (affectedResources != null && affectedResources.size() > 0) {
159+
if (hasResources()) {
160+
// don't want to have a new chunk & nested iterator for this, so we start the object here
162161
builder.startObject("affected_resources");
163162
}
164163
return builder;
165-
}), resourcesIterator, Iterators.single((builder, params) -> {
166-
if (affectedResources != null && affectedResources.size() > 0) {
167-
builder.endObject();
168-
}
169-
builder.endObject();
170-
return builder;
171-
}));
164+
}),
165+
hasResources()
166+
? Iterators.flatMap(affectedResources.iterator(), s -> s.toXContentChunked(outerParams))
167+
: Collections.emptyIterator(),
168+
ChunkedToXContentHelper.chunk((b, p) -> {
169+
if (hasResources()) {
170+
b.endObject();
171+
}
172+
return b.endObject();
173+
})
174+
);
172175
}
173176
}

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

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
package org.elasticsearch.health;
1111

1212
import org.elasticsearch.common.collect.Iterators;
13+
import org.elasticsearch.common.xcontent.ChunkedToXContentHelper;
1314
import org.elasticsearch.common.xcontent.ChunkedToXContentObject;
1415
import org.elasticsearch.xcontent.ToXContent;
1516

@@ -25,15 +26,14 @@ public record HealthIndicatorResult(
2526
List<HealthIndicatorImpact> impacts,
2627
List<Diagnosis> diagnosisList
2728
) implements ChunkedToXContentObject {
29+
30+
private boolean hasDiagnosis() {
31+
return diagnosisList != null && diagnosisList.isEmpty() == false;
32+
}
33+
2834
@Override
2935
public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerParams) {
30-
final Iterator<? extends ToXContent> diagnosisIterator;
31-
if (diagnosisList == null) {
32-
diagnosisIterator = Collections.emptyIterator();
33-
} else {
34-
diagnosisIterator = Iterators.flatMap(diagnosisList.iterator(), s -> s.toXContentChunked(outerParams));
35-
}
36-
return Iterators.concat(Iterators.single((ToXContent) (builder, params) -> {
36+
return Iterators.concat(ChunkedToXContentHelper.chunk((builder, params) -> {
3737
builder.startObject();
3838
builder.field("status", status.xContentValue());
3939
builder.field("symptom", symptom);
@@ -43,16 +43,21 @@ public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params outerP
4343
if (impacts != null && impacts.isEmpty() == false) {
4444
builder.field("impacts", impacts);
4545
}
46-
if (diagnosisList != null && diagnosisList.isEmpty() == false) {
46+
if (hasDiagnosis()) {
47+
// don't want to have a new chunk & nested iterator for this, so we start the object here
4748
builder.startArray("diagnosis");
4849
}
4950
return builder;
50-
}), diagnosisIterator, Iterators.single((builder, params) -> {
51-
if (diagnosisList != null && diagnosisList.isEmpty() == false) {
52-
builder.endArray();
53-
}
54-
builder.endObject();
55-
return builder;
56-
}));
51+
}),
52+
hasDiagnosis()
53+
? Iterators.flatMap(diagnosisList.iterator(), s -> s.toXContentChunked(outerParams))
54+
: Collections.emptyIterator(),
55+
ChunkedToXContentHelper.chunk((b, p) -> {
56+
if (hasDiagnosis()) {
57+
b.endArray();
58+
}
59+
return b.endObject();
60+
})
61+
);
5762
}
5863
}

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

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import org.elasticsearch.xcontent.XContentBuilder;
1818

1919
import java.io.IOException;
20+
import java.util.Arrays;
2021
import java.util.Collections;
2122
import java.util.Map;
2223
import java.util.Objects;
@@ -47,6 +48,13 @@ public static ScriptCacheStats read(StreamInput in) throws IOException {
4748
return new ScriptCacheStats(context);
4849
}
4950

51+
private Map.Entry<String, ScriptStats>[] sortedContextStats() {
52+
@SuppressWarnings("unchecked")
53+
Map.Entry<String, ScriptStats>[] stats = context.entrySet().toArray(Map.Entry[]::new);
54+
Arrays.sort(stats, Map.Entry.comparingByKey());
55+
return stats;
56+
}
57+
5058
@Override
5159
public void writeTo(StreamOutput out) throws IOException {
5260
if (general != null) {
@@ -57,38 +65,36 @@ public void writeTo(StreamOutput out) throws IOException {
5765

5866
out.writeBoolean(true);
5967
out.writeInt(context.size());
60-
for (String name : context.keySet().stream().sorted().toList()) {
61-
out.writeString(name);
62-
context.get(name).writeTo(out);
68+
for (Map.Entry<String, ScriptStats> stats : sortedContextStats()) {
69+
out.writeString(stats.getKey());
70+
stats.getValue().writeTo(out);
6371
}
6472
}
6573

74+
private static void scriptStatsToXContent(ScriptStats s, XContentBuilder builder) throws IOException {
75+
builder.field(ScriptStats.Fields.COMPILATIONS, s.getCompilations());
76+
builder.field(ScriptStats.Fields.CACHE_EVICTIONS, s.getCacheEvictions());
77+
builder.field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, s.getCompilationLimitTriggered());
78+
}
79+
6680
@Override
6781
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
6882
builder.startObject(Fields.SCRIPT_CACHE_STATS);
6983
builder.startObject(Fields.SUM);
7084
if (general != null) {
71-
builder.field(ScriptStats.Fields.COMPILATIONS, general.getCompilations());
72-
builder.field(ScriptStats.Fields.CACHE_EVICTIONS, general.getCacheEvictions());
73-
builder.field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, general.getCompilationLimitTriggered());
85+
scriptStatsToXContent(general, builder);
7486
builder.endObject().endObject();
7587
return builder;
7688
}
7789

78-
ScriptStats sum = sum();
79-
builder.field(ScriptStats.Fields.COMPILATIONS, sum.getCompilations());
80-
builder.field(ScriptStats.Fields.CACHE_EVICTIONS, sum.getCacheEvictions());
81-
builder.field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, sum.getCompilationLimitTriggered());
90+
scriptStatsToXContent(sum(), builder);
8291
builder.endObject();
8392

8493
builder.startArray(Fields.CONTEXTS);
85-
for (String name : context.keySet().stream().sorted().toList()) {
86-
ScriptStats stats = context.get(name);
94+
for (Map.Entry<String, ScriptStats> stats : sortedContextStats()) {
8795
builder.startObject();
88-
builder.field(Fields.CONTEXT, name);
89-
builder.field(ScriptStats.Fields.COMPILATIONS, stats.getCompilations());
90-
builder.field(ScriptStats.Fields.CACHE_EVICTIONS, stats.getCacheEvictions());
91-
builder.field(ScriptStats.Fields.COMPILATION_LIMIT_TRIGGERED, stats.getCompilationLimitTriggered());
96+
builder.field(Fields.CONTEXT, stats.getKey());
97+
scriptStatsToXContent(stats.getValue(), builder);
9298
builder.endObject();
9399
}
94100
builder.endArray();

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/StreamingChatCompletionResults.java‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -82,11 +82,7 @@ public record Result(String delta) implements ChunkedToXContent {
8282

8383
@Override
8484
public Iterator<? extends ToXContent> toXContentChunked(ToXContent.Params params) {
85-
return Iterators.concat(
86-
ChunkedToXContentHelper.startObject(),
87-
ChunkedToXContentHelper.field(RESULT, delta),
88-
ChunkedToXContentHelper.endObject()
89-
);
85+
return ChunkedToXContentHelper.chunk((b, p) -> b.startObject().field(RESULT, delta).endObject());
9086
}
9187
}
9288
}

0 commit comments

Comments
 (0)