Add task cancellation check in aggregation code paths by kaushalmahi12 · Pull Request #18386 · opensearch-project/OpenSearch

kaushalmahi12 · 2025-05-28T19:00:37Z

Description

This change adds cancellation checks for nested and bucket aggregations. This change will solve the long running cancelled queries stuck in aggregation code paths both at shard and co-ordinator level.

Related Issues

Resolves #15413

Check List

Functionality includes testing.
API changes companion pull request created, if applicable.
Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

github-actions · 2025-05-28T19:51:17Z

❌ Gradle check result for 5fe20f9: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

jainankitk

@kaushalmahi12 - Thanks for raising this PR! This will significantly harden the cancellation mechanism for aggregation queries, especially at the coordinator level. I have added few comments to ensure maintainability and keep performance overhead minimal

jainankitk · 2025-05-28T21:06:12Z

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

                buckets = new LongObjectPagedHashMap<>(grid.buckets.size(), reduceContext.bigArrays());
            }
            for (Object obj : grid.buckets) {
+                checkCancelled(reduceContext);


Does checkingCancelled for every bucket adds too much overhead? How about sampling after specific number of iterations or something?

jainankitk · 2025-05-28T21:07:14Z

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


    @Override
-    public void postCollection() throws IOException {
+    public void postCollection(Runnable checkCancelled) throws IOException {


This is something that might be overridden by subclass in plugins. We need to ensure this is not breaking those

jainankitk · 2025-05-28T21:09:35Z

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

+            BooleanSupplier isRequestCancelled = () -> {
+                if (task instanceof CancellableTask) {
+                    return ((CancellableTask) task).isCancelled();
+                }
+                return false;
+            };


This could be static lambda somewhere? Doesn't have any logic specific to this class?

It is capturing the task instance from the enclosing scope

jainankitk · 2025-05-28T21:14:42Z

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

+        Runnable checkCancelled = () -> {
+            if (isCancelled.getAsBoolean()) {
+                throw new OpenSearchRejectedExecutionException("Search was cancelled while post collection phase");
+            }
+        };


Again, this can be common Runnable for rejecting requests post cancellation

jainankitk · 2025-05-28T21:16:00Z

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

+        if (reduceContext.isTaskCancelled()) {
+            throw new OpenSearchRejectedExecutionException("search is cancelled");
+        }
+        reduceContext.consumeBucketsAndMaybeBreak(0);


Why do we need this as part of checkCancelled?

this triggers the check for memory based CB underneath after every X calls to consumeBucketsAndMaybeBreak

jainankitk · 2025-05-28T21:18:38Z

server/src/main/java/org/opensearch/search/aggregations/bucket/terms/DoubleTerms.java

    @Override
    public InternalAggregation reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
        boolean promoteToDouble = false;
+        checkCancelled(reduceContext);


This check looks redundant to me

jainankitk · 2025-05-28T21:20:23Z

...er/src/main/java/org/opensearch/search/aggregations/bucket/terms/NumericTermsAggregator.java


        @Override
-        void collectZeroDocEntriesIfNeeded(long owningBucketOrd) throws IOException {}
+        void collectZeroDocEntriesIfNeeded(long owningBucketOrd, Runnable checkCancelled) throws IOException {}


Do we really need to pass this checkCancelled runnable around? We should be able to use class level object or something?

jainankitk · 2025-05-28T21:20:36Z

...src/main/java/org/opensearch/search/aggregations/metrics/AbstractInternalHDRPercentiles.java


    @Override
    public AbstractInternalHDRPercentiles reduce(List<InternalAggregation> aggregations, ReduceContext reduceContext) {
+        checkCancelled(reduceContext);


Again, this check is redundant. Seeing this at few other places as well, which we can probably get rid of

kaushalmahi12 · 2025-06-03T20:12:00Z

Closing this PR in favor of : #18426

kaushalmahi12 added 19 commits May 28, 2025 10:33

add in-memory rule processing service

3ecf0f1

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

add missing javadoc

62f9eb2

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

fix generics error

981792d

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

remove stubs

606a836

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

move generic logic to lib

79e075d

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

fix javadoc error

297ca5a

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

fix javadoc error

660623e

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

expose feature level attribute value store init method

1d930b3

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

address comments

1f36b3d

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

use constructors over static methods

84f45ef

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

make member var final in InMemoryRuleProcessingService

d73dbb6

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

add task cancellation check in aggregations

7cdcf08

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

add cancellation checks

c39afed

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

add cancellation check in major agg reduce phases

968e96f

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

fix compilation error

a4e315e

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

add more checks

c498cdf

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

refactor code

b34734b

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

fix UTs

f39d96f

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

add deleted test

9ca1bc4

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

kaushalmahi12 requested review from Bukhtawar, CEHENKLE, anasalkouz, andrross, ashking94, cwperks, dbwiddis, gbbafna, jed326, kotwanikunal and mch2 as code owners May 28, 2025 19:00

kaushalmahi12 requested review from a team, Rishikesh1159, VachaShah, msfroh, owaiskazi19, reta, sachinpkale, saratvemulapalli, shwetathareja and sohami as code owners May 28, 2025 19:00

github-actions bot added bug Something isn't working Search:Aggregations labels May 28, 2025

add CHANGELOG

5fe20f9

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>

kaushalmahi12 added backport 2.x Backport to 2.x branch backport 3.0 labels May 28, 2025

jainankitk reviewed May 28, 2025

View reviewed changes

This was referenced May 22, 2025

[AUTOCUT] Gradle Check Flaky Test Report for SimpleSecureNetty4TransportTests #18071

Open

[AUTOCUT] Gradle Check Flaky Test Report for SecureReactorNetty4HttpServerTransportTests #17486

Closed

kaushalmahi12 closed this Jun 3, 2025

Conversation

kaushalmahi12 commented May 28, 2025

Description

Related Issues

Check List

Uh oh!

github-actions bot commented May 28, 2025

Uh oh!

jainankitk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kaushalmahi12 commented Jun 3, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants