Throttle SSE & MSE Tasks if Server heap usage is above a threshold #16271
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Contributor
There was a problem hiding this comment.
Pull Request Overview
Adds a new mechanism to stop submitting multi-stage execution (MSE) tasks when heap usage exceeds a configured critical threshold.
- Introduce
ThrottleOnCriticalHeapUsageExecutorthat throws whenthrottleQuerySubmission()is true. - Add a default
throttleQuerySubmissioninThreadResourceUsageAccountantand implement it inPerQueryCPUMemAccountantFactory. - Wire the new executor through server startup (
BaseServerStarter→ServerInstance→WorkerQueryServer→QueryRunner) and expose a config flag.
Reviewed Changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java | Add test for throttling logic |
| pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java | Define new config key and default |
| pinot-spi/src/main/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutor.java | Implement executor wrapper that checks heap usage threshold |
| pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java | Provide a default throttleQuerySubmission() |
| pinot-server/src/main/java/org/apache/pinot/server/worker/WorkerQueryServer.java | Pass ThreadResourceUsageAccountant into QueryRunner |
| pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java | Inject ThreadResourceUsageAccountant when creating ServerInstance |
| pinot-server/src/main/java/org/apache/pinot/server/starter/ServerInstance.java | Update constructor and WorkerQueryServer instantiation to accept accountant |
| pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java | Update test to supply a default accountant instance |
| pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java | Extend init signature and wrap executor with ThrottleOnCriticalHeapUsageExecutor based on config |
| pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java | Implement throttleQuerySubmission() and expose current heap usage |
Comments suppressed due to low confidence (1)
pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java:213
Server.CONFIG_OF_ENABLE_MSE_THROTTLING_ON_CRITICAL_HEAP_USAGEis not in scope; it should be referenced asCommonConstants.Server.CONFIG_OF_ENABLE_MSE_THROTTLING_ON_CRITICAL_HEAP_USAGE.
if (serverConf.getProperty(Server.CONFIG_OF_ENABLE_MSE_THROTTLING_ON_CRITICAL_HEAP_USAGE,
|
|
||
| @Override | ||
| public boolean throttleQuerySubmission() { | ||
| return _numCalls.getAndIncrement() > 1; |
There was a problem hiding this comment.
The stubbed throttleQuerySubmission() returns true only after two submissions, but the test expects throttling on the second submission. Change the condition to > 0 to block the second task as intended.
Suggested change
| return _numCalls.getAndIncrement() > 1; | |
| return _numCalls.getAndIncrement() > 0; |
Contributor
Author
There was a problem hiding this comment.
This is fine. The function is called twice for each task in the executor
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #16271 +/- ##
============================================
+ Coverage 62.90% 63.21% +0.31%
+ Complexity 1386 1365 -21
============================================
Files 2867 2963 +96
Lines 163354 171523 +8169
Branches 24952 26277 +1325
============================================
+ Hits 102755 108431 +5676
- Misses 52847 54848 +2001
- Partials 7752 8244 +492
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
vrajat
changed the title
gortiz
approved these changes
Jul 3, 2025
PerQueryCPUMemAccountant defines critical heap usage based on its configuration. This PR adds the ability to disallow anymore MSE tasks if heap usage is above the critical level. A new configuration parameter is added `pinot.server.query.executor.mse.enableThrottlingOnCriticalHeapUsage`. Default is false. If enabled, the executor checks with `ThreadResourceUsageAccountant` if heap level is critical. In the default implementation, `false` is returned. `PerQueryCPUMemAccountant` compares the configuration with current heap usage. Heap usage is updated every time watcher task runs. By default, it runs every 30ms. Closes apache#16270
vrajat
force-pushed
the
rv-tracing-throttle
branch
from
68c3862 to
77c05dd
Compare
vrajat
changed the title
gortiz
approved these changes
Jul 4, 2025
Jackie-Jiang
added
documentation
release-notes
| Tracing.getThreadAccountant().getClusterConfigChangeListener()); | ||
| } | ||
|
|
||
| SendStatsPredicate sendStatsPredicate = SendStatsPredicate.create(_serverConf, _helixManager); |
Contributor
There was a problem hiding this comment.
We should not move this behind thread accountant initialization since server metrics need to be initialized.
cc @priyen-stripe who helped find the issue
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
PerQueryCPUMemAccountant defines critical heap usage based on its configuration. This PR adds the ability to disallow anymore SSE & MSE tasks if heap usage is above the critical level.
A new configuration parameter is added
pinot.server.query.executor.enableThrottlingOnHeapUsage. Default is false.If enabled, the executor checks with
ThreadResourceUsageAccountantif heap level is alarm. In the default implementation,falseis returned.PerQueryCPUMemAccountantcompares the configuration with current heap usage. Heap usage is updated every time watcher task runs. By default, it runs every 30ms.Closes #16270