Fix pipeline agg serialization for ccs (backport of #54282) (#54468)

This fixes pipeline aggregations used in cross cluster search from an older
version of Elasticsearch to a newer version of Elasticsearch. I broke
this in #53730 when I was too aggressive in shutting off serialization
of pipeline aggs. In particular, this comes up when the coordinating
node is pre-7.8.0 and the gateway node is on or after 7.8.0.

The fix is another step down the line to remove pipeline aggregators
from the aggregation tree. Sort of. It create a new
`List<PipelineAggregator>` member in `InternalAggregation` *but* it is
only used for bwc serialization and it is fed by the mechanism
established in #53730 to read the pipelines from the

Author: nik9000

server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregation.java

@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.IntConsumer;
 import java.util.function.Supplier;
@@ -147,6 +148,7 @@ public void consumeBucketsAndMaybeBreak(int size) {
     protected final Map<String, Object> metadata;
 
     private final List<PipelineAggregator> pipelineAggregators;
+    private List<PipelineAggregator> pipelineAggregatorsForBwcSerialization;
 
     /**
      * Constructs an aggregation result with a given name.
@@ -159,16 +161,25 @@ protected InternalAggregation(String name, List<PipelineAggregator> pipelineAggr
         this.metadata = metadata;
     }
 
+    /**
+     * Merge a {@linkplain PipelineAggregator.PipelineTree} into this
+     * aggregation result tree before serializing to a node older than
+     * 7.8.0.
+     */
+    public final void mergePipelineTreeForBWCSerialization(PipelineAggregator.PipelineTree pipelineTree) {
+        pipelineAggregatorsForBwcSerialization = pipelineTree.aggregators();
+        forEachBucket(bucketAggs -> bucketAggs.mergePipelineTreeForBWCSerialization(pipelineTree));
+    }
+
     /**
      * Read from a stream.
      */
     protected InternalAggregation(StreamInput in) throws IOException {
         name = in.readString();
         metadata = in.readMap();
+        pipelineAggregators = emptyList();
         if (in.getVersion().before(Version.V_7_8_0)) {
-            pipelineAggregators = in.readNamedWriteableList(PipelineAggregator.class);
-        } else {
-            pipelineAggregators = emptyList();
+            in.readNamedWriteableList(PipelineAggregator.class);
         }
     }
 
@@ -177,7 +188,9 @@ public final void writeTo(StreamOutput out) throws IOException {
         out.writeString(name);
         out.writeGenericValue(metadata);
         if (out.getVersion().before(Version.V_7_8_0)) {
-            out.writeNamedWriteableList(pipelineAggregators);
+            assert pipelineAggregatorsForBwcSerialization != null :
+                "serializing to pre-7.8.0 versions should have called mergePipelineTreeForBWCSerialization";
+            out.writeNamedWriteableList(pipelineAggregatorsForBwcSerialization);
         }
         doWriteTo(out);
     }
@@ -212,6 +225,11 @@ public InternalAggregation copyWithRewritenBuckets(Function<InternalAggregations
                 "Aggregation [" + getName() + "] must be a bucket aggregation but was [" + getWriteableName() + "]");
     }
 
+    /**
+     * Run a {@linkplain Consumer} over all buckets in this aggregation.
+     */
+    public void forEachBucket(Consumer<InternalAggregations> consumer) {}
+
     /**
      * Creates the output from all pipeline aggs that this aggregation is associated with.  Should only
      * be called after all aggregations have been fully reduced
@@ -278,10 +296,23 @@ public Map<String, Object> getMetadata() {
         return metadata;
     }
 
+    /**
+     * @deprecated soon to be removed because it is not longer needed
+     */
+    @Deprecated
     public List<PipelineAggregator> pipelineAggregators() {
         return pipelineAggregators;
     }
 
+    /**
+     * The {@linkplain PipelineAggregator}s sent to older versions of Elasticsearch.
+     * @deprecated only use these for serializing to older Elasticsearch versions
+     */
+    @Deprecated
+    public List<PipelineAggregator> pipelineAggregatorsForBwcSerialization() {
+        return pipelineAggregatorsForBwcSerialization;
+    }
+
     @Override
     public String getType() {
         return getWriteableName();

server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregations.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.search.aggregations.InternalAggregation.ReduceContext;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.pipeline.SiblingPipelineAggregator;
+import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree;
 import org.elasticsearch.search.aggregations.support.AggregationPath;
 
 import java.io.IOException;
@@ -100,18 +101,38 @@ public InternalAggregations(StreamInput in) throws IOException {
     }
 
     @Override
-    @SuppressWarnings("unchecked")
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeNamedWriteableList((List<InternalAggregation>)aggregations);
-        if (out.getVersion().before(Version.V_7_8_0) && out.getVersion().onOrAfter(Version.V_6_7_0)) {
+        if (out.getVersion().before(Version.V_7_8_0)) {
             if (pipelineTreeForBwcSerialization == null) {
-                out.writeNamedWriteableList(emptyList());
+                mergePipelineTreeForBWCSerialization(PipelineTree.EMPTY);
+                out.writeNamedWriteableList(getInternalAggregations());
+                if (out.getVersion().onOrAfter(Version.V_6_7_0)) {
+                    out.writeNamedWriteableList(emptyList());
+                }
             } else {
-                out.writeNamedWriteableList(pipelineTreeForBwcSerialization.get().aggregators());
+                PipelineAggregator.PipelineTree pipelineTree = pipelineTreeForBwcSerialization.get();
+                mergePipelineTreeForBWCSerialization(pipelineTree);
+                out.writeNamedWriteableList(getInternalAggregations());
+                if (out.getVersion().onOrAfter(Version.V_6_7_0)) {
+                    out.writeNamedWriteableList(emptyList());
+                }
             }
+        } else {
+            out.writeNamedWriteableList(getInternalAggregations());
         }
     }
 
+    /**
+     * Merge a {@linkplain PipelineAggregator.PipelineTree} into this
+     * aggregation result tree before serializing to a node older than
+     * 7.8.0.
+     */
+    public void mergePipelineTreeForBWCSerialization(PipelineAggregator.PipelineTree pipelineTree) {
+        getInternalAggregations().stream().forEach(agg -> {
+            agg.mergePipelineTreeForBWCSerialization(pipelineTree.subTree(agg.getName()));
+        });
+    }
+
     /**
      * Make a mutable copy of the aggregation results.
      * <p>
@@ -135,6 +156,14 @@ public List<SiblingPipelineAggregator> getTopLevelPipelineAggregators() {
                 .collect(toList());
     }
 
+    /**
+     * Get the transient pipeline tree used to serialize pipeline aggregators to old nodes.
+     */
+    @Deprecated
+    Supplier<PipelineAggregator.PipelineTree> getPipelineTreeForBwcSerialization() {
+        return pipelineTreeForBwcSerialization;
+    }
+
     @SuppressWarnings("unchecked")
     private List<InternalAggregation> getInternalAggregations() {
         return (List<InternalAggregation>) aggregations;

server/src/main/java/org/elasticsearch/search/aggregations/InternalMultiBucketAggregation.java

@@ -30,6 +30,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 public abstract class InternalMultiBucketAggregation<A extends InternalMultiBucketAggregation,
@@ -145,7 +146,7 @@ public static int countInnerBucket(Aggregation agg) {
     }
 
     /**
-     * Amulti-bucket agg needs to first reduce the buckets and *their* pipelines
+     * A multi-bucket agg needs to first reduce the buckets and *their* pipelines
      * before allowing sibling pipelines to materialize.
      */
     @Override
@@ -173,6 +174,13 @@ public InternalAggregation copyWithRewritenBuckets(Function<InternalAggregations
         return modified ? create(newBuckets) : this;
     }
 
+    @Override
+    public void forEachBucket(Consumer<InternalAggregations> consumer) {
+        for (B bucket : getBuckets()) {
+            consumer.accept((InternalAggregations) bucket.getAggregations());
+        }
+    }
+
     private List<B> reducePipelineBuckets(ReduceContext reduceContext, PipelineTree pipelineTree) {
         List<B> reducedBuckets = new ArrayList<>();
         for (B bucket : getBuckets()) {

server/src/main/java/org/elasticsearch/search/aggregations/bucket/InternalSingleBucketAggregation.java

@@ -34,6 +34,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 /**
@@ -181,6 +182,11 @@ public InternalAggregation copyWithRewritenBuckets(Function<InternalAggregations
         return create(rewritten);
     }
 
+    @Override
+    public void forEachBucket(Consumer<InternalAggregations> consumer) {
+        consumer.accept(aggregations);
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (this == obj) return true;

server/src/test/java/org/elasticsearch/search/aggregations/InternalAggregationsTests.java

@@ -88,6 +88,10 @@ private InternalAggregation.ReduceContextBuilder maxBucketReduceContext() {
     }
 
     public static InternalAggregations createTestInstance() throws Exception {
+        return createTestInstance(randomPipelineTree());
+    }
+
+    public static InternalAggregations createTestInstance(PipelineAggregator.PipelineTree pipelineTree) throws Exception {
         List<InternalAggregation> aggsList = new ArrayList<>();
         if (randomBoolean()) {
             StringTermsTests stringTermsTests = new StringTermsTests();
@@ -104,7 +108,7 @@ public static InternalAggregations createTestInstance() throws Exception {
             InternalSimpleValueTests simpleValueTests = new InternalSimpleValueTests();
             aggsList.add(simpleValueTests.createTestInstance());
         }
-        return new InternalAggregations(aggsList);
+        return new InternalAggregations(aggsList, () -> pipelineTree);
     }
 
     private static PipelineAggregator.PipelineTree randomPipelineTree() {
@@ -129,11 +133,9 @@ public void testSerialization() throws Exception {
     }
 
     public void testGetTopLevelPipelineAggregators() throws Exception {
-        InternalAggregations orig = createTestInstance();
-        PipelineAggregator.PipelineTree tree = randomPipelineTree();
-        InternalAggregations withPipelines = new InternalAggregations(orig.copyResults(), () -> tree);
-        assertThat(withPipelines.aggregations, equalTo(orig.aggregations));
-        assertThat(withPipelines.getTopLevelPipelineAggregators(), equalTo(tree.aggregators()));
+        PipelineAggregator.PipelineTree pipelineTree = randomPipelineTree();
+        InternalAggregations aggs = createTestInstance(pipelineTree);
+        assertThat(aggs.getTopLevelPipelineAggregators(), equalTo(pipelineTree.aggregators()));
     }
 
     private void writeToAndReadFrom(InternalAggregations aggregations, int iteration) throws IOException {
@@ -146,8 +148,14 @@ private void writeToAndReadFrom(InternalAggregations aggregations, int iteration
                 InternalAggregations deserialized = new InternalAggregations(in);
                 assertEquals(aggregations.aggregations, deserialized.aggregations);
                 if (iteration < 2) {
-                    //serialize this enough times to make sure that we are able to write again what we read
-                    writeToAndReadFrom(deserialized, iteration + 1);
+                    /*
+                     * Add the pipeline tree for bwc serialization just like we
+                     * do when we merge the aggregation. Without that we can't
+                     * properly serialize to older versions.
+                     */
+                    InternalAggregations asThoughReduced = new InternalAggregations(
+                        deserialized.copyResults(), aggregations.getPipelineTreeForBwcSerialization());
+                    writeToAndReadFrom(asThoughReduced, iteration + 1);
                 }
             }
         }

server/src/test/java/org/elasticsearch/search/aggregations/bucket/significant/SignificanceHeuristicTests.java

@@ -102,6 +102,9 @@ public void testStreamResponse() throws Exception {
         ByteArrayOutputStream outBuffer = new ByteArrayOutputStream();
         OutputStreamStreamOutput out = new OutputStreamStreamOutput(outBuffer);
         out.setVersion(version);
+        if (version.before(Version.V_7_8_0)) {
+            sigTerms.mergePipelineTreeForBWCSerialization(PipelineAggregator.PipelineTree.EMPTY);
+        }
         out.writeNamedWriteable(sigTerms);
 
         // read

test/framework/src/main/java/org/elasticsearch/test/InternalAggregationTestCase.java

@@ -126,10 +126,12 @@
 import org.elasticsearch.search.aggregations.metrics.TopHitsAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.ValueCountAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.WeightedAvgAggregationBuilder;
+import org.elasticsearch.search.aggregations.pipeline.AvgBucketPipelineAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.DerivativePipelineAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.ExtendedStatsBucketPipelineAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.InternalBucketMetricValue;
 import org.elasticsearch.search.aggregations.pipeline.InternalSimpleValue;
+import org.elasticsearch.search.aggregations.pipeline.MaxBucketPipelineAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.ParsedBucketMetricValue;
 import org.elasticsearch.search.aggregations.pipeline.ParsedDerivative;
 import org.elasticsearch.search.aggregations.pipeline.ParsedExtendedStatsBucket;
@@ -140,6 +142,7 @@
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator.PipelineTree;
 import org.elasticsearch.search.aggregations.pipeline.StatsBucketPipelineAggregationBuilder;
+import org.elasticsearch.search.aggregations.pipeline.SumBucketPipelineAggregationBuilder;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -153,6 +156,7 @@
 import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyList;
+import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonMap;
 import static org.elasticsearch.common.xcontent.XContentHelper.toXContent;
 import static org.elasticsearch.search.aggregations.InternalMultiBucketAggregation.countInnerBucket;
@@ -384,6 +388,52 @@ public final void testFromXContentWithRandomFields() throws IOException {
         assertFromXContent(aggregation, parsedAggregation);
     }
 
+    public void testMergePipelineTreeForBWCSerialization() {
+        T agg = createTestInstance();
+        PipelineAggregator.PipelineTree pipelineTree = randomPipelineTree(agg);
+        agg.mergePipelineTreeForBWCSerialization(pipelineTree);
+        assertMergedPipelineTreeForBWCSerialization(agg, pipelineTree);
+    }
+
+    public static PipelineAggregator.PipelineTree randomPipelineTree(InternalAggregation aggregation) {
+        Map<String, PipelineTree> subTree = new HashMap<>();
+        aggregation.forEachBucket(bucketAggs -> {
+            for (Aggregation subAgg : bucketAggs) {
+                if (subTree.containsKey(subAgg.getName())) {
+                    continue;
+                }
+                subTree.put(subAgg.getName(), randomPipelineTree((InternalAggregation) subAgg));
+            }
+        });
+        return new PipelineAggregator.PipelineTree(emptyMap(), randomPipelineAggregators());
+    }
+
+    public static List<PipelineAggregator> randomPipelineAggregators() {
+        List<PipelineAggregator> pipelines = new ArrayList<>();
+        if (randomBoolean()) {
+            if (randomBoolean()) {
+                pipelines.add(new MaxBucketPipelineAggregationBuilder("name1", "bucket1").create());
+            }
+            if (randomBoolean()) {
+                pipelines.add(new AvgBucketPipelineAggregationBuilder("name2", "bucket2").create());
+            }
+            if (randomBoolean()) {
+                pipelines.add(new SumBucketPipelineAggregationBuilder("name3", "bucket3").create());
+            }
+        }
+        return pipelines;
+    }
+
+    @SuppressWarnings("deprecation")
+    private void assertMergedPipelineTreeForBWCSerialization(InternalAggregation agg, PipelineAggregator.PipelineTree pipelineTree) {
+        assertThat(agg.pipelineAggregatorsForBwcSerialization(), equalTo(pipelineTree.aggregators()));
+        agg.forEachBucket(bucketAggs -> {
+            for (Aggregation subAgg : bucketAggs) {
+                assertMergedPipelineTreeForBWCSerialization((InternalAggregation) subAgg, pipelineTree.subTree(subAgg.getName()));
+            }
+        });
+    }
+
     protected abstract void assertFromXContent(T aggregation, ParsedAggregation parsedAggregation) throws IOException;
 
     @SuppressWarnings("unchecked")