handle all shards unassigned

Author: idegtiarenko

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

@@ -195,27 +195,30 @@ private void trySendingRequestsForPendingShards(TargetShards targetShards, Compu
                             var failure = shardFailures.get(shardId);
                             if (failure != null && failure.fatal == false && failure.failure instanceof NoShardAvailableActionException) {
                                 pendingRetries.add(shardId);
-                            } else {
-                                shardFailures.compute(
-                                    shardId,
-                                    (k, v) -> new ShardFailure(
-                                        true,
-                                        v == null ? new NoShardAvailableActionException(shardId, "no shard copies found") : v.failure
-                                    )
-                                );
                             }
                         }
                     }
+                    if (pendingRetries.isEmpty() == false && remainingUnavailableShardResolutionAttempts.decrementAndGet() >= 0) {
+                        for (var entry : resolveShards(pendingRetries).entrySet()) {
+                            targetShards.getShard(entry.getKey()).remainingNodes.addAll(entry.getValue());
+                        }
+                    }
+                    for (ShardId shardId : pendingShardIds) {
+                        if (targetShards.getShard(shardId).remainingNodes.isEmpty()) {
+                            shardFailures.compute(
+                                shardId,
+                                (k, v) -> new ShardFailure(
+                                    true,
+                                    v == null ? new NoShardAvailableActionException(shardId, "no shard copies found") : v.failure
+                                )
+                            );
+                        }
+                    }
                     if (reportedFailure
                         || (allowPartialResults == false && shardFailures.values().stream().anyMatch(shardFailure -> shardFailure.fatal))) {
                         reportedFailure = true;
                         reportFailures(computeListener);
                     } else {
-                        if (pendingRetries.isEmpty() == false && remainingUnavailableShardResolutionAttempts.decrementAndGet() >= 0) {
-                            for (var entry : resolveShards(pendingRetries).entrySet()) {
-                                targetShards.getShard(entry.getKey()).remainingNodes.addAll(entry.getValue());
-                            }
-                        }
                         for (NodeRequest request : selectNodeRequests(targetShards)) {
                             sendOneNodeRequest(targetShards, computeListener, request);
                         }

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

@@ -501,6 +501,46 @@ public void testRetryOnlyMovedShards() {
         assertThat("Must retry only affected shards", resolvedShards, contains(shard2));
     }
 
+    public void testRetryUnassignedShardWithoutPartialResults() {
+        var attempt = new AtomicInteger(0);
+        var future = sendRequests(false, -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(DriverCompletionInfo.EMPTY, Map.of())
+                        : new DataNodeComputeResponse(DriverCompletionInfo.EMPTY, Map.of(shard1, new ShardNotFoundException(shard1)))
+                )
+            )
+
+        );
+        expectThrows(NoShardAvailableActionException.class, containsString("no such shard"), future::actionGet);
+    }
+
+    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(DriverCompletionInfo.EMPTY, Map.of())
+                            : new DataNodeComputeResponse(DriverCompletionInfo.EMPTY, 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));
+    }
+
     static DataNodeRequestSender.TargetShard targetShard(ShardId shardId, DiscoveryNode... nodes) {
         return new DataNodeRequestSender.TargetShard(shardId, new ArrayList<>(Arrays.asList(nodes)), null);
     }