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/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/index/IndexService.java

@@ -834,6 +834,10 @@ public ScriptService getScriptService() {
         return scriptService;
     }
 
+    public CircuitBreakerService breakerService() {
+        return circuitBreakerService;
+    }
+
     List<IndexingOperationListener> getIndexOperationListeners() { // pkg private for testing
         return indexingOperationListeners;
     }

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

@@ -23,6 +23,7 @@
 import org.apache.lucene.util.NumericUtils;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.cluster.routing.IndexRouting;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.lucene.search.Queries;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
@@ -103,6 +104,7 @@ final class DefaultSearchContext extends SearchContext {
     private final IndexShard indexShard;
     private final IndexService indexService;
     private final ContextIndexSearcher searcher;
+    private final long memoryAccountingBufferSize;
     private DfsSearchResult dfsResult;
     private QuerySearchResult queryResult;
     private RankFeatureResult rankFeatureResult;
@@ -168,7 +170,8 @@ final class DefaultSearchContext extends SearchContext {
         Executor executor,
         SearchService.ResultsType resultsType,
         boolean enableQueryPhaseParallelCollection,
-        int minimumDocsPerSlice
+        int minimumDocsPerSlice,
+        long memoryAccountingBufferSize
     ) throws IOException {
         this.readerContext = readerContext;
         this.request = request;
@@ -179,6 +182,7 @@ final class DefaultSearchContext extends SearchContext {
             this.shardTarget = shardTarget;
             this.indexService = readerContext.indexService();
             this.indexShard = readerContext.indexShard();
+            this.memoryAccountingBufferSize = memoryAccountingBufferSize;
 
             Engine.Searcher engineSearcher = readerContext.acquireSearcher("search");
             int maximumNumberOfSlices = determineMaximumNumberOfSlices(
@@ -902,6 +906,16 @@ public Profilers getProfilers() {
         return profilers;
     }
 
+    @Override
+    public CircuitBreaker circuitBreaker() {
+        return indexService.breakerService().getBreaker(CircuitBreaker.REQUEST);
+    }
+
+    @Override
+    public long memAccountingBufferSize() {
+        return memoryAccountingBufferSize;
+    }
+
     public void setProfilers(Profilers profilers) {
         this.profilers = profilers;
     }

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);
     }
@@ -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,20 @@ 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(1, ByteSizeUnit.MB),
+        ByteSizeValue.of(1, ByteSizeUnit.MB),
+        ByteSizeValue.ofBytes(Long.MAX_VALUE),
+        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 +305,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 +325,8 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
 
     private volatile boolean enableRewriteAggsToFilterByFilter;
 
+    private volatile long memoryAccountingBufferSize;
+
     private final Cancellable keepAliveReaper;
 
     private final AtomicLong idGenerator = new AtomicLong();
@@ -341,11 +358,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 +404,10 @@ 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).getBytes();
+        clusterService.getClusterSettings()
+            .addSettingsUpdateConsumer(MEMORY_ACCOUNTING_BUFFER_SIZE, newValue -> this.memoryAccountingBufferSize = newValue.getBytes());
     }
 
     private void setEnableSearchWorkerThreads(boolean enableSearchWorkerThreads) {
@@ -1195,7 +1213,8 @@ private DefaultSearchContext createSearchContext(
                 searchExecutor,
                 resultsType,
                 enableQueryPhaseParallelCollection,
-                minimumDocsPerSlice
+                minimumDocsPerSlice,
+                memoryAccountingBufferSize
             );
             // we clone the query shard context here just for rewriting otherwise we
             // might end up with incorrect state since we are using now() or script services

server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java

@@ -13,6 +13,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.TotalHits;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.index.fieldvisitor.LeafStoredFieldLoader;
 import org.elasticsearch.index.fieldvisitor.StoredFieldLoader;
 import org.elasticsearch.index.mapper.IdLoader;
@@ -129,7 +130,7 @@ private SearchHits buildSearchHits(SearchContext context, int[] docIdsToLoad, Pr
         StoredFieldLoader storedFieldLoader = profiler.storedFields(StoredFieldLoader.fromSpec(storedFieldsSpec));
         IdLoader idLoader = context.newIdLoader();
         boolean requiresSource = storedFieldsSpec.requiresSource();
-
+        final int[] accumulatedBytesInLeaf = new int[1];
         NestedDocuments nestedDocuments = context.getSearchExecutionContext().getNestedDocuments();
 
         FetchPhaseDocsIterator docsIterator = new FetchPhaseDocsIterator() {
@@ -162,6 +163,11 @@ protected SearchHit nextDoc(int doc) throws IOException {
                 if (context.isCancelled()) {
                     throw new TaskCancelledException("cancelled");
                 }
+                if (context.checkRealMemoryCB(accumulatedBytesInLeaf[0], "fetch source")) {
+                    // if we checked the real memory breaker, we restart our local accounting
+                    accumulatedBytesInLeaf[0] = 0;
+                }
+
                 HitContext hit = prepareHitContext(
                     context,
                     requiresSource,
@@ -181,6 +187,11 @@ protected SearchHit nextDoc(int doc) throws IOException {
                     for (FetchSubPhaseProcessor processor : processors) {
                         processor.process(hit);
                     }
+
+                    BytesReference sourceRef = hit.hit().getSourceRef();
+                    if (sourceRef != null) {
+                        accumulatedBytesInLeaf[0] += sourceRef.length();
+                    }
                     success = true;
                     return hit.hit();
                 } finally {
@@ -291,7 +302,7 @@ private static HitContext prepareNonNestedHitContext(
             Source source = Source.lazy(lazyStoredSourceLoader(profiler, subReaderContext, subDocId));
             return new HitContext(hit, subReaderContext, subDocId, Map.of(), source, rankDoc);
         } else {
-            SearchHit hit = new SearchHit(docId, id);
+            SearchHit hit = new SearchHit(docId, id, null);
             Source source;
             if (requiresSource) {
                 Timer timer = profiler.startLoadingSource();
@@ -369,8 +380,8 @@ private static HitContext prepareNestedHitContext(
         assert nestedIdentity != null;
         Source nestedSource = nestedIdentity.extractSource(rootSource);
 
-        SearchHit hit = new SearchHit(topDocId, rootId, nestedIdentity);
-        return new HitContext(hit, subReaderContext, nestedInfo.doc(), childFieldLoader.storedFields(), nestedSource, rankDoc);
+        SearchHit nestedHit = new SearchHit(topDocId, rootId, nestedIdentity);
+        return new HitContext(nestedHit, subReaderContext, nestedInfo.doc(), childFieldLoader.storedFields(), nestedSource, rankDoc);
     }
 
     interface Profiler {

server/src/main/java/org/elasticsearch/search/fetch/FetchPhaseDocsIterator.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
+import org.elasticsearch.common.breaker.CircuitBreakingException;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
 import org.elasticsearch.search.SearchShardTarget;
@@ -100,6 +101,9 @@ public final SearchHit[] iterate(
             }
         } catch (SearchTimeoutException e) {
             throw e;
+        } catch (CircuitBreakingException e) {
+            purgeSearchHits(searchHits);
+            throw e;
         } catch (Exception e) {
             purgeSearchHits(searchHits);
             throw new FetchPhaseExecutionException(shardTarget, "Error running fetch phase for doc [" + currentDoc + "]", e);

server/src/main/java/org/elasticsearch/search/internal/FilteredSearchContext.java

@@ -13,6 +13,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.action.search.SearchType;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
 import org.elasticsearch.index.mapper.IdLoader;
@@ -460,4 +461,14 @@ public SourceLoader newSourceLoader() {
     public IdLoader newIdLoader() {
         return in.newIdLoader();
     }
+
+    @Override
+    public CircuitBreaker circuitBreaker() {
+        return in.circuitBreaker();
+    }
+
+    @Override
+    public long memAccountingBufferSize() {
+        return in.memAccountingBufferSize();
+    }
 }

server/src/main/java/org/elasticsearch/search/internal/SearchContext.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.action.search.SearchType;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.core.Assertions;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Releasable;
@@ -369,6 +370,31 @@ public Query rewrittenQuery() {
      */
     public abstract Profilers getProfilers();
 
+    /**
+     * The circuit breaker used to account for the search operation.
+     */
+    public abstract CircuitBreaker circuitBreaker();
+
+    /**
+     * Return the amount of memory to buffer locally before accounting for it in the breaker.
+     */
+    public abstract long memAccountingBufferSize();
+
+    /**
+     * Checks if the accumulated bytes are greater than the buffer size and if so, checks the available memory in the parent breaker
+     * (the real memory breaker).
+     * @param locallyAccumulatedBytes the number of bytes accumulated locally
+     * @param label the label to use in the breaker
+     * @return true if the real memory breaker is called and false otherwise
+     */
+    public final boolean checkRealMemoryCB(int locallyAccumulatedBytes, String label) {
+        if (locallyAccumulatedBytes >= memAccountingBufferSize()) {
+            circuitBreaker().addEstimateBytesAndMaybeBreak(0, label);
+            return true;
+        }
+        return false;
+    }
+
     /**
      * Adds a releasable that will be freed when this context is closed.
      */

server/src/main/java/org/elasticsearch/search/rank/RankSearchContext.java

@@ -13,6 +13,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHits;
 import org.elasticsearch.action.search.SearchType;
+import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
 import org.elasticsearch.index.mapper.IdLoader;
@@ -203,6 +204,16 @@ public Profilers getProfilers() {
         return null;
     }
 
+    @Override
+    public CircuitBreaker circuitBreaker() {
+        return parent.circuitBreaker();
+    }
+
+    @Override
+    public long memAccountingBufferSize() {
+        return parent.memAccountingBufferSize();
+    }
+
     @Override
     public long getRelativeTimeInMillis() {
         return parent.getRelativeTimeInMillis();