[9.0] Fix missing removal of query cancellation callback in QueryPhase (#130279)

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

@@ -786,6 +786,17 @@ private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, SearchSh
             }
             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

@@ -190,6 +190,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));