[sandbox] Fix AnalyticsQueryTask onCancel callback race

[bowenlan-amzn]

Description

When the analytics-engine transport action runs on the coordinator, there's a time gap between when the task becomes available to cancel and when the cancel callback is wired up:

1. TransportAction.execute:101 — taskManager.register(...) puts AnalyticsQueryTask into cancellableTasks. From this moment on, _tasks/_cancel can find and fire it.
2. DefaultPlanExecutor.doExecute:414 — searchExecutor.execute(...) queues the lambda containing planning + scheduler setup; the transport thread returns immediately.
3. The lambda is picked up by a SEARCH worker, runs Calcite planning, builds the DAG, and finally hits QueryScheduler.setCancellationCallback:139 which calls setOnCancelCallback(...).

The gap (1) → (3) is determined by SEARCH pool queue depth and Calcite planning latency (cold path with schema fetch). Sub-millisecond on a quiet cluster, noticeably reachable under load.

A cancel landing in that gap — server-side timeout, HTTP disconnect via RestCancellableNodeClient (fires inline on the Netty event loop), or parent-task cascade — runs onCancelled() while onCancelCallback is still null. The current code returns silently; the task is marked cancelled but the analytics query keeps running to completion.

This regresses the disconnect-cancel guarantee #22229 + opensearch-project/sql#5563 set up.

Fix

Mirror the pattern already used by AnalyticsShardTask.setCancellationListener on the data-node side: after the install-time CAS, re-check isCancelled() and run the callback inline if it's already set. Switch consumption to getAndSet(null) so the install-side and onCancelled() paths can't both fire it.

Test

testCallbackFiresImmediatelyIfAlreadyCancelled is the regression guard — empirically fails on the pre-fix AnalyticsQueryTask (callCount=0) and passes after this change. Other four tests confirm existing semantics unchanged: callback fires once on normal cancel, idempotent under repeated cancel, no-op when no callback installed, second-set throws.

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.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 40da719)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Race Condition Between lines 85-87, isCancelled() is checked after the CAS succeeds, but onCancelled() could be invoked by another thread between the CAS at line 81 and the isCancelled() check at line 85. If onCancelled() runs in that window, it will call runCallbackOnce() and null out the callback via getAndSet(null) at line 96. Then line 86 will call runCallbackOnce() again, find null, and the callback never fires. This breaks the guarantee that the callback runs exactly once when cancellation occurs before installation. public void setOnCancelCallback(Runnable callback) { if (onCancelCallback.compareAndSet(null, callback) == false) { throw new IllegalStateException("onCancelCallback already set for AnalyticsQueryTask " + queryId); } // Cancel may have arrived before this install — fire inline if so. Mirrors AnalyticsShardTask.setCancellationListener. if (isCancelled()) { runCallbackOnce(); } }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 38769dc
Explore these optional code suggestions:

Category	Suggestion	Impact
General	Log callback execution failures The callback exception is caught but not logged or handled, which can silently swallow errors during cancellation. Add logging to track when callback execution fails, as this could indicate issues in cleanup logic that need investigation. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsQueryTask.java [98-103] private void runCallbackOnce() { Runnable cb = onCancelCallback.getAndSet(null); if (cb != null) { try { cb.run(); } catch (Exception e) { + logger.warn("Exception during cancellation callback for task " + queryId, e); + } + } +} Suggestion importance[1-10]: 6 __ Why: Adding logging for caught exceptions is a reasonable improvement for debugging and monitoring. However, the suggestion assumes a logger field exists without verifying it in the PR diff, and the impact is moderate since the exception is already caught and won't crash the system.	Low
Possible issue	Fix callback race condition There's a race condition: after compareAndSet succeeds but before isCancelled() check, onCancelled() could execute and call getAndSet(null), nullifying the callback. Then isCancelled() returns true and calls runCallbackOnce() again, but the callback is already null and won't execute. Reorder to check cancellation before setting the callback. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsQueryTask.java [80-90] public void setOnCancelCallback(Runnable callback) { + boolean alreadyCancelled = isCancelled(); if (onCancelCallback.compareAndSet(null, callback) == false) { throw new IllegalStateException("onCancelCallback already set for AnalyticsQueryTask " + queryId); } - // Cancel can race the install: the framework registers the task before doExecute forks - // to the search executor, so a cancel from a timeout, HTTP disconnect, or parent cascade - // can fire onCancelled() before this method runs and find no callback. Re-check and - // fire inline. Mirrors AnalyticsShardTask.setCancellationListener. - if (isCancelled()) { + if (alreadyCancelled) { runCallbackOnce(); } } Suggestion importance[1-10]: 2 __ Why: The suggested race condition is unlikely because onCancelled() calls runCallbackOnce() which uses getAndSet(null), ensuring the callback runs exactly once. The test testCallbackFiresImmediatelyIfAlreadyCancelled validates the current implementation works correctly. The proposed reordering doesn't meaningfully improve thread safety.	Low

---

Previous suggestions

Suggestions up to commit 5faecbe

Category	Suggestion	Impact
Possible issue	Fix callback installation race condition There's a race condition: after compareAndSet succeeds but before isCancelled() is checked, onCancelled() could execute and call getAndSet(null), nullifying the callback. Then isCancelled() returns true and calls runCallbackOnce() again, but the callback is already null and won't execute. Reverse the order: check cancellation first, then set the callback atomically. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/task/AnalyticsQueryTask.java [80-91] public void setOnCancelCallback(Runnable callback) { + boolean wasCancelled = isCancelled(); if (onCancelCallback.compareAndSet(null, callback) == false) { throw new IllegalStateException("onCancelCallback already set for AnalyticsQueryTask " + queryId); } - // Cancel can race the install: the framework registers the task before doExecute forks - // to the search executor, so a cancel from a timeout, HTTP disconnect, or parent cascade - // can fire onCancelled() before this method runs and find no callback. Re-check and - // fire inline. Mirrors AnalyticsShardTask.setCancellationListener. - if (isCancelled()) { + if (wasCancelled) { runCallbackOnce(); } } Suggestion importance[1-10]: 2 __ Why: The suggested race condition is unlikely because onCancelled() calls runCallbackOnce() which uses getAndSet(null), ensuring the callback runs exactly once. The current implementation with compareAndSet followed by isCancelled() check correctly handles the race where cancellation happens before callback installation. The suggested reordering doesn't meaningfully improve the logic and the concern raised is not a real issue given the atomic operations used.	Low

[github-actions]

✅ Gradle check result for 5faecbe: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.33%. Comparing base (b53b8fe) to head (40da719).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22231      +/-   ##
============================================
- Coverage     73.43%   73.33%   -0.11%     
+ Complexity    75965    75858     -107     
============================================
  Files          6070     6070              
  Lines        344903   344903              
  Branches      49625    49625              
============================================
- Hits         253285   252936     -349     
- Misses        71493    71790     +297     
- Partials      20125    20177      +52     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 38769dc

[github-actions]

Persistent review updated to latest commit 40da719

[github-actions]

✅ Gradle check result for 40da719: SUCCESS