diff --git a/docs/changelog/130279.yaml b/docs/changelog/130279.yaml new file mode 100644 index 0000000000000..d3d51e5f88bf6 --- /dev/null +++ b/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] diff --git a/server/src/main/java/org/elasticsearch/search/SearchService.java b/server/src/main/java/org/elasticsearch/search/SearchService.java index d485b53e7e409..e29cb4f87fc76 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchService.java +++ b/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 { diff --git a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java b/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java index d64eb663b409c..df9e70ba49e1a 100644 --- a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java +++ b/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(); } diff --git a/server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java b/server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java index fe07cbf8efdfd..ce521dcc48ba5 100644 --- a/server/src/test/java/org/elasticsearch/search/internal/ContextIndexSearcherTests.java +++ b/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));