[Query Resource Isolation] Fix Refresh message by praveenc7 · Pull Request #16636 · apache/pinot

praveenc7 · 2025-08-19T18:59:57Z

Summary

Dual registration of MessageHandlerFactory leads to a No-op for the second registration. Moving the QueryWorkloadRefresh to the existing MessageHandler
Fix the check for invalid subtypes

Testing done

Quick starter test to check if message if received on config change

pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java

codecov-commenter · 2025-08-19T20:56:39Z

Codecov Report

❌ Patch coverage is 4.16667% with 46 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.34%. Comparing base (6d1b8d6) to head (c252590).
⚠️ Report is 6 commits behind head on master.

Files with missing lines	Patch %	Lines
...er/starter/helix/SegmentMessageHandlerFactory.java	0.00%	23 Missing ⚠️
.../helix/BrokerUserDefinedMessageHandlerFactory.java	0.00%	16 Missing ⚠️
...he/pinot/spi/accounting/WorkloadBudgetManager.java	25.00%	6 Missing ⚠️
...t/common/messages/QueryWorkloadRefreshMessage.java	0.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #16636      +/-   ##
============================================
- Coverage     63.41%   63.34%   -0.07%     
- Complexity     1379     1380       +1     
============================================
  Files          3036     3035       -1     
  Lines        176910   176938      +28     
  Branches      27139    27144       +5     
============================================
- Hits         112179   112076     -103     
- Misses        56118    56253     +135     
+ Partials       8613     8609       -4

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (ø)`
integration	`100.00% <ø> (ø)`
integration1	`100.00% <ø> (ø)`
integration2	`0.00% <ø> (ø)`
java-11	`63.32% <4.16%> (-0.06%)`	⬇️
java-21	`63.31% <4.16%> (-0.02%)`	⬇️
temurin	`63.34% <4.16%> (-0.07%)`	⬇️
unittests	`63.33% <4.16%> (-0.07%)`	⬇️
unittests1	`56.49% <22.22%> (-0.01%)`	⬇️
unittests2	`33.12% <0.00%> (-0.06%)`	⬇️

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.

vvivekiyer · 2025-08-27T17:38:28Z

pinot-spi/src/main/java/org/apache/pinot/spi/accounting/WorkloadBudgetManager.java

@@ -168,7 +173,7 @@ private void startBudgetResetTask() {
      // Also print the budget used in the last enforcement window.
      _workloadBudgets.forEach((workload, budget) -> {
        BudgetStats stats = budget.getStats();


Can we change all loggers here to debug?

The Updates/deletes logs are spare, will update the rest

vvivekiyer

Looks good. Please make the suggested changes.

vvivekiyer · 2025-08-27T17:39:19Z

pinot-spi/src/main/java/org/apache/pinot/spi/accounting/WorkloadBudgetManager.java

      _workloadBudgets.forEach((workload, budget) -> {
        BudgetStats stats = budget.getStats();
-        LOGGER.debug("Workload: {} -> CPU: {}ns, Memory: {} bytes", workload, stats._cpuRemaining,
+        LOGGER.info("Workload: {} -> CPU: {}ns, Memory: {} bytes", workload, stats._cpuRemaining,


Can we please keep this as debug? This will log profusely during every refresh if enforcementWindow is small (say 1s).

(Unrelated) can you remove the TODO at the beginning of the function?

I see some value in quickly debugging. But given we have metrics soon, we can keep it a debug if we anticipate aggresive enforcementWindow

vvivekiyer · 2025-08-27T17:49:38Z

pinot-spi/src/main/java/org/apache/pinot/spi/accounting/WorkloadBudgetManager.java

  }

+  public void deleteWorkload(String workload) {
+    _workloadBudgets.remove(workload);


Can we initialize _workloadBudgets only if _isEnabled is true?
Also handle isEnabled in this function before removing.

was initially thinking of keeping delete lenient. For example, if the feature is disabled and we try to delete a workload, the operation wouldn’t succeed. However, since the state is in-memory, it would eventually get cleaned up on restart.

On second thought, it’s probably better to support delete only when the feature is enabled — it’s more useful when the feature is toggled on and off and on.

Can you also move the initialization of _workloadBudgets such that it's done after the isEnabled check?

praveenc7 · 2025-08-28T17:42:37Z

The failure is due to unrelated change : OfflineClusterMemBasedServerQueryKillingTest.testDigestOOM4MSE:397 [{"errorCode":450,"message":"Received 1 error from stage 2 on Server_localhost_22001: Found an error on stage 2 while reading from a child stage"}] expected [true] but found [false]

* [Query Resource Isolation] Workload Configs (apache#15109) * Workload Configs * workload config * Add API * config * Change config structure * Propagation strategy * Fix style check * Cost spliting on update * Table addition propagation * perf * Tests * test * test 2 * Review comments 1 * review comments 3 * review comments 3 * name change * review comments 4 * Fix TableDoesNotExistError for hybrid tables in MSE queries in controller API (apache#16102) * Make ThreadResourceUsageProvider a Helper/Utility Class. (apache#16051) * ThreadResourceUsageProvider is a helper class. ThreadResourceContext tracks resource usage. Fix updateConcurrently * Rename to ThreadResourceSnapshot * Clean up * Add javadoc * Done use auto closeable * Checkstyle * Fix compilation error * Add back removed functions in SPI * Remove private constructor because japicmp complains. * Add setThreadResourceUsageProvider because of backward-incompatible checks * Add setThreadResourceUsageProvider because of backward-incompatible checks * Fix test * Fix ThreadResourceSnapshot usage and tests * Store cpu sample in nanoseconds. * Reduce logs and improve logging when queries are terminated due to OOM. (apache#16172) * Dynamic PerQueryCPUMemAccountant Config on Servers (apache#16219) * Checkpoint * Register change handler * Fix bugs. Manually tested * Checkstyle * Tests * Add pre-check that values are default * Undo typo fix * Update QueryRunner to make use of window function overflow handling server configurations (apache#16108) * Add multistage thread limiting configs at the broker and server level (apache#16080) * Adding changes for supporting RLS (apache#16043) * Use stats cache on error instead of the chained mechanism (apache#15992) * Improve broker error messaging when broker is the one reporting the failure (apache#16076) * Introduce MSE active and passive timeouts (apache#16075) * Throttle SSE & MSE Tasks if Server heap usage is above a threshold (apache#16271) * Fix QueryScheduler constructor using class name. (apache#16280) * Fix QueryScheduler constructor using class name. * Fix test * [Query Resource Isolation] WorkloadBudgetManager and Host enforcement (apache#15798) * QRI - WorkloadBudgetManager implementation * Address review comments * Remove singleton & signature fix * Fix compatibility checker * Review comments * Move WorkloadBudgetManager to core. --------- Co-authored-by: praveenc7 <praveenkchaganlal@gmail.com> * Eliminate duplicate cancel attempts in PerQueryCPUMemAccountant (apache#16299) * Add basic 1 query tests * Add more tests * Add ability to remember cancel queries. * Clean up if conditions in killMostExpensiveQuery * Fix test failures. * Address review comments. * Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant (apache#16142) * Remove all calls to System.gc() in PerQueryCPUMemAccountantFactory (apache#16374) * Initialize thread accountant just before serving queries (apache#16326) * Allow Reset of ThreadResourceUsageAccountant in Tracing.java (apache#16360) * Queries now self terminate if in panic mode. (apache#16380) * Queries now self terminate if in panic mode. * Add config test * Hard kill on critical level. * Fix configs * Separate anchor thread interruption. * Checkstyle * Review comments * remove code for critical level --------- Co-authored-by: Rajat Venkatesh <vrajat@users.noreply.github.com> * [Query Resource Isolation] Additonal Sampling for Broker and Server (apache#16164) * fix * sampling * Broker sampling * revert integ-test * Fix test failures * review comments * remove MSE * broker auth * remove per pruner & planner sample * Use Broker's accountant to sample in the request handler. (apache#16439) * [Query Resource Isolation] Workload Scheduler (apache#16018) * QRI - WorkloadBudgetManager implementation * Address review comments * scheduler * unit test * review comments: metrics, secondary, resource-manager * remove broker admission * Remove default budget --------- Co-authored-by: Vivek Iyer Vaidyanathan Iyer <vvaidyanathan@linkedin.com> * Cleanup deprecated methods in ThreadResourceUsageAccountant (apache#16479) * Remove unnecessary methods and config for ThreadResourceUsageAccountant (apache#16490) * Add tests for OOM Termination of MSE queries. (apache#16514) * Fix a flaky assert when testing OOM Cancellation of MSE Queries (apache#16533) * Disable Flaky Tests (apache#16554) This is a follow-up to apache#16533 The fix for a flaky test did not work. This PR disables these tests temporarily. * Use correlation ID instead of request id in PerQueryCpuMemAccountant (apache#16040) * [Query Resource Isolation]Interface for Workload Stats Collection (apache#16340) * Interface for Stats Collection * Additional comments * inherit * additional class comments * [Query Resource Isolation] Fix Refresh message (apache#16636) * Fix Refresh message * delete queryworkload message handler * info -> debug logs * reduce logging (apache#16698) * style check * [Query Workload Isolation] Cost-split support (apache#16672) * splits * Cost split * test * propagation entity change & java doc * Propagation scheme review comments * empty commit to trigger build * Reduce log for PerQueryCPUMemResourceUsageAccountant (apache#16642) --------- Co-authored-by: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Co-authored-by: Yash Mayya <yash.mayya@gmail.com> Co-authored-by: Satwik Pachigolla <40644097+satwik-pachigolla@users.noreply.github.com> Co-authored-by: 9aman <35227405+9aman@users.noreply.github.com> Co-authored-by: Gonzalo Ortiz Jaureguizar <gortiz@users.noreply.github.com> Co-authored-by: Vivek Iyer Vaidyanathan <vvivekiyer@gmail.com> Co-authored-by: Xiaotian (Jackie) Jiang <17555551+Jackie-Jiang@users.noreply.github.com> Co-authored-by: Rajat Venkatesh <vrajat@users.noreply.github.com> Co-authored-by: Vivek Iyer Vaidyanathan Iyer <vvaidyanathan@linkedin.com>

* [Query Resource Isolation] Workload Configs (apache#15109) * Workload Configs * workload config * Add API * config * Change config structure * Propagation strategy * Fix style check * Cost spliting on update * Table addition propagation * perf * Tests * test * test 2 * Review comments 1 * review comments 3 * review comments 3 * name change * review comments 4 * Fix TableDoesNotExistError for hybrid tables in MSE queries in controller API (apache#16102) * Make ThreadResourceUsageProvider a Helper/Utility Class. (apache#16051) * ThreadResourceUsageProvider is a helper class. ThreadResourceContext tracks resource usage. Fix updateConcurrently * Rename to ThreadResourceSnapshot * Clean up * Add javadoc * Done use auto closeable * Checkstyle * Fix compilation error * Add back removed functions in SPI * Remove private constructor because japicmp complains. * Add setThreadResourceUsageProvider because of backward-incompatible checks * Add setThreadResourceUsageProvider because of backward-incompatible checks * Fix test * Fix ThreadResourceSnapshot usage and tests * Store cpu sample in nanoseconds. * Reduce logs and improve logging when queries are terminated due to OOM. (apache#16172) * Dynamic PerQueryCPUMemAccountant Config on Servers (apache#16219) * Checkpoint * Register change handler * Fix bugs. Manually tested * Checkstyle * Tests * Add pre-check that values are default * Undo typo fix * Update QueryRunner to make use of window function overflow handling server configurations (apache#16108) * Add multistage thread limiting configs at the broker and server level (apache#16080) * Adding changes for supporting RLS (apache#16043) * Use stats cache on error instead of the chained mechanism (apache#15992) * Improve broker error messaging when broker is the one reporting the failure (apache#16076) * Introduce MSE active and passive timeouts (apache#16075) * Throttle SSE & MSE Tasks if Server heap usage is above a threshold (apache#16271) * Fix QueryScheduler constructor using class name. (apache#16280) * Fix QueryScheduler constructor using class name. * Fix test * [Query Resource Isolation] WorkloadBudgetManager and Host enforcement (apache#15798) * QRI - WorkloadBudgetManager implementation * Address review comments * Remove singleton & signature fix * Fix compatibility checker * Review comments * Move WorkloadBudgetManager to core. --------- Co-authored-by: praveenc7 <praveenkchaganlal@gmail.com> * Eliminate duplicate cancel attempts in PerQueryCPUMemAccountant (apache#16299) * Add basic 1 query tests * Add more tests * Add ability to remember cancel queries. * Clean up if conditions in killMostExpensiveQuery * Fix test failures. * Address review comments. * Use QueryCancelCallback to cancel queries from ThreadResourceUsageAccountant (apache#16142) * Remove all calls to System.gc() in PerQueryCPUMemAccountantFactory (apache#16374) * Initialize thread accountant just before serving queries (apache#16326) * Allow Reset of ThreadResourceUsageAccountant in Tracing.java (apache#16360) * Queries now self terminate if in panic mode. (apache#16380) * Queries now self terminate if in panic mode. * Add config test * Hard kill on critical level. * Fix configs * Separate anchor thread interruption. * Checkstyle * Review comments * remove code for critical level --------- Co-authored-by: Rajat Venkatesh <vrajat@users.noreply.github.com> * [Query Resource Isolation] Additonal Sampling for Broker and Server (apache#16164) * fix * sampling * Broker sampling * revert integ-test * Fix test failures * review comments * remove MSE * broker auth * remove per pruner & planner sample * Use Broker's accountant to sample in the request handler. (apache#16439) * [Query Resource Isolation] Workload Scheduler (apache#16018) * QRI - WorkloadBudgetManager implementation * Address review comments * scheduler * unit test * review comments: metrics, secondary, resource-manager * remove broker admission * Remove default budget --------- Co-authored-by: Vivek Iyer Vaidyanathan Iyer <vvaidyanathan@linkedin.com> * Cleanup deprecated methods in ThreadResourceUsageAccountant (apache#16479) * Remove unnecessary methods and config for ThreadResourceUsageAccountant (apache#16490) * Add tests for OOM Termination of MSE queries. (apache#16514) * Fix a flaky assert when testing OOM Cancellation of MSE Queries (apache#16533) * Disable Flaky Tests (apache#16554) This is a follow-up to apache#16533 The fix for a flaky test did not work. This PR disables these tests temporarily. * Use correlation ID instead of request id in PerQueryCpuMemAccountant (apache#16040) * [Query Resource Isolation]Interface for Workload Stats Collection (apache#16340) * Interface for Stats Collection * Additional comments * inherit * additional class comments * [Query Resource Isolation] Fix Refresh message (apache#16636) * Fix Refresh message * delete queryworkload message handler * info -> debug logs * reduce logging (apache#16698) * style check * [Query Workload Isolation] Cost-split support (apache#16672) * splits * Cost split * test * propagation entity change & java doc * Propagation scheme review comments * empty commit to trigger build * Reduce log for PerQueryCPUMemResourceUsageAccountant (apache#16642) * [refactor] Switching to RoutingManager for broker request handlers (apache#16442) Co-authored-by: Shaurya Chaturvedi <shauryachats@uber.com> * Fix broker request id generator to avoid generating same id (apache#16661) * Introduce QueryExecutionContext to manage query life cycle (apache#16728) * Introduce QueryExecutionContext to manage query life cycle 2 (apache#16728) --------- Co-authored-by: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Co-authored-by: Yash Mayya <yash.mayya@gmail.com> Co-authored-by: Satwik Pachigolla <40644097+satwik-pachigolla@users.noreply.github.com> Co-authored-by: 9aman <35227405+9aman@users.noreply.github.com> Co-authored-by: Gonzalo Ortiz Jaureguizar <gortiz@users.noreply.github.com> Co-authored-by: Vivek Iyer Vaidyanathan <vvivekiyer@gmail.com> Co-authored-by: Xiaotian (Jackie) Jiang <17555551+Jackie-Jiang@users.noreply.github.com> Co-authored-by: Rajat Venkatesh <vrajat@users.noreply.github.com> Co-authored-by: Vivek Iyer Vaidyanathan Iyer <vvaidyanathan@linkedin.com> Co-authored-by: Shaurya Chaturvedi <shauryachats@gmail.com> Co-authored-by: Shaurya Chaturvedi <shauryachats@uber.com>

praveenc7 added 2 commits August 19, 2025 11:50

Fix Refresh message

1453d85

delete queryworkload message handler

e0bc425

praveenc7 commented Aug 19, 2025

View reviewed changes

pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java Outdated Show resolved Hide resolved

praveenc7 marked this pull request as ready for review August 19, 2025 20:08

shauryachats mentioned this pull request Aug 20, 2025

[refactor] Refactor TableCache to interface with ZkTableCache implementation #16646

Merged

Merge branch 'master' into teststats

4ac07e6

praveenc7 force-pushed the teststats branch from 985eada to 4ac07e6 Compare August 25, 2025 21:19

vvivekiyer reviewed Aug 27, 2025

View reviewed changes

vvivekiyer approved these changes Aug 27, 2025

View reviewed changes

info -> debug logs

90991b9

Merge branch 'master' into teststats

c252590

praveenc7 force-pushed the teststats branch from 599dd12 to c252590 Compare August 28, 2025 17:46

vvivekiyer approved these changes Aug 28, 2025

View reviewed changes

vvivekiyer merged commit fb06abc into apache:master Aug 28, 2025
18 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Query Resource Isolation] Fix Refresh message#16636

[Query Resource Isolation] Fix Refresh message#16636
vvivekiyer merged 5 commits intoapache:masterfrom
praveenc7:teststats

praveenc7 commented Aug 19, 2025 •

edited

Loading

Uh oh!

Uh oh!

codecov-commenter commented Aug 19, 2025 •

edited

Loading

Uh oh!

vvivekiyer Aug 27, 2025

Uh oh!

praveenc7 Aug 28, 2025

Uh oh!

vvivekiyer left a comment

Uh oh!

vvivekiyer Aug 27, 2025

Uh oh!

praveenc7 Aug 28, 2025

Uh oh!

vvivekiyer Aug 27, 2025

Uh oh!

praveenc7 Aug 28, 2025

Uh oh!

vvivekiyer Aug 28, 2025

Uh oh!

praveenc7 commented Aug 28, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

praveenc7 commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing done

Uh oh!

Uh oh!

codecov-commenter commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

vvivekiyer Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

praveenc7 Aug 28, 2025

Choose a reason for hiding this comment

Uh oh!

vvivekiyer left a comment

Choose a reason for hiding this comment

Uh oh!

vvivekiyer Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

praveenc7 Aug 28, 2025

Choose a reason for hiding this comment

Uh oh!

vvivekiyer Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

praveenc7 Aug 28, 2025

Choose a reason for hiding this comment

Uh oh!

vvivekiyer Aug 28, 2025

Choose a reason for hiding this comment

Uh oh!

praveenc7 commented Aug 28, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

praveenc7 commented Aug 19, 2025 •

edited

Loading

codecov-commenter commented Aug 19, 2025 •

edited

Loading