Add a checkCancelled public method to ContextIndexSearcher

Author: kvanerum

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

@@ -191,6 +191,12 @@ public boolean hasCancellations() {
         return this.cancellable.isEnabled();
     }
 
+    public void checkCancelled() {
+        if (this.cancellable.isEnabled()) {
+            this.cancellable.checkCancelled();
+        }
+    }
+
     public void setAggregatedDfs(AggregatedDfs aggregatedDfs) {
         this.aggregatedDfs = aggregatedDfs;
     }

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 = () -> {