aho135 commented on code in PR #20314: URL: https://github.com/apache/druid/pull/20314#discussion_r3984328011
########## 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: Good catch, fixed in a6a99ae075. The scheduler is now built directly as a `ScheduledThreadPoolExecutor` with `setRemoveOnCancelPolicy(true)`, so a cancelled watchdog future is removed from the delay queue immediately on `future.cancel(false)` rather than lingering until `maxPlanningTimeMs` elapses. (I couldn't set the policy on the `Execs.scheduledSingleThreaded` result because `Executors.newSingleThreadScheduledExecutor` wraps the pool in a `DelegatedScheduledExecutorService` that hides it, so I construct it directly with `Execs.makeThreadFactory`, matching the pattern in `SegmentLoadDropHandler`.) -- 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]
