Remove SearchPhaseContext (#116471) (#117396)

* Stop instantiating RankFeaturePhase unnecessarily (#115724)

We should not create the phase instance when we know we won't be doing any rank feature
execution up-front. An instance of these isn't free and entails creating an array of searched_shard_count size
which along is non-trivial. Also, this needlessly obscured the threading logic for fetch which has already led to a bug
before.

* Remove SearchPhaseContext (#116471)

The only production implementation of this thing is
AbstractSearchAsyncAction, no need to keep a separate interface around.
This makes the logic a lot more obvious in terms of the lifeycle of
"context" and how it's essentially just the "main" search phase.
Plus it outright saves a lot of code, even though it adds a little on
the test side.

* Cleanup some redundancies around DfsQueryPhase (#116057)

No need to add the result consumer to teh context another time, it's
already added to it in the constructor of `SearchDfsQueryThenFetchAsyncAction`.
Also, no need to feed `this` and `this.results` to `getNextPhase`
explicitly, there's only a single call to this method so we can safely
clean up the redundant arguments.

Author: original-brownbear

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

@@ -69,7 +69,7 @@
  * The fan out and collect algorithm is traditionally used as the initial phase which can either be a query execution or collection of
  * distributed frequencies
  */
-abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> extends SearchPhase implements SearchPhaseContext {
+abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> extends SearchPhase {
     private static final float DEFAULT_INDEX_BOOST = 1.0f;
     private final Logger logger;
     private final NamedWriteableRegistry namedWriteableRegistry;
@@ -107,7 +107,8 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
     private final boolean throttleConcurrentRequests;
     private final AtomicBoolean requestCancelled = new AtomicBoolean();
 
-    private final List<Releasable> releasables = new ArrayList<>();
+    // protected for tests
+    protected final List<Releasable> releasables = new ArrayList<>();
 
     AbstractSearchAsyncAction(
         String name,
@@ -195,7 +196,9 @@ protected void notifyListShards(
         );
     }
 
-    @Override
+    /**
+     * Registers a {@link Releasable} that will be closed when the search request finishes or fails.
+     */
     public void addReleasable(Releasable releasable) {
         releasables.add(releasable);
     }
@@ -368,8 +371,12 @@ protected abstract void executePhaseOnShard(
         SearchActionListener<Result> listener
     );
 
-    @Override
-    public final void executeNextPhase(SearchPhase currentPhase, Supplier<SearchPhase> nextPhaseSupplier) {
+    /**
+     * Processes the phase transition from on phase to another. This method handles all errors that happen during the initial run execution
+     * of the next phase. If there are no successful operations in the context when this method is executed the search is aborted and
+     * a response is returned to the user indicating that all shards have failed.
+     */
+    protected void executeNextPhase(SearchPhase currentPhase, Supplier<SearchPhase> nextPhaseSupplier) {
         /* This is the main search phase transition where we move to the next phase. If all shards
          * failed or if there was a failure and partial results are not allowed, then we immediately
          * fail. Otherwise we continue to the next phase.
@@ -505,8 +512,7 @@ protected void onShardGroupFailure(int shardIndex, SearchShardTarget shardTarget
      * @param shardTarget the shard target for this failure
      * @param e the failure reason
      */
-    @Override
-    public final void onShardFailure(final int shardIndex, SearchShardTarget shardTarget, Exception e) {
+    void onShardFailure(final int shardIndex, SearchShardTarget shardTarget, Exception e) {
         if (TransportActions.isShardNotAvailableException(e)) {
             // Groups shard not available exceptions under a generic exception that returns a SERVICE_UNAVAILABLE(503)
             // temporary error.
@@ -603,32 +609,45 @@ private void successfulShardExecution(SearchShardIterator shardsIt) {
         }
     }
 
-    @Override
+    /**
+     * Returns the total number of shards to the current search across all indices
+     */
     public final int getNumShards() {
         return results.getNumShards();
     }
 
-    @Override
+    /**
+    * Returns a logger for this context to prevent each individual phase to create their own logger.
+    */
     public final Logger getLogger() {
         return logger;
     }
 
-    @Override
+    /**
+      * Returns the currently executing search task
+    */
     public final SearchTask getTask() {
         return task;
     }
 
-    @Override
+    /**
+      * Returns the currently executing search request
+      */
     public final SearchRequest getRequest() {
         return request;
     }
 
-    @Override
+    /**
+      * Returns the targeted {@link OriginalIndices} for the provided {@code shardIndex}.
+      */
     public OriginalIndices getOriginalIndices(int shardIndex) {
         return shardIterators[shardIndex].getOriginalIndices();
     }
 
-    @Override
+    /**
+      * Checks if the given context id is part of the point in time of this search (if exists).
+      * We should not release search contexts that belong to the point in time during or after searches.
+    */
     public boolean isPartOfPointInTime(ShardSearchContextId contextId) {
         final PointInTimeBuilder pointInTimeBuilder = request.pointInTimeBuilder();
         if (pointInTimeBuilder != null) {
@@ -665,7 +684,12 @@ boolean buildPointInTimeFromSearchResults() {
         return false;
     }
 
-    @Override
+    /**
+      * Builds and sends the final search response back to the user.
+      *
+      * @param internalSearchResponse the internal search response
+      * @param queryResults           the results of the query phase
+      */
     public void sendSearchResponse(SearchResponseSections internalSearchResponse, AtomicArray<SearchPhaseResult> queryResults) {
         ShardSearchFailure[] failures = buildShardFailures();
         Boolean allowPartialResults = request.allowPartialSearchResults();
@@ -690,8 +714,14 @@ public void sendSearchResponse(SearchResponseSections internalSearchResponse, At
         }
     }
 
-    @Override
-    public final void onPhaseFailure(SearchPhase phase, String msg, Throwable cause) {
+    /**
+     * This method will communicate a fatal phase failure back to the user. In contrast to a shard failure
+     * will this method immediately fail the search request and return the failure to the issuer of the request
+     * @param phase the phase that failed
+     * @param msg an optional message
+     * @param cause the cause of the phase failure
+     */
+    public void onPhaseFailure(SearchPhase phase, String msg, Throwable cause) {
         raisePhaseFailure(new SearchPhaseExecutionException(phase.getName(), msg, cause, buildShardFailures()));
     }
 
@@ -718,16 +748,32 @@ private void raisePhaseFailure(SearchPhaseExecutionException exception) {
         listener.onFailure(exception);
     }
 
+    /**
+      * Releases a search context with the given context ID on the node the given connection is connected to.
+      * @see org.elasticsearch.search.query.QuerySearchResult#getContextId()
+      * @see org.elasticsearch.search.fetch.FetchSearchResult#getContextId()
+      *
+      */
+    void sendReleaseSearchContext(ShardSearchContextId contextId, Transport.Connection connection, OriginalIndices originalIndices) {
+        assert isPartOfPointInTime(contextId) == false : "Must not release point in time context [" + contextId + "]";
+        if (connection != null) {
+            searchTransportService.sendFreeContext(connection, contextId, originalIndices);
+        }
+    }
+
     /**
      * Executed once all shard results have been received and processed
      * @see #onShardFailure(int, SearchShardTarget, Exception)
      * @see #onShardResult(SearchPhaseResult, SearchShardIterator)
      */
     final void onPhaseDone() {  // as a tribute to @kimchy aka. finishHim()
-        executeNextPhase(this, () -> getNextPhase(results, this));
+        executeNextPhase(this, this::getNextPhase);
     }
 
-    @Override
+    /**
+     * Returns a connection to the node if connected otherwise and {@link org.elasticsearch.transport.ConnectTransportException} will be
+     * thrown.
+     */
     public final Transport.Connection getConnection(String clusterAlias, String nodeId) {
         Transport.Connection conn = nodeIdToConnection.apply(clusterAlias, nodeId);
         Version minVersion = request.minCompatibleShardNode();
@@ -737,18 +783,21 @@ public final Transport.Connection getConnection(String clusterAlias, String node
         return conn;
     }
 
-    @Override
-    public final SearchTransportService getSearchTransport() {
+    /**
+     * Returns the {@link SearchTransportService} to send shard request to other nodes
+     */
+    public SearchTransportService getSearchTransport() {
         return searchTransportService;
     }
 
-    @Override
     public final void execute(Runnable command) {
         executor.execute(command);
     }
 
-    @Override
-    public final void onFailure(Exception e) {
+    /**
+      * Notifies the top-level listener of the provided exception
+    */
+    public void onFailure(Exception e) {
         listener.onFailure(e);
     }
 
@@ -786,11 +835,8 @@ protected final ShardSearchRequest buildShardSearchRequest(SearchShardIterator s
 
     /**
      * Returns the next phase based on the results of the initial search phase
-     * @param results the results of the initial search phase. Each non null element in the result array represent a successfully
-     *                executed shard request
-     * @param context the search context for the next phase
      */
-    protected abstract SearchPhase getNextPhase(SearchPhaseResults<Result> results, SearchPhaseContext context);
+    protected abstract SearchPhase getNextPhase();
 
     private static final class PendingExecutions {
         private final Semaphore semaphore;

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

@@ -22,9 +22,9 @@ final class CountedCollector<R extends SearchPhaseResult> {
     private final SearchPhaseResults<R> resultConsumer;
     private final CountDown counter;
     private final Runnable onFinish;
-    private final SearchPhaseContext context;
+    private final AbstractSearchAsyncAction<?> context;
 
-    CountedCollector(SearchPhaseResults<R> resultConsumer, int expectedOps, Runnable onFinish, SearchPhaseContext context) {
+    CountedCollector(SearchPhaseResults<R> resultConsumer, int expectedOps, Runnable onFinish, AbstractSearchAsyncAction<?> context) {
         this.resultConsumer = resultConsumer;
         this.counter = new CountDown(expectedOps);
         this.onFinish = onFinish;
@@ -50,7 +50,7 @@ void onResult(R result) {
     }
 
     /**
-     * Escalates the failure via {@link SearchPhaseContext#onShardFailure(int, SearchShardTarget, Exception)}
+     * Escalates the failure via {@link AbstractSearchAsyncAction#onShardFailure(int, SearchShardTarget, Exception)}
      * and then runs {@link #countDown()}
      */
     void onFailure(final int shardIndex, @Nullable SearchShardTarget shardTarget, Exception e) {

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

@@ -44,7 +44,7 @@ final class DfsQueryPhase extends SearchPhase {
     private final AggregatedDfs dfs;
     private final List<DfsKnnResults> knnResults;
     private final Function<SearchPhaseResults<SearchPhaseResult>, SearchPhase> nextPhaseFactory;
-    private final SearchPhaseContext context;
+    private final AbstractSearchAsyncAction<?> context;
     private final SearchTransportService searchTransportService;
     private final SearchProgressListener progressListener;
 
@@ -54,7 +54,7 @@ final class DfsQueryPhase extends SearchPhase {
         List<DfsKnnResults> knnResults,
         SearchPhaseResults<SearchPhaseResult> queryResult,
         Function<SearchPhaseResults<SearchPhaseResult>, SearchPhase> nextPhaseFactory,
-        SearchPhaseContext context
+        AbstractSearchAsyncAction<?> context
     ) {
         super("dfs_query");
         this.progressListener = context.getTask().getProgressListener();
@@ -65,10 +65,6 @@ final class DfsQueryPhase extends SearchPhase {
         this.nextPhaseFactory = nextPhaseFactory;
         this.context = context;
         this.searchTransportService = context.getSearchTransport();
-
-        // register the release of the query consumer to free up the circuit breaker memory
-        // at the end of the search
-        context.addReleasable(queryResult);
     }
 
     @Override

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

@@ -31,11 +31,11 @@
  * forwards to the next phase immediately.
  */
 final class ExpandSearchPhase extends SearchPhase {
-    private final SearchPhaseContext context;
+    private final AbstractSearchAsyncAction<?> context;
     private final SearchHits searchHits;
     private final Supplier<SearchPhase> nextPhase;
 
-    ExpandSearchPhase(SearchPhaseContext context, SearchHits searchHits, Supplier<SearchPhase> nextPhase) {
+    ExpandSearchPhase(AbstractSearchAsyncAction<?> context, SearchHits searchHits, Supplier<SearchPhase> nextPhase) {
         super("expand");
         this.context = context;
         this.searchHits = searchHits;

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

@@ -33,11 +33,15 @@
  * @see org.elasticsearch.index.mapper.LookupRuntimeFieldType
  */
 final class FetchLookupFieldsPhase extends SearchPhase {
-    private final SearchPhaseContext context;
+    private final AbstractSearchAsyncAction<?> context;
     private final SearchResponseSections searchResponse;
     private final AtomicArray<SearchPhaseResult> queryResults;
 
-    FetchLookupFieldsPhase(SearchPhaseContext context, SearchResponseSections searchResponse, AtomicArray<SearchPhaseResult> queryResults) {
+    FetchLookupFieldsPhase(
+        AbstractSearchAsyncAction<?> context,
+        SearchResponseSections searchResponse,
+        AtomicArray<SearchPhaseResult> queryResults
+    ) {
         super("fetch_lookup_fields");
         this.context = context;
         this.searchResponse = searchResponse;

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

@@ -36,7 +36,7 @@
 final class FetchSearchPhase extends SearchPhase {
     private final AtomicArray<SearchPhaseResult> searchPhaseShardResults;
     private final BiFunction<SearchResponseSections, AtomicArray<SearchPhaseResult>, SearchPhase> nextPhaseFactory;
-    private final SearchPhaseContext context;
+    private final AbstractSearchAsyncAction<?> context;
     private final Logger logger;
     private final SearchProgressListener progressListener;
     private final AggregatedDfs aggregatedDfs;
@@ -47,7 +47,7 @@ final class FetchSearchPhase extends SearchPhase {
     FetchSearchPhase(
         SearchPhaseResults<SearchPhaseResult> resultConsumer,
         AggregatedDfs aggregatedDfs,
-        SearchPhaseContext context,
+        AbstractSearchAsyncAction<?> context,
         @Nullable SearchPhaseController.ReducedQueryPhase reducedQueryPhase
     ) {
         this(
@@ -66,7 +66,7 @@ final class FetchSearchPhase extends SearchPhase {
     FetchSearchPhase(
         SearchPhaseResults<SearchPhaseResult> resultConsumer,
         AggregatedDfs aggregatedDfs,
-        SearchPhaseContext context,
+        AbstractSearchAsyncAction<?> context,
         @Nullable SearchPhaseController.ReducedQueryPhase reducedQueryPhase,
         BiFunction<SearchResponseSections, AtomicArray<SearchPhaseResult>, SearchPhase> nextPhaseFactory
     ) {