Fix race condition when resolving new location for multiple shards at once (elastic#128062)

(cherry picked from commit f4b6086)

Author: idegtiarenko

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

@@ -144,8 +144,7 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
                     var pendingRetries = new HashSet<ShardId>();
                     for (ShardId shardId : pendingShardIds) {
                         if (targetShards.getShard(shardId).remainingNodes.isEmpty()) {
-                            var failure = shardFailures.get(shardId);
-                            if (failure != null && failure.fatal == false && failure.failure instanceof NoShardAvailableActionException) {
+                            if (isRetryableFailure(shardFailures.get(shardId))) {
                                 pendingRetries.add(shardId);
                             }
                         }
@@ -156,7 +155,8 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
                         }
                     }
                     for (ShardId shardId : pendingShardIds) {
-                        if (targetShards.getShard(shardId).remainingNodes.isEmpty()) {
+                        if (targetShards.getShard(shardId).remainingNodes.isEmpty()
+                            && (isRetryableFailure(shardFailures.get(shardId)) == false || pendingRetries.contains(shardId))) {
                             shardFailures.compute(
                                 shardId,
                                 (k, v) -> new ShardFailure(
@@ -327,6 +327,10 @@ record NodeRequest(DiscoveryNode node, List<ShardId> shardIds, Map<Index, AliasF
 
     private record ShardFailure(boolean fatal, Exception failure) {}
 
+    private static boolean isRetryableFailure(ShardFailure failure) {
+        return failure != null && failure.fatal == false && failure.failure instanceof NoShardAvailableActionException;
+    }
+
     /**
      * Selects the next nodes to send requests to. Limits to at most one outstanding request per node.
      * If there is already a request in-flight to a node, another request will not be sent to the same node

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

@@ -51,8 +51,8 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
-import java.util.stream.Collectors;
 
+import static java.util.stream.Collectors.toMap;
 import static org.elasticsearch.xpack.esql.plugin.DataNodeRequestSender.NodeRequest;
 import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.containsString;
@@ -402,6 +402,32 @@ public void testRetryMovedShard() {
         assertThat(attempt.get(), equalTo(3));
     }
 
+    public void testRetryMultipleMovedShards() {
+        var attempt = new AtomicInteger(0);
+        var response = safeGet(
+            sendRequests(
+                randomBoolean(),
+                -1,
+                List.of(targetShard(shard1, node1), targetShard(shard2, node2), targetShard(shard3, node3)),
+                shardIds -> shardIds.stream().collect(toMap(Function.identity(), shardId -> List.of(randomFrom(node1, node2, node3)))),
+                (node, shardIds, aliasFilters, listener) -> runWithDelay(
+                    () -> listener.onResponse(
+                        attempt.incrementAndGet() <= 6
+                            ? new DataNodeComputeResponse(
+                                List.of(),
+                                shardIds.stream().collect(toMap(Function.identity(), ShardNotFoundException::new))
+                            )
+                            : new DataNodeComputeResponse(List.of(), Map.of())
+                    )
+                )
+            )
+        );
+        assertThat(response.totalShards, equalTo(3));
+        assertThat(response.successfulShards, equalTo(3));
+        assertThat(response.skippedShards, equalTo(0));
+        assertThat(response.failedShards, equalTo(0));
+    }
+
     public void testDoesNotRetryMovedShardIndefinitely() {
         var attempt = new AtomicInteger(0);
         var response = safeGet(sendRequests(true, -1, List.of(targetShard(shard1, node1)), shardIds -> {
@@ -463,28 +489,28 @@ public void testRetryUnassignedShardWithoutPartialResults() {
 
         );
         expectThrows(NoShardAvailableActionException.class, containsString("no such shard"), future::actionGet);
+        assertThat(attempt.get(), equalTo(1));
     }
 
     public void testRetryUnassignedShardWithPartialResults() {
-        var response = safeGet(
-            sendRequests(
-                true,
-                -1,
-                List.of(targetShard(shard1, node1), targetShard(shard2, node2)),
-                shardIds -> Map.of(shard1, List.of()),
-                (node, shardIds, aliasFilters, listener) -> runWithDelay(
-                    () -> listener.onResponse(
-                        Objects.equals(shardIds, List.of(shard2))
-                            ? new DataNodeComputeResponse(List.of(), Map.of())
-                            : new DataNodeComputeResponse(List.of(), Map.of(shard1, new ShardNotFoundException(shard1)))
-                    )
+        var attempt = new AtomicInteger(0);
+        var response = safeGet(sendRequests(true, -1, List.of(targetShard(shard1, node1), targetShard(shard2, node2)), shardIds -> {
+            attempt.incrementAndGet();
+            return Map.of(shard1, List.of());
+        },
+            (node, shardIds, aliasFilters, listener) -> runWithDelay(
+                () -> listener.onResponse(
+                    Objects.equals(shardIds, List.of(shard2))
+                        ? new DataNodeComputeResponse(List.of(), Map.of())
+                        : new DataNodeComputeResponse(List.of(), Map.of(shard1, new ShardNotFoundException(shard1)))
                 )
             )
-        );
+        ));
         assertThat(response.totalShards, equalTo(2));
         assertThat(response.successfulShards, equalTo(1));
         assertThat(response.skippedShards, equalTo(0));
         assertThat(response.failedShards, equalTo(1));
+        assertThat(attempt.get(), equalTo(1));
     }
 
     static DataNodeRequestSender.TargetShard targetShard(ShardId shardId, DiscoveryNode... nodes) {
@@ -553,11 +579,7 @@ PlainActionFuture<ComputeResponse> sendRequests(
             void searchShards(Set<String> concreteIndices, ActionListener<TargetShards> listener) {
                 runWithDelay(
                     () -> listener.onResponse(
-                        new TargetShards(
-                            shards.stream().collect(Collectors.toMap(TargetShard::shardId, Function.identity())),
-                            shards.size(),
-                            0
-                        )
+                        new TargetShards(shards.stream().collect(toMap(TargetShard::shardId, Function.identity())), shards.size(), 0)
                     )
                 );
             }