Validate missing shards after the coordinator rewrite (#116382)

The coordinate rewrite can skip searching shards when the query filters
on `@timestamp`, event.ingested  or the _tier field.

We currently check for missing shards across all the indices that are
the query is running against however,  some shards/indices might not
play a role in the query at all after the coordinator rewrite.

This moves the check for missing shards **after** we've run the
coordinator rewrite so we validate only the  shards that will be
searched by the query.

Author: andreidan

docs/changelog/116382.yaml

@@ -0,0 +1,5 @@
+pr: 116382
+summary: Validate missing shards after the coordinator rewrite
+area: Search
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhase.java

@@ -131,7 +131,6 @@ private static boolean assertSearchCoordinationThread() {
     @Override
     public void run() {
         assert assertSearchCoordinationThread();
-        checkNoMissingShards();
         runCoordinatorRewritePhase();
     }
 
@@ -175,7 +174,10 @@ private void runCoordinatorRewritePhase() {
         if (matchedShardLevelRequests.isEmpty()) {
             finishPhase();
         } else {
-            new Round(new GroupShardsIterator<>(matchedShardLevelRequests)).run();
+            GroupShardsIterator<SearchShardIterator> matchingShards = new GroupShardsIterator<>(matchedShardLevelRequests);
+            // verify missing shards only for the shards that we hit for the query
+            checkNoMissingShards(matchingShards);
+            new Round(matchingShards).run();
         }
     }
 
@@ -185,9 +187,9 @@ private void consumeResult(boolean canMatch, ShardSearchRequest request) {
         results.consumeResult(result, () -> {});
     }
 
-    private void checkNoMissingShards() {
+    private void checkNoMissingShards(GroupShardsIterator<SearchShardIterator> shards) {
         assert assertSearchCoordinationThread();
-        doCheckNoMissingShards(getName(), request, shardsIts);
+        doCheckNoMissingShards(getName(), request, shards);
     }
 
     private Map<SendingTarget, List<SearchShardIterator>> groupByNode(GroupShardsIterator<SearchShardIterator> shards) {

server/src/test/java/org/elasticsearch/action/search/CanMatchPreFilterSearchPhaseTests.java

@@ -28,6 +28,7 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.Tuple;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.mapper.DateFieldMapper;
@@ -68,6 +69,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
@@ -77,6 +79,7 @@
 import static org.elasticsearch.core.Types.forciblyCast;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.mockito.Mockito.mock;
 
@@ -1087,6 +1090,137 @@ public void testCanMatchFilteringOnCoordinatorThatCanBeSkippedTsdb() throws Exce
         );
     }
 
+    public void testCanMatchFilteringOnCoordinatorWithMissingShards() throws Exception {
+        // we'll test that we're executing _tier coordinator rewrite for indices (data stream backing or regular) without any @timestamp
+        // or event.ingested fields
+        // for both data stream backing and regular indices we'll have one index in hot and one UNASSIGNED (targeting warm though).
+        // the warm indices will be skipped as our queries will filter based on _tier: hot and the can match phase will not report error the
+        // missing index even if allow_partial_search_results is false (because the warm index would've not been part of the search anyway)
+
+        Map<Index, Settings.Builder> indexNameToSettings = new HashMap<>();
+        ClusterState state = ClusterState.EMPTY_STATE;
+
+        String dataStreamName = randomAlphaOfLengthBetween(10, 20);
+        Index warmDataStreamIndex = new Index(DataStream.getDefaultBackingIndexName(dataStreamName, 1), UUIDs.base64UUID());
+        indexNameToSettings.put(
+            warmDataStreamIndex,
+            settings(IndexVersion.current()).put(IndexMetadata.SETTING_INDEX_UUID, warmDataStreamIndex.getUUID())
+                .put(DataTier.TIER_PREFERENCE, "data_warm,data_hot")
+        );
+        Index hotDataStreamIndex = new Index(DataStream.getDefaultBackingIndexName(dataStreamName, 2), UUIDs.base64UUID());
+        indexNameToSettings.put(
+            hotDataStreamIndex,
+            settings(IndexVersion.current()).put(IndexMetadata.SETTING_INDEX_UUID, hotDataStreamIndex.getUUID())
+                .put(DataTier.TIER_PREFERENCE, "data_hot")
+        );
+        DataStream dataStream = DataStreamTestHelper.newInstance(dataStreamName, List.of(warmDataStreamIndex, hotDataStreamIndex));
+
+        Index warmRegularIndex = new Index("warm-index", UUIDs.base64UUID());
+        indexNameToSettings.put(
+            warmRegularIndex,
+            settings(IndexVersion.current()).put(IndexMetadata.SETTING_INDEX_UUID, warmRegularIndex.getUUID())
+                .put(DataTier.TIER_PREFERENCE, "data_warm,data_hot")
+        );
+        Index hotRegularIndex = new Index("hot-index", UUIDs.base64UUID());
+        indexNameToSettings.put(
+            hotRegularIndex,
+            settings(IndexVersion.current()).put(IndexMetadata.SETTING_INDEX_UUID, hotRegularIndex.getUUID())
+                .put(DataTier.TIER_PREFERENCE, "data_hot")
+        );
+
+        List<Index> allIndices = new ArrayList<>(4);
+        allIndices.addAll(dataStream.getIndices());
+        allIndices.add(warmRegularIndex);
+        allIndices.add(hotRegularIndex);
+
+        List<Index> hotIndices = List.of(hotRegularIndex, hotDataStreamIndex);
+        List<Index> warmIndices = List.of(warmRegularIndex, warmDataStreamIndex);
+
+        for (Index index : allIndices) {
+            IndexMetadata.Builder indexMetadataBuilder = IndexMetadata.builder(index.getName())
+                .settings(indexNameToSettings.get(index))
+                .numberOfShards(1)
+                .numberOfReplicas(0);
+            Metadata.Builder metadataBuilder = Metadata.builder(state.metadata()).put(indexMetadataBuilder);
+            state = ClusterState.builder(state).metadata(metadataBuilder).build();
+        }
+
+        ClusterState finalState = state;
+        CoordinatorRewriteContextProvider coordinatorRewriteContextProvider = new CoordinatorRewriteContextProvider(
+            parserConfig(),
+            mock(Client.class),
+            System::currentTimeMillis,
+            () -> finalState,
+            (index) -> null
+        );
+
+        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery()
+            .filter(QueryBuilders.termQuery(CoordinatorRewriteContext.TIER_FIELD_NAME, "data_hot"));
+
+        {
+            // test that a search doesn't fail if the query filters out the unassigned shards
+            // via _tier (coordinator rewrite will eliminate the shards that don't match)
+            assignShardsAndExecuteCanMatchPhase(
+                List.of(dataStream),
+                List.of(hotRegularIndex, warmRegularIndex),
+                coordinatorRewriteContextProvider,
+                boolQueryBuilder,
+                List.of(),
+                null,
+                warmIndices,
+                false,
+                (updatedSearchShardIterators, requests) -> {
+                    var skippedShards = updatedSearchShardIterators.stream().filter(SearchShardIterator::skip).toList();
+                    var nonSkippedShards = updatedSearchShardIterators.stream()
+                        .filter(searchShardIterator -> searchShardIterator.skip() == false)
+                        .toList();
+
+                    boolean allSkippedShardAreFromWarmIndices = skippedShards.stream()
+                        .allMatch(shardIterator -> warmIndices.contains(shardIterator.shardId().getIndex()));
+                    assertThat(allSkippedShardAreFromWarmIndices, equalTo(true));
+                    boolean allNonSkippedShardAreHotIndices = nonSkippedShards.stream()
+                        .allMatch(shardIterator -> hotIndices.contains(shardIterator.shardId().getIndex()));
+                    assertThat(allNonSkippedShardAreHotIndices, equalTo(true));
+                    boolean allRequestMadeToHotIndices = requests.stream()
+                        .allMatch(request -> hotIndices.contains(request.shardId().getIndex()));
+                    assertThat(allRequestMadeToHotIndices, equalTo(true));
+                }
+            );
+        }
+
+        {
+            // test that a search does fail if the query does NOT filter ALL the
+            // unassigned shards
+            CountDownLatch latch = new CountDownLatch(1);
+            Tuple<CanMatchPreFilterSearchPhase, List<ShardSearchRequest>> canMatchPhaseAndRequests = getCanMatchPhaseAndRequests(
+                List.of(dataStream),
+                List.of(hotRegularIndex, warmRegularIndex),
+                coordinatorRewriteContextProvider,
+                boolQueryBuilder,
+                List.of(),
+                null,
+                List.of(hotRegularIndex, warmRegularIndex, warmDataStreamIndex),
+                false,
+                new ActionListener<>() {
+                    @Override
+                    public void onResponse(GroupShardsIterator<SearchShardIterator> searchShardIterators) {
+                        fail(null, "unexpected success with result [%s] while expecting to handle failure with [%s]", searchShardIterators);
+                        latch.countDown();
+                    }
+
+                    @Override
+                    public void onFailure(Exception e) {
+                        assertThat(e, instanceOf(SearchPhaseExecutionException.class));
+                        latch.countDown();
+                    }
+                }
+            );
+
+            canMatchPhaseAndRequests.v1().start();
+            latch.await(10, TimeUnit.SECONDS);
+        }
+    }
+
     private void assertAllShardsAreQueried(List<SearchShardIterator> updatedSearchShardIterators, List<ShardSearchRequest> requests) {
         int skippedShards = (int) updatedSearchShardIterators.stream().filter(SearchShardIterator::skip).count();
 
@@ -1111,6 +1245,69 @@ private void assignShardsAndExecuteCanMatchPhase(
         SuggestBuilder suggest,
         BiConsumer<List<SearchShardIterator>, List<ShardSearchRequest>> canMatchResultsConsumer
     ) throws Exception {
+        assignShardsAndExecuteCanMatchPhase(
+            dataStreams,
+            regularIndices,
+            contextProvider,
+            query,
+            aggregations,
+            suggest,
+            List.of(),
+            true,
+            canMatchResultsConsumer
+        );
+    }
+
+    private void assignShardsAndExecuteCanMatchPhase(
+        List<DataStream> dataStreams,
+        List<Index> regularIndices,
+        CoordinatorRewriteContextProvider contextProvider,
+        QueryBuilder query,
+        List<AggregationBuilder> aggregations,
+        SuggestBuilder suggest,
+        List<Index> unassignedIndices,
+        boolean allowPartialResults,
+        BiConsumer<List<SearchShardIterator>, List<ShardSearchRequest>> canMatchResultsConsumer
+    ) throws Exception {
+        AtomicReference<GroupShardsIterator<SearchShardIterator>> result = new AtomicReference<>();
+        CountDownLatch latch = new CountDownLatch(1);
+        Tuple<CanMatchPreFilterSearchPhase, List<ShardSearchRequest>> canMatchAndShardRequests = getCanMatchPhaseAndRequests(
+            dataStreams,
+            regularIndices,
+            contextProvider,
+            query,
+            aggregations,
+            suggest,
+            unassignedIndices,
+            allowPartialResults,
+            ActionTestUtils.assertNoFailureListener(iter -> {
+                result.set(iter);
+                latch.countDown();
+            })
+        );
+
+        canMatchAndShardRequests.v1().start();
+        latch.await();
+
+        List<SearchShardIterator> updatedSearchShardIterators = new ArrayList<>();
+        for (SearchShardIterator updatedSearchShardIterator : result.get()) {
+            updatedSearchShardIterators.add(updatedSearchShardIterator);
+        }
+
+        canMatchResultsConsumer.accept(updatedSearchShardIterators, canMatchAndShardRequests.v2());
+    }
+
+    private Tuple<CanMatchPreFilterSearchPhase, List<ShardSearchRequest>> getCanMatchPhaseAndRequests(
+        List<DataStream> dataStreams,
+        List<Index> regularIndices,
+        CoordinatorRewriteContextProvider contextProvider,
+        QueryBuilder query,
+        List<AggregationBuilder> aggregations,
+        SuggestBuilder suggest,
+        List<Index> unassignedIndices,
+        boolean allowPartialResults,
+        ActionListener<GroupShardsIterator<SearchShardIterator>> canMatchActionListener
+    ) {
         Map<String, Transport.Connection> lookup = new ConcurrentHashMap<>();
         DiscoveryNode primaryNode = DiscoveryNodeUtils.create("node_1");
         DiscoveryNode replicaNode = DiscoveryNodeUtils.create("node_2");
@@ -1136,23 +1333,31 @@ private void assignShardsAndExecuteCanMatchPhase(
                 // and none is assigned, the phase is considered as failed meaning that the next phase won't be executed
                 boolean withAssignedPrimaries = randomBoolean() || atLeastOnePrimaryAssigned == false;
                 int numShards = randomIntBetween(1, 6);
-                originalShardIters.addAll(
-                    getShardsIter(dataStreamIndex, originalIndices, numShards, false, withAssignedPrimaries ? primaryNode : null, null)
-                );
-                atLeastOnePrimaryAssigned |= withAssignedPrimaries;
+                if (unassignedIndices.contains(dataStreamIndex)) {
+                    originalShardIters.addAll(getShardsIter(dataStreamIndex, originalIndices, numShards, false, null, null));
+                } else {
+                    originalShardIters.addAll(
+                        getShardsIter(dataStreamIndex, originalIndices, numShards, false, withAssignedPrimaries ? primaryNode : null, null)
+                    );
+                    atLeastOnePrimaryAssigned |= withAssignedPrimaries;
+                }
             }
         }
 
         for (Index regularIndex : regularIndices) {
-            originalShardIters.addAll(
-                getShardsIter(regularIndex, originalIndices, randomIntBetween(1, 6), randomBoolean(), primaryNode, replicaNode)
-            );
+            if (unassignedIndices.contains(regularIndex)) {
+                originalShardIters.addAll(getShardsIter(regularIndex, originalIndices, randomIntBetween(1, 6), false, null, null));
+            } else {
+                originalShardIters.addAll(
+                    getShardsIter(regularIndex, originalIndices, randomIntBetween(1, 6), randomBoolean(), primaryNode, replicaNode)
+                );
+            }
         }
         GroupShardsIterator<SearchShardIterator> shardsIter = GroupShardsIterator.sortAndCreate(originalShardIters);
 
         final SearchRequest searchRequest = new SearchRequest();
         searchRequest.indices(indices);
-        searchRequest.allowPartialSearchResults(true);
+        searchRequest.allowPartialSearchResults(allowPartialResults);
 
         final AliasFilter aliasFilter;
         if (aggregations.isEmpty() == false || randomBoolean()) {
@@ -1212,35 +1417,24 @@ public void sendCanMatch(
         );
 
         AtomicReference<GroupShardsIterator<SearchShardIterator>> result = new AtomicReference<>();
-        CountDownLatch latch = new CountDownLatch(1);
-        CanMatchPreFilterSearchPhase canMatchPhase = new CanMatchPreFilterSearchPhase(
-            logger,
-            searchTransportService,
-            (clusterAlias, node) -> lookup.get(node),
-            aliasFilters,
-            Collections.emptyMap(),
-            threadPool.executor(ThreadPool.Names.SEARCH_COORDINATION),
-            searchRequest,
-            shardsIter,
-            timeProvider,
-            null,
-            true,
-            contextProvider,
-            ActionTestUtils.assertNoFailureListener(iter -> {
-                result.set(iter);
-                latch.countDown();
-            })
+        return new Tuple<>(
+            new CanMatchPreFilterSearchPhase(
+                logger,
+                searchTransportService,
+                (clusterAlias, node) -> lookup.get(node),
+                aliasFilters,
+                Collections.emptyMap(),
+                threadPool.executor(ThreadPool.Names.SEARCH_COORDINATION),
+                searchRequest,
+                shardsIter,
+                timeProvider,
+                null,
+                true,
+                contextProvider,
+                canMatchActionListener
+            ),
+            requests
         );
-
-        canMatchPhase.start();
-        latch.await();
-
-        List<SearchShardIterator> updatedSearchShardIterators = new ArrayList<>();
-        for (SearchShardIterator updatedSearchShardIterator : result.get()) {
-            updatedSearchShardIterators.add(updatedSearchShardIterator);
-        }
-
-        canMatchResultsConsumer.accept(updatedSearchShardIterators, requests);
     }
 
     static class StaticCoordinatorRewriteContextProviderBuilder {

x-pack/plugin/searchable-snapshots/src/internalClusterTest/java/org/elasticsearch/xpack/searchablesnapshots/SearchableSnapshotsCanMatchOnCoordinatorIntegTests.java

@@ -1029,6 +1029,9 @@ public void testCanMatchSkipsPartiallyMountedIndicesWhenFrozenNodesUnavailable()
             TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("_tier", "data_content");
             List<String> indicesToSearch = List.of(regularIndex, partiallyMountedIndex);
             SearchRequest request = new SearchRequest().indices(indicesToSearch.toArray(new String[0]))
+                // we randomise the partial search results because if shards that do NOT match the query are unavailable
+                // the search is not partial
+                .allowPartialSearchResults(randomBoolean())
                 .source(new SearchSourceBuilder().query(termQueryBuilder));
 
             assertResponse(client().search(request), searchResponse -> {
@@ -1045,6 +1048,7 @@ public void testCanMatchSkipsPartiallyMountedIndicesWhenFrozenNodesUnavailable()
             TermsQueryBuilder termsQueryBuilder = QueryBuilders.termsQuery("_tier", "data_hot", "data_content");
             List<String> indicesToSearch = List.of(regularIndex, partiallyMountedIndex);
             SearchRequest request = new SearchRequest().indices(indicesToSearch.toArray(new String[0]))
+                .allowPartialSearchResults(randomBoolean())
                 .source(new SearchSourceBuilder().query(termsQueryBuilder));
 
             assertResponse(client().search(request), searchResponse -> {
@@ -1061,6 +1065,7 @@ public void testCanMatchSkipsPartiallyMountedIndicesWhenFrozenNodesUnavailable()
             BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery().mustNot(QueryBuilders.termQuery("_tier", "data_frozen"));
             List<String> indicesToSearch = List.of(regularIndex, partiallyMountedIndex);
             SearchRequest request = new SearchRequest().indices(indicesToSearch.toArray(new String[0]))
+                .allowPartialSearchResults(randomBoolean())
                 .source(new SearchSourceBuilder().query(boolQueryBuilder));
 
             assertResponse(client().search(request), searchResponse -> {
@@ -1078,6 +1083,7 @@ public void testCanMatchSkipsPartiallyMountedIndicesWhenFrozenNodesUnavailable()
                 .mustNot(randomFrom(QueryBuilders.wildcardQuery("_tier", "dat*ozen"), QueryBuilders.prefixQuery("_tier", "data_fro")));
             List<String> indicesToSearch = List.of(regularIndex, partiallyMountedIndex);
             SearchRequest request = new SearchRequest().indices(indicesToSearch.toArray(new String[0]))
+                .allowPartialSearchResults(randomBoolean())
                 .source(new SearchSourceBuilder().query(boolQueryBuilder));
 
             assertResponse(client().search(request), searchResponse -> {