Move the mem accounting buffer size in aggregationcontext

Author: andreidan

benchmarks/src/main/java/org/elasticsearch/benchmark/search/aggregations/AggConstructionContentionBenchmark.java

@@ -366,6 +366,11 @@ public Set<String> sourcePath(String fullName) {
             return Set.of(fullName);
         }
 
+        @Override
+        public long memoryAccountingBufferSize() {
+            return 1024 * 1024;
+        }
+
         @Override
         public void close() {
             List<Releasable> releaseMe = new ArrayList<>(this.releaseMe);

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -1441,7 +1441,8 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
                 context::isCancelled,
                 context::buildFilteredQuery,
                 enableRewriteAggsToFilterByFilter,
-                source.aggregations().isInSortOrderExecutionRequired()
+                source.aggregations().isInSortOrderExecutionRequired(),
+                memoryAccountingBufferSize
             );
             context.addQuerySearchResultReleasable(aggContext);
             try {

server/src/main/java/org/elasticsearch/search/aggregations/metrics/TopHitsAggregator.java

@@ -35,7 +35,6 @@
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
-import org.elasticsearch.search.SearchService;
 import org.elasticsearch.search.aggregations.AggregationExecutionContext;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.InternalAggregation;
@@ -57,8 +56,6 @@
 
 class TopHitsAggregator extends MetricsAggregator {
 
-    private final long memAccountingBufferSize;
-
     private static class Collectors {
         public final TopDocsCollector<?> topDocsCollector;
         public final MaxScoreCollector maxScoreCollector;
@@ -90,7 +87,6 @@ private static class Collectors {
         this.subSearchContext = subSearchContext;
         this.topDocsCollectors = new LongObjectPagedHashMap<>(1, bigArrays);
         this.fetchProfiles = context.profiling() ? new ArrayList<>() : null;
-        this.memAccountingBufferSize = context.getClusterSettings().get(SearchService.MEMORY_ACCOUNTING_BUFFER_SIZE).getBytes();
     }
 
     @Override
@@ -201,7 +197,12 @@ public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOE
         for (int i = 0; i < topDocs.scoreDocs.length; i++) {
             docIdsToLoad[i] = topDocs.scoreDocs[i].doc;
         }
-        FetchSearchResult fetchResult = runFetchPhase(subSearchContext, docIdsToLoad, context.breaker(), memAccountingBufferSize);
+        FetchSearchResult fetchResult = runFetchPhase(
+            subSearchContext,
+            docIdsToLoad,
+            context.breaker(),
+            context.memoryAccountingBufferSize()
+        );
         if (fetchProfiles != null) {
             fetchProfiles.add(fetchResult.profileResult());
         }

server/src/main/java/org/elasticsearch/search/aggregations/support/AggregationContext.java

@@ -308,6 +308,11 @@ public final AggregationUsageService getUsageService() {
 
     public abstract Set<String> sourcePath(String fullName);
 
+    /**
+     * Return the amount of memory to buffer locally before accounting for it in the breaker.
+     */
+    public abstract long memoryAccountingBufferSize();
+
     /**
      * Does this index have a {@code _doc_count} field in any segment?
      */
@@ -355,6 +360,7 @@ public static class ProductionAggregationContext extends AggregationContext {
         private final AnalysisRegistry analysisRegistry;
 
         private final List<Aggregator> releaseMe = new ArrayList<>();
+        private final long memoryAccountingBufferSize;
 
         public ProductionAggregationContext(
             AnalysisRegistry analysisRegistry,
@@ -372,7 +378,8 @@ public ProductionAggregationContext(
             Supplier<Boolean> isCancelled,
             Function<Query, Query> filterQuery,
             boolean enableRewriteToFilterByFilter,
-            boolean inSortOrderExecutionRequired
+            boolean inSortOrderExecutionRequired,
+            long memoryAccountingBufferSize
         ) {
             this.analysisRegistry = analysisRegistry;
             this.context = context;
@@ -407,6 +414,7 @@ public ProductionAggregationContext(
             this.filterQuery = filterQuery;
             this.enableRewriteToFilterByFilter = enableRewriteToFilterByFilter;
             this.inSortOrderExecutionRequired = inSortOrderExecutionRequired;
+            this.memoryAccountingBufferSize = memoryAccountingBufferSize;
         }
 
         @Override
@@ -612,6 +620,11 @@ public Set<String> sourcePath(String fullName) {
             return context.sourcePath(fullName);
         }
 
+        @Override
+        public long memoryAccountingBufferSize() {
+            return memoryAccountingBufferSize;
+        }
+
         @Override
         public void close() {
             /*

test/framework/src/main/java/org/elasticsearch/index/mapper/MapperServiceTestCase.java

@@ -675,6 +675,11 @@ public Set<String> sourcePath(String fullName) {
                 return Set.of(fullName);
             }
 
+            @Override
+            public long memoryAccountingBufferSize() {
+                return 1024 * 1024;
+            }
+
             @Override
             public void close() {
                 throw new UnsupportedOperationException();

test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java

@@ -68,6 +68,8 @@
 import org.elasticsearch.common.network.NetworkAddress;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.unit.ByteSizeUnit;
+import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.MockBigArrays;
 import org.elasticsearch.common.util.MockPageCacheRecycler;
@@ -444,7 +446,8 @@ public Iterable<MappedFieldType> dimensionFields() {
             isCancelled,
             q -> q,
             true,
-            isInSortOrderExecutionRequired
+            isInSortOrderExecutionRequired,
+            ByteSizeValue.of(1, ByteSizeUnit.MB).getBytes()
         );
         return context;
     }