Inline aggregation logic into QueryPhaseResultConsumer

This refactoring removes some needless indirection but more importantly,
it sets up a significant improvement to aggs: passing them to the reducer as they
become available and without retaining them in a list (serialized or materialized)
any longer than necessary.
I chose to do this step first to make that change shorter and more focussed as it will
involve a bit of non-trivial concurrency.

Author: original-brownbear

server/src/main/java/org/elasticsearch/action/search/QueryPhaseResultConsumer.java

@@ -24,13 +24,12 @@
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.core.Nullable;
-import org.elasticsearch.core.Releasable;
-import org.elasticsearch.core.Releasables;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.aggregations.AggregationReduceContext;
+import org.elasticsearch.search.aggregations.AggregatorsReducer;
 import org.elasticsearch.search.aggregations.InternalAggregations;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.query.QuerySearchResult;
@@ -246,7 +245,10 @@ public SearchPhaseController.ReducedQueryPhase reduce() throws Exception {
             if (aggsList != null) {
                 // Add an estimate of the final reduce size
                 breakerSize = addEstimateAndMaybeBreak(estimateRamBytesUsedForReduce(breakerSize));
-                aggs = aggregate(buffer.iterator(), new Iterator<>() {
+                var reduceContext = performFinalReduce
+                    ? aggReduceContextBuilder.forFinalReduction()
+                    : aggReduceContextBuilder.forPartialReduction();
+                aggs = InternalAggregations.maybeExecuteFinalReduce(reduceContext, aggregate(buffer.iterator(), new Iterator<>() {
                     @Override
                     public boolean hasNext() {
                         return aggsList.isEmpty() == false;
@@ -256,10 +258,7 @@ public boolean hasNext() {
                     public InternalAggregations next() {
                         return aggsList.pollFirst();
                     }
-                },
-                    resultSize,
-                    performFinalReduce ? aggReduceContextBuilder.forFinalReduction() : aggReduceContextBuilder.forPartialReduction()
-                );
+                }, resultSize, reduceContext));
             } else {
                 aggs = null;
             }
@@ -391,39 +390,59 @@ private static InternalAggregations aggregate(
         int resultSetSize,
         AggregationReduceContext reduceContext
     ) {
-        interface ReleasableIterator extends Iterator<InternalAggregations>, Releasable {}
-        try (var aggsIter = new ReleasableIterator() {
-
-            private Releasable toRelease;
-
-            @Override
-            public void close() {
-                Releasables.close(toRelease);
+        if (resultSetSize == 1) {
+            if (partialResults.hasNext()) {
+                return InternalAggregations.reduce(partialResults.next(), reduceContext);
             }
-
-            @Override
-            public boolean hasNext() {
-                return toConsume.hasNext();
+            try (var delayable = toConsume.next().consumeAggs()) {
+                return InternalAggregations.reduce(delayable.expand(), reduceContext);
             }
-
-            @Override
-            public InternalAggregations next() {
-                var res = toConsume.next().consumeAggs();
-                Releasables.close(toRelease);
-                toRelease = res;
-                return res.expand();
+        }
+        try {
+            // general case
+            if (partialResults.hasNext()) {
+                return consumeAggResults(partialResults, toConsume, createReducer(resultSetSize, reduceContext, partialResults.next()));
             }
-        }) {
-            return InternalAggregations.topLevelReduce(
-                partialResults.hasNext() ? Iterators.concat(partialResults, aggsIter) : aggsIter,
-                resultSetSize,
-                reduceContext
-            );
+            AggregatorsReducer reducer;
+            try (var delayable = toConsume.next().consumeAggs()) {
+                reducer = createReducer(resultSetSize, reduceContext, delayable.expand());
+            }
+            return consumeAggResults(partialResults, toConsume, reducer);
         } finally {
             toConsume.forEachRemaining(QuerySearchResult::releaseAggs);
         }
     }
 
+    private static AggregatorsReducer createReducer(int resultSetSize, AggregationReduceContext reduceContext, InternalAggregations first) {
+        boolean success = false;
+        var reducer = new AggregatorsReducer(first, reduceContext, resultSetSize);
+        try {
+            reducer.accept(first);
+            success = true;
+            return reducer;
+        } finally {
+            if (success == false) {
+                reducer.close();
+            }
+        }
+    }
+
+    private static InternalAggregations consumeAggResults(
+        Iterator<InternalAggregations> partialResults,
+        Iterator<QuerySearchResult> toConsume,
+        AggregatorsReducer reducer
+    ) {
+        try (reducer) {
+            partialResults.forEachRemaining(reducer::accept);
+            while (toConsume.hasNext()) {
+                try (var delayable = toConsume.next().consumeAggs()) {
+                    reducer.accept(delayable.expand());
+                }
+            }
+            return reducer.get();
+        }
+    }
+
     public int getNumReducePhases() {
         return numReducePhases;
     }

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

@@ -177,17 +177,7 @@ public SortValue sortValue(AggregationPath.PathElement head, Iterator<Aggregatio
         return aggregation.sortValue(Optional.ofNullable(head.key()).orElse(head.metric()));
     }
 
-    /**
-     * Equivalent to {@link #topLevelReduce(List, AggregationReduceContext)} but it takes an iterator and a count.
-     */
-    public static InternalAggregations topLevelReduce(Iterator<InternalAggregations> aggs, int count, AggregationReduceContext context) {
-        if (count == 0) {
-            return null;
-        }
-        return maybeExecuteFinalReduce(context, count == 1 ? reduce(aggs.next(), context) : reduce(aggs, count, context));
-    }
-
-    private static InternalAggregations maybeExecuteFinalReduce(AggregationReduceContext context, InternalAggregations reduced) {
+    public static InternalAggregations maybeExecuteFinalReduce(AggregationReduceContext context, InternalAggregations reduced) {
         if (reduced == null) {
             return null;
         }
@@ -242,16 +232,6 @@ public static InternalAggregations reduce(List<InternalAggregations> aggregation
         }
     }
 
-    private static InternalAggregations reduce(Iterator<InternalAggregations> aggsIterator, int count, AggregationReduceContext context) {
-        // general case
-        var first = aggsIterator.next();
-        try (AggregatorsReducer reducer = new AggregatorsReducer(first, context, count)) {
-            reducer.accept(first);
-            aggsIterator.forEachRemaining(reducer::accept);
-            return reducer.get();
-        }
-    }
-
     public static InternalAggregations reduce(InternalAggregations aggregations, AggregationReduceContext context) {
         final List<InternalAggregation> internalAggregations = aggregations.asList();
         int size = internalAggregations.size();