move logic to DataNodeComputeHandler

Author: idegtiarenko

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.RemoteException;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.breaker.CircuitBreaker;
@@ -383,6 +384,7 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         mock(SearchService.class),
         null,
         mock(ClusterService.class),
+        mock(ProjectResolver.class),
         mock(IndexNameExpressionResolver.class),
         null,
         mockInferenceRunner()

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlResolveNodesAction.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.ShardSearchFailure;
 import org.elasticsearch.cluster.RemoteException;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.concurrent.RunOnce;
@@ -128,6 +129,7 @@ public class ComputeService {
     private final LookupFromIndexService lookupFromIndexService;
     private final InferenceRunner inferenceRunner;
     private final ClusterService clusterService;
+    private final ProjectResolver projectResolver;
     private final AtomicLong childSessionIdGenerator = new AtomicLong();
     private final DataNodeComputeHandler dataNodeComputeHandler;
     private final ClusterComputeHandler clusterComputeHandler;
@@ -155,7 +157,16 @@ public ComputeService(
         this.lookupFromIndexService = lookupFromIndexService;
         this.inferenceRunner = transportActionServices.inferenceRunner();
         this.clusterService = transportActionServices.clusterService();
-        this.dataNodeComputeHandler = new DataNodeComputeHandler(this, searchService, transportService, exchangeService, esqlExecutor);
+        this.projectResolver = transportActionServices.projectResolver();
+        this.dataNodeComputeHandler = new DataNodeComputeHandler(
+            this,
+            clusterService,
+            projectResolver,
+            searchService,
+            transportService,
+            exchangeService,
+            esqlExecutor
+        );
         this.clusterComputeHandler = new ClusterComputeHandler(
             this,
             exchangeService,

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/ComputeService.java

@@ -16,6 +16,8 @@
 import org.elasticsearch.action.support.ChannelActionListener;
 import org.elasticsearch.action.support.RefCountingRunnable;
 import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.project.ProjectResolver;
+import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.compute.operator.DriverProfile;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
 import org.elasticsearch.compute.operator.exchange.ExchangeSink;
@@ -66,19 +68,25 @@
 final class DataNodeComputeHandler implements TransportRequestHandler<DataNodeRequest> {
     private final ComputeService computeService;
     private final SearchService searchService;
+    private final ClusterService clusterService;
+    private final ProjectResolver projectResolver;
     private final TransportService transportService;
     private final ExchangeService exchangeService;
     private final Executor esqlExecutor;
     private final ThreadPool threadPool;
 
     DataNodeComputeHandler(
         ComputeService computeService,
+        ClusterService clusterService,
+        ProjectResolver projectResolver,
         SearchService searchService,
         TransportService transportService,
         ExchangeService exchangeService,
         Executor esqlExecutor
     ) {
         this.computeService = computeService;
+        this.clusterService = clusterService;
+        this.projectResolver = projectResolver;
         this.searchService = searchService;
         this.transportService = transportService;
         this.exchangeService = exchangeService;
@@ -102,6 +110,8 @@ void startComputeOnDataNodes(
         Integer maxConcurrentNodesPerCluster = PlanConcurrencyCalculator.INSTANCE.calculateNodesConcurrency(dataNodePlan, configuration);
 
         new DataNodeRequestSender(
+            clusterService,
+            projectResolver,
             transportService,
             esqlExecutor,
             parentTask,

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeComputeHandler.java

@@ -19,7 +19,10 @@
 import org.elasticsearch.action.support.TransportActions;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.cluster.node.DiscoveryNodeRole;
+import org.elasticsearch.cluster.project.ProjectResolver;
+import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.breaker.CircuitBreakingException;
+import org.elasticsearch.common.util.Maps;
 import org.elasticsearch.common.util.concurrent.ConcurrentCollections;
 import org.elasticsearch.compute.operator.DriverProfile;
 import org.elasticsearch.compute.operator.FailureCollector;
@@ -33,9 +36,6 @@
 import org.elasticsearch.transport.TransportException;
 import org.elasticsearch.transport.TransportRequestOptions;
 import org.elasticsearch.transport.TransportService;
-import org.elasticsearch.xpack.esql.action.EsqlResolveNodesAction;
-import org.elasticsearch.xpack.esql.action.EsqlResolveNodesAction.ResolveNodesRequest;
-import org.elasticsearch.xpack.esql.action.EsqlResolveNodesAction.ResolveNodesResponse;
 import org.elasticsearch.xpack.esql.action.EsqlSearchShardsAction;
 
 import java.util.ArrayList;
@@ -78,6 +78,8 @@ abstract class DataNodeRequestSender {
         DiscoveryNodeRole.DATA_FROZEN_NODE_ROLE.roleName()
     );
 
+    private final ClusterService clusterService;
+    private final ProjectResolver projectResolver;
     private final TransportService transportService;
     private final Executor esqlExecutor;
     private final CancellableTask rootTask;
@@ -100,6 +102,8 @@ abstract class DataNodeRequestSender {
     private final AtomicInteger remainingTargetShardSearchAttempts = new AtomicInteger(10);
 
     DataNodeRequestSender(
+        ClusterService clusterService,
+        ProjectResolver projectResolver,
         TransportService transportService,
         Executor esqlExecutor,
         CancellableTask rootTask,
@@ -109,6 +113,8 @@ abstract class DataNodeRequestSender {
         boolean allowPartialResults,
         int concurrentRequests
     ) {
+        this.clusterService = clusterService;
+        this.projectResolver = projectResolver;
         this.transportService = transportService;
         this.esqlExecutor = esqlExecutor;
         this.rootTask = rootTask;
@@ -490,13 +496,23 @@ void searchShards(Set<String> concreteIndices, ActionListener<TargetShards> list
     }
 
     void resolveShards(Set<ShardId> shardIds, ActionListener<Map<ShardId, List<DiscoveryNode>>> listener) {
-        transportService.sendChildRequest(
-            transportService.getLocalNode(),
-            EsqlResolveNodesAction.TYPE.name(),
-            new ResolveNodesRequest(shardIds),
-            rootTask,
-            TransportRequestOptions.EMPTY,
-            new ActionListenerResponseHandler<>(listener.map(ResolveNodesResponse::nodes), ResolveNodesResponse::new, esqlExecutor)
-        );
+        ActionListener.completeWith(listener, () -> doResolveShards(shardIds));
+    }
+
+    private Map<ShardId, List<DiscoveryNode>> doResolveShards(Set<ShardId> shardIds) {
+        var project = projectResolver.getProjectState(clusterService.state());
+        var nodes = Maps.<ShardId, List<DiscoveryNode>>newMapWithExpectedSize(shardIds.size());
+        for (var shardId : shardIds) {
+            nodes.put(
+                shardId,
+                project.routingTable()
+                    .shardRoutingTable(shardId)
+                    .allShards()
+                    .filter(shard -> shard.active() && shard.isSearchable())
+                    .map(shard -> project.cluster().nodes().get(shard.currentNodeId()))
+                    .toList()
+            );
+        }
+        return nodes;
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSender.java

@@ -58,7 +58,6 @@
 import org.elasticsearch.xpack.esql.action.EsqlQueryAction;
 import org.elasticsearch.xpack.esql.action.EsqlQueryRequestBuilder;
 import org.elasticsearch.xpack.esql.action.EsqlResolveFieldsAction;
-import org.elasticsearch.xpack.esql.action.EsqlResolveNodesAction;
 import org.elasticsearch.xpack.esql.action.EsqlSearchShardsAction;
 import org.elasticsearch.xpack.esql.action.RestEsqlAsyncQueryAction;
 import org.elasticsearch.xpack.esql.action.RestEsqlDeleteAsyncResultAction;
@@ -240,7 +239,6 @@ public List<ActionHandler> getActions() {
             new ActionHandler(XPackUsageFeatureAction.ESQL, EsqlUsageTransportAction.class),
             new ActionHandler(XPackInfoFeatureAction.ESQL, EsqlInfoTransportAction.class),
             new ActionHandler(EsqlResolveFieldsAction.TYPE, EsqlResolveFieldsAction.class),
-            new ActionHandler(EsqlResolveNodesAction.TYPE, EsqlResolveNodesAction.class),
             new ActionHandler(EsqlSearchShardsAction.TYPE, EsqlSearchShardsAction.class),
             new ActionHandler(EsqlAsyncStopAction.INSTANCE, TransportEsqlAsyncStopAction.class),
             new ActionHandler(EsqlListQueriesAction.INSTANCE, TransportEsqlListQueriesAction.class),

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.plugin;
 
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
 import org.elasticsearch.search.SearchService;
@@ -20,6 +21,7 @@ public record TransportActionServices(
     SearchService searchService,
     ExchangeService exchangeService,
     ClusterService clusterService,
+    ProjectResolver projectResolver,
     IndexNameExpressionResolver indexNameExpressionResolver,
     UsageService usageService,
     InferenceRunner inferenceRunner

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportActionServices.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.action.support.HandledTransportAction;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.project.ProjectResolver;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -94,6 +95,7 @@ public TransportEsqlQueryAction(
         SearchService searchService,
         ExchangeService exchangeService,
         ClusterService clusterService,
+        ProjectResolver projectResolver,
         ThreadPool threadPool,
         BigArrays bigArrays,
         BlockFactoryProvider blockFactoryProvider,
@@ -149,6 +151,7 @@ public TransportEsqlQueryAction(
             searchService,
             exchangeService,
             clusterService,
+            projectResolver,
             indexNameExpressionResolver,
             usageService,
             new InferenceRunner(client)

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java

@@ -34,7 +34,6 @@
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.transport.TransportService;
 import org.elasticsearch.xpack.esql.plugin.DataNodeRequestSender.NodeListener;
-import org.elasticsearch.xpack.esql.plugin.DataNodeRequestSender.TargetShards;
 import org.junit.After;
 import org.junit.Before;
 
@@ -569,6 +568,8 @@ PlainActionFuture<ComputeResponse> sendRequests(
             Collections.emptyMap()
         );
         new DataNodeRequestSender(
+            null,
+            null,
             transportService,
             executor,
             task,