Skip search shards with INDEX_REFRESH_BLOCK

server/src/internalClusterTest/java/org/elasticsearch/search/SearchWithIndexBlocksIT.java

@@ -0,0 +1,153 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search;
+
+import org.elasticsearch.action.index.IndexRequestBuilder;
+import org.elasticsearch.action.search.ClosePointInTimeRequest;
+import org.elasticsearch.action.search.OpenPointInTimeRequest;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.TransportClosePointInTimeAction;
+import org.elasticsearch.action.search.TransportOpenPointInTimeAction;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.block.ClusterBlocks;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.metadata.ProjectId;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.search.builder.PointInTimeBuilder;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.test.ESIntegTestCase;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+
+import static org.elasticsearch.cluster.block.ClusterBlocks.EMPTY_CLUSTER_BLOCK;
+import static org.elasticsearch.test.ClusterServiceUtils.setState;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
+
+public class SearchWithIndexBlocksIT extends ESIntegTestCase {
+
+    public void testSearchIndicesWithIndexRefreshBlocks() {
+        List<String> indices = createIndices();
+        Map<String, Integer> numDocsPerIndex = indexDocuments(indices);
+        List<String> unblockedIndices = addIndexRefreshBlockToSomeIndices(indices);
+
+        int expectedHits = 0;
+        for (String index : unblockedIndices) {
+            expectedHits += numDocsPerIndex.get(index);
+        }
+
+        assertHitCount(prepareSearch().setQuery(QueryBuilders.matchAllQuery()), expectedHits);
+    }
+
+    public void testOpenPITOnIndicesWithIndexRefreshBlocks() {
+        List<String> indices = createIndices();
+        Map<String, Integer> numDocsPerIndex = indexDocuments(indices);
+        List<String> unblockedIndices = addIndexRefreshBlockToSomeIndices(indices);
+
+        int expectedHits = 0;
+        for (String index : unblockedIndices) {
+            expectedHits += numDocsPerIndex.get(index);
+        }
+
+        BytesReference pitId = null;
+        try {
+            OpenPointInTimeRequest openPITRequest = new OpenPointInTimeRequest(indices.toArray(new String[0])).keepAlive(
+                TimeValue.timeValueSeconds(10)
+            ).allowPartialSearchResults(true);
+            pitId = client().execute(TransportOpenPointInTimeAction.TYPE, openPITRequest).actionGet().getPointInTimeId();
+            SearchRequest searchRequest = new SearchRequest().source(
+                new SearchSourceBuilder().pointInTimeBuilder(new PointInTimeBuilder(pitId).setKeepAlive(TimeValue.timeValueSeconds(10)))
+            );
+            assertHitCount(client().search(searchRequest), expectedHits);
+        } finally {
+            if (pitId != null) {
+                client().execute(TransportClosePointInTimeAction.TYPE, new ClosePointInTimeRequest(pitId)).actionGet();
+            }
+        }
+    }
+
+    public void testMultiSearchIndicesWithIndexRefreshBlocks() {
+        List<String> indices = createIndices();
+        Map<String, Integer> numDocsPerIndex = indexDocuments(indices);
+        List<String> unblockedIndices = addIndexRefreshBlockToSomeIndices(indices);
+
+        int expectedHits = 0;
+        for (String index : unblockedIndices) {
+            expectedHits += numDocsPerIndex.get(index);
+        }
+
+        final long expectedHitsL = expectedHits;
+        assertResponse(
+            client().prepareMultiSearch()
+                .add(prepareSearch().setQuery(QueryBuilders.matchAllQuery()))
+                .add(prepareSearch().setQuery(QueryBuilders.termQuery("field", "blah"))),
+            response -> {
+                assertHitCount(Objects.requireNonNull(response.getResponses()[0].getResponse()), expectedHitsL);
+                assertHitCount(Objects.requireNonNull(response.getResponses()[1].getResponse()), 0);
+            }
+        );
+    }
+
+    private List<String> createIndices() {
+        int numIndices = randomIntBetween(1, 3);
+        List<String> indices = new ArrayList<>();
+        for (int i = 0; i < numIndices; i++) {
+            indices.add("test" + i);
+            createIndex("test" + i);
+        }
+        return indices;
+    }
+
+    private Map<String, Integer> indexDocuments(List<String> indices) {
+        Map<String, Integer> numDocsPerIndex = new HashMap<>();
+        List<IndexRequestBuilder> indexRequests = new ArrayList<>();
+        for (String index : indices) {
+            int numDocs = randomIntBetween(0, 10);
+            numDocsPerIndex.put(index, numDocs);
+            for (int i = 0; i < numDocs; i++) {
+                indexRequests.add(prepareIndex(index).setId(String.valueOf(i)).setSource("field", "value"));
+            }
+        }
+        indexRandom(true, indexRequests);
+
+        return numDocsPerIndex;
+    }
+
+    private List<String> addIndexRefreshBlockToSomeIndices(List<String> indices) {
+        List<String> unblockedIndices = new ArrayList<>();
+        var blocksBuilder = ClusterBlocks.builder().blocks(EMPTY_CLUSTER_BLOCK);
+        for (String index : indices) {
+            boolean blockIndex = randomBoolean();
+            if (blockIndex) {
+                blocksBuilder.addIndexBlock(ProjectId.DEFAULT, index, IndexMetadata.INDEX_REFRESH_BLOCK);
+            } else {
+                unblockedIndices.add(index);
+            }
+        }
+
+        var dataNodes = clusterService().state().getNodes().getAllNodes();
+        for (DiscoveryNode dataNode : dataNodes) {
+            ClusterService clusterService = internalCluster().getInstance(ClusterService.class, dataNode.getName());
+            ClusterState currentState = clusterService.state();
+            ClusterState newState = ClusterState.builder(currentState).blocks(blocksBuilder).build();
+            setState(clusterService, newState);
+        }
+
+        return unblockedIndices;
+    }
+}

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

@@ -589,7 +589,7 @@ public void onFailure(Exception e) {}
         );
     }
 
-    static void adjustSearchType(SearchRequest searchRequest, boolean singleShard) {
+    static void adjustSearchType(SearchRequest searchRequest, boolean oneOrZeroShards) {
         // if there's a kNN search, always use DFS_QUERY_THEN_FETCH
         if (searchRequest.hasKnnSearch()) {
             searchRequest.searchType(DFS_QUERY_THEN_FETCH);
@@ -604,7 +604,7 @@ static void adjustSearchType(SearchRequest searchRequest, boolean singleShard) {
         }
 
         // optimize search type for cases where there is only one shard group to search on
-        if (singleShard) {
+        if (oneOrZeroShards) {
             // if we only have one group, then we always want Q_T_F, no need for DFS, and no need to do THEN since we hit one shard
             searchRequest.searchType(QUERY_THEN_FETCH);
         }
@@ -1305,7 +1305,7 @@ private void executeSearch(
 
         Map<String, Float> concreteIndexBoosts = resolveIndexBoosts(searchRequest, projectState.cluster());
 
-        adjustSearchType(searchRequest, shardIterators.size() == 1);
+        adjustSearchType(searchRequest, shardIterators.size() <= 1);
 
         final DiscoveryNodes nodes = projectState.cluster().nodes();
         BiFunction<String, String, Transport.Connection> connectionLookup = buildConnectionLookup(
@@ -1864,6 +1864,7 @@ List<SearchShardIterator> getLocalShardsIterator(
         Set<ResolvedExpression> indicesAndAliases,
         String[] concreteIndices
     ) {
+        concreteIndices = ignoreBlockedIndices(projectState, concreteIndices);
         var routingMap = indexNameExpressionResolver.resolveSearchRouting(
             projectState.metadata(),
             searchRequest.routing(),
@@ -1896,6 +1897,20 @@ List<SearchShardIterator> getLocalShardsIterator(
         return Arrays.asList(list);
     }
 
+    static String[] ignoreBlockedIndices(ProjectState projectState, String[] concreteIndices) {
+        // optimization: mostly we do not have any blocks so there's no point in the expensive per-index checking
+        boolean hasIndexBlocks = projectState.blocks().indices(projectState.projectId()).isEmpty() == false;
+        if (hasIndexBlocks) {
+            return Arrays.stream(concreteIndices)
+                .filter(
+                    index -> projectState.blocks()
+                        .hasIndexBlock(projectState.projectId(), index, IndexMetadata.INDEX_REFRESH_BLOCK) == false
+                )
+                .toArray(String[]::new);
+        }
+        return concreteIndices;
+    }
+
     private interface TelemetryListener {
         void setRemotes(int count);
 

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

@@ -1812,4 +1812,35 @@ public void onFailure(Exception ex) {
             assertTrue(ESTestCase.terminate(threadPool));
         }
     }
+
+    public void testIgnoreBlockedIndices() {
+        int numIndices = randomIntBetween(1, 10);
+        String[] concreteIndices = new String[numIndices];
+        for (int i = 0; i < numIndices; i++) {
+            concreteIndices[i] = "index" + i;
+        }
+
+        List<String> shuffledIndices = Arrays.asList(concreteIndices);
+        Collections.shuffle(shuffledIndices, random());
+        concreteIndices = shuffledIndices.toArray(new String[0]);
+
+        final ProjectId projectId = randomProjectIdOrDefault();
+        ClusterBlocks.Builder blocksBuilder = ClusterBlocks.builder();
+        int numBlockedIndices = randomIntBetween(0, numIndices);
+        for (int i = 0; i < numBlockedIndices; i++) {
+            blocksBuilder.addIndexBlock(projectId, concreteIndices[i], IndexMetadata.INDEX_REFRESH_BLOCK);
+        }
+        final ClusterState clusterState = ClusterState.builder(ClusterName.DEFAULT)
+            .putProjectMetadata(ProjectMetadata.builder(projectId).build())
+            .blocks(blocksBuilder)
+            .build();
+        final ProjectState projectState = clusterState.projectState(projectId);
+
+        String[] actual = TransportSearchAction.ignoreBlockedIndices(projectState, concreteIndices);
+        String[] expected = Arrays.stream(concreteIndices)
+            .filter(index -> clusterState.blocks().hasIndexBlock(projectId, index, IndexMetadata.INDEX_REFRESH_BLOCK) == false)
+            .toArray(String[]::new);
+
+        assertThat(Arrays.asList(actual), containsInAnyOrder(expected));
+    }
 }