Periodically check the available memory when fetching search hits source

docs/changelog/121920.yaml

@@ -0,0 +1,6 @@
+pr: 121920
+summary: Account for the `SearchHit` source in circuit breaker
+area: Search
+type: enhancement
+issues:
+ - 89656

server/src/internalClusterTest/java/org/elasticsearch/action/search/TransportSearchIT.java

@@ -78,6 +78,7 @@
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.either;
 import static org.hamcrest.Matchers.equalTo;
 
 public class TransportSearchIT extends ESIntegTestCase {
@@ -487,7 +488,10 @@ public void onFailure(Exception e) {
                     Exception.class,
                     client.prepareSearch("test").addAggregation(new TestAggregationBuilder("test"))
                 );
-                assertThat(exc.getCause().getMessage(), containsString("<reduce_aggs>"));
+                assertThat(
+                    exc.getCause().getMessage(),
+                    either(containsString("<reduce_aggs>")).or(containsString("fetch phase source loader"))
+                );
             });
 
             final AtomicArray<Exception> exceptions = new AtomicArray<>(10);
@@ -514,7 +518,10 @@ public void onFailure(Exception exc) {
             latch.await();
             assertThat(exceptions.asList().size(), equalTo(10));
             for (Exception exc : exceptions.asList()) {
-                assertThat(exc.getCause().getMessage(), containsString("<reduce_aggs>"));
+                assertThat(
+                    exc.getCause().getMessage(),
+                    either(containsString("<reduce_aggs>")).or(containsString("fetch phase source loader"))
+                );
             }
             assertBusy(() -> assertThat(requestBreakerUsed(), equalTo(0L)));
         } finally {

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.common;
+
+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 int bytes;
+    private final CircuitBreaker breaker;
+
+    private MemoryAccountingBytesRefCounted(int bytes, CircuitBreaker breaker) {
+        this.bytes = bytes;
+        this.breaker = breaker;
+    }
+
+    public static MemoryAccountingBytesRefCounted create(CircuitBreaker breaker) {
+        return new MemoryAccountingBytesRefCounted(0, breaker);
+    }
+
+    public void account(int bytes, String label) {
+        breaker.addEstimateBytesAndMaybeBreak(bytes, label);
+        this.bytes += bytes;
+    }
+
+    @Override
+    protected void closeInternal() {
+        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/SearchHit.java

@@ -111,7 +111,6 @@ public final class SearchHit implements Writeable, ToXContentObject, RefCounted
 
     private final RefCounted refCounted;
 
-    // used only in tests
     public SearchHit(int docId) {
         this(docId, null);
     }
@@ -124,7 +123,7 @@ public SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity) {
         this(nestedTopDocId, id, nestedIdentity, null);
     }
 
-    private SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity, @Nullable RefCounted refCounted) {
+    public SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity, @Nullable RefCounted refCounted) {
         this(
             nestedTopDocId,
             DEFAULT_SCORE,
@@ -293,6 +292,7 @@ public static SearchHit unpooled(int docId, String id) {
     }
 
     public static SearchHit unpooled(int nestedTopDocId, String id, NestedIdentity nestedIdentity) {
+        // always referenced search hits do NOT call #deallocate
         return new SearchHit(nestedTopDocId, id, nestedIdentity, ALWAYS_REFERENCED);
     }
 

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];
@@ -291,6 +303,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
     private final BigArrays bigArrays;
 
     private final FetchPhase fetchPhase;
+    private final CircuitBreaker circuitBreaker;
     private volatile Executor searchExecutor;
     private volatile boolean enableQueryPhaseParallelCollection;
 
@@ -310,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();
@@ -341,11 +356,8 @@ public SearchService(
         this.scriptService = scriptService;
         this.bigArrays = bigArrays;
         this.fetchPhase = fetchPhase;
-        this.multiBucketConsumerService = new MultiBucketConsumerService(
-            clusterService,
-            settings,
-            circuitBreakerService.getBreaker(CircuitBreaker.REQUEST)
-        );
+        circuitBreaker = circuitBreakerService.getBreaker(CircuitBreaker.REQUEST);
+        this.multiBucketConsumerService = new MultiBucketConsumerService(clusterService, settings, circuitBreaker);
         this.executorSelector = executorSelector;
         this.tracer = tracer;
 
@@ -390,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) {
@@ -404,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(
@@ -792,7 +811,7 @@ public void executeRankFeaturePhase(RankFeatureShardRequest request, SearchShard
                     return searchContext.rankFeatureResult();
                 }
                 RankFeatureShardPhase.prepareForFetch(searchContext, request);
-                fetchPhase.execute(searchContext, docIds, null);
+                fetchPhase.execute(searchContext, docIds, null, circuitBreaker, memoryAccountingBufferSize.getBytes());
                 RankFeatureShardPhase.processFetch(searchContext);
                 var rankFeatureResult = searchContext.rankFeatureResult();
                 rankFeatureResult.incRef();
@@ -810,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);
+            fetchPhase.execute(context, shortcutDocIdsToLoad(context), null, circuitBreaker, memoryAccountingBufferSize.getBytes());
             if (reader.singleSession()) {
                 freeReaderContext(reader.id());
             }
@@ -976,7 +995,13 @@ public void executeFetchPhase(ShardFetchRequest request, CancellableTask task, A
                             System.nanoTime()
                         )
                     ) {
-                        fetchPhase.execute(searchContext, request.docIds(), request.getRankDocks());
+                        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

@@ -25,6 +25,7 @@
 import org.apache.lucene.search.TopScoreDocCollectorManager;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.search.MaxScoreCollector;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.lucene.Lucene;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
 import org.elasticsearch.common.util.BigArrays;
@@ -34,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;
@@ -55,6 +57,8 @@
 
 class TopHitsAggregator extends MetricsAggregator {
 
+    private final long memAccountingBufferSize;
+
     private static class Collectors {
         public final TopDocsCollector<?> topDocsCollector;
         public final MaxScoreCollector maxScoreCollector;
@@ -86,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
@@ -196,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);
+        FetchSearchResult fetchResult = runFetchPhase(subSearchContext, docIdsToLoad, context.breaker(), memAccountingBufferSize);
         if (fetchProfiles != null) {
             fetchProfiles.add(fetchResult.profileResult());
         }
@@ -220,7 +225,12 @@ public InternalAggregation buildAggregation(long owningBucketOrdinal) throws IOE
         );
     }
 
-    private static FetchSearchResult runFetchPhase(SubSearchContext subSearchContext, int[] docIdsToLoad) {
+    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) {
@@ -229,7 +239,7 @@ public SearchExecutionContext getSearchExecutionContext() {
                 return searchExecutionContext;
             }
         };
-        fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null);
+        fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null, breaker, memAccountingBufferSize);
         return fetchSubSearchContext.fetchResult();
     }