Ref count search response bytes

docs/changelog/103763.yaml

@@ -0,0 +1,6 @@
+pr: 103763
+summary: Ref count search response bytes
+area: Search
+type: enhancement
+issues:
+ - 102657

modules/ingest-geoip/src/test/java/org/elasticsearch/ingest/geoip/DatabaseNodeServiceTests.java

@@ -318,7 +318,7 @@ private String mockSearches(String databaseName, int firstChunk, int lastChunk)
         Map<String, ActionFuture<SearchResponse>> requestMap = new HashMap<>();
         for (int i = firstChunk; i <= lastChunk; i++) {
             byte[] chunk = data.get(i - firstChunk);
-            SearchHit hit = new SearchHit(i);
+            SearchHit hit = SearchHit.unpooled(i);
             try (XContentBuilder builder = XContentBuilder.builder(XContentType.SMILE.xContent())) {
                 builder.map(Map.of("data", chunk));
                 builder.flush();
@@ -328,7 +328,7 @@ private String mockSearches(String databaseName, int firstChunk, int lastChunk)
                 throw new UncheckedIOException(ex);
             }
 
-            SearchHits hits = new SearchHits(new SearchHit[] { hit }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1f);
+            SearchHits hits = SearchHits.unpooled(new SearchHit[] { hit }, new TotalHits(1, TotalHits.Relation.EQUAL_TO), 1f);
             SearchResponse searchResponse = new SearchResponse(hits, null, null, false, null, null, 0, null, 1, 1, 0, 1L, null, null);
             toRelease.add(searchResponse::decRef);
             @SuppressWarnings("unchecked")

modules/percolator/src/main/java/org/elasticsearch/percolator/PercolatorHighlightSubFetchPhase.java

@@ -97,7 +97,7 @@ public void process(HitContext hit) throws IOException {
                             BytesReference document = percolateQuery.getDocuments().get(slot);
                             leafStoredFields.advanceTo(slot);
                             HitContext subContext = new HitContext(
-                                new SearchHit(slot, "unknown"),
+                                SearchHit.unpooled(slot, "unknown"),
                                 percolatorLeafReaderContext,
                                 slot,
                                 leafStoredFields.storedFields(),

modules/percolator/src/test/java/org/elasticsearch/percolator/PercolatorMatchedSlotSubFetchPhaseTests.java

@@ -56,7 +56,7 @@ public void testHitsExecute() throws Exception {
                 LeafReaderContext context = reader.leaves().get(0);
                 // A match:
                 {
-                    HitContext hit = new HitContext(new SearchHit(0), context, 0, Map.of(), Source.empty(null));
+                    HitContext hit = new HitContext(SearchHit.unpooled(0), context, 0, Map.of(), Source.empty(null));
                     PercolateQuery.QueryStore queryStore = ctx -> docId -> new TermQuery(new Term("field", "value"));
                     MemoryIndex memoryIndex = new MemoryIndex();
                     memoryIndex.addField("field", "value", new WhitespaceAnalyzer());
@@ -87,7 +87,7 @@ public void testHitsExecute() throws Exception {
 
                 // No match:
                 {
-                    HitContext hit = new HitContext(new SearchHit(0), context, 0, Map.of(), Source.empty(null));
+                    HitContext hit = new HitContext(SearchHit.unpooled(0), context, 0, Map.of(), Source.empty(null));
                     PercolateQuery.QueryStore queryStore = ctx -> docId -> new TermQuery(new Term("field", "value"));
                     MemoryIndex memoryIndex = new MemoryIndex();
                     memoryIndex.addField("field", "value1", new WhitespaceAnalyzer());
@@ -117,7 +117,7 @@ public void testHitsExecute() throws Exception {
 
                 // No query:
                 {
-                    HitContext hit = new HitContext(new SearchHit(0), context, 0, Map.of(), Source.empty(null));
+                    HitContext hit = new HitContext(SearchHit.unpooled(0), context, 0, Map.of(), Source.empty(null));
                     PercolateQuery.QueryStore queryStore = ctx -> docId -> null;
                     MemoryIndex memoryIndex = new MemoryIndex();
                     memoryIndex.addField("field", "value", new WhitespaceAnalyzer());

modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/RatedSearchHit.java

@@ -33,12 +33,12 @@ public class RatedSearchHit implements Writeable, ToXContentObject {
     private final OptionalInt rating;
 
     public RatedSearchHit(SearchHit searchHit, OptionalInt rating) {
-        this.searchHit = searchHit;
+        this.searchHit = searchHit.asUnpooled();
         this.rating = rating;
     }
 
     RatedSearchHit(StreamInput in) throws IOException {
-        this(SearchHit.readFrom(in), in.readBoolean() ? OptionalInt.of(in.readVInt()) : OptionalInt.empty());
+        this(SearchHit.readFrom(in, false), in.readBoolean() ? OptionalInt.of(in.readVInt()) : OptionalInt.empty());
     }
 
     @Override

modules/reindex/src/internalClusterTest/java/org/elasticsearch/index/reindex/BulkByScrollUsesAllScrollDocumentsAfterConflictsIntegTests.java

@@ -209,7 +209,7 @@ public void testDeleteByQuery() throws Exception {
                 .addSort(SORTING_FIELD, SortOrder.DESC),
             response -> {
                 // Modify a subset of the target documents concurrently
-                final List<SearchHit> originalDocs = Arrays.asList(response.getHits().getHits());
+                final List<SearchHit> originalDocs = Arrays.asList(response.getHits().asUnpooled().getHits());
                 docsModifiedConcurrently.addAll(randomSubsetOf(finalConflictingOps, originalDocs));
             }
         );

modules/reindex/src/test/java/org/elasticsearch/reindex/AsyncBulkByScrollActionTests.java

@@ -567,8 +567,8 @@ protected RequestWrapper<?> buildRequest(Hit doc) {
         action.start();
 
         // create a simulated response.
-        SearchHit hit = new SearchHit(0, "id").sourceRef(new BytesArray("{}"));
-        SearchHits hits = new SearchHits(
+        SearchHit hit = SearchHit.unpooled(0, "id").sourceRef(new BytesArray("{}"));
+        SearchHits hits = SearchHits.unpooled(
             IntStream.range(0, 100).mapToObj(i -> hit).toArray(SearchHit[]::new),
             new TotalHits(0, TotalHits.Relation.EQUAL_TO),
             0

modules/reindex/src/test/java/org/elasticsearch/reindex/ClientScrollableHitSourceTests.java

@@ -162,8 +162,8 @@ public void testScrollKeepAlive() {
 
     private SearchResponse createSearchResponse() {
         // create a simulated response.
-        SearchHit hit = new SearchHit(0, "id").sourceRef(new BytesArray("{}"));
-        SearchHits hits = new SearchHits(
+        SearchHit hit = SearchHit.unpooled(0, "id").sourceRef(new BytesArray("{}"));
+        SearchHits hits = SearchHits.unpooled(
             IntStream.range(0, randomIntBetween(0, 20)).mapToObj(i -> hit).toArray(SearchHit[]::new),
             new TotalHits(0, TotalHits.Relation.EQUAL_TO),
             0

server/src/internalClusterTest/java/org/elasticsearch/search/functionscore/RandomScoreFunctionIT.java

@@ -112,7 +112,7 @@ public void testConsistentHitsWithSameSeed() throws Exception {
                             CoreMatchers.equalTo(0)
                         );
                         final int hitCount = response.getHits().getHits().length;
-                        final SearchHit[] currentHits = response.getHits().getHits();
+                        final SearchHit[] currentHits = response.getHits().asUnpooled().getHits();
                         ArrayUtil.timSort(currentHits, (o1, o2) -> {
                             // for tie-breaking we have to resort here since if the score is
                             // identical we rely on collection order which might change.

server/src/main/java/org/elasticsearch/action/search/AbstractSearchAsyncAction.java

@@ -172,7 +172,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
         this.results = resultConsumer;
         // register the release of the query consumer to free up the circuit breaker memory
         // at the end of the search
-        addReleasable(resultConsumer::decRef);
+        addReleasable(resultConsumer);
         this.clusters = clusters;
     }
 

server/src/main/java/org/elasticsearch/action/search/ArraySearchPhaseResults.java

@@ -9,11 +9,11 @@
 package org.elasticsearch.action.search;
 
 import org.elasticsearch.common.util.concurrent.AtomicArray;
-import org.elasticsearch.core.AbstractRefCounted;
-import org.elasticsearch.core.RefCounted;
+import org.elasticsearch.core.Releasable;
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.transport.LeakTracker;
 
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Stream;
 
 /**
@@ -22,7 +22,13 @@
 class ArraySearchPhaseResults<Result extends SearchPhaseResult> extends SearchPhaseResults<Result> {
     final AtomicArray<Result> results;
 
-    private final RefCounted refCounted = LeakTracker.wrap(AbstractRefCounted.of(this::doClose));
+    private final AtomicBoolean closed = new AtomicBoolean(false);
+
+    private final Releasable releasable = LeakTracker.wrap(() -> {
+        for (Result result : getAtomicArray().asList()) {
+            result.decRef();
+        }
+    });
 
     ArraySearchPhaseResults(int size) {
         super(size);
@@ -41,12 +47,16 @@ void consumeResult(Result result, Runnable next) {
         next.run();
     }
 
-    protected void doClose() {
-        for (Result result : getAtomicArray().asList()) {
-            result.decRef();
+    @Override
+    public final void close() {
+        if (closed.compareAndSet(false, true)) {
+            releasable.close();
+            doClose();
         }
     }
 
+    protected void doClose() {}
+
     boolean hasResult(int shardIndex) {
         return results.get(shardIndex) != null;
     }
@@ -55,24 +65,4 @@ boolean hasResult(int shardIndex) {
     AtomicArray<Result> getAtomicArray() {
         return results;
     }
-
-    @Override
-    public void incRef() {
-        refCounted.incRef();
-    }
-
-    @Override
-    public boolean tryIncRef() {
-        return refCounted.tryIncRef();
-    }
-
-    @Override
-    public boolean decRef() {
-        return refCounted.decRef();
-    }
-
-    @Override
-    public boolean hasReferences() {
-        return refCounted.hasReferences();
-    }
 }

server/src/main/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhase.java

@@ -482,24 +482,7 @@ Stream<CanMatchShardResponse> getSuccessfulResults() {
         }
 
         @Override
-        public void incRef() {
-
-        }
-
-        @Override
-        public boolean tryIncRef() {
-            return false;
-        }
-
-        @Override
-        public boolean decRef() {
-            return false;
-        }
-
-        @Override
-        public boolean hasReferences() {
-            return false;
-        }
+        public void close() {}
     }
 
     private GroupShardsIterator<SearchShardIterator> getIterator(

server/src/main/java/org/elasticsearch/action/search/CountOnlyQueryPhaseResultConsumer.java

@@ -104,20 +104,5 @@ AtomicArray<SearchPhaseResult> getAtomicArray() {
     }
 
     @Override
-    public void incRef() {}
-
-    @Override
-    public boolean tryIncRef() {
-        return true;
-    }
-
-    @Override
-    public boolean decRef() {
-        return true;
-    }
-
-    @Override
-    public boolean hasReferences() {
-        return false;
-    }
+    public void close() {}
 }

server/src/main/java/org/elasticsearch/action/search/CountedCollector.java

@@ -25,7 +25,6 @@ final class CountedCollector<R extends SearchPhaseResult> {
 
     CountedCollector(SearchPhaseResults<R> resultConsumer, int expectedOps, Runnable onFinish, SearchPhaseContext context) {
         this.resultConsumer = resultConsumer;
-        resultConsumer.incRef();
         this.counter = new CountDown(expectedOps);
         this.onFinish = onFinish;
         this.context = context;
@@ -38,11 +37,7 @@ final class CountedCollector<R extends SearchPhaseResult> {
     void countDown() {
         assert counter.isCountedDown() == false : "more operations executed than specified";
         if (counter.countDown()) {
-            try {
-                onFinish.run();
-            } finally {
-                resultConsumer.decRef();
-            }
+            onFinish.run();
         }
     }
 

server/src/main/java/org/elasticsearch/action/search/DfsQueryPhase.java

@@ -66,7 +66,7 @@ final class DfsQueryPhase extends SearchPhase {
 
         // register the release of the query consumer to free up the circuit breaker memory
         // at the end of the search
-        context.addReleasable(queryResult::decRef);
+        context.addReleasable(queryResult);
     }
 
     @Override

server/src/main/java/org/elasticsearch/action/search/ExpandSearchPhase.java

@@ -96,6 +96,7 @@ public void run() {
                             hit.setInnerHits(Maps.newMapWithExpectedSize(innerHitBuilders.size()));
                         }
                         hit.getInnerHits().put(innerHitBuilder.getName(), innerHits);
+                        innerHits.mustIncRef();
                     }
                 }
                 onPhaseDone();

server/src/main/java/org/elasticsearch/action/search/FetchSearchPhase.java

@@ -38,11 +38,16 @@ final class FetchSearchPhase extends SearchPhase {
     private final AggregatedDfs aggregatedDfs;
 
     FetchSearchPhase(SearchPhaseResults<SearchPhaseResult> resultConsumer, AggregatedDfs aggregatedDfs, SearchPhaseContext context) {
-        this(resultConsumer, aggregatedDfs, context, (response, queryPhaseResults) -> {
-            response.mustIncRef();
-            context.addReleasable(response::decRef);
-            return new ExpandSearchPhase(context, response.hits, () -> new FetchLookupFieldsPhase(context, response, queryPhaseResults));
-        });
+        this(
+            resultConsumer,
+            aggregatedDfs,
+            context,
+            (response, queryPhaseResults) -> new ExpandSearchPhase(
+                context,
+                response.hits,
+                () -> new FetchLookupFieldsPhase(context, response, queryPhaseResults)
+            )
+        );
     }
 
     FetchSearchPhase(
@@ -61,7 +66,7 @@ final class FetchSearchPhase extends SearchPhase {
             );
         }
         this.fetchResults = new ArraySearchPhaseResults<>(resultConsumer.getNumShards());
-        context.addReleasable(fetchResults::decRef);
+        context.addReleasable(fetchResults);
         this.queryResults = resultConsumer.getAtomicArray();
         this.aggregatedDfs = aggregatedDfs;
         this.nextPhaseFactory = nextPhaseFactory;
@@ -225,10 +230,8 @@ private void moveToNextPhase(
         AtomicArray<? extends SearchPhaseResult> fetchResultsArr
     ) {
         var resp = SearchPhaseController.merge(context.getRequest().scroll() != null, reducedQueryPhase, fetchResultsArr);
-        try {
-            context.executeNextPhase(this, nextPhaseFactory.apply(resp, queryResults));
-        } finally {
-            resp.decRef();
-        }
+        context.addReleasable(resp::decRef);
+        fetchResults.close();
+        context.executeNextPhase(this, nextPhaseFactory.apply(resp, queryResults));
     }
 }

server/src/main/java/org/elasticsearch/action/search/QueryPhaseResultConsumer.java

@@ -105,11 +105,7 @@ public QueryPhaseResultConsumer(
 
     @Override
     protected void doClose() {
-        try {
-            super.doClose();
-        } finally {
-            pendingMerges.close();
-        }
+        pendingMerges.close();
     }
 
     @Override

server/src/main/java/org/elasticsearch/action/search/SearchDfsQueryThenFetchAsyncAction.java

@@ -68,7 +68,7 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
             clusters
         );
         this.queryPhaseResultConsumer = queryPhaseResultConsumer;
-        addReleasable(queryPhaseResultConsumer::decRef);
+        addReleasable(queryPhaseResultConsumer);
         this.progressListener = task.getProgressListener();
         // don't build the SearchShard list (can be expensive) if the SearchProgressListener won't use it
         if (progressListener != SearchProgressListener.NOOP) {
@@ -95,7 +95,6 @@ protected SearchPhase getNextPhase(final SearchPhaseResults<DfsSearchResult> res
         final List<DfsSearchResult> dfsSearchResults = results.getAtomicArray().asList();
         final AggregatedDfs aggregatedDfs = SearchPhaseController.aggregateDfs(dfsSearchResults);
         final List<DfsKnnResults> mergedKnnResults = SearchPhaseController.mergeKnnResults(getRequest(), dfsSearchResults);
-        queryPhaseResultConsumer.incRef();
         return new DfsQueryPhase(
             dfsSearchResults,
             aggregatedDfs,

server/src/main/java/org/elasticsearch/action/search/SearchPhaseController.java

@@ -364,11 +364,15 @@ public static SearchResponseSections merge(
         }
         ScoreDoc[] sortedDocs = reducedQueryPhase.sortedTopDocs.scoreDocs;
         var fetchResults = fetchResultsArray.asList();
-        SearchHits hits = getHits(reducedQueryPhase, ignoreFrom, fetchResultsArray);
-        if (reducedQueryPhase.suggest != null && fetchResults.isEmpty() == false) {
-            mergeSuggest(reducedQueryPhase, fetchResultsArray, hits, sortedDocs);
+        final SearchHits hits = getHits(reducedQueryPhase, ignoreFrom, fetchResultsArray);
+        try {
+            if (reducedQueryPhase.suggest != null && fetchResults.isEmpty() == false) {
+                mergeSuggest(reducedQueryPhase, fetchResultsArray, hits, sortedDocs);
+            }
+            return reducedQueryPhase.buildResponse(hits, fetchResults);
+        } finally {
+            hits.decRef();
         }
-        return reducedQueryPhase.buildResponse(hits, fetchResults);
     }
 
     private static void mergeSuggest(
@@ -462,6 +466,7 @@ private static SearchHits getHits(
                     searchHit.score(shardDoc.score);
                 }
                 hits.add(searchHit);
+                searchHit.incRef();
             }
         }
         return new SearchHits(