FrankChen021 commented on code in PR #20314:
URL: https://github.com/apache/druid/pull/20314#discussion_r4062164238
##########
sql/src/main/java/org/apache/druid/sql/DirectStatement.java:
##########
@@ -196,20 +199,61 @@ public ResultSet plan()
}
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 the wall-clock time spent planning this query. A non-positive
timeout disables this. The budget is
+ // measured from planningStartNanos (above), so any time already spent
constructing the planner counts against
+ // it and a query cannot get a fresh full budget after an expensive
planner/schema setup.
+ final long maxPlanningTimeMs =
planner.getPlannerContext().getPlannerConfig().getMaxPlanningTimeMs();
+ // The budget is measured from planningStartNanos (above), so time
already spent constructing the planner
+ // counts against it. Compute the remaining budget once and, if the
timeout is enabled but already exhausted,
+ // fail immediately: arming with a non-positive budget would return a
disabled (no-op) watchdog, letting a
+ // query that crossed the deadline keep the planning thread busy until
it happens to return on its own.
+ 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()
+ )) {
+ // Share this query's cancel flag with any nested planners created
during view expansion, so that the same
+ // planning timeout governs the whole planning session (see
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;
+ }
+ // Planning may have finished after the deadline: the watchdog
fired but planning was in a
+ // non-cancellable section (or caught the interrupt and
returned), or the watchdog callback was
+ // delayed past the deadline under scheduler jitter. Re-check
the wall-clock deadline as well as the
+ // flag, and reject a late plan rather than executing it.
+ if (timeout.isTimedOut() ||
planningDeadlineExceeded(maxPlanningTimeMs, planningStartNanos)) {
Review Comment:
[P2] Recheck the watchdog after the final success check
**Finding:** After this success-path check, the lambda returns and the
try-with-resources block closes SqlPlanningTimeout. The scheduler can acquire
the watchdog lock in that interval, set timedOut, and interrupt the planning
thread; close() then clears that interrupt, but plan() never rechecks timedOut
and returns the plan. A plan that crosses the configured deadline in this last
window can therefore be accepted and executed, weakening the timeout guardrail.
**Suggestion:** Make the final timeout decision and watchdog close
race-safe, for example by having close report whether the watchdog fired or by
performing the deadline check after closing, and add a deterministic race test.
--
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]
