gortiz opened a new pull request, #19867: URL: https://github.com/apache/pinot/pull/19867
## Summary When a single-stage query is terminated (for example, killed by the resource accountant under heap pressure), the combine operator keeps waiting until the query deadline. During that wait, the main query thread keeps all the merged results reachable (for a group-by, the full `IndexedTable`). As a result, killing the most expensive query does not free its memory, and the server can still fail with an `OutOfMemoryError`. This PR makes the combine wait aware of termination. The main thread now wakes up every 100 ms while it waits for the worker tasks. If the query is terminated, the combine operator returns the termination error immediately. ## Background `QueryExecutionContext.terminate()` sets the terminate exception and calls `cancel(true)` on the registered tasks. The combine worker tasks are registered, because the query scheduler gives them a context-aware executor. Two facts keep the main thread waiting: 1. The main query thread is not registered for cancellation, so `terminate()` does not wake it. The combine waits on `_operatorLatch.await(...)` (group-by) or `_blockingQueue.poll(...)` (all other combine operators). Neither wait checks the terminate exception. 2. A worker task that is still queued in the executor when `cancel(true)` runs never executes. Its `finally` block in `BaseCombineOperator.startProcess()` does not run, so it never counts down the latch and never adds a block to the queue. With a saturated worker pool, this is the normal case. The two facts together make the main thread wait until the deadline. Worker tasks that already run stop correctly, because they call `checkTerminationAndSampleUsagePeriodically`. ## Change `BaseCombineOperator` gets two helpers: - `pollResultsBlock(queue, endTimeMs)` replaces the `BlockingQueue.poll(...)` waits. - `awaitLatch(latch, endTimeMs)` replaces the `CountDownLatch.await(...)` waits. Both helpers wait in slices of `TERMINATION_CHECK_INTERVAL_MS` (100 ms). Before each slice, they read `QueryThreadContext.getTerminateException()`. If the query is terminated, they throw the `TerminationException`. The existing exception handling of each combine operator converts it into the termination error block. The timeout behavior at the deadline does not change. The change applies to all the combine operators that wait for worker results: - `BaseSingleBlockCombineOperator` (selection, aggregation, distinct) - `GroupByCombineOperator`, `SortedGroupByCombineOperator`, `SequentialSortedGroupByCombineOperator` - `MinMaxValueBasedSelectionOrderByCombineOperator` - `BaseStreamingCombineOperator`, `StreamingGroupByCombineOperator`, `StreamingDistinctCombineOperator` No separate change is necessary for the queued tasks. Only `terminate()` cancels a task before it starts (`stopProcess()` cancels tasks only after the merge). When the main thread observes the termination, it does not need a signal from the cancelled tasks. The `Phaser` in `stopProcess()` is already correct, because a task that never starts never registers with it. ### Why not interrupt the main thread The main thread is a pool thread that lives longer than the query. An interrupt that arrives after the wait ends hits unrelated work, such as the final merge, the response serialization, or the segment release. An interrupt during I/O on an `InterruptibleChannel` closes the channel. A safe interrupt needs a register/deregister window under a lock shared with `terminate()`, plus a flag reset. The bounded 100 ms check gives the same result with no change to the threading model and no SPI change. ## Testing `CombineTerminationTest` (new) fills the worker pool with blocking tasks, so that the combine tasks stay queued. Then it terminates the query and checks that the combine operator returns `QUERY_CANCELLATION` in less than 5 seconds, with a query deadline of 60 seconds. It covers the selection-only, selection order-by, aggregation, group-by, sorted group-by, sequential sorted group-by, and streaming selection-only combine operators. Without the fix, the cases time out. With the fix, all seven cases pass in less than 1 second in total. The existing combine and streaming tests pass (`CombineSlowOperatorsTest`, `CombineErrorOperatorsTest`, `SelectionCombineOperatorTest`, `SortedGroupByCombineOperatorsTest`, `MixedLayoutGroupByCombineOperatorTest`, `StreamingGroupByCombineOperatorTest`, `StreamingDistinctCombineOperatorTest`, and more). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
