Always create search context for scroll queries (#52078)

We need to either exclude null responses from the scroll search response
or always create a search context for every target shards, although that
scroll query can be written to match_no_docs. Otherwise, we won't find
search_context for subsequent scroll requests.

This commit implements the latter option as it's less error-prone.

Relates #51708

Author: dnhatn

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

@@ -610,7 +610,8 @@ public final ShardSearchRequest buildShardSearchRequest(SearchShardIterator shar
         // if we already received a search result we can inform the shard that it
         // can return a null response if the request rewrites to match none rather
         // than creating an empty response in the search thread pool.
-        shardRequest.canReturnNullResponseIfMatchNoDocs(hasShardResponse.get());
+        // Note that, we have to disable this shortcut for scroll queries.
+        shardRequest.canReturnNullResponseIfMatchNoDocs(hasShardResponse.get() && request.scroll() == null);
         return shardRequest;
     }
 

server/src/main/java/org/elasticsearch/search/SearchService.java

@@ -372,6 +372,7 @@ public void executeQueryPhase(ShardSearchRequest request, SearchShardTask task,
                     if (rewritten.canReturnNullResponseIfMatchNoDocs()
                             && canRewriteToMatchNone(rewritten.source())
                             && rewritten.source().query() instanceof MatchNoneQueryBuilder) {
+                        assert request.scroll() == null : "must always create search context for scroll requests";
                         onMatchNoDocs(context, listener);
                     } else {
                         // fork the execution in the search thread pool and wraps the searcher

server/src/test/java/org/elasticsearch/search/scroll/SearchScrollIT.java

@@ -35,6 +35,7 @@
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.index.query.RangeQueryBuilder;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.sort.FieldSortBuilder;
@@ -53,6 +54,7 @@
 import static org.elasticsearch.index.query.QueryBuilders.termQuery;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoSearchHits;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchHits;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertSearchResponse;
@@ -616,6 +618,43 @@ public void testInvalidScrollKeepAlive() throws IOException {
         assertThat(illegalArgumentException.getMessage(), containsString("Keep alive for scroll (3h) is too large"));
     }
 
+    /**
+     * Ensures that we always create and retain search contexts on every target shards for a scroll request
+     * regardless whether that query can be written to match_no_docs on some target shards or not.
+     */
+    public void testScrollRewrittenToMatchNoDocs() {
+        final int numShards = randomIntBetween(3, 5);
+        assertAcked(
+            client().admin().indices().prepareCreate("test")
+                .setSettings(Settings.builder().put(IndexMetaData.SETTING_NUMBER_OF_SHARDS, numShards))
+                .addMapping("_doc", "created_date", "type=date,format=yyyy-MM-dd"));
+        client().prepareIndex("test", "_doc").setId("1").setSource("created_date", "2020-01-01").get();
+        client().prepareIndex("test", "_doc").setId("2").setSource("created_date", "2020-01-02").get();
+        client().prepareIndex("test", "_doc").setId("3").setSource("created_date", "2020-01-03").get();
+        client().admin().indices().prepareRefresh("test").get();
+        SearchResponse resp = null;
+        try {
+            int totalHits = 0;
+            resp = client().prepareSearch("test")
+                .setQuery(new RangeQueryBuilder("created_date").gte("2020-01-02").lte("2020-01-03"))
+                .setMaxConcurrentShardRequests(randomIntBetween(1, 3)) // sometimes fan out shard requests one by one
+                .setSize(randomIntBetween(1, 2))
+                .setScroll(TimeValue.timeValueMinutes(1))
+                .get();
+            assertNoFailures(resp);
+            while (resp.getHits().getHits().length > 0) {
+                totalHits += resp.getHits().getHits().length;
+                resp = client().prepareSearchScroll(resp.getScrollId()).setScroll(TimeValue.timeValueMinutes(1)).get();
+                assertNoFailures(resp);
+            }
+            assertThat(totalHits, equalTo(2));
+        } finally {
+            if (resp != null && resp.getScrollId() != null) {
+                client().prepareClearScroll().addScrollId(resp.getScrollId()).get();
+            }
+        }
+    }
+
     private void assertToXContentResponse(ClearScrollResponse response, boolean succeed, int numFreed) throws IOException {
         XContentBuilder builder = XContentFactory.jsonBuilder();
         response.toXContent(builder, ToXContent.EMPTY_PARAMS);