Query Execution Threads self-terminate when heap usage is at critical level

[vrajat]

When heap usage is at critical level, query execution threads check if the tasks belong to the most expensive query and then self-terminate.

If accounting.thread.self.terminate is true, then the watcher task does not terminate queries. Therefore this config can be used to choose between the two modes of termination.

The unit test infrastructure has also been updated.

• TestResourceUsageAccountant has setter/getter to set specific values of heap usage.
• BasePerQueryCPUMemAccountantTest creates threads that simulate SSE query task threads.
• PerQueryCPUMemHardCancelTest tests hard cancel in critical and panic level.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 71.64179% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.37%. Comparing base (058f7aa) to head (fe0d589).
⚠️ Report is 21 commits behind head on master.

Files with missing lines	Patch %	Lines
...re/accounting/PerQueryCPUMemAccountantFactory.java	71.64%	8 Missing and 11 partials ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #16445      +/-   ##
============================================
+ Coverage     63.34%   63.37%   +0.03%     
- Complexity     1379     1380       +1     
============================================
  Files          3015     3027      +12     
  Lines        174795   176468    +1673     
  Branches      26788    27086     +298     
============================================
+ Hits         110722   111837    +1115     
- Misses        55612    56067     +455     
- Partials       8461     8564     +103     

Flag	Coverage Δ
custom-integration1	100.00% <ø> (ø)
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-11	63.33% <71.64%> (+0.02%)	⬆️
java-21	63.28% <70.14%> (-0.02%)	⬇️
temurin	63.37% <71.64%> (+0.03%)	⬆️
unittests	63.37% <71.64%> (+0.03%)	⬆️
unittests1	56.47% <71.64%> (+0.06%)	⬆️
unittests2	33.18% <0.00%> (-0.22%)	⬇️

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 implements query execution thread self-termination when heap usage reaches critical levels, providing an alternative to the watcher task terminating queries. The implementation adds logic for threads to check if they belong to the most expensive query and self-terminate when heap usage is critical.

Key changes:

• Enhanced isQueryTerminated() method to check heap usage levels and identify expensive queries for self-termination
• Updated query cancellation logic to use QueryResourceTracker instead of just query ID strings
• Refactored test infrastructure to support thread simulation and heap usage control

Reviewed Changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
PerQueryCPUMemAccountantFactory.java	Implements self-termination logic and updates query cancellation methods
TestResourceAccountant.java	Adds heap usage control and test-specific watcher task implementation
BasePerQueryCPUMemAccountantTest.java	Creates base test class with thread simulation utilities
PerQueryCPUMemAccountantTest.java	Refactors existing tests to use new base class and thread simulation
PerQueryCPUMemAccountCancelTest.java	Updates cancellation tests to use new infrastructure
PerQueryCPUMemAccountHardCancelTest.java	Adds new tests for hard cancellation at critical and panic levels
PerQueryCPUMemAccountantTest.java (runtime)	Updates method signature for query cancellation

[Copilot]

The comment incorrectly states 'task id = 2 (2500 bytes)' but the original test setup shows task id 2 has 2000 bytes and task id 2 (the third task with index 2) has 2500 bytes. The comment should clarify which task is being referenced.

Suggested change

	// Replace task id = 2 (2500 bytes) with a new task id 5 (1500 bytes)
	// Replace task id = 2 (2000 bytes) with a new task id 5 (1500 bytes)

[Copilot]

Iterating over _threadEntriesMap without synchronization while other threads may be modifying it could lead to ConcurrentModificationException. Consider using appropriate synchronization or a thread-safe iteration approach.

[Copilot]

The _cancelSentQueries.add() check-and-act pattern is not atomic. Multiple threads could pass the condition simultaneously before any can add the queryId, potentially leading to duplicate termination logging or processing.

[vrajat]

_cancelSentQueries is created using ConcurrentHashMap.newKeySet(). So it is backed by a ConcurrentHashMap. Is add still not atomic ?

State your sources.

[gortiz]

Probably copilot is not smart enough to know the runtime implementation used for the set. It may help if you change the static type from Set to ConcurrentHashMap.KeySetView. But yes, it should be atomic

[gortiz]

I think we already discuss about that. IIRC the anchor thread interruption is what is used to kill SSE queries. I think it would be better (easier to read and maintain) if we move that anchor thread interruption into a callback we register for SSE queries in the same way we have a custom callback for MSE

[vrajat]

I hadnt realized hard-cancel (this PR) is so easy to implement. I want to delete callbacks, anchor threads and reduce watcher task to only aggregate data. Do not try to interrupt threads. Let them self-terminate as in this code.

[gortiz]

This is almost the same code that we have in cancelQuery. I think it would be better to extract the code into an internal cancel query.

[vrajat]

I did try to consolidate. The differences are:

• other instances use QueryResourceTracker while this loop uses TaskEntry (thread metadata vs query metadata)
• Logging is done outside the loop here. The code should not log for every thread.

Both of these difference made it hard and I gave up.

[gortiz]

Each time we modify this class, I have to spend a long time remembering which thread touches each attribute. Can we add a javadoc on the attributes indicating the thread constraint?

[vrajat]

I've added a javadoc to explain what the the variable stores as well as who writes and reads it.

[Jackie-Jiang]

(MAJOR) Without interrupt, query might block on certain operations, such as polling a blocking queue. We probably need to do both interrupt and self terminate. How do we verify the current approach is working or not?

[vrajat]

True. This is an enhancement though. Note that the current approach of interrupting the anchor thread does not interrupt I/O ops in query execution threads. It still depends on the a thread exiting the I/O operation and then checking the state of the anchor thread.

For pre-emption from I/O ops, a cancel callback is required.

re: testing - I have added basic unit tests as a starting point.

[Jackie-Jiang]

What I meant is that we stop sending interrupt when self terminate is enabled. This could cause thread blocking on I/O forever right?

[vrajat]

Interrupt is only sent to the anchor thread. Refer: https://github.com/apache/pinot/blob/master/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java#L529

Anchor Thread does not interrupt other threads in SSE. So there is no code yet for interrupting query execution thread from PerQuery....

Refer:

pinot/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QueryScheduler.java

Line 147 in e78c836

Tracing.ThreadAccountantOps.setupRunner(QueryThreadContext.getCid(), workloadName);

For MSE, Cancel Callback was added recently. The callback does call interrupt on all threads. However there isn't high confidence if OpChainSchedulerService.cancel works.

Refer: https://github.com/apache/pinot/blob/master/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java#L120

If you plan to add code to also cancel query execution threads - then yes - both self-terminate and interrupt of query execution threads are required. One option is to use cancel callbacks. The better approach is for QueryScheduler or QueryRunner to proactively cancel all threads of a query when one of them returns an error. Its control flow smell for one query execution thread to take on the responsibility of interrupting other query execution threads.

[vrajat]

Flaky test. More info #16554