[8.19] Remove early phase failure in batched (#136889) (#138209)

#121885 attempted to shortcut a phase failure caused by a reduction
failure on the data node by failing the query phase in the batched
query action response listener. Before batching the query phase, we
did not fail the phase immediately upon a reduction failure. We held
on to the failure and continued querying all shards, only failing
during final reduction at the beginning of the fetch phase.

I can't think of anything inherently wrong with this approach,
besides the fact that the phase cannot be failed multiple times
(#134151). However certain cleanup aspects of the code (specifically
releasing reader contexts and query search results, see: #130821,
#122707) rely on the assumption that all shards are queried before
failing the phase.

This commit reworks batched requests to fail in the same way: only
after all shards are queried. To do this, we must include results in
transport response even when a reduction failure occurred.

Author: benchaplin

docs/changelog/136889.yaml

@@ -0,0 +1,6 @@
+pr: 136889
+summary: Remove early phase failure in batched
+area: Search
+type: bug
+issues:
+ - 134151

muted-tests.yml

@@ -361,9 +361,6 @@ tests:
   - class: org.elasticsearch.repositories.blobstore.testkit.analyze.GCSRepositoryAnalysisRestIT
     method: testRepositoryAnalysis
     issue: https://github.com/elastic/elasticsearch/issues/125668
-  - class: org.elasticsearch.search.basic.SearchWithRandomDisconnectsIT
-    method: testSearchWithRandomDisconnects
-    issue: https://github.com/elastic/elasticsearch/issues/122707
   - class: org.elasticsearch.docker.test.DockerYmlTestSuiteIT
     method: test {p0=/11_nodes/Test cat nodes output}
     issue: https://github.com/elastic/elasticsearch/issues/125906
@@ -409,4 +406,3 @@ tests:
   - class: org.elasticsearch.xpack.inference.InferenceRestIT
     method: test {p0=inference/70_text_similarity_rank_retriever/Text similarity reranker with min_score zero includes all docs}
     issue: https://github.com/elastic/elasticsearch/issues/137732
-

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

@@ -80,6 +80,9 @@ public class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<S
     private static final Logger logger = LogManager.getLogger(SearchQueryThenFetchAsyncAction.class);
 
     private static final TransportVersion BATCHED_QUERY_PHASE_VERSION = TransportVersion.fromName("batched_query_phase_version");
+    private static final TransportVersion BATCHED_RESPONSE_MIGHT_INCLUDE_REDUCTION_FAILURE = TransportVersion.fromName(
+        "batched_response_might_include_reduction_failure"
+    );
 
     private final SearchProgressListener progressListener;
 
@@ -218,18 +221,27 @@ public static final class NodeQueryResponse extends TransportResponse {
         private final RefCounted refCounted = LeakTracker.wrap(new SimpleRefCounted());
 
         private final Object[] results;
+        private final Exception reductionFailure;
         private final SearchPhaseController.TopDocsStats topDocsStats;
         private final QueryPhaseResultConsumer.MergeResult mergeResult;
 
         NodeQueryResponse(StreamInput in) throws IOException {
             this.results = in.readArray(i -> i.readBoolean() ? new QuerySearchResult(i) : i.readException(), Object[]::new);
-            this.mergeResult = QueryPhaseResultConsumer.MergeResult.readFrom(in);
-            this.topDocsStats = SearchPhaseController.TopDocsStats.readFrom(in);
+            if (in.getTransportVersion().supports(BATCHED_RESPONSE_MIGHT_INCLUDE_REDUCTION_FAILURE) && in.readBoolean()) {
+                this.reductionFailure = in.readException();
+                this.mergeResult = null;
+                this.topDocsStats = null;
+            } else {
+                this.reductionFailure = null;
+                this.mergeResult = QueryPhaseResultConsumer.MergeResult.readFrom(in);
+                this.topDocsStats = SearchPhaseController.TopDocsStats.readFrom(in);
+            }
         }
 
         NodeQueryResponse(
-            QueryPhaseResultConsumer.MergeResult mergeResult,
             Object[] results,
+            Exception reductionFailure,
+            QueryPhaseResultConsumer.MergeResult mergeResult,
             SearchPhaseController.TopDocsStats topDocsStats
         ) {
             this.results = results;
@@ -238,6 +250,7 @@ public static final class NodeQueryResponse extends TransportResponse {
                     r.incRef();
                 }
             }
+            this.reductionFailure = reductionFailure;
             this.mergeResult = mergeResult;
             this.topDocsStats = topDocsStats;
             assert Arrays.stream(results).noneMatch(Objects::isNull) : Arrays.toString(results);
@@ -248,6 +261,10 @@ public Object[] getResults() {
             return results;
         }
 
+        Exception getReductionFailure() {
+            return reductionFailure;
+        }
+
         @Override
         public void writeTo(StreamOutput out) throws IOException {
             out.writeArray((o, v) -> {
@@ -260,8 +277,19 @@ public void writeTo(StreamOutput out) throws IOException {
                     ((QuerySearchResult) v).writeTo(o);
                 }
             }, results);
-            mergeResult.writeTo(out);
-            topDocsStats.writeTo(out);
+            if (out.getTransportVersion().supports(BATCHED_RESPONSE_MIGHT_INCLUDE_REDUCTION_FAILURE)) {
+                boolean hasReductionFailure = reductionFailure != null;
+                out.writeBoolean(hasReductionFailure);
+                if (hasReductionFailure) {
+                    out.writeException(reductionFailure);
+                } else {
+                    mergeResult.writeTo(out);
+                    topDocsStats.writeTo(out);
+                }
+            } else {
+                mergeResult.writeTo(out);
+                topDocsStats.writeTo(out);
+            }
         }
 
         @Override
@@ -495,7 +523,12 @@ public Executor executor() {
                     @Override
                     public void handleResponse(NodeQueryResponse response) {
                         if (results instanceof QueryPhaseResultConsumer queryPhaseResultConsumer) {
-                            queryPhaseResultConsumer.addBatchedPartialResult(response.topDocsStats, response.mergeResult);
+                            Exception reductionFailure = response.getReductionFailure();
+                            if (reductionFailure != null) {
+                                queryPhaseResultConsumer.failure.compareAndSet(null, reductionFailure);
+                            } else {
+                                queryPhaseResultConsumer.addBatchedPartialResult(response.topDocsStats, response.mergeResult);
+                            }
                         }
                         for (int i = 0; i < response.results.length; i++) {
                             var s = request.shards.get(i);
@@ -515,6 +548,21 @@ public void handleResponse(NodeQueryResponse response) {
 
                     @Override
                     public void handleException(TransportException e) {
+                        if (connection.getTransportVersion().supports(BATCHED_RESPONSE_MIGHT_INCLUDE_REDUCTION_FAILURE) == false) {
+                            bwcHandleException(e);
+                            return;
+                        }
+                        Exception cause = (Exception) ExceptionsHelper.unwrapCause(e);
+                        logger.debug("handling node search exception coming from [" + nodeId + "]", cause);
+                        onNodeQueryFailure(e, request, routing);
+                    }
+
+                    /**
+                     * This code is strictly for _snapshot_ backwards compatibility. The feature flag
+                     * {@link SearchService#BATCHED_QUERY_PHASE_FEATURE_FLAG} was not turned on when the transport version
+                     * {@link SearchQueryThenFetchAsyncAction#BATCHED_RESPONSE_MIGHT_INCLUDE_REDUCTION_FAILURE} was introduced.
+                     */
+                    private void bwcHandleException(TransportException e) {
                         Exception cause = (Exception) ExceptionsHelper.unwrapCause(e);
                         logger.debug("handling node search exception coming from [" + nodeId + "]", cause);
                         if (e instanceof SendRequestTransportException || cause instanceof TaskCancelledException) {
@@ -786,11 +834,93 @@ void onShardDone() {
             if (countDown.countDown() == false) {
                 return;
             }
+            if (channel.getVersion().supports(BATCHED_RESPONSE_MIGHT_INCLUDE_REDUCTION_FAILURE) == false) {
+                bwcRespond();
+                return;
+            }
+            var channelListener = new ChannelActionListener<>(channel);
+            NodeQueryResponse nodeQueryResponse;
+            try (queryPhaseResultConsumer) {
+                Exception reductionFailure = queryPhaseResultConsumer.failure.get();
+                if (reductionFailure == null) {
+                    nodeQueryResponse = getSuccessfulResponse();
+                } else {
+                    nodeQueryResponse = getReductionFailureResponse(reductionFailure);
+                }
+            } catch (IOException e) {
+                releaseAllResultsContexts();
+                channelListener.onFailure(e);
+                return;
+            }
+            ActionListener.respondAndRelease(channelListener, nodeQueryResponse);
+        }
+
+        private NodeQueryResponse getSuccessfulResponse() throws IOException {
+            final QueryPhaseResultConsumer.MergeResult mergeResult;
+            try {
+                mergeResult = Objects.requireNonNullElse(
+                    queryPhaseResultConsumer.consumePartialMergeResultDataNode(),
+                    EMPTY_PARTIAL_MERGE_RESULT
+                );
+            } catch (Exception e) {
+                return getReductionFailureResponse(e);
+            }
+            // translate shard indices to those on the coordinator so that it can interpret the merge result without adjustments,
+            // also collect the set of indices that may be part of a subsequent fetch operation here so that we can release all other
+            // indices without a roundtrip to the coordinating node
+            final BitSet relevantShardIndices = new BitSet(searchRequest.shards.size());
+            if (mergeResult.reducedTopDocs() != null) {
+                for (ScoreDoc scoreDoc : mergeResult.reducedTopDocs().scoreDocs) {
+                    final int localIndex = scoreDoc.shardIndex;
+                    scoreDoc.shardIndex = searchRequest.shards.get(localIndex).shardIndex;
+                    relevantShardIndices.set(localIndex);
+                }
+            }
+            final Object[] results = new Object[queryPhaseResultConsumer.getNumShards()];
+            for (int i = 0; i < results.length; i++) {
+                var result = queryPhaseResultConsumer.results.get(i);
+                if (result == null) {
+                    results[i] = failures.get(i);
+                } else {
+                    // free context id and remove it from the result right away in case we don't need it anymore
+                    maybeFreeContext(result, relevantShardIndices, namedWriteableRegistry);
+                    results[i] = result;
+                }
+                assert results[i] != null;
+            }
+            return new NodeQueryResponse(results, null, mergeResult, queryPhaseResultConsumer.topDocsStats);
+        }
+
+        private NodeQueryResponse getReductionFailureResponse(Exception reductionFailure) throws IOException {
+            try {
+                final Object[] results = new Object[queryPhaseResultConsumer.getNumShards()];
+                for (int i = 0; i < results.length; i++) {
+                    var result = queryPhaseResultConsumer.results.get(i);
+                    if (result == null) {
+                        results[i] = failures.get(i);
+                    } else {
+                        results[i] = result;
+                    }
+                    assert results[i] != null;
+                }
+                return new NodeQueryResponse(results, reductionFailure, null, null);
+            } finally {
+                releaseAllResultsContexts();
+            }
+        }
+
+        /**
+         * This code is strictly for _snapshot_ backwards compatibility. The feature flag
+         * {@link SearchService#BATCHED_QUERY_PHASE_FEATURE_FLAG} was not turned on when the transport version
+         * {@link SearchQueryThenFetchAsyncAction#BATCHED_RESPONSE_MIGHT_INCLUDE_REDUCTION_FAILURE} was introduced.
+         */
+        void bwcRespond() {
             var channelListener = new ChannelActionListener<>(channel);
             try (queryPhaseResultConsumer) {
                 var failure = queryPhaseResultConsumer.failure.get();
                 if (failure != null) {
-                    handleMergeFailure(failure, channelListener, namedWriteableRegistry);
+                    releaseAllResultsContexts();
+                    channelListener.onFailure(failure);
                     return;
                 }
                 final QueryPhaseResultConsumer.MergeResult mergeResult;
@@ -800,7 +930,8 @@ void onShardDone() {
                         EMPTY_PARTIAL_MERGE_RESULT
                     );
                 } catch (Exception e) {
-                    handleMergeFailure(e, channelListener, namedWriteableRegistry);
+                    releaseAllResultsContexts();
+                    channelListener.onFailure(e);
                     return;
                 }
                 // translate shard indices to those on the coordinator so that it can interpret the merge result without adjustments,
@@ -839,16 +970,30 @@ && isPartOfPIT(searchRequest.searchRequest, q.getContextId(), namedWriteableRegi
 
                 ActionListener.respondAndRelease(
                     channelListener,
-                    new NodeQueryResponse(mergeResult, results, queryPhaseResultConsumer.topDocsStats)
+                    new NodeQueryResponse(results, null, mergeResult, queryPhaseResultConsumer.topDocsStats)
                 );
             }
         }
 
-        private void handleMergeFailure(
-            Exception e,
-            ChannelActionListener<TransportResponse> channelListener,
+        private void maybeFreeContext(
+            SearchPhaseResult result,
+            BitSet relevantShardIndices,
             NamedWriteableRegistry namedWriteableRegistry
         ) {
+            if (result instanceof QuerySearchResult q
+                && q.getContextId() != null
+                && relevantShardIndices.get(q.getShardIndex()) == false
+                && q.hasSuggestHits() == false
+                && q.getRankShardResult() == null
+                && searchRequest.searchRequest.scroll() == null
+                && isPartOfPIT(searchRequest.searchRequest, q.getContextId(), namedWriteableRegistry) == false) {
+                if (dependencies.searchService.freeReaderContext(q.getContextId())) {
+                    q.clearContextId();
+                }
+            }
+        }
+
+        private void releaseAllResultsContexts() {
             queryPhaseResultConsumer.getSuccessfulResults()
                 .forEach(
                     searchPhaseResult -> releaseLocalContext(
@@ -858,7 +1003,6 @@ private void handleMergeFailure(
                         namedWriteableRegistry
                     )
                 );
-            channelListener.onFailure(e);
         }
 
         void consumeResult(QuerySearchResult queryResult) {

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

@@ -286,7 +286,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
         Property.NodeScope
     );
 
-    private static final boolean BATCHED_QUERY_PHASE_FEATURE_FLAG = new FeatureFlag("batched_query_phase").isEnabled();
+    public static final boolean BATCHED_QUERY_PHASE_FEATURE_FLAG = new FeatureFlag("batched_query_phase").isEnabled();
 
     public static final int DEFAULT_SIZE = 10;
     public static final int DEFAULT_FROM = 0;

server/src/main/java/org/elasticsearch/search/query/QuerySearchResult.java

@@ -491,7 +491,6 @@ public void writeToNoId(StreamOutput out) throws IOException {
                 out.writeBoolean(true);
                 writeTopDocs(out, topDocsAndMaxScore);
             } else {
-                assert isPartiallyReduced();
                 out.writeBoolean(false);
             }
         } else {

server/src/main/resources/transport/definitions/referable/batched_response_might_include_reduction_failure.csv

@@ -0,0 +1 @@
+9213000,9185007,9112012,8841074

server/src/main/resources/transport/upper_bounds/8.19.csv

@@ -1 +1 @@
-initial_8.19.8,8841073
+batched_response_might_include_reduction_failure,8841074