Fix missing removal of query cancellation callback in QueryPhase (#130279) (#131809)

When a search targets a single shard, the fetch phase is executed in the same roundtrip as the query phase. In that case the search context is reused across phases, and the cancellation runnables registered to the searcher are reused. This is incorrect as the search timeout check should not be applied to the fetch phase, or it will make fetching partial results unfeasible. For instance if the query phase times out, the fetch phase will time out as well and there won't be partial results. If the query phase does not time out, but the fetch phase does, the latter will return a different set of results compared to those expected which will lead to unexpected situations like an array index out of bound exception.

This commit clears the cancellation runnables ahead of running the fetch phase in the same rountrip as query and performs their registration once again, without the timeout checks which are never applied to the fetch phase.

Closes #130071

Co-authored-by: jessepeixoto <[email protected]>

Author: piergm

docs/changelog/130279.yaml

@@ -0,0 +1,5 @@
+pr: 130279
+summary: Fix missing removal of query cancellation callback in QueryPhase
+area: Search
+type: bug
+issues: [130071]

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

@@ -889,6 +889,17 @@ private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, Cancella
             }
             if (request.numberOfShards() == 1 && (request.source() == null || request.source().rankBuilder() == null)) {
                 // we already have query results, but we can run fetch at the same time
+                // in this case we reuse the search context across search and fetch phase, hence we need to clear the cancellation
+                // checks that were applied by the query phase before running fetch. Note that the timeout checks are not applied
+                // to the fetch phase, while the cancellation checks are.
+                context.searcher().clearQueryCancellations();
+                if (context.lowLevelCancellation()) {
+                    context.searcher().addQueryCancellation(() -> {
+                        if (task != null) {
+                            task.ensureNotCancelled();
+                        }
+                    });
+                }
                 context.addFetchResult();
                 return executeFetchPhase(readerContext, context, afterQueryTime);
             } else {

server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java

@@ -187,6 +187,11 @@ public void close() {
         this.cancellable.clear();
     }
 
+    // clear all registered cancellation callbacks to prevent them from leaking into other phases
+    public void clearQueryCancellations() {
+        this.cancellable.clear();
+    }
+
     public boolean hasCancellations() {
         return this.cancellable.isEnabled();
     }

server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java

@@ -400,6 +400,29 @@ private static void assertSlices(LeafSlice[] slices, int numDocs, int numThreads
         assertThat(sumDocs, equalTo(numDocs));
     }
 
+    public void testClearQueryCancellations() throws IOException {
+        Directory dir = newDirectory();
+        RandomIndexWriter w = new RandomIndexWriter(random(), dir);
+        w.addDocument(new Document());
+        DirectoryReader reader = w.getReader();
+        ContextIndexSearcher searcher = new ContextIndexSearcher(
+            reader,
+            IndexSearcher.getDefaultSimilarity(),
+            IndexSearcher.getDefaultQueryCache(),
+            IndexSearcher.getDefaultQueryCachingPolicy(),
+            true
+        );
+
+        assertFalse(searcher.hasCancellations());
+        searcher.addQueryCancellation(() -> {});
+        assertTrue(searcher.hasCancellations());
+
+        searcher.clearQueryCancellations();
+        assertFalse(searcher.hasCancellations());
+
+        IOUtils.close(reader, w, dir);
+    }
+
     public void testExitableTermsMinAndMax() throws IOException {
         Directory dir = newDirectory();
         IndexWriter w = new IndexWriter(dir, newIndexWriterConfig(null));