Local accounting up to 32kb

Author: andreidan

server/src/main/java/org/elasticsearch/common/MemoryAccountingBytesRefCounted.java

@@ -12,29 +12,31 @@
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.core.AbstractRefCounted;
 
+/**
+ * A ref counted object that accounts for memory usage in bytes and releases the
+ * accounted memory from the circuit breaker when the reference count reaches zero.
+ */
 public final class MemoryAccountingBytesRefCounted extends AbstractRefCounted {
 
-    private long bytes;
+    private int bytes;
     private final CircuitBreaker breaker;
 
-    private MemoryAccountingBytesRefCounted(long bytes, CircuitBreaker breaker) {
+    private MemoryAccountingBytesRefCounted(int bytes, CircuitBreaker breaker) {
         this.bytes = bytes;
         this.breaker = breaker;
     }
 
     public static MemoryAccountingBytesRefCounted create(CircuitBreaker breaker) {
-        return new MemoryAccountingBytesRefCounted(0L, breaker);
+        return new MemoryAccountingBytesRefCounted(0, breaker);
     }
 
-    public void account(long bytes, String label) {
+    public void account(int bytes, String label) {
         breaker.addEstimateBytesAndMaybeBreak(bytes, label);
         this.bytes += bytes;
     }
 
     @Override
     protected void closeInternal() {
-        if (bytes != 0L) {
-            breaker.addWithoutBreaking(-bytes);
-        }
+        breaker.addWithoutBreaking(-bytes);
     }
 }

server/src/main/java/org/elasticsearch/common/settings/ClusterSettings.java

@@ -535,6 +535,7 @@ public void apply(Settings value, Settings current, Settings previous) {
         SearchModule.SCRIPTED_METRICS_AGG_ALLOWED_STORED_SCRIPTS,
         SearchService.SEARCH_WORKER_THREADS_ENABLED,
         SearchService.QUERY_PHASE_PARALLEL_COLLECTION_ENABLED,
+        SearchService.MEMORY_ACCOUNTING_BUFFER_SIZE,
         ThreadPool.ESTIMATED_TIME_INTERVAL_SETTING,
         ThreadPool.LATE_TIME_INTERVAL_WARN_THRESHOLD_SETTING,
         ThreadPool.SLOW_SCHEDULER_TASK_WARN_THRESHOLD_SETTING,

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

@@ -274,6 +274,18 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         Property.NodeScope
     );
 
+    /**
+     * The size of the buffer used for memory accounting.
+     * This buffer is used to locally track the memory accummulate during the executiong of
+     * a search request before submitting the accumulated value to the circuit breaker.
+     */
+    public static final Setting<ByteSizeValue> MEMORY_ACCOUNTING_BUFFER_SIZE = Setting.byteSizeSetting(
+        "search.memory_accounting_buffer_size",
+        ByteSizeValue.of(32, ByteSizeUnit.KB),
+        Property.Dynamic,
+        Property.NodeScope
+    );
+
     public static final int DEFAULT_SIZE = 10;
     public static final int DEFAULT_FROM = 0;
     private static final StackTraceElement[] EMPTY_STACK_TRACE_ARRAY = new StackTraceElement[0];
@@ -311,6 +323,8 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
 
     private volatile boolean enableRewriteAggsToFilterByFilter;
 
+    private volatile ByteSizeValue memoryAccountingBufferSize;
+
     private final Cancellable keepAliveReaper;
 
     private final AtomicLong idGenerator = new AtomicLong();
@@ -388,6 +402,9 @@ public SearchService(
         enableQueryPhaseParallelCollection = QUERY_PHASE_PARALLEL_COLLECTION_ENABLED.get(settings);
         clusterService.getClusterSettings()
             .addSettingsUpdateConsumer(QUERY_PHASE_PARALLEL_COLLECTION_ENABLED, this::setEnableQueryPhaseParallelCollection);
+
+        memoryAccountingBufferSize = MEMORY_ACCOUNTING_BUFFER_SIZE.get(settings);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(MEMORY_ACCOUNTING_BUFFER_SIZE, this::setMemoryAccountingBufferSize);
     }
 
     private void setEnableSearchWorkerThreads(boolean enableSearchWorkerThreads) {
@@ -402,6 +419,10 @@ private void setEnableQueryPhaseParallelCollection(boolean enableQueryPhaseParal
         this.enableQueryPhaseParallelCollection = enableQueryPhaseParallelCollection;
     }
 
+    private void setMemoryAccountingBufferSize(ByteSizeValue memoryAccountingBufferSize) {
+        this.memoryAccountingBufferSize = memoryAccountingBufferSize;
+    }
+
     private static void validateKeepAlives(TimeValue defaultKeepAlive, TimeValue maxKeepAlive) {
         if (defaultKeepAlive.millis() > maxKeepAlive.millis()) {
             throw new IllegalArgumentException(
@@ -790,7 +811,7 @@ public void executeRankFeaturePhase(RankFeatureShardRequest request, SearchShard
                     return searchContext.rankFeatureResult();
                 }
                 RankFeatureShardPhase.prepareForFetch(searchContext, request);
-                fetchPhase.execute(searchContext, docIds, null, circuitBreaker);
+                fetchPhase.execute(searchContext, docIds, null, circuitBreaker, memoryAccountingBufferSize.getBytes());
                 RankFeatureShardPhase.processFetch(searchContext);
                 var rankFeatureResult = searchContext.rankFeatureResult();
                 rankFeatureResult.incRef();
@@ -808,7 +829,7 @@ private QueryFetchSearchResult executeFetchPhase(ReaderContext reader, SearchCon
             Releasable scope = tracer.withScope(context.getTask());
             SearchOperationListenerExecutor executor = new SearchOperationListenerExecutor(context, true, afterQueryTime)
         ) {
-            fetchPhase.execute(context, shortcutDocIdsToLoad(context), null, circuitBreaker);
+            fetchPhase.execute(context, shortcutDocIdsToLoad(context), null, circuitBreaker, memoryAccountingBufferSize.getBytes());
             if (reader.singleSession()) {
                 freeReaderContext(reader.id());
             }
@@ -974,7 +995,13 @@ public void executeFetchPhase(ShardFetchRequest request, CancellableTask task, A
                             System.nanoTime()
                         )
                     ) {
-                        fetchPhase.execute(searchContext, request.docIds(), request.getRankDocks(), circuitBreaker);
+                        fetchPhase.execute(
+                            searchContext,
+                            request.docIds(),
+                            request.getRankDocks(),
+                            circuitBreaker,
+                            memoryAccountingBufferSize.getBytes()
+                        );
                         if (readerContext.singleSession()) {
                             freeReaderContext(request.contextId());
                         }

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

@@ -35,6 +35,7 @@
 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;
@@ -56,6 +57,8 @@
 
 class TopHitsAggregator extends MetricsAggregator {
 
+    private final long memAccountingBufferSize;
+
     private static class Collectors {
         public final TopDocsCollector<?> topDocsCollector;
         public final MaxScoreCollector maxScoreCollector;
@@ -87,6 +90,7 @@ 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
@@ -197,7 +201,7 @@ 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());
+        FetchSearchResult fetchResult = runFetchPhase(subSearchContext, docIdsToLoad, context.breaker(), memAccountingBufferSize);
         if (fetchProfiles != null) {
             fetchProfiles.add(fetchResult.profileResult());
         }
@@ -221,7 +225,12 @@ public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOE
         );
     }
 
-    private static FetchSearchResult runFetchPhase(SubSearchContext subSearchContext, int[] docIdsToLoad, CircuitBreaker breaker) {
+    private static FetchSearchResult runFetchPhase(
+        SubSearchContext subSearchContext,
+        int[] docIdsToLoad,
+        CircuitBreaker breaker,
+        long memAccountingBufferSize
+    ) {
         // Fork the search execution context for each slice, because the fetch phase does not support concurrent execution yet.
         SearchExecutionContext searchExecutionContext = new SearchExecutionContext(subSearchContext.getSearchExecutionContext());
         SubSearchContext fetchSubSearchContext = new SubSearchContext(subSearchContext) {
@@ -230,7 +239,7 @@ public SearchExecutionContext getSearchExecutionContext() {
                 return searchExecutionContext;
             }
         };
-        fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null, breaker);
+        fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null, breaker, memAccountingBufferSize);
         return fetchSubSearchContext.fetchResult();
     }