FrankChen021 commented on code in PR #20314:
URL: https://github.com/apache/druid/pull/20314#discussion_r4067902861
##########
sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java:
##########
@@ -568,15 +574,107 @@ private SqlStatementFactory buildSqlStatementFactory()
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);
+ }
+
+ /**
+ * A query whose planning exceeds the configured {@code maxPlanningTimeMs}
should fail with a
+ * {@link QueryTimeoutException} rather than occupying the planning thread
indefinitely. Planning is simulated as a
+ * CPU-bound loop that honours the Calcite cancel flag (as Calcite's planner
does), so we verify the watchdog trips
+ * that flag and the failure is surfaced as a timeout.
+ */
+ @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();
+ // Busy-wait like a CPU-bound Calcite planning phase that periodically
checks for cancellation.
+ while (!cancelFlag.isCancelRequested() &&
!Thread.currentThread().isInterrupted()) {
+ // spin until the planning-timeout watchdog aborts us
+ }
+ // Calcite throws when it observes a tripped cancel flag; emulate that
here.
+ throw new RuntimeException("Preparation aborted");
+ }
+ };
+
+ try {
Review Comment:
use `assertThrows` to replace the try-catch
##########
sql/src/test/java/org/apache/druid/sql/SqlStatementTest.java:
##########
@@ -568,15 +574,107 @@ private SqlStatementFactory buildSqlStatementFactory()
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);
+ }
+
+ /**
+ * A query whose planning exceeds the configured {@code maxPlanningTimeMs}
should fail with a
+ * {@link QueryTimeoutException} rather than occupying the planning thread
indefinitely. Planning is simulated as a
+ * CPU-bound loop that honours the Calcite cancel flag (as Calcite's planner
does), so we verify the watchdog trips
+ * that flag and the failure is surfaced as a timeout.
+ */
+ @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();
+ // Busy-wait like a CPU-bound Calcite planning phase that periodically
checks for cancellation.
+ while (!cancelFlag.isCancelRequested() &&
!Thread.currentThread().isInterrupted()) {
+ // spin until the planning-timeout watchdog aborts us
+ }
+ // Calcite throws when it observes a tripped cancel flag; emulate that
here.
+ throw new RuntimeException("Preparation aborted");
+ }
+ };
+
+ try {
+ stmt.plan();
+ fail("Expected planning to time out");
+ }
+ catch (QueryTimeoutException e) {
+ Assertions.assertTrue(
+ e.getMessage().contains("exceeded the configured maximum planning
time"),
+ "Unexpected message: " + e.getMessage()
+ );
+ }
+ finally {
+ stmt.close();
+ }
+ }
+
+ /**
+ * The planning budget covers planner construction too: if {@link
DirectStatement#createPlanner()} (schema/planner
+ * setup) alone exhausts {@code maxPlanningTimeMs}, planning must fail with
a {@link QueryTimeoutException} before any
+ * further work, rather than getting a fresh budget once the watchdog is
armed.
+ */
+ @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 {
Review Comment:
same as above
--
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]
