This is an automated email from the ASF dual-hosted git repository.
yashmayya 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 1ecbd5c99db Fail op chains of terminated MSE queries instead of
dropping them (#19854)
1ecbd5c99db is described below
commit 1ecbd5c99db783075b04aa365f5a302fa11b3691
Author: Yash Mayya <[email protected]>
AuthorDate: Fri Oct 16 09:28:14 2026 -0700
Fail op chains of terminated MSE queries instead of dropping them (#19854)
---
.../apache/pinot/query/runtime/QueryRunner.java | 9 +-
.../runtime/executor/OpChainSchedulerService.java | 131 ++++++++++----
.../pinot/query/runtime/plan/OpChainConverter.java | 6 +-
.../plan/pipeline/PipelineBreakerOperator.java | 12 ++
.../pinot/query/service/server/QueryServer.java | 9 +-
.../apache/pinot/query/QueryServerEnclosure.java | 9 +-
.../HardLimitExecutorOpChainRegressionTest.java | 13 +-
.../executor/OpChainSchedulerServiceTest.java | 175 +++++++++++++++----
.../plan/OpChainConverterDispatcherTest.java | 5 +-
.../plan/pipeline/PipelineBreakerExecutorTest.java | 51 ++++++
.../queries/GatedExecutorServicePlugin.java | 142 +++++++++++++++
.../query/runtime/queries/QueryRunnerTestBase.java | 10 +-
.../QueryTerminatedBeforeStageStartsTest.java | 190 +++++++++++++++++++++
.../apache/pinot/spi/query/QueryThreadContext.java | 10 +-
14 files changed, 678 insertions(+), 94 deletions(-)
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
index d632202afe8..c65756cda57 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/QueryRunner.java
@@ -332,11 +332,10 @@ public class QueryRunner {
// This can fail if the executor rejects the task.
_opChainScheduler.register(opChain);
} catch (RuntimeException e) {
- // register() can fail after the op chain was built — the query was
cancelled/terminated before scheduling
- // (checkTermination / cancelled-query cache), or the executor rejected
the task (e.g. under load shedding). In
- // those cases the op chain is never scheduled, so the scheduler's
FutureCallback, which normally close()s the
- // chain on completion, never fires — leaking any operator resources
that are released only in close(). Close the
- // built chain here (idempotent; guarded so a close failure cannot mask
the original error).
+ // register() can fail after the op chain was built, when the executor
rejects the task (e.g. under load
+ // shedding). In that case the op chain is never scheduled, so the
scheduler's FutureCallback, which normally
+ // close()s the chain on completion, never fires — leaking any operator
resources that are released only in
+ // close(). Close the built chain here (idempotent; guarded so a close
failure cannot mask the original error).
if (opChain != null) {
try {
opChain.close();
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java
index 24cbdef09ec..694d54ebd41 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerService.java
@@ -21,6 +21,7 @@ package org.apache.pinot.query.runtime.executor;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
+import com.google.common.util.concurrent.AbstractFuture;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFutureTask;
@@ -35,6 +36,7 @@ import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.common.metrics.MseMeter;
import org.apache.pinot.common.metrics.MseMetrics;
@@ -46,7 +48,6 @@ import org.apache.pinot.query.runtime.operator.OpChain;
import org.apache.pinot.query.runtime.operator.OpChainId;
import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.spi.env.PinotConfiguration;
-import org.apache.pinot.spi.exception.QueryCancelledException;
import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.exception.TerminationException;
import org.apache.pinot.spi.metrics.PinotMeter;
@@ -112,44 +113,38 @@ public class OpChainSchedulerService {
.build();
}
+ /// Schedules the op chain on the executor.
+ ///
+ /// If the query is terminated, or was cancelled on this server (which
terminates it), the op chain does not run: it
+ /// is cancelled with the termination error on the calling thread, which
sends that error to its receivers. The op
+ /// chain must not be dropped, because its receivers would then wait for
data until the query deadline.
+ ///
+ /// @throws RuntimeException if the executor refuses the op chain
public void register(OpChain operatorChain) {
QueryExecutionContext executionContext =
QueryThreadContext.get().getExecutionContext();
- // Check if query is already terminated before acquiring the read lock.
- checkTermination(operatorChain, executionContext);
// Acquire read lock for the query to ensure that the query is not
cancelled while scheduling the operator chain.
long requestId = operatorChain.getId().getRequestId();
Lock readLock = getQueryLock(requestId).readLock();
+ TerminationException terminateException;
readLock.lock();
try {
- // Check if query is already terminated again after acquiring the read
lock.
- checkTermination(operatorChain, executionContext);
- // Do not schedule the operator chain if the query has been cancelled.
if (_cancelledQueryCache.getIfPresent(requestId) != null) {
LOGGER.debug("({}): Query has been cancelled", operatorChain);
executionContext.terminate(QueryErrorCode.QUERY_CANCELLATION,
"Cancelled on: " + _instanceId);
- throw new QueryCancelledException(
- "Query has been cancelled before op-chain: " +
operatorChain.getId() + " being scheduled");
- } else {
+ }
+ terminateException = executionContext.getTerminateException();
+ if (terminateException == null) {
registerInternal(operatorChain, executionContext);
+ return;
}
} finally {
readLock.unlock();
}
- }
-
- private void checkTermination(OpChain operatorChain, QueryExecutionContext
executionContext) {
- TerminationException terminateException =
executionContext.getTerminateException();
- if (terminateException != null) {
- LOGGER.debug("({}): Query has been terminated", operatorChain,
terminateException);
- if (terminateException.getErrorCode() ==
QueryErrorCode.QUERY_CANCELLATION) {
- throw new QueryCancelledException(
- "Query has been cancelled before op-chain: " +
operatorChain.getId() + " being scheduled");
- } else {
- throw new QueryCancelledException(
- "Query has been terminated before op-chain: " +
operatorChain.getId() + " being scheduled: "
- + terminateException.getErrorCode() + " - " +
terminateException.getMessage(), terminateException);
- }
- }
+ // Not scheduled, so that the error does not depend on the executor
accepting the op chain
+ LOGGER.debug("({}): Query has been terminated", operatorChain,
terminateException);
+ notifyCompletionListener(operatorChain.getId(), operatorChain, null,
terminateException);
+ operatorChain.cancel(terminateException);
+ operatorChain.close();
}
private void registerInternal(OpChain operatorChain, QueryExecutionContext
executionContext) {
@@ -173,11 +168,17 @@ public class OpChainSchedulerService {
// that by passing null.
AtomicReference<MultiStageQueryStats> statsRef = new AtomicReference<>();
- // Create a ListenableFutureTask to ensure the opChain is cancelled even
if the task is not scheduled
+ // The FutureCallback below runs on the thread that runs the op chain,
once the op chain completes
ListenableFutureTask<Void> listenableFutureTask =
ListenableFutureTask.create(new TraceRunnable() {
@Override
public void runJob() {
_metrics.onOpChainStarted();
+ // The query was terminated while the op chain waited for a thread.
Fail without doing any work: the
+ // FutureCallback then cancels the op chain with the termination
error, which sends that error to the receivers.
+ TerminationException terminateException =
executionContext.getTerminateException();
+ if (terminateException != null) {
+ throw terminateException;
+ }
LOGGER.trace("({}): Executing", operatorChain);
MseBlock result = rootOperator.nextBlock();
while (result.isData()) {
@@ -228,19 +229,25 @@ public class OpChainSchedulerService {
}
}, MoreExecutors.directExecutor());
+ OpChainTask opChainTask = new OpChainTask(listenableFutureTask);
try {
- _executorService.submit(listenableFutureTask);
+ // Use execute() and not submit(): submit() registers its Future with
the execution context, and terminating the
+ // query would cancel that Future, which drops the op chain if it has
not started yet. See OpChainTask.
+ _executorService.execute(opChainTask);
} catch (RuntimeException e) {
- // The MSE executor is wrapped in HardLimitExecutor + heap throttling,
so submit() can throw
- // RejectedExecutionException. When it does, the task never runs, so the
directExecutor FutureCallback above
- // (which decrements the active-opchain counter and removes the
per-request context entry) never fires. Back
- // that bookkeeping out here so the per-request context entry does not
leak until a later cancel, then rethrow
- // so the caller propagates the failure as a stage error. The caller,
QueryRunner#processQueryBlocking,
- // close()s the op chain on this rethrow, releasing operator resources
that the never-fired FutureCallback
- // would otherwise have closed.
+ // The MSE executor is wrapped in HardLimitExecutor + heap throttling,
so execute() can throw (e.g.
+ // IllegalStateException when the limit is reached). When it does, the
task never runs, so the directExecutor
+ // FutureCallback above (which decrements the active-opchain counter and
removes the per-request context entry)
+ // never fires. Back that bookkeeping out here so the per-request
context entry does not leak until a later
+ // cancel, then rethrow so the caller propagates the failure as a stage
error. The caller,
+ // QueryRunner#processQueryBlocking, close()s the op chain on this
rethrow, releasing operator resources that the
+ // never-fired FutureCallback would otherwise have closed.
decrementActiveOpChains(requestId);
throw e;
}
+ // Registered after execute(), so that a task the executor refuses is not
kept by the context. If the query is
+ // terminated in between, addTask() cancels the task right away.
+ executionContext.addTask(opChainTask);
}
private void decrementActiveOpChains(long requestId) {
@@ -302,12 +309,12 @@ public class OpChainSchedulerService {
/// opchain finishes after being interrupted. In legacy mode, cancel-path
stats are not collected.
///
/// If no opchains are currently registered (pre-registration cancel race),
the requestId is marked in the
- /// cancelled-query cache so that any future registrations for that
requestId are immediately rejected.
+ /// cancelled-query cache so that the op chains registered later for that
requestId fail without running.
public void cancel(long requestId) {
// Acquire the write lock BEFORE eviction so that register()'s readLock
cannot overlap with the window
// between the context eviction and the cancelled-query cache write.
Without the write lock first, a
// concurrent register() could pass the cache check (cache not yet
written) after the context was already
- // evicted, slipping through and starting an opchain that should have been
rejected.
+ // evicted, slipping through and running an opchain of a cancelled query.
AtomicReference<QueryExecutionContext> ctxRef = new AtomicReference<>();
Lock writeLock = getQueryLock(requestId).writeLock();
writeLock.lock();
@@ -325,8 +332,8 @@ public class OpChainSchedulerService {
}
QueryExecutionContext context = ctxRef.get();
if (context != null) {
- // terminate() interrupts all registered tasks (via addTask) and sets
the termination flag so that
- // checkTermination() blocks any subsequent registrations for the same
requestId.
+ // terminate() interrupts all registered tasks (via addTask) and sets
the termination flag, so that op chains
+ // that have not started yet, and the ones registered later for the same
requestId, fail without running.
context.terminate(QueryErrorCode.QUERY_CANCELLATION, "Cancelled on: " +
_instanceId);
}
}
@@ -335,6 +342,56 @@ public class OpChainSchedulerService {
return _queryLocks[(int) (requestId & QUERY_LOCK_MASK)];
}
+ /// Runs an op chain. It is registered with the [QueryExecutionContext] of
the query, so that terminating the query
+ /// interrupts the op chain.
+ ///
+ /// [QueryExecutionContext#terminate] cancels the registered tasks. A
`Future` from [ExecutorService#submit] that is
+ /// cancelled before its task starts never runs the task. An op chain
dropped this way never sends an EOS or an error
+ /// to its receivers, so they wait for data until the query deadline.
Cancelling this task only interrupts the op
+ /// chain while it runs. An op chain that has not started yet still runs,
and fails on its first check: the context
+ /// sets the termination error before it cancels the tasks.
+ ///
+ /// `_lock` makes sure that an interrupt reaches the thread only while it
runs the op chain, and never the next task
+ /// of that thread.
+ private static final class OpChainTask extends AbstractFuture<Void>
implements Runnable {
+ private final Runnable _opChainRunner;
+ private final Object _lock = new Object();
+ @GuardedBy("_lock")
+ @Nullable
+ private Thread _thread;
+
+ OpChainTask(Runnable opChainRunner) {
+ _opChainRunner = opChainRunner;
+ }
+
+ @Override
+ public void run() {
+ synchronized (_lock) {
+ _thread = Thread.currentThread();
+ }
+ try {
+ _opChainRunner.run();
+ } finally {
+ synchronized (_lock) {
+ _thread = null;
+ }
+ // _thread is null, so interruptTask() can no longer interrupt this
thread. Clear an interrupt that it sent
+ // earlier, after the op chain stopped checking for interrupts.
+ Thread.interrupted();
+ set(null);
+ }
+ }
+
+ @Override
+ protected void interruptTask() {
+ synchronized (_lock) {
+ if (_thread != null) {
+ _thread.interrupt();
+ }
+ }
+ }
+ }
+
private static class Metrics {
private final MseMetrics _mseMetrics = MseMetrics.get();
private final PinotMeter _startedOpchains =
_mseMetrics.getMeteredValue(MseMeter.OPCHAINS_STARTED);
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainConverter.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainConverter.java
index cd87d58f60d..1f3cf546d2e 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainConverter.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/OpChainConverter.java
@@ -29,7 +29,6 @@ import org.apache.pinot.query.runtime.operator.ErrorOperator;
import org.apache.pinot.query.runtime.operator.MailboxSendOperator;
import org.apache.pinot.query.runtime.operator.MultiStageOperator;
import org.apache.pinot.query.runtime.operator.OpChain;
-import org.apache.pinot.spi.exception.QueryErrorCode;
/// Extension point for alternative plan-to-[OpChain] converters. The active
implementation is chosen by
@@ -57,7 +56,8 @@ public interface OpChainConverter {
/// can keep track of this mapping.
OpChain convert(PlanNode node, OpChainExecutionContext context,
BiConsumer<PlanNode, MultiStageOperator> tracker);
- /// Builds a minimal op-chain that sends an already-detected stage error to
downstream receivers.
+ /// Builds a minimal op-chain that sends an already-detected stage error,
with its error code, to downstream
+ /// receivers.
///
/// This is used, for example, when a pipeline breaker detects an error and
needs to propagate it to downstream
/// stages.
@@ -71,7 +71,7 @@ public interface OpChainConverter {
}
MailboxSendNode rootNode = (MailboxSendNode) stagePlan.getRootNode();
String message = extractErrorMessage(errorBlock);
- MultiStageOperator errorOp = new ErrorOperator(context,
QueryErrorCode.QUERY_EXECUTION, message);
+ MultiStageOperator errorOp = new ErrorOperator(context,
errorBlock.getMainErrorCode(), message);
MultiStageOperator sendOp = new MailboxSendOperator(context, errorOp,
rootNode);
return new OpChain(context, sendOp);
}
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
index a9f10d7a0d9..318ee51c9e2 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerOperator.java
@@ -101,6 +101,18 @@ public class PipelineBreakerOperator extends
MultiStageOperator {
return _errorBlock;
}
+ /// Keeps the cancellation reason as the error, so that a pipeline breaker
cancelled before it collected its results
+ /// (e.g. the query was terminated before it started) does not look like it
collected no rows.
+ @Override
+ public void cancel(Throwable e) {
+ if (_errorBlock == null) {
+ // Not only an Exception: the scheduler also cancels an op chain that
failed with an Error, e.g. an
+ // OutOfMemoryError while the pipeline breaker collects its results
+ _errorBlock = ErrorMseBlock.fromException(e instanceof Exception ?
(Exception) e : new RuntimeException(e));
+ }
+ super.cancel(e);
+ }
+
@Override
public String toExplainString() {
return EXPLAIN_NAME;
diff --git
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
index bd5bde1ee96..5388a26a92e 100644
---
a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
+++
b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java
@@ -755,11 +755,10 @@ public class QueryServer extends
PinotQueryWorkerGrpc.PinotQueryWorkerImplBase {
} catch (Throwable t) {
LOGGER.warn("Caught exception cancelling request {} via
SubmitWithStream", cancel.getRequestId(), t);
}
- // Once this stream's request is cancelled, no further meaningful
OpChainComplete reports will arrive — in
- // particular, opchains cancelled before they register never fire the
completion listener, so _completedOpChains
- // would never reach _expectedOpChains and ServerDone would never be
sent. Close the stream promptly so the
- // broker stops waiting out its full drain window. Gated on the ack so
the broker still sees submit_ack before
- // ServerDone; cleanupListener and sendDoneAndComplete are both
idempotent, so racing onOpChainComplete is safe.
+ // Once this stream's request is cancelled, no further meaningful
OpChainComplete reports will arrive. Close the
+ // stream promptly so the broker stops waiting out its full drain
window. Gated on the ack so the broker still
+ // sees submit_ack before ServerDone; cleanupListener and
sendDoneAndComplete are both idempotent, so racing
+ // onOpChainComplete is safe.
if (cancel.getRequestId() == _requestId) {
_ackSentFuture.thenRun(() -> {
cleanupListener();
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java
index 42510b40f1a..706df8fe8d8 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/QueryServerEnclosure.java
@@ -82,7 +82,14 @@ public class QueryServerEnclosure {
public CompletableFuture<Void> processQuery(WorkerMetadata workerMetadata,
StagePlan stagePlan,
Map<String, String> requestMetadataMap) {
- QueryExecutionContext executionContext =
QueryExecutionContext.forMseServerRequest(requestMetadataMap, "serverId");
+ return processQuery(workerMetadata, stagePlan, requestMetadataMap,
+ QueryExecutionContext.forMseServerRequest(requestMetadataMap,
"serverId"));
+ }
+
+ /// Like [#processQuery(WorkerMetadata, StagePlan, Map)], but runs the stage
in the given execution context. Like the
+ /// query server, callers can share one context between all the stages of a
query on this server.
+ public CompletableFuture<Void> processQuery(WorkerMetadata workerMetadata,
StagePlan stagePlan,
+ Map<String, String> requestMetadataMap, QueryExecutionContext
executionContext) {
QueryThreadContext.MseWorkerInfo mseWorkerInfo =
new
QueryThreadContext.MseWorkerInfo(stagePlan.getStageMetadata().getStageId(),
workerMetadata.getWorkerId());
try (QueryThreadContext ignore = QueryThreadContext.open(executionContext,
mseWorkerInfo,
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/HardLimitExecutorOpChainRegressionTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/HardLimitExecutorOpChainRegressionTest.java
index ad5b8008286..2b458cf5c68 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/HardLimitExecutorOpChainRegressionTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/HardLimitExecutorOpChainRegressionTest.java
@@ -24,7 +24,6 @@ import java.util.Map;
import java.util.concurrent.AbstractExecutorService;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.apache.pinot.common.datatable.StatMap;
import org.apache.pinot.query.mailbox.MailboxService;
@@ -54,7 +53,9 @@ import static org.testng.Assert.assertTrue;
/// decrements the per-request counter never fired, and the exception was
swallowed by the `FutureTask` created by
/// `AbstractExecutorService.submit`, so nothing was logged either.
///
-/// Now the limit is checked only once, on submit. A task that passes that
check always runs.
+/// Now the limit is checked only once, on submit. A task that passes that
check always runs. The scheduler also hands
+/// op chains to the executor with `execute()` now, not `submit()`, so an
exception from the executor is no longer
+/// swallowed by a `FutureTask` either: the test runs the queued task
directly, and such an exception fails it.
public class HardLimitExecutorOpChainRegressionTest {
/// Executor that queues tasks instead of running them so the test controls
when each one starts.
@@ -139,11 +140,9 @@ public class HardLimitExecutorOpChainRegressionTest {
threadA.start();
assertTrue(aStarted.await(10, TimeUnit.SECONDS), "A did not start");
- // Start B: it passed the submit-time check, so it must run even though A
pushed the counter to the limit.
- Runnable taskB = deferred._queued.get(1);
- taskB.run();
- assertTrue(taskB instanceof Future, "AbstractExecutorService.submit wraps
the task in a FutureTask");
- ((Future<?>) taskB).get(10, TimeUnit.SECONDS);
+ // Start B: it passed the submit-time check, so it must run even though A
pushed the counter to the limit. It runs
+ // on this thread, so an exception from the executor would fail the test
here.
+ deferred._queued.get(1).run();
verify(operatorB).nextBlock();
verify(operatorB).close();
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java
index a1d21d8d151..eab550f6e99 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/executor/OpChainSchedulerServiceTest.java
@@ -20,6 +20,7 @@ package org.apache.pinot.query.runtime.executor;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -39,9 +40,11 @@ import
org.apache.pinot.query.runtime.operator.MultiStageOperator;
import org.apache.pinot.query.runtime.operator.OpChain;
import org.apache.pinot.query.runtime.plan.MultiStageQueryStats;
import org.apache.pinot.query.runtime.plan.OpChainExecutionContext;
-import org.apache.pinot.spi.exception.QueryCancelledException;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.executor.DecoratorExecutorService;
import org.apache.pinot.spi.executor.ExecutorServiceUtils;
import org.apache.pinot.spi.query.QueryThreadContext;
+import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;
import org.testng.annotations.AfterClass;
@@ -53,8 +56,9 @@ import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertSame;
import static org.testng.Assert.assertTrue;
-import static org.testng.Assert.fail;
public class OpChainSchedulerServiceTest {
@@ -203,6 +207,127 @@ public class OpChainSchedulerServiceTest {
Mockito.verify(_operatorA, Mockito.times(1)).cancel(Mockito.any());
}
+ /// An op chain that is still waiting for a thread when the query is
terminated must still fail with the termination
+ /// reason. Otherwise it never sends an EOS or an error to its receivers,
and they wait until the query deadline.
+ @Test
+ public void shouldCancelOpChainWhenQueryIsTerminatedBeforeItStarts()
+ throws InterruptedException {
+ ExecutorService singleThreadExecutor =
+ Executors.newSingleThreadExecutor(new
NamedThreadFactory("single_worker_on_" + getClass().getSimpleName()));
+ try {
+ // Keep the only thread busy, so that the op chain waits in the executor
queue until the query is terminated.
+ CountDownLatch releaseThread = new CountDownLatch(1);
+ singleThreadExecutor.execute(() -> {
+ try {
+ releaseThread.await();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ });
+ CountDownLatch cancelLatch = new CountDownLatch(1);
+ AtomicReference<Throwable> cancelCause = new AtomicReference<>();
+ Mockito.doAnswer(inv -> {
+ cancelCause.set(inv.getArgument(0));
+ cancelLatch.countDown();
+ return null;
+ }).when(_operatorA).cancel(Mockito.any());
+
+ OpChainSchedulerService schedulerService =
+ new
OpChainSchedulerService(QueryThreadContext.contextAwareExecutorService(singleThreadExecutor));
+ try (QueryThreadContext threadContext =
QueryThreadContext.openForMseTest()) {
+ schedulerService.register(getChain(_operatorA));
+ threadContext.getExecutionContext()
+ .terminate(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED, "OOM
killed on SERVER");
+ }
+ releaseThread.countDown();
+
+ assertTrue(cancelLatch.await(10, TimeUnit.SECONDS), "expected the op
chain to be cancelled");
+ assertEquals(QueryErrorCode.fromThrowable(cancelCause.get(),
QueryErrorCode.UNKNOWN),
+ QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED);
+ assertEquals(cancelCause.get().getMessage(), "OOM killed on SERVER");
+ Mockito.verify(_operatorA, Mockito.timeout(10_000)).close();
+ Mockito.verify(_operatorA, Mockito.never()).nextBlock();
+ } finally {
+ singleThreadExecutor.shutdownNow();
+ }
+ }
+
+ /// An op chain registered after its query is terminated does not run, and
must not be dropped either. It is
+ /// cancelled with the termination error before register() returns, without
going through the executor, so that its
+ /// receivers get that error.
+ @Test
+ public void shouldCancelOpChainRegisteredAfterQueryIsTerminated() {
+ AtomicReference<Throwable> listenerError = new AtomicReference<>();
+ OpChainSchedulerService schedulerService = new
OpChainSchedulerService(_executor);
+ schedulerService.registerCompletionListener(123L, (id, root, stats, ctx,
error) -> listenerError.set(error));
+ try (QueryThreadContext threadContext =
QueryThreadContext.openForMseTest()) {
+ threadContext.getExecutionContext()
+ .terminate(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED, "OOM
killed on SERVER");
+ schedulerService.register(getChain(_operatorA));
+ }
+
+ ArgumentCaptor<Throwable> cancelCause =
ArgumentCaptor.forClass(Throwable.class);
+ Mockito.verify(_operatorA).cancel(cancelCause.capture());
+ assertEquals(QueryErrorCode.fromThrowable(cancelCause.getValue(),
QueryErrorCode.UNKNOWN),
+ QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED);
+ assertEquals(cancelCause.getValue().getMessage(), "OOM killed on SERVER");
+ Mockito.verify(_operatorA).close();
+ Mockito.verify(_operatorA, Mockito.never()).nextBlock();
+ assertSame(listenerError.get(), cancelCause.getValue());
+ assertEquals(schedulerService.activeRequestCount(), 0);
+ }
+
+ /// Terminating the query interrupts the thread that runs the op chain. The
interrupt must not outlive the op chain,
+ /// or it would hit the next task of that thread.
+ @Test
+ public void shouldNotLeakInterruptAfterOpChainCompletes()
+ throws InterruptedException {
+ CountDownLatch registered = new CountDownLatch(1);
+ Mockito.when(_operatorA.nextBlock()).thenAnswer(inv -> {
+ registered.await();
+ return SuccessMseBlock.INSTANCE;
+ });
+ Mockito.doAnswer(inv ->
MultiStageQueryStats.emptyStats(0)).when(_operatorA).calculateStats();
+ // Records whether the thread is still interrupted when the task returns,
which the next task of a pooled thread
+ // would see
+ AtomicBoolean interruptedAfterTask = new AtomicBoolean();
+ CountDownLatch taskDone = new CountDownLatch(1);
+ ExecutorService threadPool = Executors.newCachedThreadPool();
+ ExecutorService recordingExecutor = new
DecoratorExecutorService(threadPool) {
+ @Override
+ protected <T> Callable<T> decorate(Callable<T> task) {
+ return task;
+ }
+
+ @Override
+ protected Runnable decorate(Runnable task) {
+ return () -> {
+ task.run();
+ interruptedAfterTask.set(Thread.currentThread().isInterrupted());
+ taskDone.countDown();
+ };
+ }
+ };
+ try {
+ OpChainSchedulerService schedulerService =
+ new
OpChainSchedulerService(QueryThreadContext.contextAwareExecutorService(recordingExecutor));
+ try (QueryThreadContext threadContext =
QueryThreadContext.openForMseTest()) {
+ // Terminate the query while the op chain closes, once it no longer
checks for interrupts
+ Mockito.doAnswer(inv -> threadContext.getExecutionContext()
+ .terminate(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED, "OOM
killed on SERVER"))
+ .when(_operatorA).close();
+ schedulerService.register(getChain(_operatorA));
+ }
+ registered.countDown();
+
+ assertTrue(taskDone.await(10, TimeUnit.SECONDS), "expected the op chain
to complete");
+ Mockito.verify(_operatorA).close();
+ assertFalse(interruptedAfterTask.get(), "The interrupt of the op chain
leaked into its thread");
+ } finally {
+ threadPool.shutdownNow();
+ }
+ }
+
@Test
public void shouldFireCompletionListenerOnSuccess()
throws InterruptedException {
@@ -294,28 +419,23 @@ public class OpChainSchedulerServiceTest {
Mockito.verify(_operatorA, Mockito.times(1)).cancel(Mockito.any());
}
+ /// The cancel of a query can reach the server before its stages. The op
chains registered after that do not run:
+ /// they are cancelled with the cancellation error.
@Test
- public void
shouldThrowQueryCancelledExceptionWhenRegisteringOpChainAfterQueryCancellation()
{
- Mockito.when(_operatorA.nextBlock()).thenReturn(SuccessMseBlock.INSTANCE);
- Mockito.doAnswer(inv ->
MultiStageQueryStats.emptyStats(1)).when(_operatorA).calculateStats();
-
+ public void shouldCancelOpChainRegisteredAfterQueryCancellation() {
OpChainSchedulerService schedulerService = new
OpChainSchedulerService(_executor);
+ // No op chain of the query is registered yet, so only the cancelled-query
cache keeps the cancellation
+ schedulerService.cancel(123L);
try (QueryThreadContext ignore = QueryThreadContext.openForMseTest()) {
schedulerService.register(getChain(_operatorA));
-
- // First cancel the query with the same requestId (123L as defined in
getChain method)
- schedulerService.cancel(123L);
-
- // Now try to register an OpChain for the same query - should throw
QueryCancelledException
- try {
- schedulerService.register(getChain(_operatorA));
- fail("Expected QueryCancelledException to be thrown when registering
OpChain after query cancellation");
- } catch (QueryCancelledException e) {
- String message = e.getMessage();
- assertTrue(message.contains("Query has been cancelled before
op-chain"));
- assertTrue(message.contains("being scheduled"));
- }
}
+
+ ArgumentCaptor<Throwable> cancelCause =
ArgumentCaptor.forClass(Throwable.class);
+ Mockito.verify(_operatorA).cancel(cancelCause.capture());
+ assertEquals(QueryErrorCode.fromThrowable(cancelCause.getValue(),
QueryErrorCode.UNKNOWN),
+ QueryErrorCode.QUERY_CANCELLATION);
+ Mockito.verify(_operatorA).close();
+ Mockito.verify(_operatorA, Mockito.never()).nextBlock();
}
/// Registers two opchains for the same request, waits for both to complete,
and verifies that the per-request
@@ -374,9 +494,6 @@ public class OpChainSchedulerServiceTest {
try {
schedulerService.register(getChain(_operatorA));
registrationSucceeded[0] = true;
- } catch (QueryCancelledException e) {
- // Expected if cancellation happens before registration
- registrationSucceeded[0] = false;
} finally {
registrationCompleted.countDown();
}
@@ -397,13 +514,13 @@ public class OpChainSchedulerServiceTest {
// The registration should have succeeded since it started before
cancellation
assertTrue(registrationSucceeded[0], "Registration should have
succeeded");
- // Now try to register another OpChain with the same requestId - should
fail
- try {
- schedulerService.register(getChain(_operatorA));
- fail("Expected QueryCancelledException for subsequent registration
after cancellation");
- } catch (QueryCancelledException e) {
- assertTrue(e.getMessage().contains("Query has been cancelled before
op-chain"));
- }
+ // An op chain registered for the same query after the cancellation does
not run, but is cancelled
+ MultiStageOperator operatorB = Mockito.mock(MultiStageOperator.class);
+ Mockito.when(operatorB.copyStatMaps()).thenAnswer(inv -> new
StatMap<>(MailboxSendOperator.StatKey.class));
+ schedulerService.register(getChain(operatorB));
+ Mockito.verify(operatorB,
Mockito.timeout(10_000)).cancel(Mockito.argThat(
+ t -> QueryErrorCode.fromThrowable(t, QueryErrorCode.UNKNOWN) ==
QueryErrorCode.QUERY_CANCELLATION));
+ Mockito.verify(operatorB, Mockito.never()).nextBlock();
}
}
}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/OpChainConverterDispatcherTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/OpChainConverterDispatcherTest.java
index f2f23446966..c9aebfe1039 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/OpChainConverterDispatcherTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/OpChainConverterDispatcherTest.java
@@ -86,7 +86,8 @@ public class OpChainConverterDispatcherTest {
new OpChainExecutionContext(mailboxService, 123L, "cid-123",
System.currentTimeMillis() + 60_000,
System.currentTimeMillis() + 60_000, "broker", Map.of(),
stageMetadata, workerMetadata, null, false, false);
StagePlan stagePlan = new StagePlan(createMailboxSendNode(schema, stageId,
receiverStageId), stageMetadata);
- ErrorMseBlock errorBlock =
ErrorMseBlock.fromError(QueryErrorCode.QUERY_EXECUTION, "simulated pipeline
failure");
+ ErrorMseBlock errorBlock =
+ ErrorMseBlock.fromError(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED,
"simulated pipeline failure");
OpChain opChain = OpChainConverterDispatcher.sendEarlyError(context,
stagePlan, errorBlock);
try {
@@ -98,7 +99,7 @@ public class OpChainConverterDispatcherTest {
Assert.assertTrue(block instanceof ErrorMseBlock);
ErrorMseBlock emittedError = (ErrorMseBlock) block;
QueryErrorCode mainErrorCode = emittedError.getMainErrorCode();
- Assert.assertEquals(mainErrorCode, QueryErrorCode.QUERY_EXECUTION);
+ Assert.assertEquals(mainErrorCode,
QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED);
String message = emittedError.getErrorMessages().get(mainErrorCode);
Assert.assertTrue(message.contains("simulated pipeline failure"),
"Unexpected error message from early-error opchain: " + message);
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutorTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutorTest.java
index d3baa636c94..2b5524a9f8a 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutorTest.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/pipeline/PipelineBreakerExecutorTest.java
@@ -43,11 +43,13 @@ import org.apache.pinot.query.routing.SharedMailboxInfos;
import org.apache.pinot.query.routing.StageMetadata;
import org.apache.pinot.query.routing.StagePlan;
import org.apache.pinot.query.routing.WorkerMetadata;
+import org.apache.pinot.query.runtime.blocks.ErrorMseBlock;
import org.apache.pinot.query.runtime.blocks.MseBlock;
import org.apache.pinot.query.runtime.blocks.SuccessMseBlock;
import org.apache.pinot.query.runtime.executor.OpChainSchedulerService;
import org.apache.pinot.query.runtime.operator.OperatorTestUtil;
import org.apache.pinot.spi.accounting.ThreadAccountantUtils;
+import org.apache.pinot.spi.exception.QueryErrorCode;
import org.apache.pinot.spi.executor.ExecutorServiceUtils;
import org.apache.pinot.spi.query.QueryExecutionContext;
import org.apache.pinot.spi.query.QueryThreadContext;
@@ -61,6 +63,8 @@ import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.INT;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -243,6 +247,53 @@ public class PipelineBreakerExecutorTest {
latch.countDown();
}
+ /// The query is terminated before its pipeline breaker starts. The pipeline
breaker does not run, and its result
+ /// carries the termination error instead of looking like it collected no
rows.
+ @Test
+ public void
shouldReturnTerminationErrorWhenQueryIsTerminatedBeforePBStarts() {
+ MailboxReceiveNode mailboxReceiveNode = getPBReceiveNode(1);
+ StagePlan stagePlan = new StagePlan(mailboxReceiveNode, _stageMetadata);
+
when(_mailboxService.getReceivingMailbox(MAILBOX_ID_1)).thenReturn(_mailbox1);
+
+ QueryExecutionContext executionContext =
+ new QueryExecutionContext(QueryExecutionContext.QueryType.MSE, 0, "0",
Accounting.DEFAULT_WORKLOAD_NAME,
+ System.currentTimeMillis(), Long.MAX_VALUE, Long.MAX_VALUE,
"brokerId", "serverId", "");
+ executionContext.terminate(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED,
"OOM killed on SERVER");
+ PipelineBreakerResult pipelineBreakerResult;
+ try (QueryThreadContext ignore = QueryThreadContext.open(executionContext,
+ new QueryThreadContext.MseWorkerInfo(1, 2),
ThreadAccountantUtils.getNoOpAccountant())) {
+ pipelineBreakerResult =
PipelineBreakerExecutor.executePipelineBreakers(_scheduler, _mailboxService,
+ _workerMetadata, stagePlan, Map.of(), true, true);
+ }
+
+ Assert.assertNotNull(pipelineBreakerResult);
+ ErrorMseBlock errorBlock = pipelineBreakerResult.getErrorBlock();
+ Assert.assertNotNull(errorBlock);
+ Assert.assertEquals(errorBlock.getErrorMessages(),
+ Map.of(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED, "OOM killed on
SERVER"));
+ verify(_mailbox1, never()).poll();
+ }
+
+ /// The scheduler cancels an op chain that fails with an Error too. The
pipeline breaker keeps that Error as its
+ /// error instead of looking like it collected no rows.
+ @Test
+ public void shouldReturnErrorBlockWhenPBFailsWithError() {
+ MailboxReceiveNode mailboxReceiveNode = getPBReceiveNode(1);
+ StagePlan stagePlan = new StagePlan(mailboxReceiveNode, _stageMetadata);
+
when(_mailboxService.getReceivingMailbox(MAILBOX_ID_1)).thenReturn(_mailbox1);
+ when(_mailbox1.poll()).thenThrow(new OutOfMemoryError("Java heap space
(test)"));
+
+ PipelineBreakerResult pipelineBreakerResult =
+ executePipelineBreakers(_scheduler, _mailboxService, _workerMetadata,
stagePlan, Map.of(), 0, Long.MAX_VALUE);
+
+ Assert.assertNotNull(pipelineBreakerResult);
+ ErrorMseBlock errorBlock = pipelineBreakerResult.getErrorBlock();
+ Assert.assertNotNull(errorBlock);
+ String message = errorBlock.getErrorMessages().get(QueryErrorCode.UNKNOWN);
+ Assert.assertNotNull(message, "Unexpected error: " +
errorBlock.getErrorMessages());
+ Assert.assertTrue(message.contains("Java heap space (test)"), message);
+ }
+
@Test
public void shouldReturnWhenAnyPBReturnsEmpty() {
MailboxReceiveNode mailboxReceiveNode1 = getPBReceiveNode(1);
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/GatedExecutorServicePlugin.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/GatedExecutorServicePlugin.java
new file mode 100644
index 00000000000..efd33023dd3
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/GatedExecutorServicePlugin.java
@@ -0,0 +1,142 @@
+/**
+ * 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.query.runtime.queries;
+
+import com.google.auto.service.AutoService;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.AbstractExecutorService;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import org.apache.pinot.spi.executor.ExecutorServicePlugin;
+import org.apache.pinot.spi.executor.ExecutorServiceProvider;
+
+
+/// [ExecutorServicePlugin] for tests that need to act after an op chain is
scheduled but before it starts.
+///
+/// While the gate is closed, the tasks that the threads of these executors
submit are parked instead of run, until
+/// [#openGate()]. In a query runner, these threads run
`QueryRunner#processQueryBlocking`, so the parked tasks are the
+/// op chains that it schedules. The stage submissions come from other threads
and run normally.
+///
+/// The gate is static, so it is shared by all the executors of this type in
the JVM. Tests that close it must open it
+/// again when they finish.
+@AutoService(ExecutorServicePlugin.class)
+public class GatedExecutorServicePlugin implements ExecutorServicePlugin {
+ public static final String ID = "gatedForTest";
+
+ private static final Set<Thread> EXECUTOR_THREADS =
ConcurrentHashMap.newKeySet();
+ private static final AtomicInteger THREAD_COUNT = new AtomicInteger();
+ private static final List<ParkedTask> PARKED_TASKS = new ArrayList<>();
+ private static boolean _gateClosed;
+
+ @Override
+ public String id() {
+ return ID;
+ }
+
+ @Override
+ public ExecutorServiceProvider provider() {
+ return (conf, confPrefix, baseName) -> new
GatedExecutorService(Executors.newCachedThreadPool(runnable -> {
+ Thread thread = new Thread(runnable, baseName + "-gated-" +
THREAD_COUNT.incrementAndGet());
+ thread.setDaemon(true);
+ EXECUTOR_THREADS.add(thread);
+ return thread;
+ }));
+ }
+
+ public static synchronized void closeGate() {
+ _gateClosed = true;
+ }
+
+ /// Opens the gate and starts the parked tasks.
+ public static void openGate() {
+ List<ParkedTask> parkedTasks;
+ synchronized (GatedExecutorServicePlugin.class) {
+ _gateClosed = false;
+ parkedTasks = new ArrayList<>(PARKED_TASKS);
+ PARKED_TASKS.clear();
+ }
+ for (ParkedTask parkedTask : parkedTasks) {
+ parkedTask._executorService.execute(parkedTask._task);
+ }
+ }
+
+ public static synchronized int getNumParkedTasks() {
+ return PARKED_TASKS.size();
+ }
+
+ private static class ParkedTask {
+ final ExecutorService _executorService;
+ final Runnable _task;
+
+ ParkedTask(ExecutorService executorService, Runnable task) {
+ _executorService = executorService;
+ _task = task;
+ }
+ }
+
+ private static class GatedExecutorService extends AbstractExecutorService {
+ private final ExecutorService _delegate;
+
+ GatedExecutorService(ExecutorService delegate) {
+ _delegate = delegate;
+ }
+
+ @Override
+ public void execute(Runnable task) {
+ synchronized (GatedExecutorServicePlugin.class) {
+ if (_gateClosed && EXECUTOR_THREADS.contains(Thread.currentThread())) {
+ PARKED_TASKS.add(new ParkedTask(_delegate, task));
+ return;
+ }
+ }
+ _delegate.execute(task);
+ }
+
+ @Override
+ public void shutdown() {
+ _delegate.shutdown();
+ }
+
+ @Override
+ public List<Runnable> shutdownNow() {
+ return _delegate.shutdownNow();
+ }
+
+ @Override
+ public boolean isShutdown() {
+ return _delegate.isShutdown();
+ }
+
+ @Override
+ public boolean isTerminated() {
+ return _delegate.isTerminated();
+ }
+
+ @Override
+ public boolean awaitTermination(long timeout, TimeUnit unit)
+ throws InterruptedException {
+ return _delegate.awaitTermination(timeout, unit);
+ }
+ }
+}
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
index f9a9983f761..ea375c4069d 100644
---
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryRunnerTestBase.java
@@ -202,12 +202,20 @@ public abstract class QueryRunnerTestBase extends
QueryTestSet {
new StageMetadata(stageId, workerMetadataList,
dispatchableStagePlan.getCustomProperties());
StagePlan stagePlan = new
StagePlan(dispatchableStagePlan.getPlanFragment().getFragmentRoot(),
stageMetadata);
for (WorkerMetadata workerMetadata : workerMetadataList) {
- submissionStubs.add(serverEnclosure.processQuery(workerMetadata,
stagePlan, requestMetadataMap));
+ submissionStubs.add(serverEnclosure.processQuery(workerMetadata,
stagePlan, requestMetadataMap,
+ getExecutionContext(serverEnclosure, requestMetadataMap)));
}
}
return submissionStubs;
}
+ /// Returns the execution context in which the given server runs a stage
worker. By default, each worker gets its own
+ /// context, while the query server shares one context between all the
stages of a query.
+ protected QueryExecutionContext getExecutionContext(QueryServerEnclosure
server,
+ Map<String, String> requestMetadataMap) {
+ return QueryExecutionContext.forMseServerRequest(requestMetadataMap,
"serverId");
+ }
+
protected List<Object[]> queryH2(String sql)
throws Exception {
int firstSemi = sql.indexOf(';');
diff --git
a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryTerminatedBeforeStageStartsTest.java
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryTerminatedBeforeStageStartsTest.java
new file mode 100644
index 00000000000..0655d4df5e1
--- /dev/null
+++
b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/QueryTerminatedBeforeStageStartsTest.java
@@ -0,0 +1,190 @@
+/**
+ * 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.query.runtime.queries;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import org.apache.pinot.common.response.broker.QueryProcessingException;
+import org.apache.pinot.query.QueryEnvironmentTestBase;
+import org.apache.pinot.query.QueryServerEnclosure;
+import org.apache.pinot.query.mailbox.MailboxService;
+import org.apache.pinot.query.planner.physical.DispatchablePlanFragment;
+import org.apache.pinot.query.routing.QueryServerInstance;
+import org.apache.pinot.query.service.dispatch.QueryDispatcher;
+import org.apache.pinot.query.testutils.MockInstanceDataManagerFactory;
+import org.apache.pinot.query.testutils.QueryTestUtils;
+import org.apache.pinot.spi.config.instance.InstanceType;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.query.QueryExecutionContext;
+import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner;
+import org.apache.pinot.spi.utils.CommonConstants.Server;
+import org.apache.pinot.util.TestUtils;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// A server can terminate a query before some of its stages run, for example
when the resource accountant kills the
+/// query to protect the heap. These stages must still send the termination
error to their receivers. Otherwise the
+/// broker gets no error and waits for data until the query times out.
+///
+/// Like the query server, each server runs all the stages of a query in one
[QueryExecutionContext], and the test
+/// terminates the context of one server.
+public class QueryTerminatedBeforeStageStartsTest extends QueryRunnerTestBase {
+ private static final long QUERY_TIMEOUT_MS = 30_000L;
+ private static final String QUERY_OPTIONS = "SET timeoutMs = " +
QUERY_TIMEOUT_MS + "; ";
+ private static final String GROUP_BY_QUERY = QUERY_OPTIONS + "SELECT col1,
COUNT(*) FROM a GROUP BY col1";
+ /// The leaf stage of `a` runs a pipeline breaker, which collects the keys
from `b` before the stage starts.
+ private static final String DYNAMIC_BROADCAST_QUERY = QUERY_OPTIONS
+ + "SELECT /*+ joinOptions(join_strategy='dynamic_broadcast') */ col1
FROM a WHERE a.col1 IN "
+ + "(SELECT b.col2 FROM b WHERE b.col3 < 10) AND a.col3 > 0";
+ private static final String KILL_MESSAGE = "OOM killed on SERVER: test";
+ /// The broker must report the kill long before the query times out.
+ private static final long MAX_REPORT_TIME_MS = 10_000L;
+
+ private QueryServerEnclosure _killedServer;
+ private final Map<QueryServerEnclosure, QueryExecutionContext>
_executionContexts = new ConcurrentHashMap<>();
+ private volatile boolean _killBeforeScheduling;
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ MockInstanceDataManagerFactory factory1 = new
MockInstanceDataManagerFactory("server1");
+
factory1.registerTable(QueryRunnerTest.SCHEMA_BUILDER.setSchemaName("a").build(),
"a_REALTIME");
+
factory1.registerTable(QueryRunnerTest.SCHEMA_BUILDER.setSchemaName("b").build(),
"b_REALTIME");
+ factory1.addSegment("a_REALTIME", QueryRunnerTest.buildRows("a_REALTIME"));
+ factory1.addSegment("b_REALTIME", QueryRunnerTest.buildRows("b_REALTIME"));
+ MockInstanceDataManagerFactory factory2 = new
MockInstanceDataManagerFactory("server2");
+
factory2.registerTable(QueryRunnerTest.SCHEMA_BUILDER.setSchemaName("a").build(),
"a_REALTIME");
+ factory2.addSegment("a_REALTIME", QueryRunnerTest.buildRows("a_REALTIME"));
+
+ _reducerHostname = "localhost";
+ _reducerPort = QueryTestUtils.getAvailablePort();
+ Map<String, Object> reducerConfig = new HashMap<>();
+ reducerConfig.put(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME,
_reducerHostname);
+ reducerConfig.put(MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT,
_reducerPort);
+ _mailboxService =
+ new MailboxService(_reducerHostname, _reducerPort,
InstanceType.BROKER, new PinotConfiguration(reducerConfig));
+ _mailboxService.start();
+
+ // The gated executor lets a test act after the op chains are scheduled
but before they start.
+ Map<String, Object> serverConfig =
+ Map.of(Server.MULTISTAGE_EXECUTOR_CONFIG_PREFIX + ".type",
GatedExecutorServicePlugin.ID);
+ QueryServerEnclosure server1 = new QueryServerEnclosure(factory1,
serverConfig);
+ server1.start();
+ QueryServerEnclosure server2 = new QueryServerEnclosure(factory2,
serverConfig);
+ server2.start();
+ int port1 = server1.getPort();
+ int port2 = server2.getPort();
+ _servers.put(new QueryServerInstance("Server_localhost_" + port1,
"localhost", port1, port1), server1);
+ _servers.put(new QueryServerInstance("Server_localhost_" + port2,
"localhost", port2, port2), server2);
+ _killedServer = server1;
+
+ _queryEnvironment =
QueryEnvironmentTestBase.getQueryEnvironment(_reducerPort, port1, port2,
+ factory1.getRegisteredSchemaMap(),
factory1.buildTableSegmentNameMap(), factory2.buildTableSegmentNameMap(),
+ null);
+ }
+
+ @AfterClass
+ public void tearDown() {
+ for (QueryServerEnclosure server : _servers.values()) {
+ server.shutDown();
+ }
+ _mailboxService.shutdown();
+ }
+
+ @BeforeMethod
+ public void resetQueryState() {
+ _executionContexts.clear();
+ _killBeforeScheduling = false;
+ }
+
+ @AfterMethod
+ public void openGate() {
+ GatedExecutorServicePlugin.openGate();
+ }
+
+ @DataProvider(name = "queries")
+ public Object[][] queries() {
+ return new Object[][]{{GROUP_BY_QUERY}, {DYNAMIC_BROADCAST_QUERY}};
+ }
+
+ /// The server kills the query after the broker dispatched it, but before
the server scheduled any of its stages.
+ @Test(dataProvider = "queries")
+ public void testKillBeforeStagesAreScheduled(String query) {
+ _killBeforeScheduling = true;
+ long startTimeMs = System.currentTimeMillis();
+ QueryDispatcher.QueryResult result = queryRunner(query, false);
+ assertKillReported(result, startTimeMs);
+ }
+
+ /// The server kills the query after it scheduled the op chains of its
stages, but before they started.
+ @Test
+ public void testKillAfterStagesAreScheduledBeforeTheyStart()
+ throws Exception {
+ int numWorkers = 0;
+ for (DispatchablePlanFragment stagePlan :
planQuery(GROUP_BY_QUERY).getQueryPlan().getQueryStagesWithoutRoot()) {
+ numWorkers += stagePlan.getWorkerMetadataList().size();
+ }
+ int numOpChains = numWorkers;
+ GatedExecutorServicePlugin.closeGate();
+ CompletableFuture<QueryDispatcher.QueryResult> resultFuture =
+ CompletableFuture.supplyAsync(() -> queryRunner(GROUP_BY_QUERY,
false));
+ TestUtils.waitForCondition(aVoid ->
GatedExecutorServicePlugin.getNumParkedTasks() == numOpChains, 10_000L,
+ "Op chains were not scheduled");
+ long killTimeMs = System.currentTimeMillis();
+
_executionContexts.get(_killedServer).terminate(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED,
KILL_MESSAGE);
+ GatedExecutorServicePlugin.openGate();
+ assertKillReported(resultFuture.get(2 * QUERY_TIMEOUT_MS,
TimeUnit.MILLISECONDS), killTimeMs);
+ }
+
+ private static void assertKillReported(QueryDispatcher.QueryResult result,
long fromTimeMs) {
+ long elapsedMs = System.currentTimeMillis() - fromTimeMs;
+ QueryProcessingException processingException =
result.getProcessingException();
+ assertNotNull(processingException, "Expected the query to fail");
+ assertEquals(processingException.getErrorCode(),
QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED.getId(),
+ "Unexpected error after " + elapsedMs + "ms: " +
processingException.getMessage());
+ assertTrue(processingException.getMessage().contains(KILL_MESSAGE),
processingException.getMessage());
+ assertTrue(elapsedMs < MAX_REPORT_TIME_MS, "The kill was reported after "
+ elapsedMs + "ms");
+ }
+
+ /// Like the query server, each server runs all the stages of the query in
one execution context.
+ @Override
+ protected QueryExecutionContext getExecutionContext(QueryServerEnclosure
server,
+ Map<String, String> requestMetadataMap) {
+ return _executionContexts.computeIfAbsent(server, key -> {
+ QueryExecutionContext context =
QueryExecutionContext.forMseServerRequest(requestMetadataMap, "serverId");
+ if (_killBeforeScheduling && key == _killedServer) {
+ context.terminate(QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED,
KILL_MESSAGE);
+ }
+ return context;
+ });
+ }
+}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryThreadContext.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryThreadContext.java
index ce07866de43..4b6092acf9c 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryThreadContext.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/query/QueryThreadContext.java
@@ -309,16 +309,18 @@ public class QueryThreadContext implements AutoCloseable {
}
/// Returns a new [ExecutorService] whose tasks will be executed with the
[QueryThreadContext] initialized with the
- /// state of the thread submitting the tasks. Tasks are registered for
cancellation when the query is terminated.
+ /// state of the thread submitting the tasks. The [Future] of each task
passed to `submit()` is registered for
+ /// cancellation when the query is terminated. Tasks passed to `execute()`
are not registered.
public static ExecutorService contextAwareExecutorService(ExecutorService
executorService) {
return contextAwareExecutorService(executorService, true);
}
/// Returns a new [ExecutorService] whose tasks will be executed with the
[QueryThreadContext] initialized with the
/// state of the thread submitting the tasks.
- /// @param registerTaskForCancellation If true, tasks are registered with
QueryExecutionContext.addTask() and will
- /// be cancelled (via Thread.interrupt())
when the query is terminated. Set to
- /// false for executors where interrupt
could cause corruption
+ /// @param registerTaskForCancellation If true, tasks passed to `submit()`
are registered with
+ /// QueryExecutionContext.addTask() and
will be cancelled (via Thread.interrupt())
+ /// when the query is terminated. Set to
false for executors where interrupt could
+ /// cause corruption
public static ExecutorService contextAwareExecutorService(ExecutorService
executorService,
boolean registerTaskForCancellation) {
Consumer<Future<?>> onSubmit = registerTaskForCancellation
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
