FrankChen021 commented on code in PR #20314:
URL: https://github.com/apache/druid/pull/20314#discussion_r3979110376
##########
sql/src/main/java/org/apache/druid/sql/DirectStatement.java:
##########
@@ -196,20 +197,44 @@ public ResultSet plan()
}
long planningStartNanos = System.nanoTime();
try (DruidPlanner planner = createPlanner()) {
- validate(planner);
- authorize(planner, authorizer());
+ // Bound the wall-clock time spent planning this query. A non-positive
timeout disables this.
+ final long maxPlanningTimeMs =
planner.getPlannerContext().getPlannerConfig().getMaxPlanningTimeMs();
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(
+ maxPlanningTimeMs,
+ planner.getPlannerContext().getCancelFlag(),
Review Comment:
[P1] Propagate the deadline into view planners
DruidViewMacro.apply creates a new DruidPlanner and PlannerContext for each
view and runs validate/prepare on it. That nested context has its own
CancelFlag and does not inherit the watchdog armed here, so planning a query
that expands a view can spend unbounded time in the nested planner even after
the outer timeout fires; Thread.interrupt is only advisory for this CPU-bound
work. Pass the outer cancellation/deadline through view expansion or arm the
nested planner against the same deadline.
##########
sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java:
##########
@@ -220,6 +221,10 @@ public SqlConformance conformance()
if (aClass.equals(PlannerContext.class)) {
return (C) plannerContext;
}
+ if (aClass.equals(CancelFlag.class)) {
Review Comment:
[P1] Propagate the cancel flag to Hep planners
CalciteRulesManager.programs uses Programs.of(...) for the pre, reduction,
and decoupled Hep stages. In Calcite 1.42, Programs.of constructs its
HepPlanner with a null Context, so it does not unwrap this per-query CancelFlag
and instead checks a separate flag that is never tripped. A query that spends
its planning budget in those Hep rule loops can therefore ignore the watchdog
and keep the Broker planning thread busy; the shared flag or an explicit
cancellation check must reach those planners.
##########
sql/src/main/java/org/apache/druid/sql/DirectStatement.java:
##########
@@ -196,20 +197,44 @@ public ResultSet plan()
}
long planningStartNanos = System.nanoTime();
try (DruidPlanner planner = createPlanner()) {
- validate(planner);
- authorize(planner, authorizer());
+ // Bound the wall-clock time spent planning this query. A non-positive
timeout disables this.
+ final long maxPlanningTimeMs =
planner.getPlannerContext().getPlannerConfig().getMaxPlanningTimeMs();
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(
+ maxPlanningTimeMs,
+ planner.getPlannerContext().getCancelFlag(),
+ Thread.currentThread()
+ )) {
+ try {
+ validate(planner);
+ authorize(planner, authorizer());
- // Adding the statement to the lifecycle manager allows cancellation.
- // Tests cancel during this call; real clients might do so if the plan
- // or execution prep stages take too long for some unexpected reason.
- sqlToolbox.sqlLifecycleManager.add(sqlQueryId(), this);
- transition(State.PREPARED);
- resultSet = createResultSet(createPlan(planner));
- prepareResult = planner.prepareResult();
- // Double check needed by SqlResourceTest
- transition(State.PREPARED);
- reporter.planningTimeNanos(System.nanoTime() - planningStartNanos);
- return resultSet;
+ // Adding the statement to the lifecycle manager allows cancellation.
+ // Tests cancel during this call; real clients might do so if the
plan
+ // or execution prep stages take too long for some unexpected reason.
+ sqlToolbox.sqlLifecycleManager.add(sqlQueryId(), this);
+ transition(State.PREPARED);
+ resultSet = createResultSet(createPlan(planner));
+ prepareResult = planner.prepareResult();
+ // Double check needed by SqlResourceTest
+ transition(State.PREPARED);
+ reporter.planningTimeNanos(System.nanoTime() - planningStartNanos);
+ return resultSet;
Review Comment:
[P2] Reject a plan completed after the deadline
The watchdog state is only inspected in the inner catch. If createPlan or
prepareResult is in a non-cancellable section, or catches the interrupt and
returns normally after the watchdog fires, this path reports success and
returns the result; close then clears the interrupt and the late plan can
execute. Check timeout.isTimedOut() before reporting success and returning the
ResultSet, translating that case to QueryTimeoutException.
##########
sql/src/main/java/org/apache/druid/sql/SqlPlanningTimeout.java:
##########
@@ -0,0 +1,129 @@
+/*
+ * 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.druid.sql;
+
+import com.google.common.annotations.VisibleForTesting;
+import org.apache.calcite.util.CancelFlag;
+import org.apache.druid.java.util.common.concurrent.Execs;
+
+import java.io.Closeable;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Bounds the wall-clock time spent planning a single SQL query. See
+ * {@link
org.apache.druid.sql.calcite.planner.PlannerConfig#getMaxPlanningTimeMs()}.
+ *
+ * <p>When {@link #arm} is called with a positive timeout, a task is scheduled
that, on deadline, trips the query's
+ * Calcite {@link CancelFlag} (so the planner aborts at its next cancellation
checkpoint) and interrupts the planning
+ * thread. The caller plans on its own thread and then calls {@link #close()}
(ideally in a {@code finally}), which
+ * cancels the pending task and, if the watchdog fired, clears the interrupt
so it is not leaked to a pooled request
+ * thread. Check {@link #isTimedOut()} when planning throws to decide whether
to translate the failure into a timeout.
+ */
+public class SqlPlanningTimeout implements Closeable
+{
+ // Single daemon thread suffices: each task only flips a flag and interrupts
a thread, and is usually cancelled first.
+ private static final ScheduledExecutorService SCHEDULER =
+ Execs.scheduledSingleThreaded("sql-planning-timeout-%d");
+
+ // No-op instance returned when no timeout is configured, so callers need no
null checks.
+ private static final SqlPlanningTimeout DISABLED = new SqlPlanningTimeout();
+
+ private final Object lock = new Object();
+ private final ScheduledFuture<?> future;
+
+ // Whether the deadline was reached. Written under lock; read via
isTimedOut().
+ private volatile boolean timedOut;
+
+ // Whether close() has been called. Once closed, a still-running watchdog
task must not interrupt the thread.
+ private boolean closed;
+
+ private SqlPlanningTimeout()
+ {
+ this.future = null;
+ }
+
+ private SqlPlanningTimeout(long maxPlanningTimeMs, CancelFlag cancelFlag,
Thread planningThread)
+ {
+ this.future = SCHEDULER.schedule(
+ () -> fire(cancelFlag, planningThread),
+ maxPlanningTimeMs,
+ TimeUnit.MILLISECONDS
+ );
+ }
+
+ /**
+ * Arm a watchdog for {@code planningThread}. A non-positive {@code
maxPlanningTimeMs} returns a no-op instance.
+ */
+ public static SqlPlanningTimeout arm(long maxPlanningTimeMs, CancelFlag
cancelFlag, Thread planningThread)
+ {
+ if (maxPlanningTimeMs <= 0) {
+ return DISABLED;
+ }
+ return new SqlPlanningTimeout(maxPlanningTimeMs, cancelFlag,
planningThread);
+ }
+
+ private void fire(CancelFlag cancelFlag, Thread planningThread)
+ {
+ synchronized (lock) {
+ if (closed) {
+ // Planning already finished; do not interrupt a thread that may have
been recycled.
+ return;
+ }
+ timedOut = true;
+ cancelFlag.requestCancel();
+ planningThread.interrupt();
+ }
+ }
+
+ /**
+ * Whether the planning deadline was reached before {@link #close()}.
+ */
+ public boolean isTimedOut()
+ {
+ return timedOut;
+ }
+
+ @Override
+ public void close()
+ {
+ if (future == null) {
+ return;
+ }
+ boolean wasTimedOut;
+ synchronized (lock) {
+ closed = true;
+ future.cancel(false);
Review Comment:
[P1] Remove cancelled watchdogs from the queue
The scheduler comes from Executors.newSingleThreadScheduledExecutor via
Execs.scheduledSingleThreaded, whose underlying ScheduledThreadPoolExecutor has
remove-on-cancel disabled by default. Every query that finishes before a
positive deadline calls future.cancel(false) here, but the cancelled delayed
future remains in the work queue until maxPlanningTimeMs elapses. Under
sustained SQL load, especially with a large per-query timeout, the queue grows
with request rate times timeout and can consume substantial Broker heap; enable
remove-on-cancel or explicitly purge/remove cancelled futures.
##########
sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerConfig.java:
##########
@@ -205,6 +234,7 @@ public boolean equals(Object o)
&& useNativeQueryExplain == that.useNativeQueryExplain
&& forceExpressionVirtualColumns ==
that.forceExpressionVirtualColumns
&& maxNumericInFilters == that.maxNumericInFilters
+ && maxPlanningTimeMs == that.maxPlanningTimeMs
Review Comment:
[P2] Include the new setting in config round-trips
getNonDefaultAsQueryContext() still omits maxPlanningTimeMs even though this
field is now part of equals and has a query-context key. Any non-default config
such as PlannerConfig.builder().maxPlanningTimeMs(5000).build() therefore fails
the defensive equality check at lines 569-575; QueryTestRunner calls this
helper when constructing QTest cases, so such a planner config cannot be used.
Emit the new context key and add a round-trip assertion.
##########
sql/src/main/java/org/apache/druid/sql/DirectStatement.java:
##########
@@ -196,20 +197,44 @@ public ResultSet plan()
}
long planningStartNanos = System.nanoTime();
try (DruidPlanner planner = createPlanner()) {
- validate(planner);
- authorize(planner, authorizer());
+ // Bound the wall-clock time spent planning this query. A non-positive
timeout disables this.
+ final long maxPlanningTimeMs =
planner.getPlannerContext().getPlannerConfig().getMaxPlanningTimeMs();
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(
Review Comment:
[P2] Start the watchdog before planner construction
planningStartNanos is captured before createPlanner(), so
sqlQuery/planningTimeMs includes PlannerContext/root-schema creation and
rule-program setup, but the watchdog is armed only after createPlanner()
returns. If schema discovery or setup consumes the configured budget, the
request gets a fresh full timeout afterward and can exceed the advertised
maximum wall-clock planning time. Start with a deadline before construction or
pass the remaining budget into arm().
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
