Check for query cancellation during rewrite (#53166)

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

@@ -110,7 +110,24 @@ public QueryPhase() {
 
     @Override
     public void preProcess(SearchContext context) {
-        context.preProcess(true);
+        final Runnable cancellation;
+        if (context.lowLevelCancellation()) {
+            SearchShardTask task = context.getTask();
+            cancellation = context.searcher().addQueryCancellation(() -> {
+                if (task.isCancelled()) {
+                    throw new TaskCancelledException("cancelled");
+                }
+            });
+        } else {
+            cancellation = null;
+        }
+        try {
+            context.preProcess(true);
+        } finally {
+            if (cancellation != null) {
+                context.searcher().removeQueryCancellation(cancellation);
+            }
+        }
     }
 
     @Override

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();
+
+            try (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));
+            }
+        }
+    }
+
+    private static class TestSearchContextWithRewriteAndCancellation extends TestSearchContext {
 
+        private 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 {