Fix missing removal of query cancellation callback in QueryPhase

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

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