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..1c705b011b642 100644 --- a/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java +++ b/server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java @@ -191,6 +191,17 @@ public boolean hasCancellations() { return this.cancellable.isEnabled(); } + /** + * Checks if the current task has been cancelled and throws a {@link TaskCancelledException} if so. + * This method provides plugin developers (e.g., custom runtime fields) with a way to detect + * task cancellation and act accordingly, for example to release resources early. + */ + public void checkCancelled() { + if (this.cancellable.isEnabled()) { + this.cancellable.checkCancelled(); + } + } + public void setAggregatedDfs(AggregatedDfs aggregatedDfs) { this.aggregatedDfs = aggregatedDfs; } diff --git a/server/src/test/java/org/elasticsearch/search/SearchCancellationTests.java b/server/src/test/java/org/elasticsearch/search/SearchCancellationTests.java index aa2e76f512cc8..ebd59f8df27e0 100644 --- a/server/src/test/java/org/elasticsearch/search/SearchCancellationTests.java +++ b/server/src/test/java/org/elasticsearch/search/SearchCancellationTests.java @@ -120,6 +120,25 @@ public void testCancellableCollector() throws IOException { assertThat(totalHits2, equalTo(reader.numDocs())); } + public void testCheckCancelled() throws IOException { + Runnable cancellation = () -> { throw new TaskCancelledException("cancelled"); }; + ContextIndexSearcher searcher = new ContextIndexSearcher( + reader, + IndexSearcher.getDefaultSimilarity(), + IndexSearcher.getDefaultQueryCache(), + IndexSearcher.getDefaultQueryCachingPolicy(), + true + ); + + searcher.checkCancelled(); + + searcher.addQueryCancellation(cancellation); + expectThrows(TaskCancelledException.class, searcher::checkCancelled); + + searcher.removeQueryCancellation(cancellation); + searcher.checkCancelled(); + } + public void testExitableDirectoryReader() throws IOException { AtomicBoolean cancelled = new AtomicBoolean(true); Runnable cancellation = () -> {