You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by go...@apache.org on 2026/10/15 10:21:18 UTC
(pinot) branch master updated: Return promptly from SSE combine operators when the query is terminated (#19867)
This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 9c5bd72a456 Return promptly from SSE combine operators when the query is terminated (#19867)
9c5bd72a456 is described below
commit 9c5bd72a456c832c5084050e675f423326890397
Author: Gonzalo Ortiz Jaureguizar <go...@users.noreply.github.com>
AuthorDate: Thu Oct 15 12:21:11 2026 +0200
Return promptly from SSE combine operators when the query is terminated (#19867)
* Return promptly from SSE combine operators when the query is terminated
When a single-stage query is terminated (for example, killed by the
resource accountant), the combine operator kept waiting until the query
deadline and kept the merged results (such as the IndexedTable) reachable.
Killing the most expensive query did not free its memory.
Two facts caused the wait:
- The main query thread is not registered for cancellation, so
QueryExecutionContext.terminate() does not wake it. The waits on
_operatorLatch.await() and _blockingQueue.poll() did not check the
terminate exception.
- A worker task cancelled while still queued never runs, so its finally
block in BaseCombineOperator.startProcess() never counts down the latch
or offers a block.
Add pollResultsBlock() and awaitLatch() to BaseCombineOperator. They wait
in 100 ms slices and throw the TerminationException as soon as
QueryThreadContext.getTerminateException() is set. The existing exception
handling turns it into the termination error block. Use them in every
combine operator that waits for worker results, including the streaming
ones. The timeout behavior at the deadline does not change.
Add CombineTerminationTest, which saturates the worker pool, terminates
the query, and checks that the combine operator returns
QUERY_CANCELLATION promptly.
* Make CombineTerminationTest terminate only after the combine thread waits
Wait until the combine thread is in TIMED_WAITING before calling
terminate(), so that the wait loop in pollResultsBlock()/awaitLatch()
must observe the termination. Before, terminate() could run while the
main thread was still in startProcess(), and the termination check
before the wait was enough to pass. Without the fix, all seven cases
now fail, including the streaming one.
---------
Co-authored-by: Gonzalo Ortiz <19...@users.noreply.github.com>
---
.../core/operator/combine/BaseCombineOperator.java | 55 +++++
.../combine/BaseSingleBlockCombineOperator.java | 3 +-
.../operator/combine/GroupByCombineOperator.java | 6 +-
...xValueBasedSelectionOrderByCombineOperator.java | 4 +-
.../SequentialSortedGroupByCombineOperator.java | 3 +-
.../combine/SortedGroupByCombineOperator.java | 6 +-
.../streaming/BaseStreamingCombineOperator.java | 3 +-
.../StreamingDistinctCombineOperator.java | 4 +-
.../streaming/StreamingGroupByCombineOperator.java | 4 +-
.../operator/combine/CombineTerminationTest.java | 241 +++++++++++++++++++++
10 files changed, 308 insertions(+), 21 deletions(-)
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java
index f4e4f67f6c0..a9f5fc6ec0f 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java
@@ -19,12 +19,16 @@
package org.apache.pinot.core.operator.combine;
import java.util.List;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.Phaser;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nullable;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.operator.BaseOperator;
import org.apache.pinot.core.operator.ExecutionStatistics;
@@ -60,6 +64,14 @@ import org.slf4j.LoggerFactory;
public abstract class BaseCombineOperator<T extends BaseResultsBlock> extends BaseOperator<BaseResultsBlock> {
private static final Logger LOGGER = LoggerFactory.getLogger(BaseCombineOperator.class);
+ /// Interval at which the main thread wakes up while waiting for the worker tasks, to check whether the query is
+ /// terminated.
+ ///
+ /// The main thread is not registered for cancellation, so terminating the query does not interrupt its wait. Also, a
+ /// worker task cancelled before it starts never runs, so it never signals completion. Without this check, the main
+ /// thread of a terminated query waits until the query deadline and keeps the merged results reachable.
+ private static final long TERMINATION_CHECK_INTERVAL_MS = 100L;
+
protected final ResultsBlockMerger<T> _resultsBlockMerger;
protected final List<Operator> _operators;
protected final int _numOperators;
@@ -254,6 +266,49 @@ public abstract class BaseCombineOperator<T extends BaseResultsBlock> extends Ba
_phaser.awaitAdvance(_phaser.arriveAndDeregister());
}
+ /// Waits for the next results block in the queue until the end time.
+ ///
+ /// Returns `null` when the end time is reached without a results block. Throws the [TerminationException] when the
+ /// query is terminated while waiting.
+ @Nullable
+ protected BaseResultsBlock pollResultsBlock(BlockingQueue<BaseResultsBlock> blockingQueue, long endTimeMs)
+ throws InterruptedException {
+ while (true) {
+ throwIfTerminated();
+ long waitTimeMs = endTimeMs - System.currentTimeMillis();
+ BaseResultsBlock resultsBlock =
+ blockingQueue.poll(Math.min(waitTimeMs, TERMINATION_CHECK_INTERVAL_MS), TimeUnit.MILLISECONDS);
+ if (resultsBlock != null || waitTimeMs <= TERMINATION_CHECK_INTERVAL_MS) {
+ return resultsBlock;
+ }
+ }
+ }
+
+ /// Waits for the latch to reach zero until the end time.
+ ///
+ /// Returns `false` when the end time is reached before the latch reaches zero. Throws the [TerminationException]
+ /// when the query is terminated while waiting.
+ protected boolean awaitLatch(CountDownLatch latch, long endTimeMs)
+ throws InterruptedException {
+ while (true) {
+ throwIfTerminated();
+ long waitTimeMs = endTimeMs - System.currentTimeMillis();
+ if (latch.await(Math.min(waitTimeMs, TERMINATION_CHECK_INTERVAL_MS), TimeUnit.MILLISECONDS)) {
+ return true;
+ }
+ if (waitTimeMs <= TERMINATION_CHECK_INTERVAL_MS) {
+ return false;
+ }
+ }
+ }
+
+ private static void throwIfTerminated() {
+ TerminationException terminateException = QueryThreadContext.getTerminateException();
+ if (terminateException != null) {
+ throw terminateException;
+ }
+ }
+
protected ExceptionResultsBlock getTimeoutResultsBlock(int numBlocksMerged) {
String logMsg = "Timed out while polling results block, numBlocksMerged: " + numBlocksMerged + " (query: "
+ _queryContext + ")";
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseSingleBlockCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseSingleBlockCombineOperator.java
index 3a00498c08d..3ea32dee657 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseSingleBlockCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseSingleBlockCombineOperator.java
@@ -22,7 +22,6 @@ import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
-import java.util.concurrent.TimeUnit;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
@@ -146,7 +145,7 @@ public abstract class BaseSingleBlockCombineOperator<T extends BaseResultsBlock>
if (waitTimeMs <= 0) {
return getTimeoutResultsBlock(numBlocksMerged);
}
- BaseResultsBlock blockToMerge = _blockingQueue.poll(waitTimeMs, TimeUnit.MILLISECONDS);
+ BaseResultsBlock blockToMerge = pollResultsBlock(_blockingQueue, endTimeMs);
if (blockToMerge == null) {
return getTimeoutResultsBlock(numBlocksMerged);
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java
index 30ced63ebff..1d168fbdbff 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java
@@ -25,7 +25,6 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
import org.apache.pinot.common.metrics.ServerMeter;
import org.apache.pinot.common.metrics.ServerMetrics;
import org.apache.pinot.core.common.Operator;
@@ -259,8 +258,9 @@ public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<Group
@Override
public BaseResultsBlock mergeResults()
throws Exception {
- long timeoutMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
- boolean opCompleted = _operatorLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
+ long endTimeMs = _queryContext.getEndTimeMs();
+ long timeoutMs = endTimeMs - System.currentTimeMillis();
+ boolean opCompleted = awaitLatch(_operatorLatch, endTimeMs);
if (!opCompleted) {
// If this happens, the broker side should already timed out, just log the error and return
String userError = "Timed out while combining group-by order-by results after " + timeoutMs + "ms";
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/MinMaxValueBasedSelectionOrderByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/MinMaxValueBasedSelectionOrderByCombineOperator.java
index b349ed7ef6e..001f79ec36b 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/MinMaxValueBasedSelectionOrderByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/MinMaxValueBasedSelectionOrderByCombineOperator.java
@@ -21,7 +21,6 @@ package org.apache.pinot.core.operator.combine;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.pinot.common.request.context.ExpressionContext;
@@ -259,8 +258,7 @@ public class MinMaxValueBasedSelectionOrderByCombineOperator
int numBlocksMerged = 0;
long endTimeMs = _queryContext.getEndTimeMs();
while (numBlocksMerged < _numOperators) {
- BaseResultsBlock blockToMerge =
- _blockingQueue.poll(endTimeMs - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
+ BaseResultsBlock blockToMerge = pollResultsBlock(_blockingQueue, endTimeMs);
if (blockToMerge == null) {
// Query times out, skip merging the remaining results blocks
String logMsg = "Timed out while polling results block, numBlocksMerged: " + numBlocksMerged + " (query: "
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SequentialSortedGroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SequentialSortedGroupByCombineOperator.java
index 0aac1776026..221b5333f7c 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SequentialSortedGroupByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SequentialSortedGroupByCombineOperator.java
@@ -21,7 +21,6 @@ package org.apache.pinot.core.operator.combine;
import java.util.Comparator;
import java.util.List;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.data.table.Record;
@@ -127,7 +126,7 @@ public class SequentialSortedGroupByCombineOperator extends BaseSingleBlockCombi
if (waitTimeMs <= 0) {
return getTimeoutResultsBlock(numBlocksMerged);
}
- BaseResultsBlock blockToMerge = _blockingQueue.poll(waitTimeMs, TimeUnit.MILLISECONDS);
+ BaseResultsBlock blockToMerge = pollResultsBlock(_blockingQueue, endTimeMs);
if (blockToMerge == null) {
return getTimeoutResultsBlock(numBlocksMerged);
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SortedGroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SortedGroupByCombineOperator.java
index 992beff9352..0117a5c700e 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SortedGroupByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/SortedGroupByCombineOperator.java
@@ -22,7 +22,6 @@ import java.util.Comparator;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.core.common.Operator;
@@ -198,8 +197,9 @@ public class SortedGroupByCombineOperator extends BaseSingleBlockCombineOperator
@Override
public BaseResultsBlock mergeResults()
throws Exception {
- long timeoutMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
- boolean opCompleted = _operatorLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
+ long endTimeMs = _queryContext.getEndTimeMs();
+ long timeoutMs = endTimeMs - System.currentTimeMillis();
+ boolean opCompleted = awaitLatch(_operatorLatch, endTimeMs);
if (!opCompleted) {
// If this happens, the broker side should already timed out, just log the error and return
String userError = "Timed out while combining group-by order-by results after " + timeoutMs + "ms";
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/BaseStreamingCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/BaseStreamingCombineOperator.java
index ca1b5930bbe..81205702f1d 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/BaseStreamingCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/BaseStreamingCombineOperator.java
@@ -76,8 +76,7 @@ public abstract class BaseStreamingCombineOperator<T extends BaseResultsBlock> e
try {
while (!_querySatisfied && _numOperatorsFinished < _numOperators) {
QueryThreadContext.checkTermination(this::getExplainName);
- BaseResultsBlock resultsBlock =
- _blockingQueue.poll(endTimeMs - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
+ BaseResultsBlock resultsBlock = pollResultsBlock(_blockingQueue, endTimeMs);
if (resultsBlock == null) {
// Query times out, skip streaming the remaining results blocks
throw QueryErrorCode.EXECUTION_TIMEOUT.asException("Timed out while polling results block");
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingDistinctCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingDistinctCombineOperator.java
index 0ed892c8101..c2d349beb61 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingDistinctCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingDistinctCombineOperator.java
@@ -21,7 +21,6 @@ package org.apache.pinot.core.operator.streaming;
import com.google.common.base.Preconditions;
import java.util.List;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
import org.apache.pinot.core.common.Operator;
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock.EarlyTerminationReason;
@@ -132,8 +131,7 @@ public class StreamingDistinctCombineOperator extends BaseStreamingCombineOperat
try {
while (!_querySatisfied && _numOperatorsFinished < _numOperators) {
QueryThreadContext.checkTermination(this::getExplainName);
- BaseResultsBlock resultsBlock =
- _blockingQueue.poll(endTimeMs - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
+ BaseResultsBlock resultsBlock = pollResultsBlock(_blockingQueue, endTimeMs);
if (resultsBlock == null) {
throw QueryErrorCode.EXECUTION_TIMEOUT.asException("Timed out while streaming distinct results");
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java
index d76da326ed8..835d2bf6187 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/streaming/StreamingGroupByCombineOperator.java
@@ -24,7 +24,6 @@ import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.TimeUnit;
import org.apache.pinot.common.utils.DataSchema;
import org.apache.pinot.common.utils.HashUtil;
import org.apache.pinot.core.common.Operator;
@@ -126,8 +125,7 @@ public class StreamingGroupByCombineOperator extends BaseStreamingCombineOperato
try {
while (_numOperatorsFinished < _numOperators) {
QueryThreadContext.checkTermination(this::getExplainName);
- BaseResultsBlock resultsBlock =
- _blockingQueue.poll(endTimeMs - System.currentTimeMillis(), TimeUnit.MILLISECONDS);
+ BaseResultsBlock resultsBlock = pollResultsBlock(_blockingQueue, endTimeMs);
if (resultsBlock == null) {
throw QueryErrorCode.EXECUTION_TIMEOUT.asException("Timed out while streaming group-by results");
}
diff --git a/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineTerminationTest.java b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineTerminationTest.java
new file mode 100644
index 00000000000..5cea24f412d
--- /dev/null
+++ b/pinot-core/src/test/java/org/apache/pinot/core/operator/combine/CombineTerminationTest.java
@@ -0,0 +1,241 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.operator.combine;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import org.apache.pinot.core.common.Block;
+import org.apache.pinot.core.common.Operator;
+import org.apache.pinot.core.operator.BaseOperator;
+import org.apache.pinot.core.operator.ExecutionStatistics;
+import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
+import org.apache.pinot.core.operator.streaming.BaseStreamingCombineOperator;
+import org.apache.pinot.core.operator.streaming.StreamingSelectionOnlyCombineOperator;
+import org.apache.pinot.core.query.request.context.QueryContext;
+import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.exception.QueryErrorMessage;
+import org.apache.pinot.spi.query.QueryExecutionContext;
+import org.apache.pinot.spi.query.QueryThreadContext;
+import org.apache.pinot.util.TestUtils;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Tests that a combine operator returns promptly when its query is terminated, even when the worker tasks are
+/// cancelled while still queued in the executor. A cancelled queued task never runs, so it never signals completion
+/// to the main thread. The main thread must observe the termination itself instead of waiting until the query
+/// deadline while it keeps the merged results reachable.
+@SuppressWarnings("rawtypes")
+public class CombineTerminationTest {
+ private static final int NUM_OPERATORS = 10;
+ private static final int NUM_THREADS = 2;
+ private static final long QUERY_TIMEOUT_MS = 60_000L;
+ private static final long MAX_RETURN_TIME_MS = 10_000L;
+ // The combine operator checks termination every 100 ms. Leave a margin for slow CI machines.
+ private static final long MAX_TERMINATION_LATENCY_MS = 5_000L;
+
+ private ThreadPoolExecutor _workerPool;
+ private CountDownLatch _releaseBlockers;
+ private ExecutorService _combineExecutor;
+
+ @BeforeMethod
+ public void setUp()
+ throws InterruptedException {
+ // Saturate the worker pool, so that every task the combine operator submits stays queued
+ _workerPool = (ThreadPoolExecutor) Executors.newFixedThreadPool(NUM_THREADS);
+ _releaseBlockers = new CountDownLatch(1);
+ CountDownLatch blockersStarted = new CountDownLatch(NUM_THREADS);
+ for (int i = 0; i < NUM_THREADS; i++) {
+ _workerPool.submit(() -> {
+ blockersStarted.countDown();
+ _releaseBlockers.await();
+ return null;
+ });
+ }
+ assertTrue(blockersStarted.await(MAX_RETURN_TIME_MS, TimeUnit.MILLISECONDS));
+ // Avoid early finalization by not using Executors.newSingleThreadExecutor (java <= 20, JDK-8145304)
+ _combineExecutor = Executors.newFixedThreadPool(1);
+ }
+
+ @AfterMethod
+ public void tearDown() {
+ _releaseBlockers.countDown();
+ _workerPool.shutdownNow();
+ _combineExecutor.shutdownNow();
+ }
+
+ @Test
+ public void testSelectionOnlyCombineOperator()
+ throws Exception {
+ testTerminateWithQueuedTasks("SELECT * FROM testTable",
+ (args) -> new SelectionOnlyCombineOperator(args._operators, args._queryContext, args._executorService));
+ }
+
+ @Test
+ public void testSelectionOrderByCombineOperator()
+ throws Exception {
+ testTerminateWithQueuedTasks("SELECT * FROM testTable ORDER BY column",
+ (args) -> new SelectionOrderByCombineOperator(args._operators, args._queryContext, args._executorService));
+ }
+
+ @Test
+ public void testAggregationCombineOperator()
+ throws Exception {
+ testTerminateWithQueuedTasks("SELECT COUNT(*) FROM testTable",
+ (args) -> new AggregationCombineOperator(args._operators, args._queryContext, args._executorService));
+ }
+
+ @Test
+ public void testGroupByCombineOperator()
+ throws Exception {
+ testTerminateWithQueuedTasks("SELECT COUNT(*) FROM testTable GROUP BY column",
+ (args) -> new GroupByCombineOperator(args._operators, args._queryContext, args._executorService));
+ }
+
+ @Test
+ public void testSortedGroupByCombineOperator()
+ throws Exception {
+ testTerminateWithQueuedTasks(
+ "SET sortAggregateSingleThreadedNumSegmentsThreshold=1; "
+ + "SELECT column, COUNT(*) FROM testTable GROUP BY column ORDER BY column LIMIT 10",
+ (args) -> new SortedGroupByCombineOperator(args._operators, args._queryContext, args._executorService));
+ }
+
+ @Test
+ public void testSequentialSortedGroupByCombineOperator()
+ throws Exception {
+ testTerminateWithQueuedTasks("SELECT column, COUNT(*) FROM testTable GROUP BY column ORDER BY column LIMIT 10",
+ (args) -> new SequentialSortedGroupByCombineOperator(args._operators, args._queryContext,
+ args._executorService));
+ }
+
+ @Test
+ public void testStreamingSelectionOnlyCombineOperator()
+ throws Exception {
+ testTerminateWithQueuedTasks("SELECT * FROM testTable",
+ (args) -> new StreamingSelectionOnlyCombineOperator(args._operators, args._queryContext,
+ args._executorService));
+ }
+
+ private void testTerminateWithQueuedTasks(String query, Function<CombineArgs, BaseCombineOperator<?>> factory)
+ throws Exception {
+ List<Operator> operators = new ArrayList<>(NUM_OPERATORS);
+ for (int i = 0; i < NUM_OPERATORS; i++) {
+ operators.add(new NeverScheduledOperator());
+ }
+ QueryContext queryContext = QueryContextConverterUtils.getQueryContext(query);
+ queryContext.setEndTimeMs(System.currentTimeMillis() + QUERY_TIMEOUT_MS);
+ QueryExecutionContext executionContext = QueryExecutionContext.forSseTest();
+ CountDownLatch combineStarted = new CountDownLatch(1);
+ AtomicReference<Thread> combineThread = new AtomicReference<>();
+
+ // Run the combine operator like the query scheduler does: on a thread with its own query context, and with a
+ // context aware executor that registers the worker tasks for cancellation
+ Future<BaseResultsBlock> future = _combineExecutor.submit(() -> {
+ try (QueryThreadContext ignore = QueryThreadContext.open(executionContext,
+ ThreadAccountantUtils.getNoOpAccountant())) {
+ ExecutorService executorService = QueryThreadContext.contextAwareExecutorService(_workerPool);
+ BaseCombineOperator<?> combineOperator = factory.apply(new CombineArgs(operators, queryContext,
+ executorService));
+ combineThread.set(Thread.currentThread());
+ combineStarted.countDown();
+ if (combineOperator instanceof BaseStreamingCombineOperator) {
+ ((BaseStreamingCombineOperator<?>) combineOperator).start();
+ }
+ return (BaseResultsBlock) combineOperator.nextBlock();
+ }
+ });
+ assertTrue(combineStarted.await(MAX_RETURN_TIME_MS, TimeUnit.MILLISECONDS));
+ // Wait until the combine operator has submitted its tasks and its thread waits for the results, so that
+ // termination cancels the tasks while they are queued and the wait loop must observe the termination
+ TestUtils.waitForCondition(
+ aVoid -> !_workerPool.getQueue().isEmpty() && combineThread.get().getState() == Thread.State.TIMED_WAITING,
+ MAX_RETURN_TIME_MS, "Combine operator did not wait for its tasks");
+
+ long terminateTimeMs = System.currentTimeMillis();
+ assertTrue(executionContext.terminate(QueryErrorCode.QUERY_CANCELLATION, "terminated for test"));
+
+ BaseResultsBlock resultsBlock = future.get(MAX_RETURN_TIME_MS, TimeUnit.MILLISECONDS);
+ long returnTimeMs = System.currentTimeMillis() - terminateTimeMs;
+ assertTrue(returnTimeMs < MAX_TERMINATION_LATENCY_MS, "Combine operator returned after " + returnTimeMs + "ms");
+ List<QueryErrorMessage> errorMessages = resultsBlock.getErrorMessages();
+ assertNotNull(errorMessages);
+ assertEquals(errorMessages.size(), 1);
+ assertEquals(errorMessages.get(0).getErrCode(), QueryErrorCode.QUERY_CANCELLATION);
+ for (Operator operator : operators) {
+ assertFalse(((NeverScheduledOperator) operator)._called);
+ }
+ }
+
+ private static class CombineArgs {
+ final List<Operator> _operators;
+ final QueryContext _queryContext;
+ final ExecutorService _executorService;
+
+ CombineArgs(List<Operator> operators, QueryContext queryContext, ExecutorService executorService) {
+ _operators = operators;
+ _queryContext = queryContext;
+ _executorService = executorService;
+ }
+ }
+
+ /// An operator whose task is cancelled before it starts, so it must never be called.
+ private static class NeverScheduledOperator extends BaseOperator {
+ private static final String EXPLAIN_NAME = "NEVER_SCHEDULED";
+
+ volatile boolean _called;
+
+ @Override
+ protected Block getNextBlock() {
+ _called = true;
+ throw new IllegalStateException("Operator must not be called after its task is cancelled");
+ }
+
+ @Override
+ public String toExplainString() {
+ return EXPLAIN_NAME;
+ }
+
+ @Override
+ public List<Operator> getChildOperators() {
+ return List.of();
+ }
+
+ @Override
+ public ExecutionStatistics getExecutionStatistics() {
+ return new ExecutionStatistics(0, 0, 0, 0);
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org