This is an automated email from the ASF dual-hosted git repository.
aho135 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git
The following commit(s) were added to refs/heads/master by this push:
new 13749ce5336 feat: Add a timeout for SQL query planning on the Broker
(#20314)
13749ce5336 is described below
commit 13749ce5336cd52864378371df5a17358196a7f0
Author: aho135 <[email protected]>
AuthorDate: Sat Sep 26 17:50:21 2026 -0700
feat: Add a timeout for SQL query planning on the Broker (#20314)
* Add a timeout for SQL query planning on the Broker
An extremely complex SQL query (for example one with a very large IN
filter) can spend tens of seconds in the Calcite planning phase on the
Broker. A flood of such queries can occupy all of the Broker's request
threads and effectively freeze the process, degrading the whole cluster.
This adds a wall-clock timeout on the planning phase (the phase measured
by the sqlQuery/planningTimeMs metric):
- New runtime property druid.sql.planner.maxPlanningTimeMs (default 0,
meaning no timeout), also overridable per query via the
maxPlanningTimeMs query context key.
- PlannerContext now owns a Calcite CancelFlag, which PlannerFactory
wires into the framework Context so that Calcite aborts at its next
cancellation checkpoint when the flag is tripped.
- A new SqlPlanningTimeout watchdog is armed in DirectStatement.plan().
When the deadline is exceeded it trips the CancelFlag and interrupts
the planning thread; the failure is surfaced as a QueryTimeoutException
(HTTP 504). The watchdog clears the interrupt on close so it cannot
leak to a pooled request thread.
Adds unit tests for the watchdog, the new config (including the query
context override), and an end-to-end DirectStatement planning-timeout
test, plus documentation.
* Make planning-timeout code comments more concise
* Remove cancelled planning-timeout tasks from the scheduler queue
Address review feedback: build the scheduler as a
ScheduledThreadPoolExecutor
with setRemoveOnCancelPolicy(true) so that cancelled watchdog tasks (the
common
case, when planning finishes before the deadline) are removed from the queue
immediately instead of lingering until their delay elapses.
* Address review: propagate cancel flag to Hep/view planners, tighten
timeout
Addresses PR review feedback beyond the scheduler queue fix:
- Propagate the per-query CancelFlag to the Hep-based planning stages
(pre/reduction/pre-Volcano/cleanup). Calcite's Programs.of builds its
HepPlanner with a null Context, so it never observed the flag; the new
CalciteRulesManager#hepProgram helper builds the HepPlanner with a
Context carrying the flag, otherwise identical to Programs.of. This is
the stage where RexSimplify (CALCITE-3178) runs, so it must be
cancellable.
- Share the top-level query's CancelFlag with nested planners created
during view expansion (DruidViewMacro), via
PlannerContext#withInheritedCancelFlag, so one planning timeout governs
the whole session.
- Reject a plan that completes after the watchdog fired (e.g. in a
non-cancellable section that swallowed the interrupt) instead of
returning it: check isTimedOut() on the success path too.
- Measure the planning budget from planningStartNanos so time spent
constructing the planner counts against maxPlanningTimeMs and a query
cannot get a fresh full budget afterwards.
- Emit maxPlanningTimeMs from PlannerConfig#getNonDefaultAsQueryContext()
so the defensive config<->context round-trip check passes for a
non-default value; add a round-trip test.
* Enforce planning deadline across construction; drop deprecated Calcite
call
Address review follow-ups:
- Enforce the planning deadline against wall-clock time elapsed since
planningStartNanos, which includes createPlanner() (planner/schema
construction). If construction alone exhausts maxPlanningTimeMs, fail
with QueryTimeoutException before arming the watchdog instead of
scheduling a 1ms timer and letting a fast follow-on plan slip through.
Also re-check the wall-clock deadline (not just the watchdog flag) on
the success path, so a plan that completes after the deadline under
scheduler jitter is still rejected. Adds a unit test that exhausts the
budget during planner construction.
- Drop the deprecated HepPlanner#registerMetadataProviders call from the
cancel-flag-aware hepProgram helper; modern Calcite routes rule
metadata through the cluster provider (cluster.setMetadataProvider),
which the helper already sets, matching Programs.of behavior. Verified
against the decoupled/join/array/select Calcite query tests.
* Close watchdog race, cover subquery rewriting, restore Hep metadata
provider
Address the second automated review round (3x P2):
- Watchdog last-millisecond race (DirectStatement): compute the remaining
budget once and, if the timeout is enabled but already exhausted, fail
immediately. Previously the deadline check and the budget recompute were
separate steps; if time crossed the deadline in between, arm() received a
non-positive budget and returned the disabled (no-op) watchdog, so a query
past the deadline got no cancel-flag trip or interrupt and was only caught
after planning returned.
- Subquery rewriting cancellation (CalciteRulesManager): buildPreProgram
used
Programs.subQuery(...), whose HepPlanners are built with a null context
and
never see the query CancelFlag. Replaced with buildSubQueryProgram(), a
cancel-aware replica of Programs.subQuery's two Hep stages, so a query
that
spends its budget in subquery rewriting is interruptible.
- Restore Hep metadata provider (CalciteRulesManager): the previous change
dropped HepPlanner#registerMetadataProviders, which also dropped the
HepRelMetadataProvider that method prepends (needed for metadata lookups
on
HepRelVertex during Hep planning). Restored it and suppressed the
deprecation warning with a comment explaining why it must match
Programs.of.
* docs: clarify that a non-positive maxPlanningTimeMs disables the timeout
A negative value is coalesced to disabled everywhere the setting is consumed
(all guards use <= 0 / > 0), same as 0. Document that non-positive disables
the
timeout, in the runtime-property table, the query-context table, and the
PlannerConfig javadoc.
* test: use imports instead of fully-qualified class names in
SqlStatementTest
Address review: import DruidPlanner, PlannerResult, CancelFlag, and JUnit
Timeout in the planning-timeout tests rather than referencing them by
fully-qualified name inline.
* Fix merge conflict resolution in PlannerConfig
The merge of master (#19830, catalog DDL) dropped the closing brace of
isPlanningTimeLimited() and the opening Javadoc delimiter of the following
isEnableCatalogDdl() method, so PlannerConfig.java no longer compiled.
Restore
both.
* Drop druid.sql.planner.maxPlanningTimeMs in favor of default query context
Per review: since dynamic query-context configuration exists, a per-query
setting should not also be exposed as a dedicated druid.sql.planner.*
runtime
property (the existing planner properties predate that feature). Remove the
@JsonProperty binding so maxPlanningTimeMs is configured only via the query
context; a cluster-wide default is set with
druid.query.default.context.maxPlanningTimeMs
(or the Broker's dynamic queryContext config). The context key, watchdog,
and
per-query override are unchanged.
Docs: remove the runtime-property row from the Broker SQL config table,
update
the maxPlanningTimeMs query-context entry to point at "Overriding default
query
context values", and add that link to the sql-query-context "Learn more"
list.
* Trim verbose javadoc on PlannerConfig#maxPlanningTimeMs
* Trim PlannerConfig#maxPlanningTimeMs javadoc to what/why only
* test: use assertThrows in planning-timeout tests
Replace the try/plan()/fail()/catch pattern with Assertions.assertThrows in
testPlanningTimeout and testPlanningTimeoutDuringPlannerConstruction,
matching
the file's existing assertThrows usages. Behavior unchanged; still asserts
the
QueryTimeoutException message and closes the statement in a finally.
* Drop Hep cancel-flag duplication; rely on Volcano coverage for the timeout
Thread-dumping a planning-bound large-IN query (forced OR path) showed the
cost
is in VolcanoPlanner.findBestExp (Druid-convention conversion + memo digest
computation), not the Hep reduction stage as originally assumed. The Volcano
planner already receives the per-query CancelFlag via the framework
Context, so
the timeout covers that path without any custom Hep plumbing.
Revert the hepProgram/buildSubQueryProgram helpers that reimplemented
Calcite's
Programs.of/Programs.subQuery just to inject the CancelFlag into the
internally
constructed HepPlanners. Those duplicated Calcite internals (and a
deprecated
API) to cover a Hep-reduction-bound scenario that could not be reproduced.
Back
to plain Programs.of/Programs.subQuery. The framework-Context wiring and the
view-planner flag inheritance (PlannerContext#withInheritedCancelFlag)
remain.
* Trim verbose comments across the planning-timeout changes
---
docs/querying/sql-query-context.md | 2 +
.../java/org/apache/druid/sql/DirectStatement.java | 97 ++++++++++++++---
.../org/apache/druid/sql/SqlPlanningTimeout.java | 117 +++++++++++++++++++++
.../druid/sql/calcite/planner/PlannerConfig.java | 48 +++++++++
.../druid/sql/calcite/planner/PlannerContext.java | 46 ++++++++
.../druid/sql/calcite/planner/PlannerFactory.java | 5 +
.../apache/druid/sql/SqlPlanningTimeoutTest.java | 98 +++++++++++++++++
.../org/apache/druid/sql/SqlStatementTest.java | 106 +++++++++++++++++--
.../sql/calcite/planner/PlannerConfigTest.java | 51 +++++++++
9 files changed, 547 insertions(+), 23 deletions(-)
diff --git a/docs/querying/sql-query-context.md
b/docs/querying/sql-query-context.md
index 23e9f41347a..88cafe1ced7 100644
--- a/docs/querying/sql-query-context.md
+++ b/docs/querying/sql-query-context.md
@@ -56,8 +56,10 @@ The table below lists the query context parameters you can
use with Druid SQL.
|`inFunctionThreshold`| At or beyond this threshold number of values, Druid
converts SQL `IN` to [`SCALAR_IN_ARRAY`](sql-functions.md#scalar_in_array). A
threshold of 0 forces this conversion in all cases. A threshold of
`Integer.MAX_VALUE` disables this conversion. The converted function is
eligible for fewer planning-time optimizations, which speeds up planning, but
may prevent certain planning-time optimizations.| `100`|
|`inFunctionExprThreshold`|At or beyond this threshold number of values, SQL
`IN` is eligible for execution using the native function `scalar_in_array`
rather than an <code>||</code> of `==`, even if the number of values
is below `inFunctionThreshold`. This property only affects translation of SQL
`IN` to a [native expression](math-expr.md). It doesn't affect translation of
SQL `IN` to a [native filter](filters.md). This property is provided for
backwards compatibility purposes [...]
|`inSubQueryThreshold`|At or beyond this threshold number of values, Druid
converts SQL `IN` to `JOIN` on an inline table. `inFunctionThreshold` takes
priority over this setting. A threshold of 0 forces usage of an inline table in
all cases where the size of a SQL `IN` is larger than `inFunctionThreshold`. A
threshold of `2147483647` disables the rewrite of SQL `IN` to `JOIN`.
|`2147483647`|
+|`maxPlanningTimeMs`|Maximum wall-clock time, in milliseconds, allowed for
planning this query on the Broker. When planning exceeds this budget, planning
is aborted and the query fails with an HTTP 504 `Query timeout` error. Use this
to guard against pathological queries, such as one with a very large `IN`
filter, whose planning time can spike. A non-positive value disables the
timeout. To apply a cluster-wide default, set the
`druid.query.default.context.maxPlanningTimeMs` runtime prope [...]
## Learn more
- [Set query context](../querying/query-context.md) for how to set the query
context.
- [Query context reference](query-context-reference.md) for available query
context parameters.
+- [Overriding default query context
values](../configuration/index.md#overriding-default-query-context-values) for
how to set cluster-wide defaults for any query context parameter with
`druid.query.default.context.*`.
- [MSQ context
parameters](../multi-stage-query/reference.md#context-parameters) for how to
set context parameters for Multi-Stage Queries.
diff --git a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java
b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java
index ee50d0b8e18..88c9b079b79 100644
--- a/sql/src/main/java/org/apache/druid/sql/DirectStatement.java
+++ b/sql/src/main/java/org/apache/druid/sql/DirectStatement.java
@@ -27,15 +27,18 @@ import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.query.QueryException;
import org.apache.druid.query.QueryInterruptedException;
+import org.apache.druid.query.QueryTimeoutException;
import org.apache.druid.server.QueryResponse;
import org.apache.druid.server.security.ResourceAction;
import org.apache.druid.sql.SqlLifecycleManager.Cancelable;
import org.apache.druid.sql.calcite.planner.DruidPlanner;
+import org.apache.druid.sql.calcite.planner.PlannerContext;
import org.apache.druid.sql.calcite.planner.PlannerResult;
import org.apache.druid.sql.calcite.planner.PrepareResult;
import javax.annotation.Nullable;
import java.util.Set;
+import java.util.concurrent.TimeUnit;
/**
* Lifecycle for direct SQL statement execution, which means that the query
@@ -196,20 +199,52 @@ public class DirectStatement extends AbstractStatement
implements Cancelable
}
long planningStartNanos = System.nanoTime();
try (DruidPlanner planner = createPlanner()) {
- 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;
+ // Bound planning wall-clock. Budget is measured from planningStartNanos
(before createPlanner), so planner
+ // construction counts against it; if already exhausted, fail now rather
than arm a no-op (non-positive) watchdog.
+ final long maxPlanningTimeMs =
planner.getPlannerContext().getPlannerConfig().getMaxPlanningTimeMs();
+ final long remainingBudgetMs =
remainingPlanningBudgetMs(maxPlanningTimeMs, planningStartNanos);
+ if (maxPlanningTimeMs > 0 && remainingBudgetMs <= 0) {
+ throw planningTimedOut(maxPlanningTimeMs);
+ }
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(
+ remainingBudgetMs,
+ planner.getPlannerContext().getCancelFlag(),
+ Thread.currentThread()
+ )) {
+ // Nested planners created during view expansion inherit this flag
(PlannerContext#withInheritedCancelFlag).
+ return PlannerContext.withInheritedCancelFlag(
+ planner.getPlannerContext().getCancelFlag(),
+ () -> {
+ 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);
+ }
+ catch (RuntimeException | AssertionError e) {
+ // On timeout, the failure is a side effect of aborting the
planner; surface it as a timeout.
+ if (timeout.isTimedOut()) {
+ throw planningTimedOut(maxPlanningTimeMs);
+ }
+ throw e;
+ }
+ // Reject a plan that completed past the deadline
(non-cancellable section, or watchdog scheduler jitter).
+ if (timeout.isTimedOut() ||
planningDeadlineExceeded(maxPlanningTimeMs, planningStartNanos)) {
+ throw planningTimedOut(maxPlanningTimeMs);
+ }
+ return resultSet;
+ }
+ );
+ }
}
catch (RelOptPlanner.CannotPlanException e) {
// Not sure if this is even thrown here.
@@ -229,6 +264,40 @@ public class DirectStatement extends AbstractStatement
implements Cancelable
}
}
+ /**
+ * Budget left in the planning timeout, given elapsed time since {@code
planningStartNanos}. Returns
+ * {@code maxPlanningTimeMs} unchanged when non-positive (disabled); may be
non-positive if already overrun.
+ */
+ private static long remainingPlanningBudgetMs(long maxPlanningTimeMs, long
planningStartNanos)
+ {
+ if (maxPlanningTimeMs <= 0) {
+ return maxPlanningTimeMs;
+ }
+ final long elapsedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() -
planningStartNanos);
+ return maxPlanningTimeMs - elapsedMs;
+ }
+
+ /**
+ * Whether the planning deadline has been exceeded by elapsed wall-clock
time. False when disabled.
+ */
+ private static boolean planningDeadlineExceeded(long maxPlanningTimeMs, long
planningStartNanos)
+ {
+ return maxPlanningTimeMs > 0
+ && TimeUnit.NANOSECONDS.toMillis(System.nanoTime() -
planningStartNanos) >= maxPlanningTimeMs;
+ }
+
+ private QueryTimeoutException planningTimedOut(long maxPlanningTimeMs)
+ {
+ return new QueryTimeoutException(
+ StringUtils.format(
+ "Query planning for [%s] exceeded the configured maximum planning
time of [%,d] ms. This may "
+ + "indicate an overly complex query, for example one with a very
large IN filter.",
+ sqlQueryId(),
+ maxPlanningTimeMs
+ )
+ );
+ }
+
protected DruidPlanner createPlanner()
{
return sqlToolbox.plannerFactory.createPlanner(
diff --git a/sql/src/main/java/org/apache/druid/sql/SqlPlanningTimeout.java
b/sql/src/main/java/org/apache/druid/sql/SqlPlanningTimeout.java
new file mode 100644
index 00000000000..34bc00631ef
--- /dev/null
+++ b/sql/src/main/java/org/apache/druid/sql/SqlPlanningTimeout.java
@@ -0,0 +1,117 @@
+/*
+ * 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 org.apache.calcite.util.CancelFlag;
+import org.apache.druid.java.util.common.concurrent.Execs;
+
+import java.io.Closeable;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.ScheduledThreadPoolExecutor;
+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()}). On
deadline it trips the query's
+ * Calcite {@link CancelFlag} and interrupts the planning thread; {@link
#close()} cancels the watchdog and clears any
+ * leaked interrupt. Check {@link #isTimedOut()} to translate a planning
failure into a timeout.
+ */
+public class SqlPlanningTimeout implements Closeable
+{
+ private static final ScheduledThreadPoolExecutor SCHEDULER =
createScheduler();
+
+ private static ScheduledThreadPoolExecutor createScheduler()
+ {
+ // One daemon thread suffices; remove cancelled tasks (the common case)
from the queue so it can't grow under load.
+ final ScheduledThreadPoolExecutor scheduler =
+ new ScheduledThreadPoolExecutor(1,
Execs.makeThreadFactory("sql-planning-timeout-%d"));
+ scheduler.setRemoveOnCancelPolicy(true);
+ return scheduler;
+ }
+
+ // 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;
+ private volatile boolean timedOut;
+ 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 possibly-recycled
thread.
+ return;
+ }
+ timedOut = true;
+ cancelFlag.requestCancel();
+ planningThread.interrupt();
+ }
+ }
+
+ public boolean isTimedOut()
+ {
+ return timedOut;
+ }
+
+ @Override
+ public void close()
+ {
+ if (future == null) {
+ return;
+ }
+ boolean wasTimedOut;
+ synchronized (lock) {
+ closed = true;
+ future.cancel(false);
+ wasTimedOut = timedOut;
+ }
+ if (wasTimedOut) {
+ // Clear the watchdog's interrupt so it isn't leaked to a pooled request
thread (close() runs on that thread).
+ Thread.interrupted();
+ }
+ }
+}
diff --git
a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerConfig.java
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerConfig.java
index 3afdbe88ea0..bf6a3f20c0b 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerConfig.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerConfig.java
@@ -40,7 +40,12 @@ public class PlannerConfig
public static final String CTX_KEY_FORCE_EXPRESSION_VIRTUAL_COLUMNS =
"forceExpressionVirtualColumns";
public static final String CTX_MAX_NUMERIC_IN_FILTERS =
"maxNumericInFilters";
public static final String CTX_REQUIRE_TIME_CONDITION =
"requireTimeCondition";
+ public static final String CTX_KEY_MAX_PLANNING_TIME_MS =
"maxPlanningTimeMs";
public static final int NUM_FILTER_NOT_USED = -1;
+ /**
+ * Sentinel value for {@link #maxPlanningTimeMs} meaning "no planning
timeout".
+ */
+ public static final long PLANNING_TIME_NOT_LIMITED = 0;
@JsonProperty
private int maxTopNLimit = 100_000;
@@ -86,6 +91,13 @@ public class PlannerConfig
@JsonProperty
private boolean authorizeTableVisibility = true;
+ /**
+ * Maximum wall-clock time, in milliseconds, allowed for planning a SQL
query. When exceeded, planning is aborted and
+ * the query fails with a {@link
org.apache.druid.query.QueryTimeoutException}. A non-positive value (default
+ * {@link #PLANNING_TIME_NOT_LIMITED}) disables the timeout.
+ */
+ private long maxPlanningTimeMs = PLANNING_TIME_NOT_LIMITED;
+
@JsonProperty
private boolean enableCatalogDdl = false;
@@ -180,6 +192,19 @@ public class PlannerConfig
return authorizeTableVisibility;
}
+ /**
+ * See {@link #maxPlanningTimeMs}. Returns {@link
#PLANNING_TIME_NOT_LIMITED} when planning time is unbounded.
+ */
+ public long getMaxPlanningTimeMs()
+ {
+ return maxPlanningTimeMs;
+ }
+
+ public boolean isPlanningTimeLimited()
+ {
+ return maxPlanningTimeMs > PLANNING_TIME_NOT_LIMITED;
+ }
+
/**
* Whether catalog DDL statements (CREATE TABLE, ALTER TABLE) may be
executed. Off by default: the Broker's SQL
* endpoint is typically reachable by far more people than the Coordinator's
catalog API, so enabling this widens
@@ -218,6 +243,7 @@ public class PlannerConfig
&& useNativeQueryExplain == that.useNativeQueryExplain
&& forceExpressionVirtualColumns ==
that.forceExpressionVirtualColumns
&& maxNumericInFilters == that.maxNumericInFilters
+ && maxPlanningTimeMs == that.maxPlanningTimeMs
&& enableSysQueriesTable == that.enableSysQueriesTable
&& authorizeTableVisibility == that.authorizeTableVisibility
&& enableCatalogDdl == that.enableCatalogDdl
@@ -241,6 +267,7 @@ public class PlannerConfig
useNativeQueryExplain,
forceExpressionVirtualColumns,
maxNumericInFilters,
+ maxPlanningTimeMs,
nativeQuerySqlPlanningMode,
enableSysQueriesTable,
authorizeTableVisibility,
@@ -259,6 +286,7 @@ public class PlannerConfig
", requireTimeCondition=" + requireTimeCondition +
", sqlTimeZone=" + sqlTimeZone +
", useNativeQueryExplain=" + useNativeQueryExplain +
+ ", maxPlanningTimeMs=" + maxPlanningTimeMs +
", nativeQuerySqlPlanningMode=" + nativeQuerySqlPlanningMode +
", enableSysQueriesTable=" + enableSysQueriesTable +
", authorizeTableVisibility=" + authorizeTableVisibility +
@@ -296,6 +324,7 @@ public class PlannerConfig
private boolean useNativeQueryExplain;
private boolean forceExpressionVirtualColumns;
private int maxNumericInFilters;
+ private long maxPlanningTimeMs;
private String nativeQuerySqlPlanningMode;
private boolean enableSysQueriesTable;
private boolean authorizeTableVisibility;
@@ -318,6 +347,7 @@ public class PlannerConfig
useNativeQueryExplain = base.isUseNativeQueryExplain();
forceExpressionVirtualColumns = base.isForceExpressionVirtualColumns();
maxNumericInFilters = base.getMaxNumericInFilters();
+ maxPlanningTimeMs = base.getMaxPlanningTimeMs();
nativeQuerySqlPlanningMode = base.getNativeQuerySqlPlanningMode();
enableSysQueriesTable = base.isEnableSysQueriesTable();
authorizeTableVisibility = base.isAuthorizeTableVisibility();
@@ -342,6 +372,12 @@ public class PlannerConfig
return this;
}
+ public Builder maxPlanningTimeMs(long value)
+ {
+ this.maxPlanningTimeMs = value;
+ return this;
+ }
+
public Builder useApproximateCountDistinct(boolean option)
{
this.useApproximateCountDistinct = option;
@@ -469,6 +505,11 @@ public class PlannerConfig
CTX_REQUIRE_TIME_CONDITION,
requireTimeCondition
);
+ maxPlanningTimeMs = QueryContexts.parseLong(
+ queryContext,
+ CTX_KEY_MAX_PLANNING_TIME_MS,
+ maxPlanningTimeMs
+ );
return this;
}
@@ -508,6 +549,7 @@ public class PlannerConfig
config.authorizeSystemTablesDirectly = authorizeSystemTablesDirectly;
config.useNativeQueryExplain = useNativeQueryExplain;
config.maxNumericInFilters = maxNumericInFilters;
+ config.maxPlanningTimeMs = maxPlanningTimeMs;
config.forceExpressionVirtualColumns = forceExpressionVirtualColumns;
config.nativeQuerySqlPlanningMode = nativeQuerySqlPlanningMode;
config.enableSysQueriesTable = enableSysQueriesTable;
@@ -544,6 +586,12 @@ public class PlannerConfig
String.valueOf(requireTimeCondition)
);
}
+ if (def.maxPlanningTimeMs != maxPlanningTimeMs) {
+ overrides.put(
+ CTX_KEY_MAX_PLANNING_TIME_MS,
+ maxPlanningTimeMs
+ );
+ }
PlannerConfig newConfig =
PlannerConfig.builder().withOverrides(overrides).build();
if (!equals(newConfig)) {
diff --git
a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java
index e9826af5dd4..923bc1d198a 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerContext.java
@@ -30,6 +30,7 @@ import org.apache.calcite.linq4j.QueryProvider;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.schema.Table;
import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.util.CancelFlag;
import org.apache.druid.error.InvalidSqlInput;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.ISE;
@@ -80,6 +81,8 @@ import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Supplier;
/**
* Like {@link PlannerConfig}, but that has static configuration and this class
@@ -151,6 +154,18 @@ public class PlannerContext
private final PlannerHook hook;
private final Set<String> lookupsToLoad = new HashSet<>();
+ /**
+ * Cancel flag inherited by any {@link PlannerContext} created on this
thread while set, so nested view-expansion
+ * planners share the top-level query's {@link #cancelFlag}. Managed by
{@link #withInheritedCancelFlag}.
+ */
+ private static final ThreadLocal<CancelFlag> INHERITED_CANCEL_FLAG = new
ThreadLocal<>();
+
+ /**
+ * Calcite cancellation flag, wired into the planner via the framework
{@link org.apache.calcite.plan.Context} in
+ * {@link PlannerFactory}; tripping it enforces the planning timeout ({@link
PlannerConfig#getMaxPlanningTimeMs()}).
+ */
+ private final CancelFlag cancelFlag = inheritedOrNewCancelFlag();
+
private PlannerConfig plannerConfig;
/**
* Root schema created with {@link
DruidSchemaCatalogProvider#createEscalatedRootSchema()}, for use in
@@ -280,6 +295,37 @@ public class PlannerContext
return plannerConfig;
}
+ public CancelFlag getCancelFlag()
+ {
+ return cancelFlag;
+ }
+
+ private static CancelFlag inheritedOrNewCancelFlag()
+ {
+ final CancelFlag inherited = INHERITED_CANCEL_FLAG.get();
+ return inherited != null ? inherited : new CancelFlag(new
AtomicBoolean(false));
+ }
+
+ /**
+ * Runs {@code action} with {@code cancelFlag} installed as the flag that
{@link PlannerContext}s created on this
+ * thread inherit, so one planning timeout governs nested view-expansion
planners. Restores the prior value after.
+ */
+ public static <T> T withInheritedCancelFlag(final CancelFlag cancelFlag,
final Supplier<T> action)
+ {
+ final CancelFlag previous = INHERITED_CANCEL_FLAG.get();
+ INHERITED_CANCEL_FLAG.set(cancelFlag);
+ try {
+ return action.get();
+ }
+ finally {
+ if (previous != null) {
+ INHERITED_CANCEL_FLAG.set(previous);
+ } else {
+ INHERITED_CANCEL_FLAG.remove();
+ }
+ }
+ }
+
public DateTime getLocalNow()
{
return localNow;
diff --git
a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
index c232108f65f..7ae657ec454 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/planner/PlannerFactory.java
@@ -40,6 +40,7 @@ import org.apache.calcite.sql.validate.SqlConformance;
import org.apache.calcite.sql2rel.SqlToRelConverter;
import org.apache.calcite.tools.FrameworkConfig;
import org.apache.calcite.tools.Frameworks;
+import org.apache.calcite.util.CancelFlag;
import org.apache.druid.guice.annotations.Json;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.query.QueryContexts;
@@ -326,6 +327,10 @@ public class PlannerFactory extends PlannerToolbox
if (aClass.equals(PlannerContext.class)) {
return (C) plannerContext;
}
+ if (aClass.equals(CancelFlag.class)) {
+ // Hand Calcite the per-query cancellation flag so planning can
be aborted (e.g. on planning timeout).
+ return (C) plannerContext.getCancelFlag();
+ }
return null;
}
diff --git a/sql/src/test/java/org/apache/druid/sql/SqlPlanningTimeoutTest.java
b/sql/src/test/java/org/apache/druid/sql/SqlPlanningTimeoutTest.java
new file mode 100644
index 00000000000..80577ac186d
--- /dev/null
+++ b/sql/src/test/java/org/apache/druid/sql/SqlPlanningTimeoutTest.java
@@ -0,0 +1,98 @@
+/*
+ * 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 org.apache.calcite.util.CancelFlag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class SqlPlanningTimeoutTest
+{
+ private static CancelFlag newCancelFlag()
+ {
+ return new CancelFlag(new AtomicBoolean(false));
+ }
+
+ @Test
+ public void testDisabledWhenTimeoutNotPositive()
+ {
+ final CancelFlag cancelFlag = newCancelFlag();
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(0, cancelFlag,
Thread.currentThread())) {
+ assertFalse(timeout.isTimedOut());
+ assertFalse(cancelFlag.isCancelRequested());
+ }
+ assertFalse(cancelFlag.isCancelRequested());
+ assertFalse(Thread.currentThread().isInterrupted());
+ }
+
+ @Test
+ public void testDisabledWhenTimeoutNegative()
+ {
+ final CancelFlag cancelFlag = newCancelFlag();
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(-100, cancelFlag,
Thread.currentThread())) {
+ assertFalse(timeout.isTimedOut());
+ }
+ assertFalse(cancelFlag.isCancelRequested());
+ }
+
+ /**
+ * On deadline the watchdog must trip the cancel flag and interrupt the
thread, and {@link #close()} must clear it.
+ */
+ @Test
+ @Timeout(30)
+ public void testTimeoutTripsCancelFlagAndInterrupts()
+ {
+ final CancelFlag cancelFlag = newCancelFlag();
+ boolean sawInterruptOrCancel;
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(50, cancelFlag,
Thread.currentThread())) {
+ // Mimic CPU-bound planning that periodically checks for cancellation.
+ while (!cancelFlag.isCancelRequested() &&
!Thread.currentThread().isInterrupted()) {
+ // spin
+ }
+ sawInterruptOrCancel = true;
+ assertTrue(timeout.isTimedOut());
+ assertTrue(cancelFlag.isCancelRequested());
+ }
+ assertTrue(sawInterruptOrCancel);
+ assertFalse(Thread.currentThread().isInterrupted());
+ }
+
+ /**
+ * Planning that finishes before the deadline must not trip the flag or
leave the thread interrupted.
+ */
+ @Test
+ @Timeout(30)
+ public void testNoTimeoutWhenPlanningFinishesEarly() throws
InterruptedException
+ {
+ final CancelFlag cancelFlag = newCancelFlag();
+ try (SqlPlanningTimeout timeout = SqlPlanningTimeout.arm(5000, cancelFlag,
Thread.currentThread())) {
+ Thread.sleep(10);
+ assertFalse(timeout.isTimedOut());
+ assertFalse(cancelFlag.isCancelRequested());
+ }
+ assertFalse(cancelFlag.isCancelRequested());
+ assertFalse(Thread.currentThread().isInterrupted());
+ }
+}
diff --git a/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java
b/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java
index 4abb504c524..956a550f254 100644
--- a/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java
@@ -26,6 +26,7 @@ import
com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.util.CancelFlag;
import org.apache.druid.error.DruidException;
import org.apache.druid.error.DruidExceptionMatcher;
import org.apache.druid.java.util.common.ISE;
@@ -38,6 +39,7 @@ import org.apache.druid.query.DefaultQueryConfig;
import org.apache.druid.query.Query;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.QueryRunnerFactoryConglomerate;
+import org.apache.druid.query.QueryTimeoutException;
import org.apache.druid.query.policy.NoopPolicyEnforcer;
import org.apache.druid.query.policy.PolicyEnforcer;
import org.apache.druid.query.policy.RestrictAllTablesPolicyEnforcer;
@@ -58,8 +60,10 @@ import org.apache.druid.sql.calcite.BaseCalciteQueryTest;
import org.apache.druid.sql.calcite.planner.CalciteRulesManager;
import org.apache.druid.sql.calcite.planner.CatalogResolver;
import org.apache.druid.sql.calcite.planner.DruidOperatorTable;
+import org.apache.druid.sql.calcite.planner.DruidPlanner;
import org.apache.druid.sql.calcite.planner.PlannerConfig;
import org.apache.druid.sql.calcite.planner.PlannerFactory;
+import org.apache.druid.sql.calcite.planner.PlannerResult;
import org.apache.druid.sql.calcite.planner.PrepareResult;
import org.apache.druid.sql.calcite.schema.DruidSchemaCatalogProvider;
import org.apache.druid.sql.calcite.util.CalciteTests;
@@ -71,6 +75,7 @@ import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
import org.junit.jupiter.api.io.TempDir;
import javax.servlet.http.HttpServletRequest;
@@ -102,6 +107,7 @@ public class SqlStatementTest
private PolicyEnforcer policyEnforcer;
private SqlStatementFactory sqlStatementFactory;
+ private SqlToolbox sqlToolbox;
@BeforeAll
public static void setUpClass()
@@ -568,15 +574,97 @@ public class SqlStatementTest
new DruidHookDispatcher()
);
- return new SqlStatementFactory(
- new SqlToolbox(
- CalciteTests.createMockSqlEngine(walker, conglomerate),
- plannerFactory,
- new NoopServiceEmitter(),
- testRequestLogger,
- QueryStackTests.DEFAULT_NOOP_SCHEDULER,
- new SqlLifecycleManager()
- )
+ this.sqlToolbox = new SqlToolbox(
+ CalciteTests.createMockSqlEngine(walker, conglomerate),
+ plannerFactory,
+ new NoopServiceEmitter(),
+ testRequestLogger,
+ QueryStackTests.DEFAULT_NOOP_SCHEDULER,
+ new SqlLifecycleManager()
);
+ return new SqlStatementFactory(sqlToolbox);
+ }
+
+ /**
+ * Planning that exceeds {@code maxPlanningTimeMs} fails with a {@link
QueryTimeoutException}. Planning is simulated
+ * as a CPU-bound loop that honours the cancel flag (as Calcite does), so
the watchdog trips it and we time out.
+ */
+ @Test
+ @Timeout(30)
+ public void testPlanningTimeout()
+ {
+ SqlQueryPlus sqlReq = SqlQueryPlus
+ .builder("SELECT COUNT(*) AS cnt, 'foo' AS TheFoo FROM druid.foo")
+
.queryContext(ImmutableMap.of(PlannerConfig.CTX_KEY_MAX_PLANNING_TIME_MS, 100))
+ .auth(CalciteTests.REGULAR_USER_AUTH_RESULT)
+ .build();
+
+ // A DirectStatement whose planning step blocks until the query is
cancelled, mimicking a pathological query.
+ DirectStatement stmt = new DirectStatement(sqlToolbox, sqlReq, null)
+ {
+ @Override
+ protected PlannerResult createPlan(DruidPlanner planner)
+ {
+ final CancelFlag cancelFlag =
planner.getPlannerContext().getCancelFlag();
+ while (!cancelFlag.isCancelRequested() &&
!Thread.currentThread().isInterrupted()) {
+ // spin until the watchdog aborts us
+ }
+ // Calcite throws when it observes a tripped cancel flag; emulate that.
+ throw new RuntimeException("Preparation aborted");
+ }
+ };
+
+ try {
+ QueryTimeoutException e =
Assertions.assertThrows(QueryTimeoutException.class, stmt::plan);
+ Assertions.assertTrue(
+ e.getMessage().contains("exceeded the configured maximum planning
time"),
+ "Unexpected message: " + e.getMessage()
+ );
+ }
+ finally {
+ stmt.close();
+ }
+ }
+
+ /**
+ * The budget covers planner construction: if {@link
DirectStatement#createPlanner()} alone exhausts
+ * {@code maxPlanningTimeMs}, planning fails with a {@link
QueryTimeoutException} before any further work.
+ */
+ @Test
+ @Timeout(30)
+ public void testPlanningTimeoutDuringPlannerConstruction()
+ {
+ SqlQueryPlus sqlReq = SqlQueryPlus
+ .builder("SELECT COUNT(*) AS cnt, 'foo' AS TheFoo FROM druid.foo")
+
.queryContext(ImmutableMap.of(PlannerConfig.CTX_KEY_MAX_PLANNING_TIME_MS, 50))
+ .auth(CalciteTests.REGULAR_USER_AUTH_RESULT)
+ .build();
+
+ // Simulate an expensive planner/schema construction that by itself
exceeds the budget.
+ DirectStatement stmt = new DirectStatement(sqlToolbox, sqlReq, null)
+ {
+ @Override
+ protected DruidPlanner createPlanner()
+ {
+ try {
+ Thread.sleep(300);
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ return super.createPlanner();
+ }
+ };
+
+ try {
+ QueryTimeoutException e =
Assertions.assertThrows(QueryTimeoutException.class, stmt::plan);
+ Assertions.assertTrue(
+ e.getMessage().contains("exceeded the configured maximum planning
time"),
+ "Unexpected message: " + e.getMessage()
+ );
+ }
+ finally {
+ stmt.close();
+ }
}
}
diff --git
a/sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerConfigTest.java
b/sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerConfigTest.java
index 6b3442132aa..5dd6865e13a 100644
---
a/sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerConfigTest.java
+++
b/sql/src/test/java/org/apache/druid/sql/calcite/planner/PlannerConfigTest.java
@@ -19,11 +19,14 @@
package org.apache.druid.sql.calcite.planner;
+import com.google.common.collect.ImmutableMap;
import nl.jqno.equalsverifier.EqualsVerifier;
import nl.jqno.equalsverifier.Warning;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
+import java.util.Map;
+
public class PlannerConfigTest
{
@Test
@@ -35,6 +38,54 @@ public class PlannerConfigTest
Assertions.assertTrue(config.isUseApproximateCountDistinct());
}
+ @Test
+ public void testMaxPlanningTimeMsDisabledByDefault()
+ {
+ PlannerConfig config = new PlannerConfig();
+ Assertions.assertEquals(PlannerConfig.PLANNING_TIME_NOT_LIMITED,
config.getMaxPlanningTimeMs());
+ Assertions.assertFalse(config.isPlanningTimeLimited());
+ }
+
+ @Test
+ public void testMaxPlanningTimeMsBuilder()
+ {
+ PlannerConfig config = PlannerConfig.builder()
+ .maxPlanningTimeMs(5000)
+ .build();
+ Assertions.assertEquals(5000, config.getMaxPlanningTimeMs());
+ Assertions.assertTrue(config.isPlanningTimeLimited());
+ }
+
+ @Test
+ public void testMaxPlanningTimeMsQueryContextOverride()
+ {
+ PlannerConfig base =
PlannerConfig.builder().maxPlanningTimeMs(10_000).build();
+ PlannerConfig overridden = base.withOverrides(
+ ImmutableMap.of(PlannerConfig.CTX_KEY_MAX_PLANNING_TIME_MS, 2500)
+ );
+ Assertions.assertEquals(2500, overridden.getMaxPlanningTimeMs());
+ // The base config is untouched.
+ Assertions.assertEquals(10_000, base.getMaxPlanningTimeMs());
+ }
+
+ @Test
+ public void testMaxPlanningTimeMsInheritedWhenNotOverridden()
+ {
+ PlannerConfig base =
PlannerConfig.builder().maxPlanningTimeMs(10_000).build();
+ PlannerConfig overridden =
base.withOverrides(ImmutableMap.of("someOtherKey", "someValue"));
+ Assertions.assertEquals(10_000, overridden.getMaxPlanningTimeMs());
+ }
+
+ @Test
+ public void testMaxPlanningTimeMsRoundTripsThroughQueryContext()
+ {
+ // Non-default maxPlanningTimeMs must round-trip through
getNonDefaultAsQueryContext()'s config<->context check.
+ PlannerConfig config =
PlannerConfig.builder().maxPlanningTimeMs(5000).build();
+ Map<String, Object> asContext = config.getNonDefaultAsQueryContext();
+ Assertions.assertEquals(5000L, ((Number)
asContext.get(PlannerConfig.CTX_KEY_MAX_PLANNING_TIME_MS)).longValue());
+ Assertions.assertEquals(config,
PlannerConfig.builder().withOverrides(asContext).build());
+ }
+
@Test
public void testPlannerConfigBuilder()
{
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
