aho135 commented on code in PR #20314:
URL: https://github.com/apache/druid/pull/20314#discussion_r4066692458
##########
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:
Good catch — this is a real but very narrow race. It's on the success path,
so planning has already completed (the thread isn't held) and `close()` still
clears the interrupt, so there's no leak. The only effect is a query that
crosses the deadline within the final check→close window (sub-ms) getting
executed instead of returning 504 — the planning work is already spent at that
point. Given the impact is a sub-millisecond boundary imprecision and a fully
deterministic test of that exact window isn't feasible without adding a
production test hook, I'd prefer to leave it as-is rather than restructure the
guardrail for it. Happy to revisit if you feel strongly.
--
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]
