Introduce QueryExecutionContext to manage query life cycle by Jackie-Jiang · Pull Request #16728 · apache/pinot

Jackie-Jiang · 2025-09-01T11:10:52Z

Introduce QueryExecutionContext to manage the query life cycle so that all threads executing the same query shares the same query execution context. It is used for the following purpose:

Track all the running threads
Enhance query resource tracking/accounting
Easier and more robust query cancellation
Better observability on query termination

High level abstraction:

QueryThreadContext is a thread-local context for storing common query-related information associated to the current thread.
QueryExecutionContext is shared across all threads running the same query
Query cancellation can be performed through QueryExecutionContext and propagated to all threads working on the same query
ThreadResourceUsageAccountant is attached to QueryExecutionContext to avoid collision between broker and server

Incompatible

Several interfaces/classes related to query life cycle tracking and accounting are changed in SPI
In order to make the tracking work/not throw exception, query execution must be wrapped within the QueryThreadContext, executor must be context aware. See examples in the test changes

Jackie-Jiang · 2025-09-01T11:19:41Z

@siddharthteotia @vvivekiyer @chenboat @ankitsultana @jadami10 Please take a look at the SPI changes and see if you have plug-ins using them

jadami10 · 2025-09-01T19:16:03Z

it's a bit hard to evaluate due to the sheer size of the change. without trying to deploy this, this looks the affected files

pinot-spi/src/main/java/org/apache/pinot/spi/accounting/QueryResourceTracker.java
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadAccounting.java
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadExecutionContext.java
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceTracker.java
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/WorkloadBudgetManager.java
pinot-spi/src/main/java/org/apache/pinot/spi/executor/DecoratorExecutorService.java
pinot-spi/src/main/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutor.java
pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryExecutionContext.java
pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryThreadContext.java
pinot-spi/src/main/java/org/apache/pinot/spi/trace/LoggerConstants.java
pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java

i've confirmed we don't use any of these internally

ankitsultana · 2025-09-02T15:29:28Z

@Jackie-Jiang : can you confirm what features this would unlock? It's not quite clear from the description.

Copilot

Pull Request Overview

This PR introduces a new query lifecycle management system centered around QueryExecutionContext and refactors thread-local context handling. The primary goal is to provide better query cancellation capabilities and resource management across all threads working on the same query.

Introduces QueryExecutionContext as a shared context across all threads executing the same query
Refactors QueryThreadContext to work with the new execution context and simplified thread-local state management
Consolidates thread accounting utilities into a new ThreadAccounting class with separate broker/server accountant registries

Reviewed Changes

Copilot reviewed 179 out of 179 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
pinot-tools/src/main/java/org/apache/pinot/tools/MultistageEngineQuickStart.java	Simplifies config overrides by removing query cancellation configs and using Map.of()
pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java	Adds default values for query cancellation configs and fixes typo in constant name
pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java	Major refactor removing thread accounting functionality, now only handles tracer registration
pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryThreadContext.java	Complete rewrite to work with QueryExecutionContext and simplified API
pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryExecutionContext.java	New class managing shared query state and task lifecycle
pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadAccounting.java	New utility class consolidating thread accounting with separate broker/server registries

Comments suppressed due to low confidence (4)

Jackie-Jiang · 2025-09-03T18:12:56Z

@Jackie-Jiang : can you confirm what features this would unlock? It's not quite clear from the description.

This PR was introduced mainly for more robust OOM protection and query cancellation. Updated the PR description

gortiz

I've reviewed most of the important changes. I think the refactor make sense and appreciate some improvements in code style, but I have to say (mostly for other reviewers, but also for future PRs) that the number of lines changed would be significantly smaller if these cosmetic refactors (ie javadoc as markdown or import static) were applied on another PR.

pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryExecutionContext.java

pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryThreadContext.java

pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryExecutionContext.java

gortiz · 2025-09-05T09:54:28Z

pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotClientRequest.java

      Map<String, Integer> serverResponses = verbose ? new HashMap<>() : null;
      if (isClient) {
-        long reqId = _requestHandler.getRequestIdByClientId(id).orElse(-1L);
-        QueryThreadContext.setIds(reqId, id);


This is a con of this refactor. With the previous model we had the ability to partially create the thread local object. This means that any log printed after this line would include the CID, which is may be very useful.

It is not a critial problem, but something that would be great if we find a way to keep

I don't follow why we need this to be thread local. It should be shared across all threads for a given query

...rc/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java

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

codecov-commenter · 2025-09-16T01:23:51Z

Codecov Report

❌ Patch coverage is 38.16885% with 1040 lines in your changes missing coverage. Please review.
✅ Project coverage is 63.35%. Comparing base (a7cad30) to head (2966942).
⚠️ Report is 222 commits behind head on master.

Files with missing lines	Patch %	Lines
...pinot/core/accounting/QueryResourceAggregator.java	0.00%	223 Missing ⚠️
...re/accounting/PerQueryCPUMemAccountantFactory.java	18.18%	133 Missing and 2 partials ⚠️
...ore/accounting/ResourceUsageAccountantFactory.java	0.00%	117 Missing ⚠️
...ot/core/accounting/WorkloadResourceAggregator.java	0.00%	63 Missing ⚠️
...apache/pinot/query/service/server/QueryServer.java	45.16%	51 Missing ⚠️
...org/apache/pinot/spi/query/QueryThreadContext.java	57.57%	35 Missing and 7 partials ⚠️
.../apache/pinot/spi/query/QueryExecutionContext.java	40.62%	37 Missing and 1 partial ⚠️
...roker/requesthandler/BaseBrokerRequestHandler.java	39.62%	25 Missing and 7 partials ⚠️
...che/pinot/core/transport/grpc/GrpcQueryServer.java	0.00%	25 Missing ⚠️
...e/pinot/core/transport/InstanceRequestHandler.java	58.18%	17 Missing and 6 partials ⚠️
... and 56 more

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #16728      +/-   ##
============================================
- Coverage     63.49%   63.35%   -0.15%     
+ Complexity     1412     1411       -1     
============================================
  Files          3073     3071       -2     
  Lines        180784   180348     -436     
  Branches      27636    27583      -53     
============================================
- Hits         114795   114257     -538     
- Misses        57132    57325     +193     
+ Partials       8857     8766      -91

Flag	Coverage Δ
custom-integration1	`100.00% <ø> (ø)`
integration	`100.00% <ø> (ø)`
integration1	`100.00% <ø> (ø)`
integration2	`0.00% <ø> (ø)`
java-11	`63.27% <37.87%> (-0.19%)`	⬇️
java-21	`63.32% <38.16%> (-0.14%)`	⬇️
temurin	`63.35% <38.16%> (-0.15%)`	⬇️
unittests	`63.35% <38.16%> (-0.15%)`	⬇️
unittests1	`56.15% <38.86%> (-0.31%)`	⬇️
unittests2	`33.73% <3.56%> (+0.12%)`	⬆️

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.

gortiz

I'm having a hard time reviewing this PR for a second time. The number of style changes are very, very high. Can we revert them?

Also, in case we want to continue with them, I would strongly ask to not amend and force-push changes, given then github gets a bit mad and forget about all files marked as reviewed and even worse, makes it impossible to compare the new changes commit by commit.

...ker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java

Jackie-Jiang · 2025-09-19T22:48:24Z

@gortiz Reverted most of the unnecessary changes in production code. Unfortunately I still need to amend the change this time because I reset the commit to see the changes. Now most of the changes should be settled and I can try adding commit instead of amending

gortiz

I've added some comments, but nothing is actually blocking

...rc/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/InternalFunctions.java

gortiz · 2025-09-23T06:55:26Z

pinot-core/src/main/java/org/apache/pinot/core/operator/BaseOperator.java

+  protected void checkTermination() {
+    QueryThreadContext.checkTermination(this::getExplainName);
+  }
+
+  protected void checkTerminationAndSampleUsage() {
+    QueryThreadContext.checkTerminationAndSampleUsage(this::getExplainName);
+  }


Couldn't we have a version of checkTerminationAndSampleUsage and checkTermination that accepts a String instead of a supplier? That way we could call this method without allocating

I didn't add it because we can directly call the static method

What do you mean? My point is that we need to allocate the lambda each time we call the method. That lambda makes sense to calculate the name lazily, but in this case, the name is a static, so the laziness is actually making the call more expensive

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java

gortiz · 2025-09-23T06:58:45Z

pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java

+    // For exception results block, check terminate exception and use it as the results block if exists. We want to
+    // return the termination reason when query is explicitly terminated.
+    if (resultsBlock instanceof ExceptionResultsBlock) {
+      TerminationException terminateException = QueryThreadContext.getTerminateException();


Why dont we do this when the received ExceptionResultBlock is created instead of here?

This is more or less a catch all check. Ideally we should check at the creator side.

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

pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryAggregator.java

pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryResourceAggregator.java

pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkAggregateGroupByOrderByQueriesSSE.java

gortiz · 2025-09-23T09:05:28Z

...y-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java

+    _opChainCache.put(opChainId, Pair.of(rootOperator, executionContext));
+
+    // Create a ListenableFutureTask to ensure the opChain is cancelled even if the task is not scheduled
+    ListenableFutureTask<Void> listenableFutureTask = ListenableFutureTask.create(new TraceRunnable() {


nit: this can be done without Guava with CompletableFutures, but ListenableFuture is fine as well

IIUC, CompletableFutures doesn't interrupt the underlying execution, which will waste resource

* [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>

Jackie-Jiang added enhancement incompatible Indicate PR that introduces backward incompatibility bugfix refactor oom-protection labels Sep 1, 2025

Jackie-Jiang requested review from ankitsultana, chenboat, gortiz, siddharthteotia, vvivekiyer, xiangfu0 and yashmayya September 1, 2025 11:12

Jackie-Jiang force-pushed the query_execution_context branch 3 times, most recently from 685ac57 to 86144eb Compare September 1, 2025 23:31

xiangfu0 requested a review from Copilot September 3, 2025 15:24

Copilot AI reviewed Sep 3, 2025

View reviewed changes

gortiz reviewed Sep 5, 2025

View reviewed changes

gortiz mentioned this pull request Sep 10, 2025

Create a DescDocIdSetOperator and improve order by desc when columns are sorted #16789

Merged

Jackie-Jiang mentioned this pull request Sep 16, 2025

Fix OfflineClusterMemBasedServerQueryKillingTest #16628

Closed

Jackie-Jiang force-pushed the query_execution_context branch from 86144eb to 22e65ff Compare September 16, 2025 00:40

Jackie-Jiang force-pushed the query_execution_context branch 2 times, most recently from 11423b8 to 4dd537f Compare September 17, 2025 02:02

Jackie-Jiang mentioned this pull request Sep 17, 2025

Apply improvements from PerQueryCPUMemAccountantFactory to Query aggregator #16793

Closed

gortiz reviewed Sep 18, 2025

View reviewed changes

Jackie-Jiang force-pushed the query_execution_context branch 3 times, most recently from cf84e08 to 57ab14e Compare September 19, 2025 22:44

Jackie-Jiang force-pushed the query_execution_context branch 5 times, most recently from 26edf47 to 957fec0 Compare September 23, 2025 01:41

gortiz approved these changes Sep 23, 2025

View reviewed changes

Jackie-Jiang force-pushed the query_execution_context branch 8 times, most recently from ca26b0c to 331a98f Compare September 30, 2025 06:16

Introduce QueryExecutionContext to manage query life cycle

2966942

Jackie-Jiang force-pushed the query_execution_context branch from 331a98f to 2966942 Compare October 1, 2025 01:24

Jackie-Jiang merged commit 8a416da into apache:master Oct 1, 2025
16 of 18 checks passed

Jackie-Jiang deleted the query_execution_context branch October 1, 2025 05:35

pengding-stripe mentioned this pull request Feb 10, 2026

Server not responded query errors following big amount of segment refresh events after upgrade to 1.5 #17657

Closed

Conversation

Jackie-Jiang commented Sep 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Incompatible

Uh oh!

Jackie-Jiang commented Sep 1, 2025

Uh oh!

jadami10 commented Sep 1, 2025

Uh oh!

ankitsultana commented Sep 2, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Jackie-Jiang commented Sep 3, 2025

Uh oh!

gortiz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

gortiz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Jackie-Jiang commented Sep 19, 2025

Uh oh!

gortiz left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Jackie-Jiang commented Sep 1, 2025 •

edited

Loading

codecov-commenter commented Sep 16, 2025 •

edited

Loading