Use the existing refCounted field in SearchHit

Author: andreidan

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

@@ -10,30 +10,30 @@
 package org.elasticsearch.common;
 
 import org.elasticsearch.common.breaker.CircuitBreaker;
-import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.core.AbstractRefCounted;
 
 public final class MemoryAccountingBytesRefCounted extends AbstractRefCounted {
 
-    private BytesReference bytes;
-    private CircuitBreaker breaker;
+    private long bytes;
+    private final CircuitBreaker breaker;
 
-    private MemoryAccountingBytesRefCounted(BytesReference bytes, CircuitBreaker breaker) {
+    private MemoryAccountingBytesRefCounted(long bytes, CircuitBreaker breaker) {
         this.bytes = bytes;
         this.breaker = breaker;
     }
 
-    public static MemoryAccountingBytesRefCounted createAndAccountForBytes(
-        BytesReference bytes,
-        CircuitBreaker breaker,
-        String memAccountingLabel
-    ) {
-        breaker.addEstimateBytesAndMaybeBreak(bytes.length(), memAccountingLabel);
-        return new MemoryAccountingBytesRefCounted(bytes, breaker);
+    public static MemoryAccountingBytesRefCounted create(CircuitBreaker breaker) {
+        return new MemoryAccountingBytesRefCounted(0L, breaker);
+    }
+
+    public void account(long bytes, String label) {
+        breaker.addEstimateBytesAndMaybeBreak(bytes, label);
+        this.bytes += bytes;
     }
 
     @Override
     protected void closeInternal() {
-        breaker.addWithoutBreaking(-bytes.length());
+        System.err.println("REMOVING " + bytes);
+        breaker.addWithoutBreaking(-bytes);
     }
 }

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

@@ -83,8 +83,6 @@ public final class SearchHit implements Writeable, ToXContentObject, RefCounted
     private long primaryTerm;
 
     private BytesReference source;
-    @Nullable
-    private RefCounted unfilteredSource;
 
     private final Map<String, DocumentField> documentFields;
     private final Map<String, DocumentField> metaFields;
@@ -125,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,
@@ -449,17 +447,6 @@ 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 unfilteredSource(RefCounted source) {
-        this.unfilteredSource = 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.

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

@@ -40,6 +40,7 @@
 import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.search.rank.RankDocShardInfo;
 import org.elasticsearch.tasks.TaskCancelledException;
+import org.elasticsearch.transport.LeakTracker;
 import org.elasticsearch.xcontent.XContentType;
 
 import java.io.IOException;
@@ -302,23 +303,20 @@ private static HitContext prepareNonNestedHitContext(
 
         String id = idLoader.getId(subDocId);
         if (id == null) {
-            SearchHit hit = new SearchHit(docId);
+            MemoryAccountingBytesRefCounted memAccountingRefCounted = MemoryAccountingBytesRefCounted.create(circuitBreaker);
+            SearchHit hit = new SearchHit(docId, null, null, LeakTracker.wrap(memAccountingRefCounted));
             // TODO: can we use real pooled buffers here as well?
-            Source source = Source.lazy(lazyStoredSourceLoader(profiler, subReaderContext, subDocId, hit, circuitBreaker));
+            Source source = Source.lazy(lazyStoredSourceLoader(profiler, subReaderContext, subDocId, memAccountingRefCounted));
             return new HitContext(hit, subReaderContext, subDocId, Map.of(), source, rankDoc);
         } else {
-            SearchHit hit = new SearchHit(docId, id);
+            MemoryAccountingBytesRefCounted memAccountingRefCounted = MemoryAccountingBytesRefCounted.create(circuitBreaker);
+            SearchHit hit = new SearchHit(docId, id, null, LeakTracker.wrap(memAccountingRefCounted));
             Source source;
             if (requiresSource) {
                 Timer timer = profiler.startLoadingSource();
                 try {
                     source = sourceLoader.source(leafStoredFieldLoader, subDocId);
-                    MemoryAccountingBytesRefCounted sourceRef = MemoryAccountingBytesRefCounted.createAndAccountForBytes(
-                        source.internalSourceRef(),
-                        circuitBreaker,
-                        "fetch phase source loader"
-                    );
-                    hit.unfilteredSource(sourceRef);
+                    memAccountingRefCounted.account(source.internalSourceRef().length(), "fetch phase source loader");
                 } catch (CircuitBreakingException e) {
                     hit.decRef();
                     throw e;
@@ -328,7 +326,7 @@ private static HitContext prepareNonNestedHitContext(
                     }
                 }
             } else {
-                source = Source.lazy(lazyStoredSourceLoader(profiler, subReaderContext, subDocId, hit, circuitBreaker));
+                source = Source.lazy(lazyStoredSourceLoader(profiler, subReaderContext, subDocId, memAccountingRefCounted));
             }
             return new HitContext(hit, subReaderContext, subDocId, leafStoredFieldLoader.storedFields(), source, rankDoc);
         }
@@ -338,28 +336,15 @@ private static Supplier<Source> lazyStoredSourceLoader(
         Profiler profiler,
         LeafReaderContext ctx,
         int doc,
-        SearchHit hit,
-        CircuitBreaker circuitBreaker
+        MemoryAccountingBytesRefCounted memAccountingRefCounted
     ) {
         return () -> {
             StoredFieldLoader rootLoader = profiler.storedFields(StoredFieldLoader.create(true, Collections.emptySet()));
             try {
                 LeafStoredFieldLoader leafRootLoader = rootLoader.getLoader(ctx, null);
                 leafRootLoader.advanceTo(doc);
                 BytesReference source = leafRootLoader.source();
-                MemoryAccountingBytesRefCounted memAccountingSourceRef = MemoryAccountingBytesRefCounted.createAndAccountForBytes(
-                    source,
-                    circuitBreaker,
-                    "fetch phase source loader"
-                );
-                // Saving the entire source we loaded in the hit, so that we can release it entirely when the hit is released
-                // This is important for the circuit breaker accounting - note that this lazy loader can be triggered in the case of
-                // inner hits even though the top hit source is not requested (see {@link FetchPhase#prepareNestedHitContext} when
-                // the `nestedSource` is created), so we need to save the entire source on the hit - we account
-                // for the top level source via the {@link SearchHit#unfilteredSourceRef(BytesReference)} method because the
-                // {@link SearchHit#source()} method can be null when the top level source is not requested.
-                // NB all of the above also applies for source filtering and field collapsing
-                hit.unfilteredSource(memAccountingSourceRef);
+                memAccountingRefCounted.account(source.length(), "fetch phase source loader");
                 return Source.fromBytes(source);
             } catch (IOException e) {
                 throw new UncheckedIOException(e);
@@ -404,7 +389,6 @@ private static HitContext prepareNestedHitContext(
             if (requiresSource) {
                 BytesReference source = leafRootLoader.source();
                 if (source != null) {
-                    NOOP_CIRCUIT_BREAKER.addEstimateBytesAndMaybeBreak(source.length(), "fetch phase nested hit source loader");
                     rootSource = Source.fromBytes(source);
                 }
             }

x-pack/plugin/enrich/src/main/java/org/elasticsearch/xpack/enrich/action/EnrichShardMultiSearchAction.java

@@ -263,12 +263,11 @@ protected MultiSearchResponse shardOperation(Request request, ShardId shardId) t
                             }
                             return context.getFieldType(field);
                         });
-                        final SearchHit hit = new SearchHit(scoreDoc.doc, visitor.id());
+                        MemoryAccountingBytesRefCounted memAccountingRefCounted = MemoryAccountingBytesRefCounted.create(breaker);
+                        final SearchHit hit = new SearchHit(scoreDoc.doc, visitor.id(), null, memAccountingRefCounted);
                         try {
                             BytesReference sourceBytesRef = visitor.source();
-                            MemoryAccountingBytesRefCounted memAccountingSourceRef = MemoryAccountingBytesRefCounted
-                                .createAndAccountForBytes(sourceBytesRef, breaker, "enrich source");
-                            hit.unfilteredSource(memAccountingSourceRef);
+                            memAccountingRefCounted.account(sourceBytesRef.length(), "enrich source");
                             hit.sourceRef(filterSource(fetchSourceContext, sourceBytesRef));
                             hits[j] = hit;
                         } catch (CircuitBreakingException e) {