Periodically check the available memory when fetching search hits source

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

@@ -13,6 +13,8 @@
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.breaker.CircuitBreaker;
+import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.compress.CompressorFactory;
@@ -83,6 +85,8 @@ public final class SearchHit implements Writeable, ToXContentObject, RefCounted
     private long primaryTerm;
 
     private BytesReference source;
+    @Nullable
+    private BytesReference unfilteredSourceRef;
 
     private final Map<String, DocumentField> documentFields;
     private final Map<String, DocumentField> metaFields;
@@ -110,21 +114,27 @@ public final class SearchHit implements Writeable, ToXContentObject, RefCounted
     private Map<String, SearchHits> innerHits;
 
     private final RefCounted refCounted;
+    private final CircuitBreaker circuitBreaker;
 
-    // used only in tests
-    public SearchHit(int docId) {
-        this(docId, null);
+    public SearchHit(int docId, CircuitBreaker circuitBreaker) {
+        this(docId, null, circuitBreaker);
     }
 
-    public SearchHit(int docId, String id) {
-        this(docId, id, null);
+    public SearchHit(int docId, String id, CircuitBreaker circuitBreaker) {
+        this(docId, id, null, circuitBreaker);
     }
 
-    public SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity) {
-        this(nestedTopDocId, id, nestedIdentity, null);
+    public SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity, CircuitBreaker circuitBreaker) {
+        this(nestedTopDocId, id, nestedIdentity, null, circuitBreaker);
     }
 
-    private SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity, @Nullable RefCounted refCounted) {
+    private SearchHit(
+        int nestedTopDocId,
+        String id,
+        NestedIdentity nestedIdentity,
+        @Nullable RefCounted refCounted,
+        @Nullable CircuitBreaker circuitBreaker
+    ) {
         this(
             nestedTopDocId,
             DEFAULT_SCORE,
@@ -145,7 +155,8 @@ private SearchHit(int nestedTopDocId, String id, NestedIdentity nestedIdentity,
             null,
             new HashMap<>(),
             new HashMap<>(),
-            refCounted
+            refCounted,
+            circuitBreaker
         );
     }
 
@@ -169,7 +180,8 @@ public SearchHit(
         Map<String, SearchHits> innerHits,
         Map<String, DocumentField> documentFields,
         Map<String, DocumentField> metaFields,
-        @Nullable RefCounted refCounted
+        @Nullable RefCounted refCounted,
+        @Nullable CircuitBreaker circuitBreaker
     ) {
         this.docId = docId;
         this.score = score;
@@ -191,6 +203,7 @@ public SearchHit(
         this.documentFields = documentFields;
         this.metaFields = metaFields;
         this.refCounted = refCounted == null ? LeakTracker.wrap(new SimpleRefCounted()) : refCounted;
+        this.circuitBreaker = circuitBreaker;
     }
 
     public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOException {
@@ -280,7 +293,8 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept
             innerHits,
             documentFields,
             metaFields,
-            isPooled ? null : ALWAYS_REFERENCED
+            isPooled ? null : ALWAYS_REFERENCED,
+            null
         );
     }
 
@@ -293,7 +307,8 @@ public static SearchHit unpooled(int docId, String id) {
     }
 
     public static SearchHit unpooled(int nestedTopDocId, String id, NestedIdentity nestedIdentity) {
-        return new SearchHit(nestedTopDocId, id, nestedIdentity, ALWAYS_REFERENCED);
+        // always referenced search hits do NOT call #deallocate
+        return new SearchHit(nestedTopDocId, id, nestedIdentity, ALWAYS_REFERENCED, null);
     }
 
     private static final Text SINGLE_MAPPING_TYPE = new Text(MapperService.SINGLE_MAPPING_NAME);
@@ -447,6 +462,17 @@ public SearchHit sourceRef(BytesReference source) {
         return this;
     }
 
+    /**
+     * We track the unfiltered, entire, source so we can release the entire size from the
+     * circuit breakers when the hit is released.
+     * The regular source might be a subset of the unfiltered source due to either
+     * source filtering, field collapsing or inner hits.
+     */
+    public SearchHit unfilteredSourceRef(BytesReference source) {
+        this.unfilteredSourceRef = source;
+        return this;
+    }
+
     /**
      * Is the source available or not. A source with no fields will return true. This will return false if {@code fields} doesn't contain
      * {@code _source} or if source is disabled in the mapping.
@@ -724,6 +750,12 @@ private void deallocate() {
             r.decRef();
         }
         SearchHit.this.source = null;
+
+        if (unfilteredSourceRef != null && circuitBreaker != null) {
+            System.out.println(" removing source loaded by inner hit " + unfilteredSourceRef.length());
+            circuitBreaker.addWithoutBreaking(-unfilteredSourceRef.length());
+        }
+        this.unfilteredSourceRef = null;
     }
 
     @Override
@@ -758,7 +790,8 @@ public SearchHit asUnpooled() {
                 : innerHits.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> e.getValue().asUnpooled())),
             documentFields,
             metaFields,
-            ALWAYS_REFERENCED
+            ALWAYS_REFERENCED,
+            null
         );
     }
 

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

@@ -290,6 +290,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;
 
@@ -340,11 +341,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;
 
@@ -788,7 +786,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);
                 RankFeatureShardPhase.processFetch(searchContext);
                 var rankFeatureResult = searchContext.rankFeatureResult();
                 rankFeatureResult.incRef();
@@ -806,7 +804,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);
             if (reader.singleSession()) {
                 freeReaderContext(reader.id());
             }
@@ -972,7 +970,7 @@ public void executeFetchPhase(ShardFetchRequest request, SearchShardTask task, A
                             System.nanoTime()
                         )
                     ) {
-                        fetchPhase.execute(searchContext, request.docIds(), request.getRankDocks());
+                        fetchPhase.execute(searchContext, request.docIds(), request.getRankDocks(), circuitBreaker);
                         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;
@@ -196,7 +197,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());
         if (fetchProfiles != null) {
             fetchProfiles.add(fetchResult.profileResult());
         }
@@ -220,7 +221,7 @@ 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) {
         // 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 +230,7 @@ public SearchExecutionContext getSearchExecutionContext() {
                 return searchExecutionContext;
             }
         };
-        fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null);
+        fetchSubSearchContext.fetchPhase().execute(fetchSubSearchContext, docIdsToLoad, null, breaker);
         return fetchSubSearchContext.fetchResult();
     }