Limit concurrent node requests (elastic#122850) (elastic#127742)

(cherry picked from commit 8d11dd2)

Author: idegtiarenko

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

@@ -10,15 +10,16 @@
 import org.apache.lucene.tests.util.LuceneTestCase;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.admin.indices.stats.IndicesStatsRequest;
 import org.elasticsearch.action.bulk.BulkRequestBuilder;
-import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.settings.Settings;
 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.
@@ -91,11 +94,11 @@ public void setupIndices() {
                 .setMapping("user", "type=keyword", "tags", "type=keyword")
                 .get();
             BulkRequestBuilder bulk = client().prepareBulk(index).setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
-            int numDocs = between(5, 10);
+            int numDocs = between(10, 25); // every shard has at least 1 doc
             for (int d = 0; d < numDocs; d++) {
                 String user = randomFrom("u1", "u2", "u3");
                 String tag = randomFrom("java", "elasticsearch", "lucene");
-                bulk.add(new IndexRequest().source(Map.of("user", user, "tags", tag)));
+                bulk.add(client().prepareIndex().setSource(Map.of("user", user, "tags", tag)));
             }
             bulk.get();
         }
@@ -257,4 +260,43 @@ 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)) {
+                logger.info("Opening exchange on node [{}]", connection.getNode().getId());
+                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
+        } catch (AssertionError e) {
+            client().admin().indices().stats(new IndicesStatsRequest()).actionGet().asMap().forEach((shard, stats) -> {
+                logger.info(
+                    "Shard {} node {} status {} docs {}",
+                    shard.shardId(),
+                    shard.currentNodeId(),
+                    shard.state(),
+                    stats.getStats().getDocs().getCount()
+                );
+            });
+        } finally {
+            coordinatorNodeTransport.clearAllRules();
+        }
+    }
 }

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

@@ -98,13 +98,13 @@ void startComputeOnDataNodes(
         Runnable runOnTaskFailure,
         ActionListener<ComputeResponse> outListener
     ) {
-        final boolean allowPartialResults = configuration.allowPartialResults();
-        DataNodeRequestSender sender = new DataNodeRequestSender(
+        new DataNodeRequestSender(
             transportService,
             esqlExecutor,
             clusterAlias,
             parentTask,
-            allowPartialResults
+            configuration.allowPartialResults(),
+            configuration.pragmas().maxConcurrentNodesPerCluster()
         ) {
             @Override
             protected void sendRequest(
@@ -113,6 +113,11 @@ protected void sendRequest(
                 Map<Index, AliasFilter> aliasFilters,
                 NodeListener nodeListener
             ) {
+                if (exchangeSource.isFinished()) {
+                    nodeListener.onSkip();
+                    return;
+                }
+
                 final AtomicLong pagesFetched = new AtomicLong();
                 var listener = ActionListener.wrap(nodeListener::onResponse, e -> nodeListener.onFailure(e, pagesFetched.get() > 0));
                 final Transport.Connection connection;
@@ -135,7 +140,7 @@ protected void sendRequest(
                     listener.delegateFailureAndWrap((l, unused) -> {
                         final Runnable onGroupFailure;
                         final CancellableTask groupTask;
-                        if (allowPartialResults) {
+                        if (configuration.allowPartialResults()) {
                             try {
                                 groupTask = computeService.createGroupTask(
                                     parentTask,
@@ -158,7 +163,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()
@@ -190,8 +195,7 @@ protected void sendRequest(
                     })
                 );
             }
-        };
-        sender.startComputeOnDataNodes(
+        }.startComputeOnDataNodes(
             clusterAlias,
             concreteIndices,
             originalIndices,

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

@@ -48,6 +48,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantLock;
 
 /**
@@ -60,10 +61,12 @@ abstract class DataNodeRequestSender {
     private final String clusterAlias;
     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 AtomicInteger skippedShards = new AtomicInteger();
     private final AtomicBoolean changed = new AtomicBoolean();
     private boolean reportedFailure = false; // guarded by sendingLock
 
@@ -72,13 +75,15 @@ abstract class DataNodeRequestSender {
         Executor esqlExecutor,
         String clusterAlias,
         CancellableTask rootTask,
-        boolean allowPartialResults
+        boolean allowPartialResults,
+        int concurrentRequests
     ) {
         this.transportService = transportService;
         this.esqlExecutor = esqlExecutor;
         this.clusterAlias = clusterAlias;
         this.rootTask = rootTask;
         this.allowPartialResults = allowPartialResults;
+        this.concurrentRequests = concurrentRequests > 0 ? new Semaphore(concurrentRequests) : null;
     }
 
     final void startComputeOnDataNodes(
@@ -96,8 +101,8 @@ final void startComputeOnDataNodes(
                     profiles,
                     TimeValue.timeValueNanos(System.nanoTime() - startTimeInNanos),
                     targetShards.totalShards(),
-                    targetShards.totalShards() - shardFailures.size(),
-                    targetShards.skippedShards(),
+                    targetShards.totalShards() - shardFailures.size() - skippedShards.get(),
+                    targetShards.skippedShards() + skippedShards.get(),
                     shardFailures.size(),
                     selectFailures()
                 );
@@ -191,6 +196,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);
             }
@@ -219,6 +227,16 @@ public void onFailure(Exception e, boolean receivedData) {
                 }
                 onAfter(List.of());
             }
+
+            @Override
+            public void onSkip() {
+                skippedShards.incrementAndGet();
+                if (rootTask.isCancelled()) {
+                    onFailure(new TaskCancelledException("null"), true);
+                } else {
+                    onResponse(new DataNodeComputeResponse(List.of(), Map.of()));
+                }
+            }
         });
     }
 
@@ -228,6 +246,8 @@ interface NodeListener {
         void onResponse(DataNodeComputeResponse response);
 
         void onFailure(Exception e, boolean receivedData);
+
+        void onSkip();
     }
 
     private static Exception unwrapFailure(Exception e) {
@@ -288,6 +308,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);
@@ -297,31 +318,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.