Aggregations cancellation after collection

docs/changelog/120944.yaml

@@ -0,0 +1,5 @@
+pr: 120944
+summary: Aggregations cancellation after collection
+area: Aggregations
+type: bug
+issues: []

modules/aggregations/src/internalClusterTest/java/org/elasticsearch/aggregations/bucket/SearchCancellationIT.java

@@ -53,6 +53,7 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
         return List.copyOf(plugins);
     }
 
+    // NOCOMMIT - make sure this test is still verifying what it was intended to verify
     public void testCancellationDuringTimeSeriesAggregation() throws Exception {
         List<ScriptedBlockPlugin> plugins = initBlockFactory();
         int numberOfShards = between(2, 5);

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,11 @@ 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++) {
+                    if (context.isCancelled()) {
+                        throw new TaskCancelledException("cancelled");
+                    }
                     LongKeyedBucketOrds.BucketOrdsEnum ordsEnum = oldOrds.ordsEnum(owningBucketOrd);
                     Rounding.Prepared preparedRounding = preparedRoundings[roundingIndexFor(owningBucketOrd)];
                     while (ordsEnum.next()) {

server/src/internalClusterTest/java/org/elasticsearch/search/SearchCancellationIT.java

@@ -94,6 +94,7 @@ public void testCancellationDuringFetchPhase() throws Exception {
         ensureSearchWasCancelled(searchResponse);
     }
 
+    // NOCOMMIT - make sure this test is still verifying what it was intended to verify
     public void testCancellationDuringAggregation() throws Exception {
         List<ScriptedBlockPlugin> plugins = initBlockFactory();
         // This test is only meaningful with at least 2 shards to trigger reduce

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++) {

test/framework/src/main/java/org/elasticsearch/search/aggregations/AggregatorTestCase.java

@@ -129,6 +129,8 @@
 import org.elasticsearch.search.SearchModule;
 import org.elasticsearch.search.aggregations.AggregatorFactories.Builder;
 import org.elasticsearch.search.aggregations.MultiBucketConsumerService.MultiBucketConsumer;
+import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
+import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregation;
 import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregationBuilder;
 import org.elasticsearch.search.aggregations.metrics.MetricsAggregator;
 import org.elasticsearch.search.aggregations.metrics.MultiValueAggregation;
@@ -149,6 +151,7 @@
 import org.elasticsearch.search.internal.SearchContext;
 import org.elasticsearch.search.internal.ShardSearchRequest;
 import org.elasticsearch.search.internal.SubSearchContext;
+import org.elasticsearch.tasks.TaskCancelledException;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.InternalAggregationTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
@@ -251,29 +254,12 @@ protected List<SearchPlugin> getSearchPlugins() {
         return List.of();
     }
 
-    /**
-     * Deprecated - this will be made private in a future update
-     */
-    @Deprecated
-    protected <A extends Aggregator> A createAggregator(
-        AggregationBuilder aggregationBuilder,
-        IndexReader indexReader,
-        MappedFieldType... fieldTypes
-    ) throws IOException {
-        return createAggregator(aggregationBuilder, createAggregationContext(indexReader, new MatchAllDocsQuery(), fieldTypes));
-    }
-
     protected <A extends Aggregator> A createAggregator(AggregationBuilder aggregationBuilder, AggregationContext context)
         throws IOException {
         return createAggregator(new AggregatorFactories.Builder().addAggregator(aggregationBuilder), context);
     }
 
-    /**
-     * Deprecated - this will be made private in a future update
-     */
-    @Deprecated
-    protected <A extends Aggregator> A createAggregator(AggregatorFactories.Builder builder, AggregationContext context)
-        throws IOException {
+    private <A extends Aggregator> A createAggregator(AggregatorFactories.Builder builder, AggregationContext context) throws IOException {
         Aggregator[] aggregators = builder.build(context, null).createTopLevelAggregators();
         assertThat(aggregators.length, equalTo(1));
         @SuppressWarnings("unchecked")
@@ -310,11 +296,8 @@ protected AggregationContext createAggregationContext(IndexReader indexReader, Q
      * While {@linkplain AggregationContext} is {@link Releasable} the caller is
      * not responsible for releasing it. Instead, it is released automatically in
      * in {@link #cleanupReleasables()}.
-     *
-     * Deprecated - this will be made private in a future update
      */
-    @Deprecated
-    protected AggregationContext createAggregationContext(
+    private AggregationContext createAggregationContext(
         IndexReader indexReader,
         IndexSettings indexSettings,
         Query query,
@@ -346,6 +329,56 @@ private AggregationContext createAggregationContext(
         int maxBucket,
         boolean isInSortOrderExecutionRequired,
         MappedFieldType... fieldTypes
+    ) {
+        return createAggregationContext(
+            searcher,
+            indexSettings,
+            query,
+            breakerService,
+            bytesToPreallocate,
+            maxBucket,
+            isInSortOrderExecutionRequired,
+            () -> false,
+            fieldTypes
+        );
+    }
+
+    /**
+     * Creates an aggregation context that will randomly report that the query has been cancelled
+     */
+    private AggregationContext createCancellingAggregationContext(
+        IndexSearcher searcher,
+        IndexSettings indexSettings,
+        Query query,
+        CircuitBreakerService breakerService,
+        long bytesToPreallocate,
+        int maxBucket,
+        boolean isInSortOrderExecutionRequired,
+        MappedFieldType... fieldTypes
+    ) {
+        return createAggregationContext(
+            searcher,
+            indexSettings,
+            query,
+            breakerService,
+            bytesToPreallocate,
+            maxBucket,
+            isInSortOrderExecutionRequired,
+            () -> ESTestCase.random().nextInt(20) == 0,
+            fieldTypes
+        );
+    }
+
+    private AggregationContext createAggregationContext(
+        IndexSearcher searcher,
+        IndexSettings indexSettings,
+        Query query,
+        CircuitBreakerService breakerService,
+        long bytesToPreallocate,
+        int maxBucket,
+        boolean isInSortOrderExecutionRequired,
+        Supplier<Boolean> isCancelled,
+        MappedFieldType... fieldTypes
     ) {
         MappingLookup mappingLookup = MappingLookup.fromMappers(
             Mapping.EMPTY,
@@ -409,7 +442,7 @@ public Iterable<MappedFieldType> dimensionFields() {
             bitsetFilterCache,
             randomInt(),
             () -> 0L,
-            () -> false,
+            isCancelled,
             q -> q,
             true,
             isInSortOrderExecutionRequired
@@ -536,9 +569,11 @@ protected <A extends InternalAggregation> A searchAndReduce(IndexReader reader,
         IndexSettings indexSettings = createIndexSettings();
         // First run it to find circuit breaker leaks on the aggregator
         runWithCrankyCircuitBreaker(indexSettings, searcher, aggTestConfig);
-        // Second run it to the end
         CircuitBreakerService breakerService = new NoneCircuitBreakerService();
-        return searchAndReduce(indexSettings, searcher, breakerService, aggTestConfig);
+        // Next, try with random cancellations, again looking for leaks
+        runWithCancellingConfig(indexSettings, searcher, breakerService, aggTestConfig);
+        // Finally, run it to the end
+        return searchAndReduce(indexSettings, searcher, breakerService, aggTestConfig, this::createAggregationContext);
     }
 
     /**
@@ -552,20 +587,51 @@ private void runWithCrankyCircuitBreaker(IndexSettings indexSettings, IndexSearc
         CircuitBreakerService crankyService = new CrankyCircuitBreakerService();
         for (int i = 0; i < 5; i++) {
             try {
-                searchAndReduce(indexSettings, searcher, crankyService, aggTestConfig);
+                searchAndReduce(indexSettings, searcher, crankyService, aggTestConfig, this::createAggregationContext);
             } catch (CircuitBreakingException e) {
                 // Circuit breaks from the cranky breaker are expected - it randomly fails, after all
                 assertThat(e.getMessage(), equalTo(CrankyCircuitBreakerService.ERROR_MESSAGE));
             }
         }
     }
 
+    private void runWithCancellingConfig(
+        IndexSettings indexSettings,
+        IndexSearcher searcher,
+        CircuitBreakerService breakerService,
+        AggTestConfig aggTestConfig
+    ) throws IOException {
+        for (int i = 0; i < 5; i++) {
+            try {
+                searchAndReduce(indexSettings, searcher, breakerService, aggTestConfig, this::createCancellingAggregationContext);
+            } catch (TaskCancelledException e) {
+                // we don't want to expectThrows this because the randomizer might just never report cancellation,
+                // but it's also normal that it should throw here.
+            }
+        }
+    }
+
+    @FunctionalInterface
+    public interface AggregationcContextSupplier {
+        AggregationContext get(
+            IndexSearcher searcher,
+            IndexSettings indexSettings,
+            Query query,
+            CircuitBreakerService breakerService,
+            long bytesToPreallocate,
+            int maxBucket,
+            boolean isInSortOrderExecutionRequired,
+            MappedFieldType... fieldTypes
+        );
+    }
+
     @SuppressWarnings("unchecked")
     private <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(
         IndexSettings indexSettings,
         IndexSearcher searcher,
         CircuitBreakerService breakerService,
-        AggTestConfig aggTestConfig
+        AggTestConfig aggTestConfig,
+        AggregationcContextSupplier contextSupplier
     ) throws IOException {
         Query query = aggTestConfig.query();
         AggregatorFactories.Builder builder = new AggregatorFactories.Builder().addAggregator(aggTestConfig.builder());
@@ -591,7 +657,7 @@ private <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(
                 subSearchers[searcherIDX] = new ShardSearcher(leave, compCTX);
             }
             for (ShardSearcher subSearcher : subSearchers) {
-                AggregationContext context = createAggregationContext(
+                AggregationContext context = contextSupplier.get(
                     subSearcher,
                     indexSettings,
                     query,
@@ -620,7 +686,7 @@ private <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(
                 }
             }
         } else {
-            AggregationContext context = createAggregationContext(
+            AggregationContext context = contextSupplier.get(
                 searcher,
                 indexSettings,
                 query,
@@ -688,16 +754,47 @@ private <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(
                 assertRoundTrip(internalAggregation.copyResults());
             }
         }
+        /* Verify that cancellation during final reduce correctly throws.
+         * We check reduce time cancellation only when consuming buckets.
+         */
+        try {
+            // I can't remember if we mutate the InternalAggregations list, so make a defensive copy
+            List<InternalAggregations> internalAggsCopy = new ArrayList<>(internalAggs);
+            A internalAgg = doFinalReduce(maxBucket, bigArraysForReduction, builder, internalAggsCopy, true);
+            if (internalAgg instanceof MultiBucketsAggregation mb) {
+                // Empty mutli-bucket aggs are expected to return before even getting to the cancellation check
+                assertEquals("Got non-empty result for a cancelled reduction", 0, mb.getBuckets().size());
+            } // other cases?
+        } catch (TaskCancelledException e) {
+            /* We may not always honor cancellation in reduce, for example if we are returning no results, so we can't
+             * just expectThrows here.
+             */
+        }
 
         // now do the final reduce
+        A internalAgg = doFinalReduce(maxBucket, bigArraysForReduction, builder, internalAggs, false);
+        assertRoundTrip(internalAgg);
+        if (aggTestConfig.builder instanceof ValuesSourceAggregationBuilder.MetricsAggregationBuilder<?>) {
+            verifyMetricNames((ValuesSourceAggregationBuilder.MetricsAggregationBuilder<?>) aggTestConfig.builder, internalAgg);
+        }
+        return internalAgg;
+    }
+
+    private <A extends InternalAggregation> A doFinalReduce(
+        int maxBucket,
+        BigArrays bigArraysForReduction,
+        Builder builder,
+        List<InternalAggregations> internalAggs,
+        boolean cancelled
+    ) throws IOException {
         MultiBucketConsumer reduceBucketConsumer = new MultiBucketConsumer(
             maxBucket,
             new NoneCircuitBreakerService().getBreaker(CircuitBreaker.REQUEST)
         );
         AggregationReduceContext reduceContext = new AggregationReduceContext.ForFinal(
             bigArraysForReduction,
             getMockScriptService(),
-            () -> false,
+            () -> cancelled,
             builder,
             reduceBucketConsumer
         );
@@ -707,10 +804,6 @@ private <A extends InternalAggregation, C extends Aggregator> A searchAndReduce(
         assertRoundTrip(internalAgg);
 
         doAssertReducedMultiBucketConsumer(internalAgg, reduceBucketConsumer);
-        assertRoundTrip(internalAgg);
-        if (aggTestConfig.builder instanceof ValuesSourceAggregationBuilder.MetricsAggregationBuilder<?>) {
-            verifyMetricNames((ValuesSourceAggregationBuilder.MetricsAggregationBuilder<?>) aggTestConfig.builder, internalAgg);
-        }
         return internalAgg;
     }
 

x-pack/plugin/analytics/build.gradle

@@ -7,6 +7,7 @@
 
 apply plugin: 'elasticsearch.internal-es-plugin'
 apply plugin: 'elasticsearch.internal-cluster-test'
+apply plugin: 'elasticsearch.internal-java-rest-test'
 
 esplugin {
   name = 'x-pack-analytics'
@@ -18,6 +19,10 @@ base {
   archivesName = 'x-pack-analytics'
 }
 
+tasks.named('javaRestTest') {
+  usesDefaultDistribution()
+}
+
 dependencies {
   api 'org.apache.commons:commons-math3:3.6.1'
   compileOnly project(path: xpackModule('core'))