skip remaining nodes

Author: idegtiarenko

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

@@ -113,7 +113,7 @@ protected void sendRequest(
                 NodeListener nodeListener
             ) {
                 if (exchangeSource.isFinished()) {
-                    nodeListener.onSkip();
+                    nodeListener.onSkip(true);
                     return;
                 }
 

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

@@ -64,6 +64,7 @@ abstract class DataNodeRequestSender {
     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,
@@ -134,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);
                         }
@@ -200,7 +201,10 @@ public void onFailure(Exception e, boolean receivedData) {
             }
 
             @Override
-            public void onSkip() {
+            public void onSkip(boolean skipRemaining) {
+                if (skipRemaining) {
+                    DataNodeRequestSender.this.skipRemaining = true;
+                }
                 onAfter(List.of());
             }
         });
@@ -213,7 +217,7 @@ interface NodeListener {
 
         void onFailure(Exception e, boolean receivedData);
 
-        void onSkip();
+        void onSkip(boolean skipRemaining);
     }
 
     private static Exception unwrapFailure(Exception e) {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plugin/DataNodeRequestSenderTests.java

@@ -299,7 +299,7 @@ public void testLimitConcurrentNodes() {
         AtomicInteger maxConcurrentRequests = new AtomicInteger(0);
         AtomicInteger concurrentRequests = new AtomicInteger(0);
         Queue<NodeRequest> sent = ConcurrentCollections.newQueue();
-        var future = sendRequests(targetShards, randomBoolean(), concurrency, (node, shardIds, aliasFilters, listener) -> {
+        var response = safeGet(sendRequests(targetShards, randomBoolean(), concurrency, (node, shardIds, aliasFilters, listener) -> {
             concurrentRequests.incrementAndGet();
 
             while (true) {
@@ -315,10 +315,39 @@ public void testLimitConcurrentNodes() {
                 concurrentRequests.decrementAndGet();
                 listener.onResponse(new DataNodeComputeResponse(List.of(), Map.of()));
             });
-        });
-        safeGet(future);
+        }));
         assertThat(sent.size(), equalTo(5));
         assertThat(maxConcurrentRequests.get(), equalTo(concurrency));
+        assertThat(response.totalShards, equalTo(5));
+        assertThat(response.successfulShards, equalTo(5));
+        assertThat(response.failedShards, equalTo(0));
+    }
+
+    public void testDoesNotSendMoreRequestsAfterNodeIsSkipped() {
+        var targetShards = List.of(
+            targetShard(shard1, node1),
+            targetShard(shard2, node2),
+            targetShard(shard3, node3),
+            targetShard(shard4, node4),
+            targetShard(shard5, node5)
+        );
+
+        AtomicInteger processed = new AtomicInteger(0);
+        Queue<NodeRequest> sent = ConcurrentCollections.newQueue();
+        var response = safeGet(sendRequests(targetShards, randomBoolean(), 1, (node, shardIds, aliasFilters, listener) -> {
+            sent.add(new NodeRequest(node, shardIds, aliasFilters));
+            runWithDelay(() -> {
+                if (processed.incrementAndGet() == 1) {
+                    listener.onResponse(new DataNodeComputeResponse(List.of(), Map.of()));
+                } else {
+                    listener.onSkip(true);
+                }
+            });
+        }));
+        assertThat(sent.size(), equalTo(2));// onResponse() + onSkip()
+        assertThat(response.totalShards, equalTo(5));
+        assertThat(response.successfulShards, equalTo(5));
+        assertThat(response.failedShards, equalTo(0));
     }
 
     static DataNodeRequestSender.TargetShard targetShard(ShardId shardId, DiscoveryNode... nodes) {