[Enhancement] Stop leader daemons via interrupt on demotion; disconnect checkpoint HTTP on stop

Rework leader-demotion daemon shutdown from the no-interrupt + cooperative
onStopRequested model to interrupt-based fast-cancel, and make in-flight
checkpoint HTTP disconnectable so it does not block the drain.

LeaderDaemon framework:
- stopGracefully() interrupts the worker by default via a new interruptOnStop()
  hook (default true), so a cycle blocked in an interruptible primitive
  (queue poll, latch/future await, sleep, interruptible lock) unblocks at once.
- Interrupt-unsafe daemons that call BDBJE/JE directly opt out with
  interruptOnStop()=false (CheckpointController, ClusterSnapshotJobScheduler)
  and stop cooperatively.

Fix task-level interrupt mishandlers so shutdownNow() is safe:
- ExportExportingTask: an interrupt during the sub-task await no longer cancels
  a healthy job as TIMEOUT; it leaves the job EXPORTING for the next leader and
  always clears doExportingThread (finally).
- TabletScheduler.blockingAddTabletCtxToScheduler and RoutineLoadTaskScheduler
  no longer swallow InterruptedException (re-assert / propagate).
- LeaderTaskExecutor / PriorityLeaderTaskExecutor close() uses shutdownNow() to
  cancel in-flight tasks fast.

HTTP fast-fail:
- CheckpointController holds its own in-flight HttpURLConnection and disconnects
  it in onStopRequested() to break out of an otherwise-uninterruptible socket
  read; it only ever disconnects its own connection (no collateral damage).
- MetaHelper only publishes the connection handle (new Consumer<HttpURLConnection>
  overloads) and no longer owns any cancel logic; cancelInFlight() removed.

Update LeaderDaemonTest (default-interrupt + opt-out paths) and
CheckpointControllerTest (disconnect-on-stop).

Signed-off-by: gengjun-git <gengjun@starrocks.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 42d82ca)

Author: gengjun-git

fe/fe-core/src/main/java/com/starrocks/clone/TabletScheduler.java

@@ -346,7 +346,11 @@ public Pair<Boolean, Long> blockingAddTabletCtxToScheduler(TabletSchedCtx tablet
             result.first = (res == AddResult.ADDED);
         } catch (InterruptedException e) {
             // heldLock has already been re-acquired by sleepUnlocked().
-            LOG.warn("Failed to execute blockingAddTabletCtxToScheduler", e);
+            // Re-assert the interrupt so the caller (e.g. TabletChecker / ColocateTableBalancer
+            // leader daemon being stopped on demotion) unwinds its cycle promptly instead of
+            // silently swallowing the cancel.
+            Thread.currentThread().interrupt();
+            LOG.warn("Interrupted while executing blockingAddTabletCtxToScheduler", e);
         }
 
         return result;

fe/fe-core/src/main/java/com/starrocks/common/util/LeaderDaemon.java

@@ -136,7 +136,9 @@ private void requestStop(boolean interruptWorker) {
 
     /**
      * Coordinated stop for leader demotion:
-     *   1. mark stopped and wake interval waits without interrupting the current cycle;
+     *   1. mark stopped, wake interval waits, and (unless {@link #interruptOnStop()} is overridden
+     *      to {@code false}) interrupt the worker so it breaks out of any blocking primitive
+     *      (queue poll, latch/future await, sleep, interruptible lock) at once;
      *   2. join up to {@code timeoutMs};
      *   3. on timeout, invoke {@link #onJoinTimeout()} (default: terminate the JVM) and return
      *      without clearing {@code worker}/{@code isRunning} - a subsequent {@link #start()}
@@ -145,7 +147,7 @@ private void requestStop(boolean interruptWorker) {
      * Idempotent.
      */
     public final void stopGracefully(long timeoutMs) {
-        requestStop(false);
+        requestStop(interruptOnStop());
         Thread t = worker;
         if (t != null) {
             try {
@@ -231,9 +233,19 @@ protected void runOneCycle() throws InterruptedException {
 
     /**
      * The body of each iteration. Runs only after FE is ready and the captured leader lease
-     * is still valid. Subclasses must not block indefinitely; leader demotion waits for the
-     * current cycle to finish and terminates the process on timeout instead of interrupting the
-     * worker inside storage or journal code that may treat interruption as fatal.
+     * is still valid. Subclasses must not block indefinitely.
+     *
+     * By default leader demotion INTERRUPTS the worker (see {@link #stopGracefully(long)}), so
+     * subclasses should block only in interruptible primitives and must let an
+     * {@link InterruptedException} propagate (or re-check {@link #isStopped()} and return) - they
+     * MUST NOT map it to a business outcome (e.g. cancel a healthy job as "timeout"). If the cycle
+     * cannot finish within {@code leader_demotion_drain_timeout_sec} the process is terminated via
+     * {@link #onJoinTimeout()} as a last resort.
+     *
+     * A subclass that runs interrupt-unsafe work on its own thread - a direct BDBJE/JE call
+     * (interrupting it can invalidate the environment) or an uninterruptible native/socket read -
+     * must override {@link #interruptOnStop()} to return {@code false} and cooperatively bail out
+     * by polling {@link #isStopped()} and/or waking its wait in {@link #onStopRequested()}.
      */
     protected abstract void runAfterLeaseValid() throws InterruptedException;
 
@@ -255,12 +267,28 @@ protected void onStopped() {
     }
 
     /**
-     * Hook called immediately after a stop request is accepted and before {@link #stopGracefully(long)}
-     * waits for the worker thread. Subclasses should use this to wake their own cancellable waits.
+     * Optional hook called immediately after a stop request is accepted and before
+     * {@link #stopGracefully(long)} waits for the worker thread. Most daemons do not need it:
+     * the default stop interrupts the worker, which already breaks any interruptible wait.
+     * It matters only for daemons that override {@link #interruptOnStop()} to {@code false} and
+     * therefore need to cooperatively wake their own uninterruptible wait (e.g. offer a sentinel
+     * to a result queue, or disconnect an in-flight HTTP connection).
      */
     protected void onStopRequested() {
     }
 
+    /**
+     * Whether {@link #stopGracefully(long)} may interrupt the worker thread. Default {@code true}:
+     * interrupt is the fast, standard way to cancel a blocked cycle. Override to return
+     * {@code false} ONLY for daemons whose worker executes interrupt-unsafe work directly on its
+     * own thread - a raw BDBJE/JE operation (an interrupt can invalidate the environment) or an
+     * uninterruptible native/socket read - and instead bail out cooperatively via
+     * {@link #isStopped()} polling and {@link #onStopRequested()}.
+     */
+    protected boolean interruptOnStop() {
+        return true;
+    }
+
     /**
      * Invoked when the worker thread fails to exit within the stop timeout. Default:
      * {@link System#exit(int)} - a stuck worker combined with a later {@link #start()} would run

fe/fe-core/src/main/java/com/starrocks/lake/snapshot/ClusterSnapshotJobScheduler.java

@@ -126,4 +126,16 @@ protected void runAfterLeaseValid() {
             CheckpointController.exclusiveUnlock();
         }
     }
+
+    /**
+     * Interrupt-unsafe: the worker calls BDBJE/JE directly (getJournal().getMaxJournalId()) and
+     * drives a full checkpoint (journal maintenance + image push) inline, where an interrupt can
+     * invalidate the BDB environment. It stops cooperatively instead - the driven checkpoint polls
+     * isStopped() between phases, backed by the stopGracefully() join-timeout -> onJoinTimeout()
+     * process-exit backstop.
+     */
+    @Override
+    protected boolean interruptOnStop() {
+        return false;
+    }
 }

fe/fe-core/src/main/java/com/starrocks/leader/CheckpointController.java

@@ -107,6 +107,15 @@ public class CheckpointController extends LeaderDaemon {
     // save the cluster snapshot info getted from the checkpoint worker.
     private volatile ClusterSnapshotInfo clusterSnapshotInfo;
 
+    // The HTTP connection this controller's worker thread is currently blocked on (image
+    // push / download / journal_id probe), or null between calls. Held so onStopRequested()
+    // can disconnect OUR OWN connection on demotion to break out of an otherwise-uninterruptible
+    // socket read (a push can block up to PUT_TIMEOUT_SECOND). Scoped to this controller, so it
+    // never touches another caller's connection. Only one HTTP op runs at a time on the
+    // single-threaded worker, so a single reference is sufficient.
+    // Package-private so same-package tests can set it to verify the disconnect-on-stop path.
+    volatile HttpURLConnection inFlightConnection;
+
     public CheckpointController(String name, Journal journal, String subDir) {
         super(name, FeConstants.checkpoint_interval_second * 1000L);
         this.journal = journal;
@@ -139,22 +148,29 @@ protected void runAfterLeaseValid() {
      * pushImage, deleteOldJournals) are idempotent at the system level - worker FE
      * keeps writing the image regardless of this leader's lifecycle, and the next
      * leader will re-orchestrate pushImage / deleteOldJournals from scratch. So we
-     * only have to wake the current checkpoint wait, release leader-session bookkeeping,
-     * and ask MetaHelper to break out of any in-flight HTTP, then let the daemon thread exit.
+     * only have to wake the current checkpoint wait, break out of any in-flight HTTP,
+     * and let the daemon thread exit; onStopped() clears the leader-session bookkeeping.
      *
-     * Today {@link MetaHelper#cancelInFlight()} is a no-op (see TODO there), so a
-     * push/download stuck in HttpURLConnection.read will block this drain up to its
-     * own connect/read timeout (up to 3600s per node for push). The bookkeeping
-     * fields below are still cleared so the next leader does not inherit stale
-     * "pending push" / "last-failed-worker" state.
+     * Bare {@link HttpURLConnection} socket reads honor neither {@link Thread#interrupt()}
+     * nor any cooperative flag, so we disconnect the connection this controller's own worker
+     * is currently blocked on. That unblocks its read within the drain budget instead of
+     * waiting out the per-node push/download timeout (up to 3600s). We only ever disconnect
+     * OUR OWN connection, so no other caller's request is affected.
      */
     @Override
     protected void onStopRequested() {
         BlockingQueue<CheckpointCompletionStatus> currentResult = result;
         if (currentResult != null) {
             currentResult.offer(new CheckpointCompletionStatus(false, "leader checkpoint controller is stopping", null));
         }
-        MetaHelper.cancelInFlight();
+        HttpURLConnection conn = inFlightConnection;
+        if (conn != null) {
+            try {
+                conn.disconnect();
+            } catch (Throwable t) {
+                LOG.warn("failed to disconnect in-flight checkpoint connection on stop", t);
+            }
+        }
     }
 
     @Override
@@ -169,6 +185,18 @@ protected void onStopped() {
         lastFailedTime.clear();
     }
 
+    /**
+     * Interrupt-unsafe: the worker calls BDBJE/JE directly (getFinalizedJournalId /
+     * deleteJournals -> removeDatabase) where an interrupt can invalidate the environment, and
+     * blocks in uninterruptible HttpURLConnection reads. It stops cooperatively instead: the
+     * isStopped() polling between phases + onStopRequested() (waking the result queue and
+     * disconnecting the in-flight HTTP connection).
+     */
+    @Override
+    protected boolean interruptOnStop() {
+        return false;
+    }
+
     protected void runCheckpointController() {
         // ignore return value in normal checkpoint controller
         runCheckpointControllerWithIds(getImageJournalId(), getCheckpointJournalId(), false);
@@ -326,7 +354,12 @@ private void downloadImage(ImageFormatVersion imageFormatVersion, String imageDi
                 "/image?version=" + journalId
                 + "&subdir=" + subDir
                 + "&image_format_version=" + imageFormatVersion;
-        MetaHelper.downloadImageFile(url, MetaService.DOWNLOAD_TIMEOUT_SECOND * 1000, String.valueOf(journalId), dir);
+        try {
+            MetaHelper.downloadImageFile(url, MetaService.DOWNLOAD_TIMEOUT_SECOND * 1000, String.valueOf(journalId), dir,
+                    conn -> inFlightConnection = conn);
+        } finally {
+            inFlightConnection = null;
+        }
 
         // clean the old images
         MetaCleaner cleaner = new MetaCleaner(imageDir);
@@ -472,7 +505,7 @@ private void pushImage(long imageVersion) {
                     + "&for_global_state=" + belongToGlobalStateMgr
                     + "&image_format_version=" + formatVersion.toString();
             try {
-                MetaHelper.httpGet(url, PUT_TIMEOUT_SECOND * 1000);
+                MetaHelper.httpGet(url, PUT_TIMEOUT_SECOND * 1000, conn -> inFlightConnection = conn);
 
                 LOG.info("push image successfully, url = {}", url);
                 if (MetricRepo.hasInit) {
@@ -481,6 +514,8 @@ private void pushImage(long imageVersion) {
             } catch (IOException e) {
                 pushSuccess = false;
                 LOG.error("Exception when pushing image file. url = {}", url, e);
+            } finally {
+                inFlightConnection = null;
             }
             if (pushSuccess) {
                 iterator.remove();
@@ -508,6 +543,7 @@ private long getMinReplayedJournalId() {
                  */
                 idURL = new URL("http://" + NetUtils.getHostPortInAccessibleFormat(host, port) + "/journal_id?prefix=" + journal.getPrefix());
                 conn = (HttpURLConnection) idURL.openConnection();
+                inFlightConnection = conn;
                 conn.setConnectTimeout(CONNECT_TIMEOUT_SECOND * 1000);
                 conn.setReadTimeout(READ_TIMEOUT_SECOND * 1000);
                 String idString = conn.getHeaderField("id");
@@ -521,6 +557,7 @@ private long getMinReplayedJournalId() {
                 minReplayedJournalId = 0;
                 break;
             } finally {
+                inFlightConnection = null;
                 if (conn != null) {
                     conn.disconnect();
                 }

fe/fe-core/src/main/java/com/starrocks/leader/MetaHelper.java

@@ -62,6 +62,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.function.Consumer;
 import java.util.stream.Stream;
 
 public class MetaHelper {
@@ -77,26 +78,6 @@ public static int getLimit() {
         return CHECKPOINT_LIMIT_BYTES;
     }
 
-    /**
-     * Best-effort cancel for any blocking HTTP operation currently issued by MetaHelper on
-     * behalf of leader-only daemons (e.g. CheckpointController push/download). Today this is
-     * a no-op: {@link #httpGet} and {@link #downloadImageFile} use bare {@link HttpURLConnection},
-     * which honors neither {@link Thread#interrupt()} nor any other external cancel signal for
-     * socket reads. The only way a stuck call exits is its own connect/read timeout, which
-     * may be up to an hour for push and {@code MetaService.DOWNLOAD_TIMEOUT_SECOND} for
-     * download.
-     *
-     * TODO(cancel-infra): wire this up to actively disconnect the active {@link HttpURLConnection}
-     * (or switch to a cancellable HTTP client) so leader-demotion {@code onStopped()} can
-     * break out of these calls within the {@code leader_demotion_drain_timeout_sec} budget.
-     * This is tracked alongside the planned {@code Locker.lockInterruptibly} cleanup, since
-     * both block the same goal of making demotion drain deterministic for non-interruptible
-     * blocking operations.
-     */
-    public static void cancelInFlight() {
-        // intentionally empty - see TODO(cancel-infra) above
-    }
-
     // rename the .PART_SUFFIX file to filename
     public static File complete(String filename, File dir) throws IOException {
         File file = new File(dir, filename + MetaHelper.PART_SUFFIX);
@@ -109,6 +90,18 @@ public static File complete(String filename, File dir) throws IOException {
 
     public static void downloadImageFile(String urlStr, int timeout, String journalId, File destDir)
             throws IOException {
+        downloadImageFile(urlStr, timeout, journalId, destDir, null);
+    }
+
+    /**
+     * Same as {@link #downloadImageFile(String, int, String, File)} but publishes the freshly
+     * opened {@link HttpURLConnection} to {@code onConnect} so the *calling* leader daemon can
+     * hold its own reference and {@link HttpURLConnection#disconnect() disconnect} it on demotion
+     * to break out of a stuck socket read. MetaHelper itself performs no cancellation.
+     */
+    public static void downloadImageFile(String urlStr, int timeout, String journalId, File destDir,
+                                         Consumer<HttpURLConnection> onConnect)
+            throws IOException {
         HttpURLConnection conn = null;
         String checksum = null;
         String destFilename = Storage.IMAGE + "." + journalId;
@@ -117,6 +110,9 @@ public static void downloadImageFile(String urlStr, int timeout, String journalI
         try (FileOutputStream out = new FileOutputStream(partFile)) {
             URL url = new URL(urlStr);
             conn = (HttpURLConnection) url.openConnection();
+            if (onConnect != null) {
+                onConnect.accept(conn);
+            }
             conn.setConnectTimeout(timeout);
             conn.setReadTimeout(timeout);
 
@@ -169,11 +165,25 @@ public static OutputStream getOutputStream(String filename, File dir)
     }
 
     public static void httpGet(String urlStr, int timeout) throws IOException {
+        httpGet(urlStr, timeout, null);
+    }
+
+    /**
+     * Same as {@link #httpGet(String, int)} but publishes the freshly opened
+     * {@link HttpURLConnection} to {@code onConnect} so the *calling* leader daemon can hold its
+     * own reference and {@link HttpURLConnection#disconnect() disconnect} it on demotion to break
+     * out of a stuck socket read (e.g. a checkpoint push that can otherwise block up to an hour).
+     * MetaHelper itself performs no cancellation.
+     */
+    public static void httpGet(String urlStr, int timeout, Consumer<HttpURLConnection> onConnect) throws IOException {
         URL url = new URL(urlStr);
         HttpURLConnection conn = null;
 
         try {
             conn = (HttpURLConnection) url.openConnection();
+            if (onConnect != null) {
+                onConnect.accept(conn);
+            }
             conn.setConnectTimeout(timeout);
             conn.setReadTimeout(timeout);
 

fe/fe-core/src/main/java/com/starrocks/load/routineload/RoutineLoadTaskScheduler.java

@@ -206,8 +206,12 @@ private void process() throws InterruptedException {
             }
 
             submitToSchedule(routineLoadTaskInfo);
+        } catch (InterruptedException e) {
+            // Propagate so the LeaderDaemon loop breaks promptly when the scheduler is being
+            // stopped (e.g. on leader demotion) instead of swallowing the cancel and re-polling.
+            throw e;
         } catch (Exception e) {
-            LOG.warn("Taking routine load task from queue has been interrupted", e);
+            LOG.warn("Failed to take/schedule routine load task from queue", e);
             return;
         }
     }