Handle index deletion while querying in ES|QL

Author: dnhatn

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlRetryIT.java

@@ -8,6 +8,9 @@
 package org.elasticsearch.xpack.esql.action;
 
 import org.elasticsearch.action.index.IndexRequestBuilder;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.index.IndexService;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.indices.IndicesService;
@@ -20,11 +23,14 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 import static org.elasticsearch.index.shard.IndexShardTestCase.closeShardNoCheck;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.not;
 
 public class EsqlRetryIT extends AbstractEsqlIntegTestCase {
 
@@ -59,6 +65,37 @@ public void testRetryOnShardFailures() throws Exception {
         }
     }
 
+    public void testQueryWhileDeletingIndices() {
+        populateIndices();
+        CountDownLatch waitForDeletion = new CountDownLatch(1);
+        try {
+            final AtomicBoolean deleted = new AtomicBoolean();
+            for (String node : internalCluster().getNodeNames()) {
+                MockTransportService.getInstance(node)
+                    .addRequestHandlingBehavior(ComputeService.DATA_ACTION_NAME, (handler, request, channel, task) -> {
+                        if (deleted.compareAndSet(false, true)) {
+                            deleteIndexCompletely("log-index-2");
+                            waitForDeletion.countDown();
+                        } else {
+                            assertTrue(waitForDeletion.await(10, TimeUnit.SECONDS));
+                        }
+                        handler.messageReceived(request, channel, task);
+                    });
+            }
+            EsqlQueryRequest request = new EsqlQueryRequest();
+            request.query("FROM log-* | STATS COUNT(timestamp) | LIMIT 1");
+            request.allowPartialResults(true);
+            try (var resp = run(request)) {
+                assertTrue(resp.isPartial());
+                assertThat(EsqlTestUtils.getValuesList(resp).getFirst().getFirst(), equalTo(4L));
+            }
+        } finally {
+            for (String node : internalCluster().getNodeNames()) {
+                MockTransportService.getInstance(node).clearAllRules();
+            }
+        }
+    }
+
     private void populateIndices() {
         internalCluster().ensureAtLeastNumDataNodes(2);
         assertAcked(prepareCreate("log-index-1").setSettings(indexSettings(between(1, 3), 1)).setMapping("timestamp", "type=date"));
@@ -88,4 +125,36 @@ private void closeOrFailShards(String nodeName) throws Exception {
             }
         }
     }
+
+    /**
+     * Deletes the given index and ensures it is completely removed from the cluster state and from all nodes
+     */
+    private void deleteIndexCompletely(String indexName) throws Exception {
+        assertAcked(indicesAdmin().prepareDelete(indexName));
+        String[] nodeNames = internalCluster().getNodeNames();
+        assertBusy(() -> {
+            for (String nodeName : nodeNames) {
+                ClusterState clusterState = internalCluster().getInstance(ClusterService.class, nodeName).state();
+                for (IndexMetadata imd : clusterState.metadata().indicesAllProjects()) {
+                    assertThat(
+                        "Index [" + indexName + "] still exists on the cluster state on [" + nodeName + "]",
+                        imd.getIndex().getName(),
+                        not(equalTo(indexName))
+                    );
+                }
+            }
+            for (String nodeName : nodeNames) {
+                final IndicesService indicesService = internalCluster().getInstance(IndicesService.class, nodeName);
+                for (IndexService indexService : indicesService) {
+                    for (IndexShard indexShard : indexService) {
+                        assertThat(
+                            "Index [" + indexName + "] still exists on node [" + nodeName + "]",
+                            indexShard.shardId().getIndexName(),
+                            not(equalTo(indexName))
+                        );
+                    }
+                }
+            }
+        });
+    }
 }

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

@@ -27,8 +27,10 @@
 import org.elasticsearch.compute.operator.DriverCompletionInfo;
 import org.elasticsearch.compute.operator.FailureCollector;
 import org.elasticsearch.index.Index;
+import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.index.shard.ShardNotFoundException;
 import org.elasticsearch.search.SearchShardTarget;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.tasks.CancellableTask;
@@ -515,15 +517,19 @@ Map<ShardId, List<DiscoveryNode>> resolveShards(Set<ShardId> shardIds) {
         var project = projectResolver.getProjectState(clusterService.state());
         var nodes = Maps.<ShardId, List<DiscoveryNode>>newMapWithExpectedSize(shardIds.size());
         for (var shardId : shardIds) {
-            nodes.put(
-                shardId,
-                project.routingTable()
+            List<DiscoveryNode> allocatedNodes;
+            try {
+                allocatedNodes = project.routingTable()
                     .shardRoutingTable(shardId)
                     .allShards()
                     .filter(shard -> shard.active() && shard.isSearchable())
                     .map(shard -> project.cluster().nodes().get(shard.currentNodeId()))
-                    .toList()
-            );
+                    .toList();
+            } catch (Exception ex) {
+                assert ex instanceof IndexNotFoundException || ex instanceof ShardNotFoundException : new AssertionError(ex);
+                continue;
+            }
+            nodes.put(shardId, allocatedNodes);
         }
         return nodes;
     }