From dd7d5d933b15ada547a983b4ae371d5faa457891 Mon Sep 17 00:00:00 2001 From: Armin Braun Date: Fri, 25 Nov 2022 11:57:01 +0100 Subject: [PATCH 1/2] Chunked FieldUsageStatsResponse These responses can become extremely large, chunk them. --- .../stats/FieldUsageStatsResponse.java | 19 +++---- .../indices/RestFieldUsageStatsAction.java | 4 +- .../stats/FieldUsageStatsResponseTests.java | 56 +++++++++++++++++++ 3 files changed, 66 insertions(+), 13 deletions(-) create mode 100644 server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java diff --git a/server/src/main/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponse.java b/server/src/main/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponse.java index 7132449cee109..9b2efa83746e2 100644 --- a/server/src/main/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponse.java +++ b/server/src/main/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponse.java @@ -9,16 +9,17 @@ package org.elasticsearch.action.admin.indices.stats; import org.elasticsearch.action.support.DefaultShardOperationFailedException; -import org.elasticsearch.action.support.broadcast.BroadcastResponse; +import org.elasticsearch.action.support.broadcast.ChunkedBroadcastResponse; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; -import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.ToXContent; import java.io.IOException; +import java.util.Iterator; import java.util.List; import java.util.Map; -public class FieldUsageStatsResponse extends BroadcastResponse { +public class FieldUsageStatsResponse extends ChunkedBroadcastResponse { private final Map> stats; FieldUsageStatsResponse( @@ -48,19 +49,15 @@ public Map> getStats() { } @Override - protected void addCustomXContentFields(XContentBuilder builder, Params params) throws IOException { - final List>> sortedEntries = stats.entrySet() - .stream() - .sorted(Map.Entry.comparingByKey()) - .toList(); - for (Map.Entry> entry : sortedEntries) { + protected Iterator customXContentChunks(ToXContent.Params params) { + return stats.entrySet().stream().sorted(Map.Entry.comparingByKey()).map(entry -> (ToXContent) (builder, p) -> { builder.startObject(entry.getKey()); builder.startArray("shards"); for (FieldUsageShardResponse resp : entry.getValue()) { resp.toXContent(builder, params); } builder.endArray(); - builder.endObject(); - } + return builder.endObject(); + }).iterator(); } } diff --git a/server/src/main/java/org/elasticsearch/rest/action/admin/indices/RestFieldUsageStatsAction.java b/server/src/main/java/org/elasticsearch/rest/action/admin/indices/RestFieldUsageStatsAction.java index c1c20724a8c74..9a5f540139985 100644 --- a/server/src/main/java/org/elasticsearch/rest/action/admin/indices/RestFieldUsageStatsAction.java +++ b/server/src/main/java/org/elasticsearch/rest/action/admin/indices/RestFieldUsageStatsAction.java @@ -17,7 +17,7 @@ import org.elasticsearch.rest.RestHandler; import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.action.RestCancellableNodeClient; -import org.elasticsearch.rest.action.RestToXContentListener; +import org.elasticsearch.rest.action.RestChunkedToXContentListener; import java.io.IOException; import java.util.List; @@ -42,7 +42,7 @@ public BaseRestHandler.RestChannelConsumer prepareRequest(final RestRequest requ fusRequest.fields(request.paramAsStringArray("fields", fusRequest.fields())); return channel -> { final RestCancellableNodeClient cancelClient = new RestCancellableNodeClient(client, request.getHttpChannel()); - cancelClient.execute(FieldUsageStatsAction.INSTANCE, fusRequest, new RestToXContentListener<>(channel)); + cancelClient.execute(FieldUsageStatsAction.INSTANCE, fusRequest, new RestChunkedToXContentListener<>(channel)); }; } } diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java new file mode 100644 index 0000000000000..2c2abfd29a8e6 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java @@ -0,0 +1,56 @@ +/* + * 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 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 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.action.admin.indices.stats; + +import org.elasticsearch.cluster.routing.ShardRoutingState; +import org.elasticsearch.cluster.routing.TestShardRouting; +import org.elasticsearch.common.UUIDs; +import org.elasticsearch.common.util.Maps; +import org.elasticsearch.index.search.stats.FieldUsageStats; +import org.elasticsearch.index.shard.ShardId; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xcontent.ToXContent; + +import java.util.List; +import java.util.Map; + +public class FieldUsageStatsResponseTests extends ESTestCase { + + public void testToXContentChunkPerIndex() { + final int indices = randomIntBetween(1, 100); + final Map> perIndex = Maps.newMapWithExpectedSize(indices); + for (int i = 0; i < indices; i++) { + perIndex.put( + "index-" + i, + List.of( + new FieldUsageShardResponse( + "tracking_id", + TestShardRouting.newShardRouting( + new ShardId("index" + i, UUIDs.randomBase64UUID(random()), 0), + "node_id", + true, + ShardRoutingState.STARTED + ), + 0, + new FieldUsageStats() + ) + ) + ); + } + final FieldUsageStatsResponse response = new FieldUsageStatsResponse(indices, indices, 0, List.of(), perIndex); + + final var iterator = response.toXContentChunked(ToXContent.EMPTY_PARAMS); + int chunks = 0; + while (iterator.hasNext()) { + iterator.next(); + chunks++; + } + assertEquals(indices + 2, chunks); + } +} From 7a3d860f1b0294cc998fff64b6058f497c51e4f8 Mon Sep 17 00:00:00 2001 From: Armin Braun Date: Fri, 25 Nov 2022 17:56:06 +0100 Subject: [PATCH 2/2] test improvements --- .../indices/stats/FieldUsageStatsResponseTests.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java index 2c2abfd29a8e6..be98f601c5f5b 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/admin/indices/stats/FieldUsageStatsResponseTests.java @@ -16,14 +16,17 @@ import org.elasticsearch.index.shard.ShardId; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.xcontent.ToXContent; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.json.JsonXContent; +import java.io.IOException; import java.util.List; import java.util.Map; public class FieldUsageStatsResponseTests extends ESTestCase { - public void testToXContentChunkPerIndex() { - final int indices = randomIntBetween(1, 100); + public void testToXContentChunkPerIndex() throws IOException { + final int indices = randomIntBetween(0, 100); final Map> perIndex = Maps.newMapWithExpectedSize(indices); for (int i = 0; i < indices; i++) { perIndex.put( @@ -45,10 +48,11 @@ public void testToXContentChunkPerIndex() { } final FieldUsageStatsResponse response = new FieldUsageStatsResponse(indices, indices, 0, List.of(), perIndex); + final XContentBuilder builder = JsonXContent.contentBuilder(); final var iterator = response.toXContentChunked(ToXContent.EMPTY_PARAMS); int chunks = 0; while (iterator.hasNext()) { - iterator.next(); + iterator.next().toXContent(builder, ToXContent.EMPTY_PARAMS); chunks++; } assertEquals(indices + 2, chunks);