Fix DataNodeRequestSender

Author: dnhatn

muted-tests.yml

@@ -364,9 +364,6 @@ tests:
 - class: org.elasticsearch.analysis.common.CommonAnalysisClientYamlTestSuiteIT
   method: test {yaml=analysis-common/40_token_filters/stemmer_override file access}
   issue: https://github.com/elastic/elasticsearch/issues/121625
-- class: org.elasticsearch.xpack.esql.plugin.DataNodeRequestSenderTests
-  method: testDoNotRetryOnRequestLevelFailure
-  issue: https://github.com/elastic/elasticsearch/issues/121966
 - class: org.elasticsearch.xpack.searchablesnapshots.hdfs.SecureHdfsSearchableSnapshotsIT
   issue: https://github.com/elastic/elasticsearch/issues/121967
 - class: org.elasticsearch.action.search.SearchQueryThenFetchAsyncActionTests

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

@@ -34,7 +34,9 @@
 import org.elasticsearch.xpack.esql.action.EsqlSearchShardsAction;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.IdentityHashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -59,6 +61,7 @@ abstract class DataNodeRequestSender {
     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
 
     DataNodeRequestSender(TransportService transportService, Executor esqlExecutor, CancellableTask rootTask, boolean allowPartialResults) {
         this.transportService = transportService;
@@ -108,7 +111,9 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
                     if (changed.compareAndSet(true, false) == false) {
                         break;
                     }
-                    for (ShardId shardId : pendingShardIds) {
+                    final Iterator<ShardId> shardIts = pendingShardIds.iterator();
+                    while (shardIts.hasNext()) {
+                        final ShardId shardId = shardIts.next();
                         if (targetShards.getShard(shardId).remainingNodes.isEmpty()) {
                             shardFailures.compute(
                                 shardId,
@@ -117,12 +122,12 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
                                     v == null ? new NoShardAvailableActionException(shardId, "no shard copies found") : v.failure
                                 )
                             );
+                            shardIts.remove();
                         }
                     }
-                    if (shardFailures.values().stream().anyMatch(shardFailure -> shardFailure.fatal)) {
-                        for (var e : shardFailures.values()) {
-                            computeListener.acquireAvoid().onFailure(e.failure);
-                        }
+                    if (reportedFailure || shardFailures.values().stream().anyMatch(shardFailure -> shardFailure.fatal)) {
+                        reportedFailure = true;
+                        reportFailures(computeListener);
                     } else {
                         var nodeRequests = selectNodeRequests(targetShards);
                         for (NodeRequest request : nodeRequests) {
@@ -138,6 +143,20 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
         }
     }
 
+    private void reportFailures(ComputeListener computeListener) {
+        assert sendingLock.isHeldByCurrentThread();
+        assert reportedFailure;
+        Iterator<ShardFailure> it = shardFailures.values().iterator();
+        Set<Exception> seen = Collections.newSetFromMap(new IdentityHashMap<>());
+        while (it.hasNext()) {
+            ShardFailure failure = it.next();
+            if (seen.add(failure.failure)) {
+                computeListener.acquireAvoid().onFailure(failure.failure);
+            }
+            it.remove();
+        }
+    }
+
     private void sendOneNodeRequest(TargetShards targetShards, ComputeListener computeListener, NodeRequest request) {
         final ActionListener<List<DriverProfile>> listener = computeListener.acquireCompute();
         sendRequest(request.node, request.shardIds, request.aliasFilters, new NodeListener() {
@@ -150,7 +169,7 @@ void onAfter(List<DriverProfile> profiles) {
             @Override
             public void onResponse(DataNodeComputeResponse response) {
                 // remove failures of successful shards
-                for (ShardId shardId : targetShards.shardIds()) {
+                for (ShardId shardId : request.shardIds()) {
                     if (response.shardLevelFailures().containsKey(shardId) == false) {
                         shardFailures.remove(shardId);
                     }