Delay creation of the next SearchPhase in executeNextPhase (#116061) (#117369)

Delaying the creation of the next phase to only when we actually need it
makes this a lot easier to reason about and should set up further
simplications. Eager creation of the next phase forced a lot of
needlessly complicated safety logic around resources on us.
Since we never "close" the `nextPhase` on failure all its resources
need to be tracked in via `context.addReleasable`. This isn't as much of
an issue with some recent refactorings leaving very little resource
creation in the constructors but still, delaying things saves memory
and makes reasoning about failure cases far easier.

Author: original-brownbear

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

@@ -56,6 +56,7 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.core.Strings.format;
@@ -368,7 +369,7 @@ protected abstract void executePhaseOnShard(
     );
 
     @Override
-    public final void executeNextPhase(SearchPhase currentPhase, SearchPhase nextPhase) {
+    public final 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.
@@ -412,6 +413,7 @@ public final void executeNextPhase(SearchPhase currentPhase, SearchPhase nextPha
                 }
                 return;
             }
+            var nextPhase = nextPhaseSupplier.get();
             if (logger.isTraceEnabled()) {
                 final String resultsFrom = results.getSuccessfulResults()
                     .map(r -> r.getSearchShardTarget().toString())
@@ -722,7 +724,7 @@ private void raisePhaseFailure(SearchPhaseExecutionException exception) {
      * @see #onShardResult(SearchPhaseResult, SearchShardIterator)
      */
     final void onPhaseDone() {  // as a tribute to @kimchy aka. finishHim()
-        executeNextPhase(this, getNextPhase(results, this));
+        executeNextPhase(this, () -> getNextPhase(results, this));
     }
 
     @Override

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

@@ -78,7 +78,7 @@ public void run() {
         final CountedCollector<SearchPhaseResult> counter = new CountedCollector<>(
             queryResult,
             searchResults.size(),
-            () -> context.executeNextPhase(this, nextPhaseFactory.apply(queryResult)),
+            () -> context.executeNextPhase(this, () -> nextPhaseFactory.apply(queryResult)),
             context
         );
 

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

@@ -164,6 +164,6 @@ private static SearchSourceBuilder buildExpandSearchSourceBuilder(InnerHitBuilde
     }
 
     private void onPhaseDone() {
-        context.executeNextPhase(this, nextPhase.get());
+        context.executeNextPhase(this, nextPhase);
     }
 }

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

@@ -260,9 +260,11 @@ private void moveToNextPhase(
         AtomicArray<? extends SearchPhaseResult> fetchResultsArr,
         SearchPhaseController.ReducedQueryPhase reducedQueryPhase
     ) {
-        var resp = SearchPhaseController.merge(context.getRequest().scroll() != null, reducedQueryPhase, fetchResultsArr);
-        context.addReleasable(resp::decRef);
-        context.executeNextPhase(this, nextPhaseFactory.apply(resp, searchPhaseShardResults));
+        context.executeNextPhase(this, () -> {
+            var resp = SearchPhaseController.merge(context.getRequest().scroll() != null, reducedQueryPhase, fetchResultsArr);
+            context.addReleasable(resp::decRef);
+            return nextPhaseFactory.apply(resp, searchPhaseShardResults);
+        });
     }
 
     private boolean shouldExplainRankScores(SearchRequest request) {

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

@@ -233,6 +233,6 @@ private float maxScore(ScoreDoc[] scoreDocs) {
     }
 
     void moveToNextPhase(SearchPhaseResults<SearchPhaseResult> phaseResults, SearchPhaseController.ReducedQueryPhase reducedQueryPhase) {
-        context.executeNextPhase(this, new FetchSearchPhase(phaseResults, aggregatedDfs, context, reducedQueryPhase));
+        context.executeNextPhase(this, () -> new FetchSearchPhase(phaseResults, aggregatedDfs, context, reducedQueryPhase));
     }
 }

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.transport.Transport;
 
 import java.util.concurrent.Executor;
+import java.util.function.Supplier;
 
 /**
  * This class provide contextual state and access to resources across multiple search phases.
@@ -120,7 +121,7 @@ default void sendReleaseSearchContext(
      * 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.
      */
-    void executeNextPhase(SearchPhase currentPhase, SearchPhase nextPhase);
+    void executeNextPhase(SearchPhase currentPhase, Supplier<SearchPhase> nextPhaseSupplier);
 
     /**
      * Registers a {@link Releasable} that will be closed when the search request finishes or fails.

server/src/test/java/org/elasticsearch/action/search/MockSearchPhaseContext.java

@@ -30,6 +30,7 @@
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 /**
  * SearchPhaseContext for tests
@@ -132,7 +133,8 @@ public SearchTransportService getSearchTransport() {
     }
 
     @Override
-    public void executeNextPhase(SearchPhase currentPhase, SearchPhase nextPhase) {
+    public void executeNextPhase(SearchPhase currentPhase, Supplier<SearchPhase> nextPhaseSupplier) {
+        var nextPhase = nextPhaseSupplier.get();
         try {
             nextPhase.run();
         } catch (Exception e) {