Check for query cancellation during rewrite

With ExitableDirectoryReader in place, check for query cancellation
during QueryPhase#preProcess where the query rewriting takes place.

Follows: #52822

Author: matriv

server/src/main/java/org/elasticsearch/search/query/QueryPhase.java

@@ -46,7 +46,6 @@
 import org.apache.lucene.search.TopFieldDocs;
 import org.apache.lucene.search.TotalHits;
 import org.apache.lucene.search.Weight;
-import org.elasticsearch.action.search.SearchShardTask;
 import org.elasticsearch.common.Booleans;
 import org.elasticsearch.common.CheckedConsumer;
 import org.elasticsearch.common.lucene.Lucene;
@@ -110,7 +109,23 @@ public QueryPhase() {
 
     @Override
     public void preProcess(SearchContext context) {
-        context.preProcess(true);
+        final Runnable cancellation;
+        if (context.lowLevelCancellation()) {
+            cancellation = context.searcher().addQueryCancellation(() -> {
+                if (context.getTask().isCancelled()) {
+                    throw new TaskCancelledException("cancelled");
+                }
+            });
+        } else {
+            cancellation = null;
+        }
+        try {
+            context.preProcess(true);
+        } finally {
+            if (cancellation != null) {
+                context.searcher().removeQueryCancellation(cancellation);
+            }
+        }
     }
 
     @Override
@@ -265,9 +280,8 @@ static boolean executeInternal(SearchContext searchContext) throws QueryPhaseExe
             }
 
             if (searchContext.lowLevelCancellation()) {
-                SearchShardTask task = searchContext.getTask();
                 searcher.addQueryCancellation(() -> {
-                    if (task.isCancelled()) {
+                    if (searchContext.getTask().isCancelled()) {
                         throw new TaskCancelledException("cancelled");
                     }
                 });

server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java

@@ -53,6 +53,8 @@
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
+import org.apache.lucene.search.MultiTermQuery;
+import org.apache.lucene.search.PrefixQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.ScoreDoc;
 import org.apache.lucene.search.Sort;
@@ -75,6 +77,7 @@
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.NumberFieldMapper;
 import org.elasticsearch.index.query.ParsedQuery;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.search.ESToParentBlockJoinQuery;
 import org.elasticsearch.index.shard.IndexShard;
 import org.elasticsearch.index.shard.IndexShardTestCase;
@@ -83,6 +86,7 @@
 import org.elasticsearch.search.internal.ScrollContext;
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.sort.SortAndFormats;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.test.TestSearchContext;
 
 import java.io.IOException;
@@ -825,7 +829,55 @@ public void testMinScore() throws Exception {
 
         reader.close();
         dir.close();
+    }
+
+    public void testCancellationDuringPreprocess() throws IOException {
+        try (Directory dir = newDirectory();
+             RandomIndexWriter w = new RandomIndexWriter(random(), dir, newIndexWriterConfig())) {
+
+            for (int i = 0; i < 10; i++) {
+                Document doc = new Document();
+                doc.add(new StringField("foo", "a".repeat(i), Store.NO));
+                w.addDocument(doc);
+            }
+            w.flush();
+            w.close();
+
+            IndexReader reader = DirectoryReader.open(dir);
+            TestSearchContext context = new TestSearchContextWithRewriteAndCancellation(
+                    null, indexShard, newContextSearcher(reader));
+            PrefixQuery prefixQuery = new PrefixQuery(new Term("foo", "a"));
+            prefixQuery.setRewriteMethod(MultiTermQuery.SCORING_BOOLEAN_REWRITE);
+            context.parsedQuery(new ParsedQuery(prefixQuery));
+            SearchShardTask task = mock(SearchShardTask.class);
+            when(task.isCancelled()).thenReturn(true);
+            context.setTask(task);
+            expectThrows(TaskCancelledException.class, () -> new QueryPhase().preProcess(context));
+
+            reader.close();
+        }
+    }
 
+    private static class TestSearchContextWithRewriteAndCancellation extends TestSearchContext {
+
+        public TestSearchContextWithRewriteAndCancellation(QueryShardContext queryShardContext, IndexShard indexShard,
+                                                           ContextIndexSearcher searcher) {
+            super(queryShardContext, indexShard, searcher);
+        }
+
+        @Override
+        public void preProcess(boolean rewrite) {
+            try {
+                searcher().rewrite(query());
+            } catch (IOException e) {
+                fail("IOException shouldn't be thrown");
+            }
+        }
+
+        @Override
+        public boolean lowLevelCancellation() {
+            return true;
+        }
     }
 
     private static ContextIndexSearcher newContextSearcher(IndexReader reader) throws IOException {