[Enhancement] Cancel in-flight agent tasks at the source on leader demotion

Replace the per-wait-site isLeaderDemoting() poll in TabletTaskExecutor with a
single source-level drain: on leader demotion, fail the completion latch of every
in-flight agent task in AgentTaskQueue so all latch-based waiters (create-tablet in
waitForFinished, and the other agent-task waiters) unblock at once and release their
db locks, instead of each wait site polling and every new waiter needing its own check.

- AgentTask.cancelPendingWaiter(Status): default no-op; overridden by the latch-holding
  subclasses (CreateReplicaTask, PushTask, TabletMetadataUpdateAgentTask,
  DropAutoIncrementMapTask) to countDownToZero their latch. Kept distinct from
  setFailed(boolean) (per-task status flag consumed by ALTER-job polling) and per-mark
  countdown (per-replica success) so retry/quorum are not broken.
- AgentTaskQueue.failAllPendingWaiters(Status): snapshot under the queue lock, then cancel
  outside it so the global queue monitor is not held across N latch operations.
- AgentTaskQueue.addTask(): reject new tasks once demoting, closing the TOCTOU where a
  create-tablet started just after the drain would otherwise wait out its timeout.
- GlobalStateMgr: new demotion stage failInFlightAgentTasks right after beginLeaderDemotion,
  so waiters release their locks before the leader-only daemons are stopped.
- TabletTaskExecutor.waitForFinished: drop the now-redundant per-iteration demotion poll.

Add AgentTaskTest coverage for the drain and the addTask demoting guard.

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

Author: gengjun-git

fe/fe-core/src/main/java/com/starrocks/server/GlobalStateMgr.java

@@ -2149,6 +2149,7 @@ void executeLeaderDemotionStages(FrontendNodeType targetType) {
         LOG.info("leader demotion to {} starting", targetType);
         long startMs = System.currentTimeMillis();
         runDemotionStage("beginLeaderDemotion", () -> beginLeaderDemotion(targetType));
+        runDemotionStage("failInFlightAgentTasks", this::failInFlightAgentTasks);
         runDemotionStage("sealJournalWriter", this::sealJournalWriter);
         runDemotionStage("stopLeaderOnlyDaemonThreads", this::stopLeaderOnlyDaemonThreads);
         runDemotionStage("flushLeaderSessionState", this::flushLeaderSessionState);
@@ -2170,6 +2171,18 @@ private void runDemotionStage(String stageName, Runnable stage) {
         LOG.info("leader demotion stage '{}' completed in {}ms", stageName, System.currentTimeMillis() - startMs);
     }
 
+    /**
+     * Abandon every in-flight BE agent task the (demoting) leader issued: fail their completion
+     * latches so waiters (e.g. create-tablet in TabletTaskExecutor.waitForFinished, which can run on
+     * a user connection thread that demotion never interrupts) unblock immediately and release their
+     * db locks, instead of waiting out tablet_create_timeout for an operation whose journal write is
+     * already fenced. AgentTaskQueue.addTask() rejects new tasks once demoting, closing the TOCTOU.
+     */
+    private void failInFlightAgentTasks() {
+        com.starrocks.task.AgentTaskQueue.failAllPendingWaiters(
+                new com.starrocks.common.Status(TStatusCode.CANCELLED, "leader is demoting"));
+    }
+
     @VisibleForTesting
     void sealJournalWriter() {
         if (journalWriter == null) {

fe/fe-core/src/main/java/com/starrocks/task/AgentTask.java

@@ -35,6 +35,7 @@
 package com.starrocks.task;
 
 import com.starrocks.common.Config;
+import com.starrocks.common.Status;
 import com.starrocks.thrift.TResourceInfo;
 import com.starrocks.thrift.TTaskType;
 
@@ -167,6 +168,19 @@ public void setFailed(boolean isFailed) {
         this.isFailed = isFailed;
     }
 
+    /**
+     * Release any waiter blocked on this task's completion latch, failing it with {@code status}.
+     * Default no-op: most agent tasks have no waiter. Latch-holding subclasses (create-replica,
+     * push, tablet-metadata-update, drop-auto-increment-map) override this so a leader demotion can
+     * abandon all in-flight agent tasks at once (see {@link AgentTaskQueue#failAllPendingWaiters})
+     * and unblock their waiters immediately, instead of each wait site polling for demotion.
+     * Distinct from {@link #setFailed(boolean)} (per-task status flag consumed by ALTER jobs) and
+     * from the per-mark countdown (per-replica success) - overriding those would break retry/quorum.
+     */
+    public void cancelPendingWaiter(Status status) {
+        // no-op by default
+    }
+
     public boolean shouldResend(long currentTimeMillis) {
         return createTime == -1 || currentTimeMillis - createTime > Config.agent_task_resend_wait_time_ms;
     }

fe/fe-core/src/main/java/com/starrocks/task/AgentTaskQueue.java

@@ -40,7 +40,9 @@
 import com.google.common.collect.Maps;
 import com.google.common.collect.Multimap;
 import com.google.common.collect.Table;
+import com.starrocks.common.Status;
 import com.starrocks.memory.estimate.Estimator;
+import com.starrocks.server.GlobalStateMgr;
 import com.starrocks.thrift.TPushType;
 import com.starrocks.thrift.TTaskType;
 import org.apache.logging.log4j.LogManager;
@@ -74,6 +76,14 @@ public static synchronized void addTaskList(List<AgentTask> taskList) {
     }
 
     public static synchronized boolean addTask(AgentTask task) {
+        // Source guard for leader demotion: once demoting, refuse to enqueue new agent tasks.
+        // Together with failAllPendingWaiters() (which drains what is already queued) this closes the
+        // TOCTOU where a create-tablet started just after the drain would otherwise wait out its
+        // timeout. The throw fails the doomed operation fast (its journal write would be fenced anyway).
+        if (GlobalStateMgr.getCurrentState().isLeaderDemoting()) {
+            throw new IllegalStateException(
+                    "leader is demoting, refuse to enqueue agent task: " + task);
+        }
         long backendId = task.getBackendId();
         TTaskType type = task.getTaskType();
 
@@ -261,6 +271,31 @@ public static synchronized void clearAllTasks() {
         taskNum = 0;
     }
 
+    /**
+     * Leader-demotion drain: fail the completion latch of every in-flight agent task so any waiter
+     * (e.g. TabletTaskExecutor.waitForFinished waiting on a create-tablet latch) unblocks at once
+     * with {@code status} instead of waiting out its timeout. Snapshots under the queue lock, then
+     * cancels outside it, so we do not hold the global queue monitor across N latch operations
+     * (which would stall concurrent BE-response processing). Tasks without a latch are no-ops.
+     * Pairs with the addTask() demoting guard, which prevents new tasks from being enqueued after
+     * this runs.
+     */
+    public static void failAllPendingWaiters(Status status) {
+        List<AgentTask> snapshot = Lists.newArrayList();
+        synchronized (AgentTaskQueue.class) {
+            for (Map<Long, AgentTask> signatureMap : tasks.values()) {
+                snapshot.addAll(signatureMap.values());
+            }
+        }
+        for (AgentTask task : snapshot) {
+            try {
+                task.cancelPendingWaiter(status);
+            } catch (Throwable t) {
+                LOG.warn("failed to cancel pending waiter for agent task {}", task, t);
+            }
+        }
+    }
+
     public static synchronized int getTaskNum() {
         return taskNum;
     }

fe/fe-core/src/main/java/com/starrocks/task/CreateReplicaTask.java

@@ -164,6 +164,13 @@ public void countDownToZero(String errMsg) {
         }
     }
 
+    @Override
+    public void cancelPendingWaiter(Status status) {
+        if (this.latch != null) {
+            latch.countDownToZero(status);
+        }
+    }
+
     public void setLatch(MarkedCountDownLatch<Long, Long> latch) {
         this.latch = latch;
     }

fe/fe-core/src/main/java/com/starrocks/task/DropAutoIncrementMapTask.java

@@ -59,6 +59,13 @@ public void countDownToZero(String errMsg) {
         }
     }
 
+    @Override
+    public void cancelPendingWaiter(Status status) {
+        if (this.latch != null) {
+            latch.countDownToZero(status);
+        }
+    }
+
     public void setLatch(MarkedCountDownLatch<Long, Long> latch) {
         this.latch = latch;
     }

fe/fe-core/src/main/java/com/starrocks/task/PushTask.java

@@ -229,6 +229,13 @@ public void countDownLatch(long backendId, long tabletId) {
         }
     }
 
+    @Override
+    public void cancelPendingWaiter(Status status) {
+        if (this.latch != null) {
+            latch.countDownToZero(status);
+        }
+    }
+
     public void countDownLatch(long backendId, long tabletId, String errMsg) {
         if (this.latch != null) {
             if (latch.markedCountDown(backendId, tabletId, new Status(TStatusCode.INTERNAL_ERROR, errMsg))) {

fe/fe-core/src/main/java/com/starrocks/task/TabletMetadataUpdateAgentTask.java

@@ -100,6 +100,13 @@ public void countDownToZero(String errMsg) {
         }
     }
 
+    @Override
+    public void cancelPendingWaiter(Status status) {
+        if (this.latch != null) {
+            latch.countDownToZero(status);
+        }
+    }
+
     public abstract Set<Long> getTablets();
 
     public abstract List<TTabletMetaInfo> getTTabletMetaInfoList();

fe/fe-core/src/main/java/com/starrocks/task/TabletTaskExecutor.java

@@ -543,14 +543,6 @@ static void waitForFinished(MarkedCountDownLatch<Long, Long> countDownLatch, lon
                     }
                 }
 
-                GlobalStateMgr globalStateMgr = GlobalStateMgr.getCurrentState();
-                if (globalStateMgr.isLeaderDemoting() || !globalStateMgr.isLeaderWorkAdmissionOpen()) {
-                    String errMsg = "Create tablet cancelled because leader work admission is closed";
-                    LOG.warn(errMsg);
-                    countDownLatch.countDownToZero(new Status(TStatusCode.CANCELLED, errMsg));
-                    throw new DdlException(errMsg);
-                }
-
                 timeLeft -= waitInterval;
             }
             // timed out

fe/fe-core/src/test/java/com/starrocks/task/AgentTaskTest.java

@@ -44,6 +44,7 @@
 import com.starrocks.common.AnalysisException;
 import com.starrocks.common.Pair;
 import com.starrocks.common.Range;
+import com.starrocks.common.Status;
 import com.starrocks.common.jmockit.Deencapsulation;
 import com.starrocks.common.util.concurrent.MarkedCountDownLatch;
 import com.starrocks.server.GlobalStateMgr;
@@ -58,6 +59,7 @@
 import com.starrocks.thrift.TBackend;
 import com.starrocks.thrift.TCompressionType;
 import com.starrocks.thrift.TCreateTabletReq;
+import com.starrocks.thrift.TStatusCode;
 import com.starrocks.thrift.TStorageMedium;
 import com.starrocks.thrift.TStorageType;
 import com.starrocks.thrift.TTabletRange;
@@ -491,4 +493,37 @@ public void testCreateReplicaTaskBuilderPreservesRange() throws AnalysisExceptio
         TCreateTabletReq req = task.toThrift();
         Assertions.assertTrue(req.isSetRange());
     }
+
+    @Test
+    public void testFailAllPendingWaitersReleasesLatch() {
+        AgentTaskQueue.clearAllTasks();
+        MarkedCountDownLatch<Long, Long> l = new MarkedCountDownLatch<>(1);
+        l.addMark(backendId1, tabletId1);
+        ((CreateReplicaTask) createReplicaTask).setLatch(l);
+        AgentTaskQueue.addTask(createReplicaTask);
+
+        // Leader-demotion drain: fail every in-flight agent task latch so a waiter
+        // (e.g. TabletTaskExecutor.waitForFinished) unblocks at once with the demotion reason
+        // instead of waiting out its timeout.
+        AgentTaskQueue.failAllPendingWaiters(new Status(TStatusCode.CANCELLED, "leader is demoting"));
+
+        Assertions.assertEquals(0, l.getCount());
+        Assertions.assertFalse(l.getStatus().ok());
+        AgentTaskQueue.clearAllTasks();
+    }
+
+    @Test
+    public void testAddTaskRejectedWhenLeaderDemoting() {
+        AgentTaskQueue.clearAllTasks();
+        new MockUp<GlobalStateMgr>() {
+            @Mock
+            public boolean isLeaderDemoting() {
+                return true;
+            }
+        };
+        // Source guard: once demoting, no new agent task may be enqueued, so a create-tablet
+        // started after the drain fails fast instead of hanging on a never-marked latch.
+        Assertions.assertThrows(IllegalStateException.class, () -> AgentTaskQueue.addTask(createReplicaTask));
+        Assertions.assertEquals(0, AgentTaskQueue.getTaskNum());
+    }
 }