From 092fa8d487e7f41c89a6ac314cd0fa2501752fec Mon Sep 17 00:00:00 2001
From: Alan Woodward <romseygeek@apache.org>
Date: Mon, 3 Apr 2023 14:23:25 +0100
Subject: [PATCH] Check formatting of after keys at construction time (#94979)

Composite after_keys are currently only formatted when they are serialized
to xcontent, meaning that a key that cannot be formatted will not be caught
until after all the results from an aggregation have been reduced. This
wastes CPU time, and also blocks any move to chunked xcontent handling
in search responses.

This commit moves formatting of an after_key into InternalComposite's
constructor, catching these errors earlier.

Resolves #94760
---
 .../bucket/composite/InternalComposite.java   | 26 +++++++++----------
 .../composite/InternalCompositeTests.java     | 18 +++++++++++++
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/InternalComposite.java b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/InternalComposite.java
index 12d1be0e02efa..752554deaf302 100644
--- a/server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/InternalComposite.java
+++ b/server/src/main/java/org/elasticsearch/search/aggregations/bucket/composite/InternalComposite.java
@@ -41,6 +41,7 @@ public class InternalComposite extends InternalMultiBucketAggregation<InternalCo
     private final int size;
     private final List<InternalBucket> buckets;
     private final CompositeKey afterKey;
+    private final Map<String, Object> formattedAfterKey;
     private final int[] reverseMuls;
     private final MissingOrder[] missingOrders;
     private final List<String> sourceNames;
@@ -69,6 +70,7 @@ public class InternalComposite extends InternalMultiBucketAggregation<InternalCo
         this.reverseMuls = reverseMuls;
         this.missingOrders = missingOrders;
         this.earlyTerminated = earlyTerminated;
+        this.formattedAfterKey = afterKey == null ? null : new ArrayMap(sourceNames, formats, afterKey.values());
     }
 
     public InternalComposite(StreamInput in) throws IOException {
@@ -89,6 +91,7 @@ public InternalComposite(StreamInput in) throws IOException {
         this.buckets = in.readList((input) -> new InternalBucket(input, sourceNames, formats, reverseMuls, missingOrders));
         this.afterKey = in.readOptionalWriteable(CompositeKey::new);
         this.earlyTerminated = in.getTransportVersion().onOrAfter(TransportVersion.V_7_6_0) ? in.readBoolean() : false;
+        this.formattedAfterKey = afterKey == null ? null : new ArrayMap(sourceNames, formats, afterKey.values());
     }
 
     @Override
@@ -171,10 +174,7 @@ List<DocValueFormat> getFormats() {
 
     @Override
     public Map<String, Object> afterKey() {
-        if (afterKey != null) {
-            return new ArrayMap(sourceNames, formats, afterKey.values());
-        }
-        return null;
+        return formattedAfterKey;
     }
 
     // Visible for tests
@@ -584,13 +584,16 @@ static Object formatObject(Object obj, DocValueFormat format) {
     static class ArrayMap extends AbstractMap<String, Object> implements Comparable<ArrayMap> {
         final List<String> keys;
         final Comparable<?>[] values;
-        final List<DocValueFormat> formats;
+        final List<Object> formattedValues;
 
         ArrayMap(List<String> keys, List<DocValueFormat> formats, Comparable<?>[] values) {
             assert keys.size() == values.length && keys.size() == formats.size();
             this.keys = keys;
-            this.formats = formats;
             this.values = values;
+            this.formattedValues = new ArrayList<>();
+            for (int i = 0; i < values.length; i++) {
+                this.formattedValues.add(formatObject(values[i], formats.get(i)));
+            }
         }
 
         @Override
@@ -602,7 +605,7 @@ public int size() {
         public Object get(Object key) {
             for (int i = 0; i < keys.size(); i++) {
                 if (key.equals(keys.get(i))) {
-                    return formatObject(values[i], formats.get(i));
+                    return formattedValues.get(i);
                 }
             }
             return null;
@@ -610,10 +613,10 @@ public Object get(Object key) {
 
         @Override
         public Set<Entry<String, Object>> entrySet() {
-            return new AbstractSet<Entry<String, Object>>() {
+            return new AbstractSet<>() {
                 @Override
                 public Iterator<Entry<String, Object>> iterator() {
-                    return new Iterator<Entry<String, Object>>() {
+                    return new Iterator<>() {
                         int pos = 0;
 
                         @Override
@@ -623,10 +626,7 @@ public boolean hasNext() {
 
                         @Override
                         public Entry<String, Object> next() {
-                            SimpleEntry<String, Object> entry = new SimpleEntry<>(
-                                keys.get(pos),
-                                formatObject(values[pos], formats.get(pos))
-                            );
+                            SimpleEntry<String, Object> entry = new SimpleEntry<>(keys.get(pos), formattedValues.get(pos));
                             ++pos;
                             return entry;
                         }
diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/composite/InternalCompositeTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/composite/InternalCompositeTests.java
index e95bc3a460133..291798cd25f2b 100644
--- a/server/src/test/java/org/elasticsearch/search/aggregations/bucket/composite/InternalCompositeTests.java
+++ b/server/src/test/java/org/elasticsearch/search/aggregations/bucket/composite/InternalCompositeTests.java
@@ -383,6 +383,24 @@ public void testFormatObjectChecked() {
         );
         long epoch = 1622060077L; // May 25th 2021, evening
         expectThrows(IllegalArgumentException.class, () -> InternalComposite.formatObject(epoch, weekYearMonth));
+
+        CompositeKey compositeKey = new CompositeKey(epoch);
+        Exception e = expectThrows(
+            IllegalArgumentException.class,
+            () -> new InternalComposite(
+                "name",
+                1,
+                List.of("date"),
+                List.of(weekYearMonth),
+                List.of(),
+                compositeKey,
+                new int[0],
+                new MissingOrder[0],
+                false,
+                null
+            )
+        );
+        assertThat(e.getMessage(), containsString("created output it couldn't parse"));
     }
 
     public void testFormatDateEpochTimezone() {