Reduce memory for big aggs run against many shards (#54758) (#55024)

This changes the behavior of aggregations when search is performed
against enough shards to enable "batch reduce" mode. In this case we
force always store aggregations in serialized form rather than a
traditional java reference. This should shrink the memory usage of large
aggregations at the cost of slightly slowing down aggregations where the
coordinating node is also a data node. Because we're only doing this
when there are many shards this is likely to be fairly rare.

As a side effect this lets us add logs for the memory usage of the aggs
buffer:
```
[2020-04-03T17:03:57,052][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [1320->448] max [1320]
[2020-04-03T17:03:57,089][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [1328->448] max [1328]
[2020-04-03T17:03:57,102][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [1328->448] max [1328]
[2020-04-03T17:03:57,103][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [1328->448] max [1328]
[2020-04-03T17:03:57,105][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs final reduction [888] max [1328]
```

These are useful, but you need to keep some things in mind before
trusting them:
1. The buffers are oversized ala Lucene's ArrayUtils. This means that we
   are using more space than we need, but probably not much more.
2. Before they are merged the aggregations are inflated into their
   traditional Java objects which *probably* take up a lot more space
   than the serialized form. That is, after all, the reason why we store
   them in serialized form in the first place.

And, just because I can, here is another example of the log:
```
[2020-04-03T17:06:18,731][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [147528->49176] max [147528]
[2020-04-03T17:06:18,750][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [147528->49176] max [147528]
[2020-04-03T17:06:18,809][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [147528->49176] max [147528]
[2020-04-03T17:06:18,827][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs partial reduction [147528->49176] max [147528]
[2020-04-03T17:06:18,829][TRACE][o.e.a.s.SearchPhaseController] [runTask-0] aggs final reduction [98352] max [147528]
```

I got that last one by building a ten shard index with a million docs in
it and running a `sum` in three layers of `terms` aggregations, all on
`long` fields, and with a `batched_reduce_size` of `3`.

Author: nik9000

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

@@ -19,18 +19,11 @@
 
 package org.elasticsearch.action.search;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.function.Function;
-import java.util.function.IntFunction;
-import java.util.function.Supplier;
-import java.util.stream.Collectors;
+import com.carrotsearch.hppc.IntArrayList;
+import com.carrotsearch.hppc.ObjectObjectHashMap;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.CollectionStatistics;
 import org.apache.lucene.search.FieldDoc;
@@ -44,6 +37,8 @@
 import org.apache.lucene.search.TotalHits.Relation;
 import org.apache.lucene.search.grouping.CollapseTopFieldDocs;
 import org.elasticsearch.common.collect.HppcMaps;
+import org.elasticsearch.common.io.stream.DelayableWriteable;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.SearchHit;
@@ -67,16 +62,28 @@
 import org.elasticsearch.search.suggest.Suggest.Suggestion;
 import org.elasticsearch.search.suggest.completion.CompletionSuggestion;
 
-import com.carrotsearch.hppc.IntArrayList;
-import com.carrotsearch.hppc.ObjectObjectHashMap;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.function.IntFunction;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 public final class SearchPhaseController {
+    private static final Logger logger = LogManager.getLogger(SearchPhaseController.class);
     private static final ScoreDoc[] EMPTY_DOCS = new ScoreDoc[0];
 
+    private final NamedWriteableRegistry namedWriteableRegistry;
     private final Function<SearchRequest, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder;
 
-    public SearchPhaseController(
+    public SearchPhaseController(NamedWriteableRegistry namedWriteableRegistry,
             Function<SearchRequest, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder) {
+        this.namedWriteableRegistry = namedWriteableRegistry;
         this.requestToAggReduceContextBuilder = requestToAggReduceContextBuilder;
     }
 
@@ -430,7 +437,8 @@ public ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResul
      * @see QuerySearchResult#consumeProfileResult()
      */
     private ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResult> queryResults,
-                                                List<Supplier<InternalAggregations>> bufferedAggs, List<TopDocs> bufferedTopDocs,
+                                                List<Supplier<InternalAggregations>> bufferedAggs,
+                                                List<TopDocs> bufferedTopDocs,
                                                 TopDocsStats topDocsStats, int numReducePhases, boolean isScrollRequest,
                                                 InternalAggregation.ReduceContextBuilder aggReduceContextBuilder,
                                                 boolean performFinalReduce) {
@@ -522,7 +530,7 @@ private ReducedQueryPhase reducedQueryPhase(Collection<? extends SearchPhaseResu
     private InternalAggregations reduceAggs(
         InternalAggregation.ReduceContextBuilder aggReduceContextBuilder,
         boolean performFinalReduce,
-        List<Supplier<InternalAggregations>> aggregationsList
+        List<? extends Supplier<InternalAggregations>> aggregationsList
     ) {
         /*
          * Parse the aggregations, clearing the list as we go so bits backing
@@ -617,8 +625,9 @@ public InternalSearchResponse buildResponse(SearchHits hits) {
      * iff the buffer is exhausted.
      */
     static final class QueryPhaseResultConsumer extends ArraySearchPhaseResults<SearchPhaseResult> {
+        private final NamedWriteableRegistry namedWriteableRegistry;
         private final SearchShardTarget[] processedShards;
-        private final Supplier<InternalAggregations>[] aggsBuffer;
+        private final DelayableWriteable.Serialized<InternalAggregations>[] aggsBuffer;
         private final TopDocs[] topDocsBuffer;
         private final boolean hasAggs;
         private final boolean hasTopDocs;
@@ -631,6 +640,8 @@ static final class QueryPhaseResultConsumer extends ArraySearchPhaseResults<Sear
         private final int topNSize;
         private final InternalAggregation.ReduceContextBuilder aggReduceContextBuilder;
         private final boolean performFinalReduce;
+        private long aggsCurrentBufferSize;
+        private long aggsMaxBufferSize;
 
         /**
          * Creates a new {@link QueryPhaseResultConsumer}
@@ -641,12 +652,14 @@ static final class QueryPhaseResultConsumer extends ArraySearchPhaseResults<Sear
          * @param bufferSize the size of the reduce buffer. if the buffer size is smaller than the number of expected results
          *                   the buffer is used to incrementally reduce aggregation results before all shards responded.
          */
-        private QueryPhaseResultConsumer(SearchProgressListener progressListener, SearchPhaseController controller,
+        private QueryPhaseResultConsumer(NamedWriteableRegistry namedWriteableRegistry, SearchProgressListener progressListener,
+                                         SearchPhaseController controller,
                                          int expectedResultSize, int bufferSize, boolean hasTopDocs, boolean hasAggs,
                                          int trackTotalHitsUpTo, int topNSize,
                                          InternalAggregation.ReduceContextBuilder aggReduceContextBuilder,
                                          boolean performFinalReduce) {
             super(expectedResultSize);
+            this.namedWriteableRegistry = namedWriteableRegistry;
             if (expectedResultSize != 1 && bufferSize < 2) {
                 throw new IllegalArgumentException("buffer size must be >= 2 if there is more than one expected result");
             }
@@ -661,7 +674,7 @@ private QueryPhaseResultConsumer(SearchProgressListener progressListener, Search
             this.processedShards = new SearchShardTarget[expectedResultSize];
             // no need to buffer anything if we have less expected results. in this case we don't consume any results ahead of time.
             @SuppressWarnings("unchecked")
-            Supplier<InternalAggregations>[] aggsBuffer = new Supplier[hasAggs ? bufferSize : 0];
+            DelayableWriteable.Serialized<InternalAggregations>[] aggsBuffer = new DelayableWriteable.Serialized[hasAggs ? bufferSize : 0];
             this.aggsBuffer = aggsBuffer;
             this.topDocsBuffer = new TopDocs[hasTopDocs ? bufferSize : 0];
             this.hasTopDocs = hasTopDocs;
@@ -684,15 +697,21 @@ public void consumeResult(SearchPhaseResult result) {
         private synchronized void consumeInternal(QuerySearchResult querySearchResult) {
             if (querySearchResult.isNull() == false) {
                 if (index == bufferSize) {
+                    InternalAggregations reducedAggs = null;
                     if (hasAggs) {
                         List<InternalAggregations> aggs = new ArrayList<>(aggsBuffer.length);
                         for (int i = 0; i < aggsBuffer.length; i++) {
                             aggs.add(aggsBuffer[i].get());
                             aggsBuffer[i] = null; // null the buffer so it can be GCed now.
                         }
-                        InternalAggregations reducedAggs = InternalAggregations.topLevelReduce(
-                                aggs, aggReduceContextBuilder.forPartialReduction());
-                        aggsBuffer[0] = () -> reducedAggs;
+                        reducedAggs = InternalAggregations.topLevelReduce(aggs, aggReduceContextBuilder.forPartialReduction());
+                        aggsBuffer[0] = DelayableWriteable.referencing(reducedAggs)
+                                .asSerialized(InternalAggregations::new, namedWriteableRegistry);
+                        long previousBufferSize = aggsCurrentBufferSize;
+                        aggsMaxBufferSize = Math.max(aggsMaxBufferSize, aggsCurrentBufferSize);
+                        aggsCurrentBufferSize = aggsBuffer[0].ramBytesUsed();
+                        logger.trace("aggs partial reduction [{}->{}] max [{}]",
+                                previousBufferSize, aggsCurrentBufferSize, aggsMaxBufferSize);
                     }
                     if (hasTopDocs) {
                         TopDocs reducedTopDocs = mergeTopDocs(Arrays.asList(topDocsBuffer),
@@ -705,12 +724,13 @@ private synchronized void consumeInternal(QuerySearchResult querySearchResult) {
                     index = 1;
                     if (hasAggs || hasTopDocs) {
                         progressListener.notifyPartialReduce(SearchProgressListener.buildSearchShards(processedShards),
-                            topDocsStats.getTotalHits(), hasAggs ? aggsBuffer[0].get() : null, numReducePhases);
+                            topDocsStats.getTotalHits(), reducedAggs, numReducePhases);
                     }
                 }
                 final int i = index++;
                 if (hasAggs) {
-                    aggsBuffer[i] = querySearchResult.consumeAggs();
+                    aggsBuffer[i] = querySearchResult.consumeAggs().asSerialized(InternalAggregations::new, namedWriteableRegistry);
+                    aggsCurrentBufferSize += aggsBuffer[i].ramBytesUsed();
                 }
                 if (hasTopDocs) {
                     final TopDocsAndMaxScore topDocs = querySearchResult.consumeTopDocs(); // can't be null
@@ -723,7 +743,7 @@ private synchronized void consumeInternal(QuerySearchResult querySearchResult) {
         }
 
         private synchronized List<Supplier<InternalAggregations>> getRemainingAggs() {
-            return hasAggs ? Arrays.asList(aggsBuffer).subList(0, index) : null;
+            return hasAggs ? Arrays.asList((Supplier<InternalAggregations>[]) aggsBuffer).subList(0, index) : null;
         }
 
         private synchronized List<TopDocs> getRemainingTopDocs() {
@@ -732,6 +752,8 @@ private synchronized List<TopDocs> getRemainingTopDocs() {
 
         @Override
         public ReducedQueryPhase reduce() {
+            aggsMaxBufferSize = Math.max(aggsMaxBufferSize, aggsCurrentBufferSize);
+            logger.trace("aggs final reduction [{}] max [{}]", aggsCurrentBufferSize, aggsMaxBufferSize);
             ReducedQueryPhase reducePhase = controller.reducedQueryPhase(results.asList(),
                 getRemainingAggs(), getRemainingTopDocs(), topDocsStats, numReducePhases, false,
                 aggReduceContextBuilder, performFinalReduce);
@@ -767,8 +789,8 @@ ArraySearchPhaseResults<SearchPhaseResult> newSearchPhaseResults(SearchProgressL
             if (request.getBatchedReduceSize() < numShards) {
                 int topNSize = getTopDocsSize(request);
                 // only use this if there are aggs and if there are more shards than we should reduce at once
-                return new QueryPhaseResultConsumer(listener, this, numShards, request.getBatchedReduceSize(), hasTopDocs, hasAggs,
-                    trackTotalHitsUpTo, topNSize, aggReduceContextBuilder, request.isFinalReduce());
+                return new QueryPhaseResultConsumer(namedWriteableRegistry, listener, this, numShards, request.getBatchedReduceSize(),
+                    hasTopDocs, hasAggs, trackTotalHitsUpTo, topNSize, aggReduceContextBuilder, request.isFinalReduce());
             }
         }
         return new ArraySearchPhaseResults<SearchPhaseResult>(numShards) {

server/src/main/java/org/elasticsearch/common/io/stream/DelayableWriteable.java

@@ -19,12 +19,14 @@
 
 package org.elasticsearch.common.io.stream;
 
-import java.io.IOException;
-import java.util.function.Supplier;
-
+import org.apache.lucene.util.Accountable;
+import org.apache.lucene.util.RamUsageEstimator;
 import org.elasticsearch.Version;
 import org.elasticsearch.common.bytes.BytesReference;
 
+import java.io.IOException;
+import java.util.function.Supplier;
+
 /**
  * A holder for {@link Writeable}s that can delays reading the underlying
  * {@linkplain Writeable} when it is read from a remote node.
@@ -43,12 +45,22 @@ public static <T extends Writeable> DelayableWriteable<T> referencing(T referenc
      * when {@link Supplier#get()} is called.
      */
     public static <T extends Writeable> DelayableWriteable<T> delayed(Writeable.Reader<T> reader, StreamInput in) throws IOException {
-        return new Delayed<>(reader, in);
+        return new Serialized<>(reader, in.getVersion(), in.namedWriteableRegistry(), in.readBytesReference());
     }
 
     private DelayableWriteable() {}
 
-    public abstract boolean isDelayed();
+    /**
+     * Returns a {@linkplain DelayableWriteable} that stores its contents
+     * in serialized form.
+     */
+    public abstract Serialized<T> asSerialized(Writeable.Reader<T> reader, NamedWriteableRegistry registry);
+
+    /**
+     * {@code true} if the {@linkplain Writeable} is being stored in
+     * serialized form, {@code false} otherwise.
+     */
+    abstract boolean isSerialized();
 
     private static class Referencing<T extends Writeable> extends DelayableWriteable<T> {
         private T reference;
@@ -59,11 +71,7 @@ private static class Referencing<T extends Writeable> extends DelayableWriteable
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            try (BytesStreamOutput buffer = new BytesStreamOutput()) {
-                buffer.setVersion(out.getVersion());
-                reference.writeTo(buffer);
-                out.writeBytesReference(buffer.bytes());
-            }
+            out.writeBytesReference(writeToBuffer(out.getVersion()).bytes());
         }
 
         @Override
@@ -72,27 +80,48 @@ public T get() {
         }
 
         @Override
-        public boolean isDelayed() {
+        public Serialized<T> asSerialized(Reader<T> reader, NamedWriteableRegistry registry) {
+            try {
+                return new Serialized<T>(reader, Version.CURRENT, registry, writeToBuffer(Version.CURRENT).bytes());
+            } catch (IOException e) {
+                throw new RuntimeException("unexpected error expanding aggregations", e);
+            }
+        }
+
+        @Override
+        boolean isSerialized() {
             return false;
         }
+
+        private BytesStreamOutput writeToBuffer(Version version) throws IOException {
+            try (BytesStreamOutput buffer = new BytesStreamOutput()) {
+                buffer.setVersion(version);
+                reference.writeTo(buffer);
+                return buffer;
+            }
+        }
     }
 
-    private static class Delayed<T extends Writeable> extends DelayableWriteable<T> {
+    /**
+     * A {@link Writeable} stored in serialized form.
+     */
+    public static class Serialized<T extends Writeable> extends DelayableWriteable<T> implements Accountable {
         private final Writeable.Reader<T> reader;
-        private final Version remoteVersion;
-        private final BytesReference serialized;
+        private final Version serializedAtVersion;
         private final NamedWriteableRegistry registry;
+        private final BytesReference serialized;
 
-        Delayed(Writeable.Reader<T> reader, StreamInput in) throws IOException {
+        Serialized(Writeable.Reader<T> reader, Version serializedAtVersion,
+                NamedWriteableRegistry registry, BytesReference serialized) throws IOException {
             this.reader = reader;
-            remoteVersion = in.getVersion();
-            serialized = in.readBytesReference();
-            registry = in.namedWriteableRegistry();
+            this.serializedAtVersion = serializedAtVersion;
+            this.registry = registry;
+            this.serialized = serialized;
         }
 
         @Override
         public void writeTo(StreamOutput out) throws IOException {
-            if (out.getVersion() == remoteVersion) {
+            if (out.getVersion() == serializedAtVersion) {
                 /*
                  * If the version *does* line up we can just copy the bytes
                  * which is good because this is how shard request caching
@@ -116,7 +145,7 @@ public T get() {
             try {
                 try (StreamInput in = registry == null ?
                         serialized.streamInput() : new NamedWriteableAwareStreamInput(serialized.streamInput(), registry)) {
-                    in.setVersion(remoteVersion);
+                    in.setVersion(serializedAtVersion);
                     return reader.read(in);
                 }
             } catch (IOException e) {
@@ -125,8 +154,18 @@ public T get() {
         }
 
         @Override
-        public boolean isDelayed() {
+        public Serialized<T> asSerialized(Reader<T> reader, NamedWriteableRegistry registry) {
+            return this; // We're already serialized
+        }
+
+        @Override
+        boolean isSerialized() {
             return true;
         }
+
+        @Override
+        public long ramBytesUsed() {
+            return serialized.ramBytesUsed() + RamUsageEstimator.NUM_BYTES_OBJECT_REF * 3 + RamUsageEstimator.NUM_BYTES_OBJECT_HEADER;
+        }
     }
 }

server/src/main/java/org/elasticsearch/node/Node.java

@@ -585,7 +585,8 @@ protected Node(final Environment initialEnvironment,
                     b.bind(MetadataCreateIndexService.class).toInstance(metadataCreateIndexService);
                     b.bind(SearchService.class).toInstance(searchService);
                     b.bind(SearchTransportService.class).toInstance(searchTransportService);
-                    b.bind(SearchPhaseController.class).toInstance(new SearchPhaseController(searchService::aggReduceContextBuilder));
+                    b.bind(SearchPhaseController.class).toInstance(new SearchPhaseController(
+                        namedWriteableRegistry, searchService::aggReduceContextBuilder));
                     b.bind(Transport.class).toInstance(transport);
                     b.bind(TransportService.class).toInstance(transportService);
                     b.bind(NetworkService.class).toInstance(networkService);

server/src/test/java/org/elasticsearch/action/search/DfsQueryPhaseTests.java

@@ -213,6 +213,6 @@ public void run() throws IOException {
     }
 
     private SearchPhaseController searchPhaseController() {
-        return new SearchPhaseController(request -> InternalAggregationTestCase.emptyReduceContextBuilder());
+        return new SearchPhaseController(writableRegistry(), request -> InternalAggregationTestCase.emptyReduceContextBuilder());
     }
  }