Limit concurrent node requests

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/exchange/ExchangeSourceHandler.java

@@ -30,6 +30,7 @@
  * @see #addRemoteSink(RemoteSink, boolean, Runnable, int, ActionListener)
  */
 public final class ExchangeSourceHandler {
+
     private final ExchangeBuffer buffer;
     private final Executor fetchExecutor;
 
@@ -56,6 +57,10 @@ public ExchangeSourceHandler(int maxBufferSize, Executor fetchExecutor) {
         this.outstandingSources = new PendingInstances(() -> finishEarly(true, ActionListener.noop()));
     }
 
+    public boolean isFinished() {
+        return buffer.isFinished();
+    }
+
     private void checkFailure() {
         if (aborted) {
             throw new TaskCancelledException("remote sinks failed");

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ManyShardsIT.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.common.util.CollectionUtils;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException;
+import org.elasticsearch.common.util.iterable.Iterables;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.plugins.Plugin;
@@ -39,13 +40,15 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.hamcrest.Matchers.lessThanOrEqualTo;
 
 /**
  * Make sures that we can run many concurrent requests with large number of shards with any data_partitioning.
@@ -256,4 +259,32 @@ public void testLimitConcurrentShards() {
             }
         }
     }
+
+    public void testCancelUnnecessaryRequests() {
+        assumeTrue("Requires pragmas", canUseQueryPragmas());
+        internalCluster().ensureAtLeastNumDataNodes(3);
+
+        var coordinatingNode = internalCluster().getNodeNames()[0];
+
+        var exchanges = new AtomicInteger(0);
+        var coordinatorNodeTransport = MockTransportService.getInstance(coordinatingNode);
+        coordinatorNodeTransport.addSendBehavior((connection, requestId, action, request, options) -> {
+            if (Objects.equals(action, ExchangeService.OPEN_EXCHANGE_ACTION_NAME)) {
+                exchanges.incrementAndGet();
+            }
+            connection.sendRequest(requestId, action, request, options);
+        });
+
+        var query = EsqlQueryRequest.syncEsqlQueryRequest();
+        query.query("from test-* | LIMIT 1");
+        query.pragmas(new QueryPragmas(Settings.builder().put(QueryPragmas.MAX_CONCURRENT_NODES_PER_CLUSTER.getKey(), 1).build()));
+
+        try {
+            var result = safeExecute(client(coordinatingNode), EsqlQueryAction.INSTANCE, query);
+            assertThat(Iterables.size(result.rows()), equalTo(1L));
+            assertThat(exchanges.get(), lessThanOrEqualTo(1));// 0 if result is populated from coordinating node
+        } finally {
+            coordinatorNodeTransport.clearAllRules();
+        }
+    }
 }

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

@@ -98,15 +98,25 @@ void startComputeOnDataNodes(
         Runnable runOnTaskFailure,
         ActionListener<ComputeResponse> outListener
     ) {
-        final boolean allowPartialResults = configuration.allowPartialResults();
-        DataNodeRequestSender sender = new DataNodeRequestSender(transportService, esqlExecutor, parentTask, allowPartialResults) {
+        new DataNodeRequestSender(
+            transportService,
+            esqlExecutor,
+            parentTask,
+            configuration.allowPartialResults(),
+            configuration.pragmas().maxConcurrentNodesPerCluster()
+        ) {
             @Override
             protected void sendRequest(
                 DiscoveryNode node,
                 List<ShardId> shardIds,
                 Map<Index, AliasFilter> aliasFilters,
                 NodeListener nodeListener
             ) {
+                if (exchangeSource.isFinished()) {
+                    nodeListener.onSkip(true);
+                    return;
+                }
+
                 final AtomicLong pagesFetched = new AtomicLong();
                 var listener = ActionListener.wrap(nodeListener::onResponse, e -> nodeListener.onFailure(e, pagesFetched.get() > 0));
                 final Transport.Connection connection;
@@ -129,7 +139,7 @@ protected void sendRequest(
                     listener.delegateFailureAndWrap((l, unused) -> {
                         final Runnable onGroupFailure;
                         final CancellableTask groupTask;
-                        if (allowPartialResults) {
+                        if (configuration.allowPartialResults()) {
                             try {
                                 groupTask = computeService.createGroupTask(
                                     parentTask,
@@ -152,7 +162,7 @@ protected void sendRequest(
                             final var remoteSink = exchangeService.newRemoteSink(groupTask, childSessionId, transportService, connection);
                             exchangeSource.addRemoteSink(
                                 remoteSink,
-                                allowPartialResults == false,
+                                configuration.allowPartialResults() == false,
                                 pagesFetched::incrementAndGet,
                                 queryPragmas.concurrentExchangeClients(),
                                 computeListener.acquireAvoid()
@@ -184,8 +194,7 @@ protected void sendRequest(
                     })
                 );
             }
-        };
-        sender.startComputeOnDataNodes(
+        }.startComputeOnDataNodes(
             clusterAlias,
             concreteIndices,
             originalIndices,

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

@@ -57,18 +57,27 @@ abstract class DataNodeRequestSender {
     private final Executor esqlExecutor;
     private final CancellableTask rootTask;
     private final boolean allowPartialResults;
+    private final Semaphore concurrentRequests;
     private final ReentrantLock sendingLock = new ReentrantLock();
     private final Queue<ShardId> pendingShardIds = ConcurrentCollections.newQueue();
     private final Map<DiscoveryNode, Semaphore> nodePermits = new HashMap<>();
     private final Map<ShardId, ShardFailure> shardFailures = ConcurrentCollections.newConcurrentMap();
     private final AtomicBoolean changed = new AtomicBoolean();
     private boolean reportedFailure = false; // guarded by sendingLock
+    private volatile boolean skipRemaining = false;
 
-    DataNodeRequestSender(TransportService transportService, Executor esqlExecutor, CancellableTask rootTask, boolean allowPartialResults) {
+    DataNodeRequestSender(
+        TransportService transportService,
+        Executor esqlExecutor,
+        CancellableTask rootTask,
+        boolean allowPartialResults,
+        int concurrentRequests
+    ) {
         this.transportService = transportService;
         this.esqlExecutor = esqlExecutor;
         this.rootTask = rootTask;
         this.allowPartialResults = allowPartialResults;
+        this.concurrentRequests = concurrentRequests > 0 ? new Semaphore(concurrentRequests) : null;
     }
 
     final void startComputeOnDataNodes(
@@ -126,7 +135,7 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
                         || (allowPartialResults == false && shardFailures.values().stream().anyMatch(shardFailure -> shardFailure.fatal))) {
                         reportedFailure = true;
                         reportFailures(computeListener);
-                    } else {
+                    } else if (skipRemaining == false) {
                         for (NodeRequest request : selectNodeRequests(targetShards)) {
                             sendOneNodeRequest(targetShards, computeListener, request);
                         }
@@ -159,6 +168,9 @@ private void sendOneNodeRequest(TargetShards targetShards, ComputeListener compu
         sendRequest(request.node, request.shardIds, request.aliasFilters, new NodeListener() {
             void onAfter(List<DriverProfile> profiles) {
                 nodePermits.get(request.node).release();
+                if (concurrentRequests != null) {
+                    concurrentRequests.release();
+                }
                 trySendingRequestsForPendingShards(targetShards, computeListener);
                 listener.onResponse(profiles);
             }
@@ -187,6 +199,14 @@ public void onFailure(Exception e, boolean receivedData) {
                 }
                 onAfter(List.of());
             }
+
+            @Override
+            public void onSkip(boolean skipRemaining) {
+                if (skipRemaining) {
+                    DataNodeRequestSender.this.skipRemaining = true;
+                }
+                onAfter(List.of());
+            }
         });
     }
 
@@ -196,6 +216,8 @@ interface NodeListener {
         void onResponse(DataNodeComputeResponse response);
 
         void onFailure(Exception e, boolean receivedData);
+
+        void onSkip(boolean skipRemaining);
     }
 
     private static Exception unwrapFailure(Exception e) {
@@ -256,6 +278,7 @@ private List<NodeRequest> selectNodeRequests(TargetShards targetShards) {
         assert sendingLock.isHeldByCurrentThread();
         final Map<DiscoveryNode, List<ShardId>> nodeToShardIds = new HashMap<>();
         final Iterator<ShardId> shardsIt = pendingShardIds.iterator();
+
         while (shardsIt.hasNext()) {
             ShardId shardId = shardsIt.next();
             ShardFailure failure = shardFailures.get(shardId);
@@ -265,31 +288,45 @@ private List<NodeRequest> selectNodeRequests(TargetShards targetShards) {
             }
             TargetShard shard = targetShards.getShard(shardId);
             Iterator<DiscoveryNode> nodesIt = shard.remainingNodes.iterator();
-            DiscoveryNode selectedNode = null;
             while (nodesIt.hasNext()) {
                 DiscoveryNode node = nodesIt.next();
-                if (nodeToShardIds.containsKey(node) || nodePermits.get(node).tryAcquire()) {
+                List<ShardId> pendingRequest = nodeToShardIds.get(node);
+                if (pendingRequest != null) {
+                    pendingRequest.add(shard.shardId);
                     nodesIt.remove();
                     shardsIt.remove();
-                    selectedNode = node;
                     break;
                 }
-            }
-            if (selectedNode != null) {
-                nodeToShardIds.computeIfAbsent(selectedNode, unused -> new ArrayList<>()).add(shard.shardId);
+
+                if (concurrentRequests == null || concurrentRequests.tryAcquire()) {
+                    if (nodePermits.get(node).tryAcquire()) {
+                        pendingRequest = new ArrayList<>();
+                        pendingRequest.add(shard.shardId);
+                        nodeToShardIds.put(node, pendingRequest);
+
+                        nodesIt.remove();
+                        shardsIt.remove();
+
+                        break;
+                    } else if (concurrentRequests != null) {
+                        concurrentRequests.release();
+                    }
+                }
             }
         }
+
         final List<NodeRequest> nodeRequests = new ArrayList<>(nodeToShardIds.size());
-        for (var e : nodeToShardIds.entrySet()) {
-            List<ShardId> shardIds = e.getValue();
+        for (var entry : nodeToShardIds.entrySet()) {
+            var node = entry.getKey();
+            var shardIds = entry.getValue();
             Map<Index, AliasFilter> aliasFilters = new HashMap<>();
             for (ShardId shardId : shardIds) {
                 var aliasFilter = targetShards.getShard(shardId).aliasFilter;
                 if (aliasFilter != null) {
                     aliasFilters.put(shardId.getIndex(), aliasFilter);
                 }
             }
-            nodeRequests.add(new NodeRequest(e.getKey(), shardIds, aliasFilters));
+            nodeRequests.add(new NodeRequest(node, shardIds, aliasFilters));
         }
         return nodeRequests;
     }

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

@@ -55,7 +55,10 @@ public final class QueryPragmas implements Writeable {
      */
     public static final Setting<TimeValue> STATUS_INTERVAL = Setting.timeSetting("status_interval", Driver.DEFAULT_STATUS_INTERVAL);
 
-    public static final Setting<Integer> MAX_CONCURRENT_SHARDS_PER_NODE = Setting.intSetting("max_concurrent_shards_per_node", 10, 1, 100);
+    public static final Setting<Integer> MAX_CONCURRENT_NODES_PER_CLUSTER = //
+        Setting.intSetting("max_concurrent_nodes_per_cluster", -1, -1, 100);
+    public static final Setting<Integer> MAX_CONCURRENT_SHARDS_PER_NODE = //
+        Setting.intSetting("max_concurrent_shards_per_node", 10, 1, 100);
 
     public static final Setting<Boolean> NODE_LEVEL_REDUCTION = Setting.boolSetting("node_level_reduction", true);
 
@@ -122,6 +125,13 @@ public int enrichMaxWorkers() {
         return ENRICH_MAX_WORKERS.get(settings);
     }
 
+    /**
+     * The maximum number of nodes to be queried at once by this query. This is safeguard to avoid overloading the cluster.
+     */
+    public int maxConcurrentNodesPerCluster() {
+        return MAX_CONCURRENT_NODES_PER_CLUSTER.get(settings);
+    }
+
     /**
      * The maximum number of shards can be executed concurrently on a single node by this query. This is a safeguard to avoid
      * opening and holding many shards (equivalent to many file descriptors) or having too many field infos created by a single query.