Implement chunked fetch streaming with circuit breaker integration

server/src/main/java/org/elasticsearch/action/ActionModule.java

@@ -423,6 +423,8 @@
 import org.elasticsearch.rest.action.synonyms.RestGetSynonymsSetsAction;
 import org.elasticsearch.rest.action.synonyms.RestPutSynonymRuleAction;
 import org.elasticsearch.rest.action.synonyms.RestPutSynonymsAction;
+import org.elasticsearch.search.fetch.chunk.ActiveFetchPhaseTasks;
+import org.elasticsearch.search.fetch.chunk.TransportFetchPhaseResponseChunkAction;
 import org.elasticsearch.snapshots.TransportUpdateSnapshotStatusAction;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.telemetry.TelemetryProvider;
@@ -761,6 +763,7 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
         actions.register(TransportMultiSearchAction.TYPE, TransportMultiSearchAction.class);
         actions.register(TransportExplainAction.TYPE, TransportExplainAction.class);
         actions.register(TransportClearScrollAction.TYPE, TransportClearScrollAction.class);
+        actions.register(TransportFetchPhaseResponseChunkAction.TYPE, TransportFetchPhaseResponseChunkAction.class);
         actions.register(RecoveryAction.INSTANCE, TransportRecoveryAction.class);
         actions.register(TransportNodesReloadSecureSettingsAction.TYPE, TransportNodesReloadSecureSettingsAction.class);
         actions.register(AutoCreateAction.INSTANCE, AutoCreateAction.TransportAction.class);
@@ -1090,6 +1093,7 @@ protected void configure() {
         bind(new TypeLiteral<RequestValidators<IndicesAliasesRequest>>() {
         }).toInstance(indicesAliasesRequestRequestValidators);
         bind(AutoCreateIndex.class).toInstance(autoCreateIndex);
+        bind(ActiveFetchPhaseTasks.class).asEagerSingleton();
 
         // register ActionType -> transportAction Map used by NodeClient
         @SuppressWarnings("rawtypes")

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

@@ -27,6 +27,7 @@
 import org.elasticsearch.search.dfs.AggregatedDfs;
 import org.elasticsearch.search.dfs.DfsKnnResults;
 import org.elasticsearch.search.dfs.DfsSearchResult;
+import org.elasticsearch.search.fetch.chunk.TransportFetchPhaseCoordinationAction;
 import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.search.query.QuerySearchRequest;
 import org.elasticsearch.search.query.QuerySearchResult;
@@ -56,18 +57,25 @@ class DfsQueryPhase extends SearchPhase {
     private final AbstractSearchAsyncAction<?> context;
     private final SearchProgressListener progressListener;
     private long phaseStartTimeInNanos;
+    private final TransportFetchPhaseCoordinationAction fetchCoordinationAction;
 
-    DfsQueryPhase(SearchPhaseResults<SearchPhaseResult> queryResult, Client client, AbstractSearchAsyncAction<?> context) {
+    DfsQueryPhase(
+        SearchPhaseResults<SearchPhaseResult> queryResult,
+        Client client,
+        AbstractSearchAsyncAction<?> context,
+        TransportFetchPhaseCoordinationAction fetchCoordinationAction
+    ) {
         super(NAME);
         this.progressListener = context.getTask().getProgressListener();
         this.queryResult = queryResult;
         this.client = client;
         this.context = context;
+        this.fetchCoordinationAction = fetchCoordinationAction;
     }
 
     // protected for testing
     protected SearchPhase nextPhase(AggregatedDfs dfs) {
-        return SearchQueryThenFetchAsyncAction.nextPhase(client, context, queryResult, dfs);
+        return SearchQueryThenFetchAsyncAction.nextPhase(client, context, queryResult, dfs, fetchCoordinationAction);
     }
 
     @SuppressWarnings("unchecked")

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

@@ -10,6 +10,8 @@
 
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.ScoreDoc;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.AtomicArray;
 import org.elasticsearch.core.Nullable;
@@ -18,6 +20,7 @@
 import org.elasticsearch.search.dfs.AggregatedDfs;
 import org.elasticsearch.search.fetch.FetchSearchResult;
 import org.elasticsearch.search.fetch.ShardFetchSearchRequest;
+import org.elasticsearch.search.fetch.chunk.TransportFetchPhaseCoordinationAction;
 import org.elasticsearch.search.internal.ShardSearchContextId;
 import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.search.rank.RankDocShardInfo;
@@ -28,6 +31,8 @@
 import java.util.List;
 import java.util.Map;
 
+import static org.elasticsearch.search.fetch.chunk.TransportFetchPhaseCoordinationAction.CHUNKED_FETCH_PHASE;
+
 /**
  * This search phase merges the query results from the previous phase together and calculates the topN hits for this search.
  * Then it reaches out to all relevant shards to fetch the topN hits.
@@ -44,12 +49,14 @@ class FetchSearchPhase extends SearchPhase {
     @Nullable
     private final SearchPhaseResults<SearchPhaseResult> resultConsumer;
     private final SearchPhaseController.ReducedQueryPhase reducedQueryPhase;
+    private final TransportFetchPhaseCoordinationAction fetchCoordinationAction;
 
     FetchSearchPhase(
         SearchPhaseResults<SearchPhaseResult> resultConsumer,
         AggregatedDfs aggregatedDfs,
         AbstractSearchAsyncAction<?> context,
-        @Nullable SearchPhaseController.ReducedQueryPhase reducedQueryPhase
+        @Nullable SearchPhaseController.ReducedQueryPhase reducedQueryPhase,
+        TransportFetchPhaseCoordinationAction fetchCoordinationAction
     ) {
         super(NAME);
         if (context.getNumShards() != resultConsumer.getNumShards()) {
@@ -67,6 +74,7 @@ class FetchSearchPhase extends SearchPhase {
         this.progressListener = context.getTask().getProgressListener();
         this.reducedQueryPhase = reducedQueryPhase;
         this.resultConsumer = reducedQueryPhase == null ? resultConsumer : null;
+        this.fetchCoordinationAction = fetchCoordinationAction;
     }
 
     // protected for tests
@@ -98,7 +106,7 @@ private void innerRun() throws Exception {
         final int numShards = context.getNumShards();
         // Usually when there is a single shard, we force the search type QUERY_THEN_FETCH. But when there's kNN, we might
         // still use DFS_QUERY_THEN_FETCH, which does not perform the "query and fetch" optimization during the query phase.
-        final boolean queryAndFetchOptimization = numShards == 1
+        boolean queryAndFetchOptimization = numShards == 1
             && context.getRequest().hasKnnSearch() == false
             && reducedQueryPhase.queryPhaseRankCoordinatorContext() == null
             && (context.getRequest().source() == null || context.getRequest().source().rankBuilder() == null);
@@ -212,6 +220,8 @@ private void executeFetch(
         final ShardSearchContextId contextId = shardPhaseResult.queryResult() != null
             ? shardPhaseResult.queryResult().getContextId()
             : shardPhaseResult.rankFeatureResult().getContextId();
+
+        // Create the listener that handles the fetch result
         var listener = new SearchActionListener<FetchSearchResult>(shardTarget, shardIndex) {
             @Override
             public void innerOnResponse(FetchSearchResult result) {
@@ -237,29 +247,43 @@ public void onFailure(Exception e) {
                 }
             }
         };
+
+        // Get connection to the target node
         final Transport.Connection connection;
         try {
             connection = context.getConnection(shardTarget.getClusterAlias(), shardTarget.getNodeId());
         } catch (Exception e) {
             listener.onFailure(e);
             return;
         }
-        context.getSearchTransport()
-            .sendExecuteFetch(
-                connection,
-                new ShardFetchSearchRequest(
-                    context.getOriginalIndices(shardPhaseResult.getShardIndex()),
-                    contextId,
-                    shardPhaseResult.getShardSearchRequest(),
-                    entry,
-                    rankDocs,
-                    lastEmittedDocForShard,
-                    shardPhaseResult.getRescoreDocIds(),
-                    aggregatedDfs
-                ),
+
+        // Create the fetch request
+        final ShardFetchSearchRequest shardFetchRequest = new ShardFetchSearchRequest(
+            context.getOriginalIndices(shardPhaseResult.getShardIndex()),
+            contextId,
+            shardPhaseResult.getShardSearchRequest(),
+            entry,
+            rankDocs,
+            lastEmittedDocForShard,
+            shardPhaseResult.getRescoreDocIds(),
+            aggregatedDfs
+        );
+
+        if (connection.getTransportVersion().supports(CHUNKED_FETCH_PHASE)) {
+            shardFetchRequest.setCoordinatingNode(context.getSearchTransport().transportService().getLocalNode());
+            shardFetchRequest.setCoordinatingTaskId(context.getTask().getId());
+
+            DiscoveryNode targetNode = connection.getNode();
+
+            // Execute via coordination action
+            fetchCoordinationAction.execute(
                 context.getTask(),
-                listener
+                new TransportFetchPhaseCoordinationAction.Request(shardFetchRequest, targetNode),
+                ActionListener.wrap(response -> listener.onResponse(response.getResult()), listener::onFailure)
             );
+        } else {
+            context.getSearchTransport().sendExecuteFetch(connection, shardFetchRequest, context.getTask(), listener);
+        }
     }
 
     private void moveToNextPhase(

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.dfs.AggregatedDfs;
+import org.elasticsearch.search.fetch.chunk.TransportFetchPhaseCoordinationAction;
 import org.elasticsearch.search.internal.ShardSearchContextId;
 import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
@@ -48,12 +49,14 @@ public class RankFeaturePhase extends SearchPhase {
     private final AggregatedDfs aggregatedDfs;
     private final SearchProgressListener progressListener;
     private final RankFeaturePhaseRankCoordinatorContext rankFeaturePhaseRankCoordinatorContext;
+    private final TransportFetchPhaseCoordinationAction fetchCoordinationAction;
 
     RankFeaturePhase(
         SearchPhaseResults<SearchPhaseResult> queryPhaseResults,
         AggregatedDfs aggregatedDfs,
         AbstractSearchAsyncAction<?> context,
-        RankFeaturePhaseRankCoordinatorContext rankFeaturePhaseRankCoordinatorContext
+        RankFeaturePhaseRankCoordinatorContext rankFeaturePhaseRankCoordinatorContext,
+        TransportFetchPhaseCoordinationAction fetchCoordinationAction
     ) {
         super(NAME);
         assert rankFeaturePhaseRankCoordinatorContext != null;
@@ -72,6 +75,7 @@ public class RankFeaturePhase extends SearchPhase {
         this.rankPhaseResults = new ArraySearchPhaseResults<>(context.getNumShards());
         context.addReleasable(rankPhaseResults);
         this.progressListener = context.getTask().getProgressListener();
+        this.fetchCoordinationAction = fetchCoordinationAction;
     }
 
     @Override
@@ -267,6 +271,9 @@ private float maxScore(ScoreDoc[] scoreDocs) {
     }
 
     void moveToNextPhase(SearchPhaseResults<SearchPhaseResult> phaseResults, SearchPhaseController.ReducedQueryPhase reducedQueryPhase) {
-        context.executeNextPhase(NAME, () -> new FetchSearchPhase(phaseResults, aggregatedDfs, context, reducedQueryPhase));
+        context.executeNextPhase(
+            NAME,
+            () -> new FetchSearchPhase(phaseResults, aggregatedDfs, context, reducedQueryPhase, fetchCoordinationAction)
+        );
     }
 }

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

@@ -18,6 +18,7 @@
 import org.elasticsearch.search.SearchPhaseResult;
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.dfs.DfsSearchResult;
+import org.elasticsearch.search.fetch.chunk.TransportFetchPhaseCoordinationAction;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.transport.Transport;
 
@@ -31,6 +32,7 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
     private final SearchPhaseResults<SearchPhaseResult> queryPhaseResultConsumer;
     private final SearchProgressListener progressListener;
     private final Client client;
+    private final TransportFetchPhaseCoordinationAction fetchCoordinationAction;
 
     SearchDfsQueryThenFetchAsyncAction(
         Logger logger,
@@ -50,7 +52,8 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
         SearchResponse.Clusters clusters,
         Client client,
         SearchResponseMetrics searchResponseMetrics,
-        Map<String, Object> searchRequestAttributes
+        Map<String, Object> searchRequestAttributes,
+        TransportFetchPhaseCoordinationAction fetchCoordinationAction
     ) {
         super(
             "dfs",
@@ -81,6 +84,7 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
             notifyListShards(progressListener, clusters, request, shardsIts);
         }
         this.client = client;
+        this.fetchCoordinationAction = fetchCoordinationAction;
     }
 
     @Override
@@ -94,7 +98,7 @@ protected void executePhaseOnShard(
 
     @Override
     protected SearchPhase getNextPhase() {
-        return new DfsQueryPhase(queryPhaseResultConsumer, client, this);
+        return new DfsQueryPhase(queryPhaseResultConsumer, client, this, fetchCoordinationAction);
     }
 
     @Override

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

@@ -42,6 +42,7 @@
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.builder.PointInTimeBuilder;
 import org.elasticsearch.search.dfs.AggregatedDfs;
+import org.elasticsearch.search.fetch.chunk.TransportFetchPhaseCoordinationAction;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.internal.ShardSearchContextId;
@@ -98,6 +99,7 @@ public class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<S
     private final Client client;
     private final boolean batchQueryPhase;
     private long phaseStartTimeNanos;
+    private final TransportFetchPhaseCoordinationAction fetchPhaseCoordinationAction;
 
     SearchQueryThenFetchAsyncAction(
         Logger logger,
@@ -118,7 +120,8 @@ public class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<S
         Client client,
         boolean batchQueryPhase,
         SearchResponseMetrics searchResponseMetrics,
-        Map<String, Object> searchRequestAttributes
+        Map<String, Object> searchRequestAttributes,
+        TransportFetchPhaseCoordinationAction fetchPhaseCoordinationAction
     ) {
         super(
             "query",
@@ -150,6 +153,7 @@ public class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<S
         if (progressListener != SearchProgressListener.NOOP) {
             notifyListShards(progressListener, clusters, request, shardsIts);
         }
+        this.fetchPhaseCoordinationAction = fetchPhaseCoordinationAction;
     }
 
     @Override
@@ -208,18 +212,19 @@ static SearchPhase nextPhase(
         Client client,
         AbstractSearchAsyncAction<?> context,
         SearchPhaseResults<SearchPhaseResult> queryResults,
-        AggregatedDfs aggregatedDfs
+        AggregatedDfs aggregatedDfs,
+        TransportFetchPhaseCoordinationAction fetchCoordinationAction
     ) {
         var rankFeaturePhaseCoordCtx = RankFeaturePhase.coordinatorContext(context.getRequest().source(), client);
         if (rankFeaturePhaseCoordCtx == null) {
-            return new FetchSearchPhase(queryResults, aggregatedDfs, context, null);
+            return new FetchSearchPhase(queryResults, aggregatedDfs, context, null, fetchCoordinationAction);
         }
-        return new RankFeaturePhase(queryResults, aggregatedDfs, context, rankFeaturePhaseCoordCtx);
+        return new RankFeaturePhase(queryResults, aggregatedDfs, context, rankFeaturePhaseCoordCtx, fetchCoordinationAction);
     }
 
     @Override
     protected SearchPhase getNextPhase() {
-        return nextPhase(client, this, results, null);
+        return nextPhase(client, this, results, null, fetchPhaseCoordinationAction);
     }
 
     /**

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

@@ -39,6 +39,8 @@
 import org.elasticsearch.search.fetch.ScrollQueryFetchSearchResult;
 import org.elasticsearch.search.fetch.ShardFetchRequest;
 import org.elasticsearch.search.fetch.ShardFetchSearchRequest;
+import org.elasticsearch.search.fetch.chunk.FetchPhaseResponseChunk;
+import org.elasticsearch.search.fetch.chunk.TransportFetchPhaseResponseChunkAction;
 import org.elasticsearch.search.internal.InternalScrollSearchRequest;
 import org.elasticsearch.search.internal.ShardSearchContextId;
 import org.elasticsearch.search.internal.ShardSearchRequest;
@@ -67,6 +69,8 @@
 import java.util.concurrent.Executor;
 import java.util.function.BiFunction;
 
+import static org.elasticsearch.search.fetch.chunk.TransportFetchPhaseCoordinationAction.CHUNKED_FETCH_PHASE;
+
 /**
  * An encapsulation of {@link SearchService} operations exposed through
  * transport.
@@ -540,8 +544,39 @@ public static void registerRequestHandler(
             namedWriteableRegistry
         );
 
-        final TransportRequestHandler<ShardFetchRequest> shardFetchRequestHandler = (request, channel, task) -> searchService
-            .executeFetchPhase(request, (SearchShardTask) task, new ChannelActionListener<>(channel));
+        final TransportRequestHandler<ShardFetchRequest> shardFetchRequestHandler = (request, channel, task) -> {
+            if (channel.getVersion().supports(CHUNKED_FETCH_PHASE)
+                && request instanceof ShardFetchSearchRequest fetchSearchReq
+                && fetchSearchReq.getCoordinatingNode() != null) {
+
+                // CHUNKED PATH
+                final FetchPhaseResponseChunk.Writer writer = new FetchPhaseResponseChunk.Writer() {
+                    final Transport.Connection conn = transportService.getConnection(fetchSearchReq.getCoordinatingNode());
+
+                    @Override
+                    public void writeResponseChunk(FetchPhaseResponseChunk responseChunk, ActionListener<Void> listener) {
+                        transportService.sendChildRequest(
+                            conn,
+                            TransportFetchPhaseResponseChunkAction.TYPE.name(),
+                            new TransportFetchPhaseResponseChunkAction.Request(fetchSearchReq.getCoordinatingTaskId(), responseChunk),
+                            task,
+                            TransportRequestOptions.EMPTY,
+                            new ActionListenerResponseHandler<>(
+                                listener.map(ignored -> null),
+                                in -> ActionResponse.Empty.INSTANCE,
+                                EsExecutors.DIRECT_EXECUTOR_SERVICE
+                            )
+                        );
+                    }
+                };
+
+                searchService.executeFetchPhase(request, (SearchShardTask) task, writer, new ChannelActionListener<>(channel));
+            } else {
+                // Normal path
+                searchService.executeFetchPhase(request, (SearchShardTask) task, new ChannelActionListener<>(channel));
+            }
+        };
+
         transportService.registerRequestHandler(
             FETCH_ID_SCROLL_ACTION_NAME,
             EsExecutors.DIRECT_EXECUTOR_SERVICE,