[9.0] Aggregations cancellation after collection (#120944)

docs/changelog/120944.yaml

@@ -0,0 +1,6 @@
+pr: 120944
+summary: Aggregations cancellation after collection
+area: Aggregations
+type: bug
+issues:
+ - 108701

modules/aggregations/src/main/java/org/elasticsearch/aggregations/bucket/histogram/AutoDateHistogramAggregator.java

@@ -35,6 +35,7 @@
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.aggregations.support.ValuesSourceConfig;
+import org.elasticsearch.tasks.TaskCancelledException;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -573,7 +574,15 @@ private void rebucket() {
                 long[] mergeMap = new long[Math.toIntExact(oldOrds.size())];
                 bucketOrds = new LongKeyedBucketOrds.FromMany(bigArrays());
                 success = true;
-                for (long owningBucketOrd = 0; owningBucketOrd <= oldOrds.maxOwningBucketOrd(); owningBucketOrd++) {
+                long maxOwning = oldOrds.maxOwningBucketOrd();
+                for (long owningBucketOrd = 0; owningBucketOrd <= maxOwning; owningBucketOrd++) {
+                    /*
+                     * Check for cancelation during this tight loop as it can take a while and the standard
+                     * cancelation checks don't run during the loop. Becuase it's a tight loop.
+                     */
+                    if (context.isCancelled()) {
+                        throw new TaskCancelledException("cancelled");
+                    }
                     LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = oldOrds.ordsEnum(owningBucketOrd);
                     Rounding.Prepared preparedRounding = preparedRoundings[roundingIndexFor(owningBucketOrd)];
                     while (ordsEnum.next()) {

modules/aggregations/src/test/java/org/elasticsearch/aggregations/metric/MatrixStatsAggregatorTests.java

@@ -36,7 +36,7 @@ public void testNoData() throws Exception {
                 MatrixStatsAggregationBuilder aggBuilder = new MatrixStatsAggregationBuilder("my_agg").fields(
                     Collections.singletonList("field")
                 );
-                InternalMatrixStats stats = searchAndReduce(reader, new AggTestConfig(aggBuilder, ft));
+                InternalMatrixStats stats = searchAndReduce(reader, new AggTestConfig(aggBuilder, ft).noReductionCancellation());
                 assertNull(stats.getStats());
                 assertEquals(0L, stats.getDocCount());
             }
@@ -54,7 +54,7 @@ public void testUnmapped() throws Exception {
                 MatrixStatsAggregationBuilder aggBuilder = new MatrixStatsAggregationBuilder("my_agg").fields(
                     Collections.singletonList("bogus")
                 );
-                InternalMatrixStats stats = searchAndReduce(reader, new AggTestConfig(aggBuilder, ft));
+                InternalMatrixStats stats = searchAndReduce(reader, new AggTestConfig(aggBuilder, ft).noReductionCancellation());
                 assertNull(stats.getStats());
                 assertEquals(0L, stats.getDocCount());
             }
@@ -88,7 +88,7 @@ public void testTwoFields() throws Exception {
                 MatrixStatsAggregationBuilder aggBuilder = new MatrixStatsAggregationBuilder("my_agg").fields(
                     Arrays.asList(fieldA, fieldB)
                 );
-                InternalMatrixStats stats = searchAndReduce(reader, new AggTestConfig(aggBuilder, ftA, ftB));
+                InternalMatrixStats stats = searchAndReduce(reader, new AggTestConfig(aggBuilder, ftA, ftB).noReductionCancellation());
                 multiPassStats.assertNearlyEqual(stats);
                 assertTrue(MatrixAggregationInspectionHelper.hasValue(stats));
             }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/BucketsAggregator.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.aggregations.support.AggregationPath;
 import org.elasticsearch.search.sort.SortOrder;
+import org.elasticsearch.tasks.TaskCancelledException;
 
 import java.io.IOException;
 import java.util.AbstractList;
@@ -163,6 +164,10 @@ protected void prepareSubAggs(LongArray ordsToCollect) throws IOException {}
      *         array of ordinals
      */
     protected final IntFunction<InternalAggregations> buildSubAggsForBuckets(LongArray bucketOrdsToCollect) throws IOException {
+        if (context.isCancelled()) {
+            throw new TaskCancelledException("not building sub-aggregations due to task cancellation");
+        }
+
         prepareSubAggs(bucketOrdsToCollect);
         InternalAggregation[][] aggregations = new InternalAggregation[subAggregators.length][];
         for (int i = 0; i < subAggregators.length; i++) {