Add task cancellation check in aggregation code paths

CHANGELOG.md

@@ -75,6 +75,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Avoid NPE if on SnapshotInfo if 'shallow' boolean not present ([#18187](https://github.com/opensearch-project/OpenSearch/issues/18187))
 - Fix 'system call filter not installed' caused when network.host: 0.0.0.0 ([#18309](https://github.com/opensearch-project/OpenSearch/pull/18309))
 - Fix MatrixStatsAggregator reuse when mode parameter changes ([#18242](https://github.com/opensearch-project/OpenSearch/issues/18242))
+- Add cancellation check in aggregation code paths at shard and co-ordinator node ([#18386](https://github.com/opensearch-project/OpenSearch/pull/18386))
 
 ### Security
 

benchmarks/src/main/java/org/opensearch/benchmark/search/aggregations/TermsReduceBenchmark.java

@@ -98,10 +98,15 @@ public class TermsReduceBenchmark {
     private final NamedWriteableRegistry namedWriteableRegistry = new NamedWriteableRegistry(searchModule.getNamedWriteables());
     private final SearchPhaseController controller = new SearchPhaseController(
         namedWriteableRegistry,
-        req -> new InternalAggregation.ReduceContextBuilder() {
+        (req, isTaskCancelled) -> new InternalAggregation.ReduceContextBuilder() {
             @Override
             public InternalAggregation.ReduceContext forPartialReduction() {
-                return InternalAggregation.ReduceContext.forPartialReduction(null, null, () -> PipelineAggregator.PipelineTree.EMPTY);
+                return InternalAggregation.ReduceContext.forPartialReduction(
+                    null,
+                    null,
+                    () -> PipelineAggregator.PipelineTree.EMPTY,
+                    isTaskCancelled
+                );
             }
 
             @Override
@@ -114,7 +119,8 @@ public InternalAggregation.ReduceContext forFinalReduction() {
                     null,
                     null,
                     bucketConsumer,
-                    PipelineAggregator.PipelineTree.EMPTY
+                    PipelineAggregator.PipelineTree.EMPTY,
+                    isTaskCancelled
                 );
             }
         }
@@ -229,7 +235,8 @@ public SearchPhaseController.ReducedQueryPhase reduceAggs(TermsList candidateLis
             SearchProgressListener.NOOP,
             namedWriteableRegistry,
             shards.size(),
-            exc -> {}
+            exc -> {},
+            () -> false
         );
         CountDownLatch latch = new CountDownLatch(shards.size());
         for (int i = 0; i < shards.size(); i++) {

modules/aggs-matrix-stats/src/main/java/org/opensearch/search/aggregations/matrix/stats/InternalMatrixStats.java

@@ -251,6 +251,7 @@ public Object getProperty(List<String> path) {
 
     @Override
     public InternalAggregation reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
+        checkCancelled(reduceContext);
         // merge stats across all shards
         List<InternalAggregation> aggs = new ArrayList<>(aggregations);
         aggs.removeIf(p -> ((InternalMatrixStats) p).stats == null);

modules/aggs-matrix-stats/src/test/java/org/opensearch/search/aggregations/matrix/stats/InternalMatrixStatsTests.java

@@ -179,7 +179,8 @@ public void testReduceRandom() {
             bigArrays,
             mockScriptService,
             b -> {},
-            PipelineTree.EMPTY
+            PipelineTree.EMPTY,
+            () -> false
         );
         InternalMatrixStats reduced = (InternalMatrixStats) shardResults.get(0).reduce(shardResults, context);
         multiPassStats.assertNearlyEqual(reduced.getResults());

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/bucket/geogrid/BaseGeoGrid.java

@@ -94,13 +94,15 @@ public List<BaseGeoGridBucket> getBuckets() {
 
     @Override
     public BaseGeoGrid reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
+        checkCancelled(reduceContext);
         LongObjectPagedHashMap<List<BaseGeoGridBucket>> buckets = null;
         for (InternalAggregation aggregation : aggregations) {
             BaseGeoGrid grid = (BaseGeoGrid) aggregation;
             if (buckets == null) {
                 buckets = new LongObjectPagedHashMap<>(grid.buckets.size(), reduceContext.bigArrays());
             }
             for (Object obj : grid.buckets) {
+                checkCancelled(reduceContext);
                 BaseGeoGridBucket bucket = (BaseGeoGridBucket) obj;
                 List<BaseGeoGridBucket> existingBuckets = buckets.get(bucket.hashAsLong());
                 if (existingBuckets == null) {
@@ -114,6 +116,7 @@ public BaseGeoGrid reduce(List<InternalAggregation> aggregations, ReduceContext
         final int size = Math.toIntExact(reduceContext.isFinalReduce() == false ? buckets.size() : Math.min(requiredSize, buckets.size()));
         BucketPriorityQueue<BaseGeoGridBucket> ordered = new BucketPriorityQueue<>(size);
         for (LongObjectPagedHashMap.Cursor<List<BaseGeoGridBucket>> cursor : buckets) {
+            checkCancelled(reduceContext);
             List<BaseGeoGridBucket> sameCellBuckets = cursor.value;
             ordered.insertWithOverflow(reduceBucket(sameCellBuckets, reduceContext));
         }

modules/geo/src/main/java/org/opensearch/geo/search/aggregations/metrics/InternalGeoBounds.java

@@ -111,6 +111,7 @@ public String getWriteableName() {
 
     @Override
     public InternalAggregation reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
+        checkCancelled(reduceContext);
         double top = Double.NEGATIVE_INFINITY;
         double bottom = Double.POSITIVE_INFINITY;
         double posLeft = Double.POSITIVE_INFINITY;
@@ -119,6 +120,7 @@ public InternalAggregation reduce(List<InternalAggregation> aggregations, Reduce
         double negRight = Double.NEGATIVE_INFINITY;
 
         for (InternalAggregation aggregation : aggregations) {
+            checkCancelled(reduceContext);
             InternalGeoBounds bounds = (InternalGeoBounds) aggregation;
 
             if (bounds.top > top) {

modules/geo/src/test/java/org/opensearch/geo/search/aggregations/bucket/geogrid/GeoGridAggregatorTestCase.java

@@ -332,10 +332,10 @@ private void testCase(
         MappedFieldType fieldType = new GeoPointFieldMapper.GeoPointFieldType(FIELD_NAME);
 
         Aggregator aggregator = createAggregator(aggregationBuilder, indexSearcher, fieldType);
-        aggregator.preCollection();
+        aggregator.preCollection(() -> {});
         indexSearcher.search(query, aggregator);
-        aggregator.postCollection();
-        verify.accept((BaseGeoGrid<T>) aggregator.buildTopLevel());
+        aggregator.postCollection(() -> {});
+        verify.accept((BaseGeoGrid<T>) aggregator.buildTopLevel(() -> {}));
 
         indexReader.close();
         directory.close();

modules/parent-join/src/main/java/org/opensearch/join/aggregations/ParentJoinAggregator.java

@@ -128,7 +128,7 @@ public void collect(int docId, long owningBucketOrd) throws IOException {
     }
 
     @Override
-    public void postCollection() throws IOException {
+    public void postCollection(Runnable checkCancelled) throws IOException {
         // Delaying until beforeBuildingBuckets
     }
 
@@ -177,7 +177,7 @@ public float score() {
                 }
             }
         }
-        super.postCollection(); // Run post collection after collecting the sub-aggs
+        super.postCollection(() -> {}); // Run post collection after collecting the sub-aggs
     }
 
     @Override

server/src/internalClusterTest/java/org/opensearch/action/search/TransportSearchIT.java

@@ -682,9 +682,9 @@ public ScoreMode scoreMode() {
         }
 
         @Override
-        public void preCollection() throws IOException {}
+        public void preCollection(Runnable r) throws IOException {}
 
         @Override
-        public void postCollection() throws IOException {}
+        public void postCollection(Runnable r) throws IOException {}
     }
 }

server/src/main/java/org/opensearch/action/search/QueryPhaseResultConsumer.java

@@ -57,6 +57,7 @@
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 
 /**
@@ -78,6 +79,7 @@ public class QueryPhaseResultConsumer extends ArraySearchPhaseResults<SearchPhas
     private final SearchProgressListener progressListener;
     private final ReduceContextBuilder aggReduceContextBuilder;
     private final NamedWriteableRegistry namedWriteableRegistry;
+    private final BooleanSupplier isTaskCancelled;
 
     private final int topNSize;
     private final boolean hasTopDocs;
@@ -99,14 +101,15 @@ public QueryPhaseResultConsumer(
         SearchProgressListener progressListener,
         NamedWriteableRegistry namedWriteableRegistry,
         int expectedResultSize,
-        Consumer<Exception> onPartialMergeFailure
+        Consumer<Exception> onPartialMergeFailure,
+        BooleanSupplier isTaskCancelled
     ) {
         super(expectedResultSize);
         this.executor = executor;
         this.circuitBreaker = circuitBreaker;
         this.controller = controller;
         this.progressListener = progressListener;
-        this.aggReduceContextBuilder = controller.getReduceContext(request);
+        this.aggReduceContextBuilder = controller.getReduceContext(request, isTaskCancelled);
         this.namedWriteableRegistry = namedWriteableRegistry;
         this.topNSize = SearchPhaseController.getTopDocsSize(request);
         this.performFinalReduce = request.isFinalReduce();
@@ -117,6 +120,7 @@ public QueryPhaseResultConsumer(
         this.hasAggs = source != null && source.aggregations() != null;
         int batchReduceSize = (hasAggs || hasTopDocs) ? Math.min(request.getBatchedReduceSize(), expectedResultSize) : expectedResultSize;
         this.pendingMerges = new PendingMerges(batchReduceSize, request.resolveTrackTotalHitsUpTo());
+        this.isTaskCancelled = isTaskCancelled;
     }
 
     @Override

server/src/main/java/org/opensearch/action/search/SearchPhaseController.java

@@ -78,8 +78,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
+import java.util.function.BiFunction;
+import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.function.IntFunction;
 import java.util.stream.Collectors;
 
@@ -92,11 +93,14 @@ public final class SearchPhaseController {
     private static final ScoreDoc[] EMPTY_DOCS = new ScoreDoc[0];
 
     private final NamedWriteableRegistry namedWriteableRegistry;
-    private final Function<SearchSourceBuilder, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder;
+    private final BiFunction<
+        SearchSourceBuilder,
+        BooleanSupplier,
+        InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder;
 
     public SearchPhaseController(
         NamedWriteableRegistry namedWriteableRegistry,
-        Function<SearchSourceBuilder, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder
+        BiFunction<SearchSourceBuilder, BooleanSupplier, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder
     ) {
         this.namedWriteableRegistry = namedWriteableRegistry;
         this.requestToAggReduceContextBuilder = requestToAggReduceContextBuilder;
@@ -744,8 +748,8 @@ public InternalSearchResponse buildResponse(SearchHits hits) {
         }
     }
 
-    InternalAggregation.ReduceContextBuilder getReduceContext(SearchRequest request) {
-        return requestToAggReduceContextBuilder.apply(request.source());
+    InternalAggregation.ReduceContextBuilder getReduceContext(SearchRequest request, BooleanSupplier isTaskCancelled) {
+        return requestToAggReduceContextBuilder.apply(request.source(), isTaskCancelled);
     }
 
     /**
@@ -757,7 +761,8 @@ QueryPhaseResultConsumer newSearchPhaseResults(
         SearchProgressListener listener,
         SearchRequest request,
         int numShards,
-        Consumer<Exception> onPartialMergeFailure
+        Consumer<Exception> onPartialMergeFailure,
+        BooleanSupplier isCancelled
     ) {
         return new QueryPhaseResultConsumer(
             request,
@@ -767,7 +772,8 @@ QueryPhaseResultConsumer newSearchPhaseResults(
             listener,
             namedWriteableRegistry,
             numShards,
-            onPartialMergeFailure
+            onPartialMergeFailure,
+            isCancelled
         );
     }
 

server/src/main/java/org/opensearch/action/search/TransportSearchAction.java

@@ -119,6 +119,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
+import java.util.function.BooleanSupplier;
 import java.util.function.Function;
 import java.util.function.LongSupplier;
 import java.util.stream.Collectors;
@@ -514,6 +515,12 @@ private ActionListener<SearchSourceBuilder> buildRewriteListener(
                 );
             }
             OriginalIndices localIndices = remoteClusterIndices.remove(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY);
+            BooleanSupplier isRequestCancelled = () -> {
+                if (task instanceof CancellableTask) {
+                    return ((CancellableTask) task).isCancelled();
+                }
+                return false;
+            };
             if (remoteClusterIndices.isEmpty()) {
                 executeLocalSearch(
                     task,
@@ -533,7 +540,7 @@ private ActionListener<SearchSourceBuilder> buildRewriteListener(
                         localIndices,
                         remoteClusterIndices,
                         timeProvider,
-                        searchService.aggReduceContextBuilder(searchRequest.source()),
+                        searchService.aggReduceContextBuilder(searchRequest.source(), isRequestCancelled),
                         remoteClusterService,
                         threadPool,
                         listener,
@@ -1265,7 +1272,8 @@ private AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction
                 task.getProgressListener(),
                 searchRequest,
                 shardIterators.size(),
-                exc -> cancelTask(task, exc)
+                exc -> cancelTask(task, exc),
+                task::isCancelled
             );
             AbstractSearchAsyncAction<? extends SearchPhaseResult> searchAsyncAction;
             switch (searchRequest.searchType()) {

server/src/main/java/org/opensearch/search/DefaultSearchContext.java

@@ -114,7 +114,8 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Executor;
-import java.util.function.Function;
+import java.util.function.BiFunction;
+import java.util.function.BooleanSupplier;
 import java.util.function.LongSupplier;
 
 import static org.opensearch.search.SearchService.AGGREGATION_REWRITE_FILTER_SEGMENT_THRESHOLD;
@@ -206,7 +207,10 @@ final class DefaultSearchContext extends SearchContext {
     private final Map<Class<?>, CollectorManager<? extends Collector, ReduceableSearchResult>> queryCollectorManagers = new HashMap<>();
     private final QueryShardContext queryShardContext;
     private final FetchPhase fetchPhase;
-    private final Function<SearchSourceBuilder, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder;
+    private final BiFunction<
+        SearchSourceBuilder,
+        BooleanSupplier,
+        InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder;
     private final String concurrentSearchMode;
     private final SetOnce<Boolean> requestShouldUseConcurrentSearch = new SetOnce<>();
     private final int maxAggRewriteFilters;
@@ -227,7 +231,7 @@ final class DefaultSearchContext extends SearchContext {
         Version minNodeVersion,
         boolean validate,
         Executor executor,
-        Function<SearchSourceBuilder, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder,
+        BiFunction<SearchSourceBuilder, BooleanSupplier, InternalAggregation.ReduceContextBuilder> requestToAggReduceContextBuilder,
         Collection<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories
     ) throws IOException {
         this.readerContext = readerContext;
@@ -1042,7 +1046,8 @@ public ReaderContext readerContext() {
 
     @Override
     public InternalAggregation.ReduceContext partialOnShard() {
-        InternalAggregation.ReduceContext rc = requestToAggReduceContextBuilder.apply(request.source()).forPartialReduction();
+        InternalAggregation.ReduceContext rc = requestToAggReduceContextBuilder.apply(request.source(), this::isCancelled)
+            .forPartialReduction();
         rc.setSliceLevel(shouldUseConcurrentSearch());
         return rc;
     }

server/src/main/java/org/opensearch/search/SearchService.java

@@ -164,6 +164,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.BooleanSupplier;
 import java.util.function.LongSupplier;
 
 import static org.opensearch.common.unit.TimeValue.timeValueHours;
@@ -1804,14 +1805,18 @@ public IndicesService getIndicesService() {
      * Returns a builder for {@link InternalAggregation.ReduceContext}. This
      * builder retains a reference to the provided {@link SearchSourceBuilder}.
      */
-    public InternalAggregation.ReduceContextBuilder aggReduceContextBuilder(SearchSourceBuilder searchSourceBuilder) {
+    public InternalAggregation.ReduceContextBuilder aggReduceContextBuilder(
+        SearchSourceBuilder searchSourceBuilder,
+        BooleanSupplier isRequestCancelled
+    ) {
         return new InternalAggregation.ReduceContextBuilder() {
             @Override
             public InternalAggregation.ReduceContext forPartialReduction() {
                 return InternalAggregation.ReduceContext.forPartialReduction(
                     bigArrays,
                     scriptService,
-                    () -> requestToPipelineTree(searchSourceBuilder)
+                    () -> requestToPipelineTree(searchSourceBuilder),
+                    isRequestCancelled
                 );
             }
 
@@ -1822,7 +1827,8 @@ public ReduceContext forFinalReduction() {
                     bigArrays,
                     scriptService,
                     multiBucketConsumerService.create(),
-                    pipelineTree
+                    pipelineTree,
+                    isRequestCancelled
                 );
             }
         };

server/src/main/java/org/opensearch/search/aggregations/AggregationCollectorManager.java

@@ -56,7 +56,8 @@ public String getCollectorReason() {
 
     @Override
     public ReduceableSearchResult reduce(Collection<Collector> collectors) throws IOException {
-        final List<InternalAggregation> internals = context.bucketCollectorProcessor().toInternalAggregations(collectors);
+        final List<InternalAggregation> internals = context.bucketCollectorProcessor()
+            .toInternalAggregations(collectors, context::isCancelled);
         assert internals.stream().noneMatch(Objects::isNull);
         context.aggregations().resetBucketMultiConsumer();
 
@@ -70,7 +71,7 @@ protected AggregationReduceableSearchResult buildAggregationResult(InternalAggre
 
     static Collector createCollector(List<Aggregator> collectors) throws IOException {
         Collector collector = MultiBucketCollector.wrap(collectors);
-        ((BucketCollector) collector).preCollection();
+        ((BucketCollector) collector).preCollection(() -> {});
         return collector;
     }
 }

server/src/main/java/org/opensearch/search/aggregations/Aggregator.java

@@ -200,8 +200,9 @@ public interface BucketComparator {
      * of another aggregation then the aggregation that contains
      * it will call {@link #buildAggregations(long[])}.
      */
-    public final InternalAggregation buildTopLevel() throws IOException {
+    public final InternalAggregation buildTopLevel(Runnable checkCancelled) throws IOException {
         assert parent() == null;
+        checkCancelled.run();
         this.internalAggregation.set(buildAggregations(new long[] { 0 })[0]);
         return internalAggregation.get();
     }