Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant

[vrajat]

PerQueryCPUMemAccountant uses callbacks to cancel queries. The main motivation is that MSE executor threads have no anchor thread.

• MSE queries register a callback at before starting execution.
• The callback calls OpchainSchedulerService.cancel.
• The callbacks are stored in a map Map<String, QueryCancelCallback>.
• The accountant uses the callback if available else falls back to interrupting the anchor thread.
• The map is regularly cleaned up by the watcher task.

Closes #16060

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 81.48148% with 5 lines in your changes missing coverage. Please review.

Project coverage is 63.12%. Comparing base (1a476de) to head (08beb9d).
Report is 439 commits behind head on master.

Files with missing lines	Patch %	Lines
...re/accounting/PerQueryCPUMemAccountantFactory.java	80.00%	4 Missing ⚠️
.../spi/accounting/ThreadResourceUsageAccountant.java	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #16142      +/-   ##
============================================
+ Coverage     62.90%   63.12%   +0.22%     
+ Complexity     1386     1363      -23     
============================================
  Files          2867     2973     +106     
  Lines        163354   172548    +9194     
  Branches      24952    26429    +1477     
============================================
+ Hits         102755   108926    +6171     
- Misses        52847    55294    +2447     
- Partials       7752     8328     +576     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-11	63.10% <81.48%> (+0.23%)	⬆️
java-21	63.09% <81.48%> (+0.27%)	⬆️
skip-bytebuffers-false	?
skip-bytebuffers-true	?
temurin	63.12% <81.48%> (+0.22%)	⬆️
unittests	63.12% <81.48%> (+0.22%)	⬆️
unittests1	56.34% <81.48%> (+0.52%)	⬆️
unittests2	33.18% <0.00%> (-0.39%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Copilot]

Pull Request Overview

This PR introduces a callback-based mechanism for cancelling queries via the ThreadResourceUsageAccountant, falling back to thread interrupts only when no callback is registered.

• Adds a registerQueryCancelCallback method to the SPI interface and implements it in PerQueryCPUMemResourceUsageAccountant.
• Hooks up callback registration in QueryRunner and QueryScheduler and wires the accountant through constructors.
• Extends tests and test utilities to verify callback registration.

Reviewed Changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
pinot-spi/.../ThreadResourceUsageAccountant.java	Added default registerQueryCancelCallback method
pinot-spi/.../QueryCancelCallback.java	Introduced new QueryCancelCallback interface
pinot-query-runtime/.../QueryRunnerAccountingTest.java	Removed unnecessary ServerMetrics.register call in test
pinot-query-runtime/.../QueryAccountantCancelTest.java	New test exercising per-query callback registration
pinot-query-runtime/.../QueryServerEnclosure.java	Overloaded constructors to accept the accountant
pinot-query-runtime/.../QueryRunner.java	Injects and registers cancel callback in processQuery
pinot-core/.../ResourceManager.java	Stores and exposes the accountant via getter
pinot-core/.../QueryScheduler.java	Registers interrupt callback on the anchor thread
pinot-core/.../PerQueryCPUMemAccountantFactory.java	Implements callback storage, cleanup, and cancelQuery

Comments suppressed due to low confidence (3)

pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java:97

• The Javadoc for registerQueryCancelCallback is missing a @param queryId description. Please add it to explain the purpose of the queryId parameter.

  default void registerQueryCancelCallback(String queryId, QueryCancelCallback queryCancelCallback) {

pinot-spi/src/main/java/org/apache/pinot/spi/accounting/QueryCancelCallback.java:23

• The Javadoc refers to a queryId parameter, but the cancelQuery() method takes no arguments. Consider updating the comment to clarify that the callback is bound to a specific query when registered.

   * Cancels the query with the given queryId.

pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerAccountingTest.java:179

• Removing the ServerMetrics.register call may lead to uninitialized metrics in this test and cause failures. If metrics are used in the accountant, re-add or mock registration here.

    HashMap<String, Object> configs = new HashMap<>();

[gortiz]

I like the idea of using an interface here, but I think we need to specify whether this method is blocking (and therefore once the call finishes the query is guarantee to be killed) or could be async

[vrajat]

It will be async. I'll add javadocs.

[gortiz]

We are using corretlations id here, right? in that case you can set the CID by using query options. See https://docs.pinot.apache.org/users/user-guide-query/query-correlation-id

[vrajat]

Not yet. The feature works with correlation id. However the test dont work with correlation id - sigh

[gortiz]

There is a conflict that has to be resolved, mark it as ready for review and add that nullable notation. Apart from these minor tasks, I think it is ready to merge

[vrajat]

There is a conflict that has to be resolved, mark it as ready for review and add that nullable notation. Apart from these minor tasks, I think it is ready to merge

Will fix. I'll get #16299 merged first and then fix conflicts and get the branch ready for merge.