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 55a760d0d64 Stop MSE query planning when the planning thread is
interrupted (#19709)
55a760d0d64 is described below
commit 55a760d0d64dfa1b3a7aafd415eda8d9a55f543e
Author: Yash Mayya <[email protected]>
AuthorDate: Mon Oct 5 13:58:48 2026 -0700
Stop MSE query planning when the planning thread is interrupted (#19709)
---
.../MultiStageBrokerRequestHandler.java | 4 +-
.../query/planner/logical/LogicalPlanner.java | 20 +++
.../pinot/query/QueryPlanningCancellationTest.java | 179 +++++++++++++++++++++
.../query/planner/logical/LogicalPlannerTest.java | 61 +++++++
4 files changed, 263 insertions(+), 1 deletion(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
index b3e8f6f18a8..f42c17adbb7 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java
@@ -288,7 +288,9 @@ public class MultiStageBrokerRequestHandler extends
BaseBrokerRequestHandler {
try (var compiled = exec.submit(() -> warmupEnv.compile("SELECT
1")).get(5, TimeUnit.SECONDS)) {
// result discarded; compile call is for JVM warmup only
} finally {
- exec.shutdownNow();
+ // Do not interrupt a warmup that takes longer than the wait. Planning
stops on an interrupt, and then the first
+ // queries pay the rest of the warmup cost.
+ exec.shutdown();
}
LOGGER.info("MSE startup warmup completed in {}ms",
System.currentTimeMillis() - startMs);
} catch (Exception e) {
diff --git
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LogicalPlanner.java
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LogicalPlanner.java
index 2e08c43162d..47242e7d1d6 100644
---
a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LogicalPlanner.java
+++
b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/LogicalPlanner.java
@@ -25,9 +25,13 @@ import org.apache.calcite.plan.RelTraitDef;
import org.apache.calcite.plan.RelTraitSet;
import org.apache.calcite.plan.hep.HepPlanner;
import org.apache.calcite.plan.hep.HepProgram;
+import org.apache.pinot.spi.exception.EarlyTerminationException;
/// The `LogicalPlanner` is an extended implementation of the Calcite's
[HepPlanner].
+///
+/// Unlike [HepPlanner], it stops before the next rule fires when the planning
thread is interrupted (see
+/// [#checkCancel()]).
public class LogicalPlanner extends HepPlanner {
private List<RelTraitDef> _traitDefs;
@@ -51,6 +55,22 @@ public class LogicalPlanner extends HepPlanner {
return _traitDefs;
}
+ /// Throws if planning must stop. Calcite calls this before it fires each
rule.
+ ///
+ /// Calcite stops planning only when the
[org.apache.calcite.util.CancelFlag] of the planner [Context] is set. It
+ /// ignores thread interrupts, but Pinot cancels planning with an interrupt:
for example, when the broker times out
+ /// waiting for the plan, or when the query is terminated. Without this
check, an interrupted thread keeps planning
+ /// to the end, which can take minutes.
+ ///
+ /// This check does not clear the interrupt status.
+ @Override
+ public void checkCancel() {
+ if (Thread.currentThread().isInterrupted()) {
+ throw new EarlyTerminationException("Interrupted while planning query");
+ }
+ super.checkCancel();
+ }
+
@Override
public RelTraitSet emptyTraitSet() {
RelTraitSet traitSet = super.emptyTraitSet();
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlanningCancellationTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlanningCancellationTest.java
new file mode 100644
index 00000000000..b8defa8788c
--- /dev/null
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlanningCancellationTest.java
@@ -0,0 +1,179 @@
+/**
+ * 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;
+
+import com.google.common.util.concurrent.Uninterruptibles;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+import org.apache.calcite.plan.RelOptRule;
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.tools.RelBuilderFactory;
+import org.apache.commons.lang3.exception.ExceptionUtils;
+import org.apache.pinot.core.routing.MockRoutingManagerFactory;
+import org.apache.pinot.query.planner.rules.DefaultRuleSetCustomizer;
+import org.apache.pinot.query.planner.rules.PinotRuleSet;
+import org.apache.pinot.query.planner.spi.Phase;
+import org.apache.pinot.spi.exception.EarlyTerminationException;
+import org.apache.pinot.spi.exception.QueryErrorCode;
+import org.apache.pinot.spi.query.QueryThreadContext;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks that query planning stops once the planning thread is interrupted.
The broker interrupts the planning
+/// thread when planning times out, and so does terminating the query.
+public class QueryPlanningCancellationTest {
+ private static final String QUERY = "SELECT col1, col3 FROM a WHERE col3 >
0";
+
+ @Test
+ public void testPlanningStopsWhenPlanningTaskIsCancelled()
+ throws Exception {
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ try {
+ // Like MultiStageBrokerRequestHandler when planning times out
+ assertPlanningStops(executor, future -> future.cancel(true));
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ public void testPlanningStopsWhenQueryIsTerminated()
+ throws Exception {
+ try (QueryThreadContext ignore = QueryThreadContext.openForMseTest()) {
+ // Like the broker's compile executor, this executor registers its tasks
with the query, so that terminating the
+ // query cancels them
+ ExecutorService executor =
QueryThreadContext.contextAwareExecutorService(Executors.newSingleThreadExecutor());
+ try {
+ assertPlanningStops(executor, future ->
QueryThreadContext.get().getExecutionContext()
+ .terminate(QueryErrorCode.QUERY_CANCELLATION, "Query is
terminated"));
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+ }
+
+ /// Plans [#QUERY] on the executor, and cancels planning with the canceller
while a rule runs. Then checks that no
+ /// other rule fires, and that planning fails.
+ private static void assertPlanningStops(ExecutorService executor,
Consumer<Future<?>> canceller)
+ throws Exception {
+ BlockingRule blockingRule = new BlockingRule();
+ QueryEnvironment queryEnvironment = createQueryEnvironment(blockingRule);
+ CompletableFuture<Throwable> planningError = new CompletableFuture<>();
+ Future<?> future = executor.submit(() -> {
+ try (QueryEnvironment.CompiledQuery ignored =
queryEnvironment.compile(QUERY)) {
+ planningError.complete(null);
+ } catch (Throwable t) {
+ planningError.complete(t);
+ }
+ });
+ try {
+ assertTrue(blockingRule._fired.await(60, TimeUnit.SECONDS),
+ "Planning did not fire the blocking rule, planning error: " +
planningError.getNow(null));
+ canceller.accept(future);
+ } finally {
+ // Release the rule also when the test fails, so that the planning
thread does not stay blocked
+ blockingRule._released.countDown();
+ }
+
+ Throwable error = planningError.get(60, TimeUnit.SECONDS);
+ assertNotNull(error, "Planning must fail after its thread is interrupted");
+ assertNotNull(ExceptionUtils.throwableOfType(error,
EarlyTerminationException.class),
+ "Planning must fail because its thread is interrupted, but failed
with: " + error);
+ // Without the interrupt, the rule fires once for each node of the plan
+ assertEquals(blockingRule._numFirings.get(), 1, "No rule must fire after
the planning thread is interrupted");
+ }
+
+ private static QueryEnvironment createQueryEnvironment(RelOptRule extraRule)
{
+ MockRoutingManagerFactory factory = new MockRoutingManagerFactory(1, 2);
+
factory.registerTable(QueryEnvironmentTestBase.TABLE_SCHEMAS.get("a_REALTIME"),
"a_REALTIME");
+ PinotRuleSet ruleSet = new PinotRuleSet(List.of(new
DefaultRuleSetCustomizer(), (phase, rules) -> {
+ if (phase == Phase.BASIC) {
+ rules.add(0, extraRule);
+ }
+ }));
+ return new QueryEnvironment(QueryEnvironment.configBuilder()
+ .requestId(-1L)
+ .database(CommonConstants.DEFAULT_DATABASE)
+ .tableCache(factory.buildTableCache())
+ .ruleSet(ruleSet)
+ .build());
+ }
+
+ /// Rule that matches every node, and that blocks the first time it fires
until the test releases it. Like Calcite
+ /// code, it ignores interrupts while it runs.
+ private static class BlockingRule extends RelRule<BlockingRule.Config> {
+ final CountDownLatch _fired = new CountDownLatch(1);
+ final CountDownLatch _released = new CountDownLatch(1);
+ final AtomicInteger _numFirings = new AtomicInteger();
+
+ BlockingRule() {
+ super(new Config(b -> b.operand(RelNode.class).anyInputs()));
+ }
+
+ @Override
+ public void onMatch(RelOptRuleCall call) {
+ if (_numFirings.incrementAndGet() == 1) {
+ _fired.countDown();
+ // Keeps the interrupt status set when it returns
+ Uninterruptibles.awaitUninterruptibly(_released);
+ }
+ }
+
+ record Config(OperandTransform operandSupplier) implements RelRule.Config {
+ @Override
+ public String description() {
+ return "BlockingRule";
+ }
+
+ @Override
+ public RelOptRule toRule() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public Config withRelBuilderFactory(RelBuilderFactory factory) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public Config withDescription(String description) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public Config withOperandSupplier(OperandTransform transform) {
+ throw new UnsupportedOperationException();
+ }
+ }
+ }
+}
diff --git
a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/LogicalPlannerTest.java
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/LogicalPlannerTest.java
new file mode 100644
index 00000000000..3bca46b0e4c
--- /dev/null
+++
b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/LogicalPlannerTest.java
@@ -0,0 +1,61 @@
+/**
+ * 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.planner.logical;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import org.apache.calcite.plan.Contexts;
+import org.apache.calcite.plan.hep.HepProgram;
+import org.apache.calcite.runtime.CalciteException;
+import org.apache.calcite.util.CancelFlag;
+import org.apache.pinot.spi.exception.EarlyTerminationException;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+
+
+/// Checks when [LogicalPlanner#checkCancel()] stops planning.
+public class LogicalPlannerTest {
+ private static final HepProgram EMPTY_PROGRAM = HepProgram.builder().build();
+
+ @Test
+ public void testCheckCancelThrowsWhenThreadIsInterrupted() {
+ LogicalPlanner planner = new LogicalPlanner(EMPTY_PROGRAM,
Contexts.empty());
+ planner.checkCancel();
+
+ Thread.currentThread().interrupt();
+ try {
+ assertThrows(EarlyTerminationException.class, planner::checkCancel);
+ assertTrue(Thread.currentThread().isInterrupted(), "Interrupt status
must stay set");
+ } finally {
+ // Clear the interrupt status so that it does not leak into other tests
+ Thread.interrupted();
+ }
+ }
+
+ @Test
+ public void testCheckCancelThrowsWhenCancelFlagIsSet() {
+ CancelFlag cancelFlag = new CancelFlag(new AtomicBoolean());
+ LogicalPlanner planner = new LogicalPlanner(EMPTY_PROGRAM,
Contexts.of(cancelFlag));
+ planner.checkCancel();
+
+ cancelFlag.requestCancel();
+ assertThrows(CalciteException.class, planner::checkCancel);
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
