[server] Move leader during reassignment when old leader stays

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorEventProcessor.java

@@ -1266,12 +1266,15 @@ private void processNotifyLeaderAndIsrResponseReceivedEvent(
             onReplicaBecomeOffline(offlineReplicas);
         }
 
-        // Try to complete rebalance tasks for the buckets in the response.
-        // This is essential for leader-only migrations to ensure they wait for the tablet
-        // server to acknowledge the leader change before proceeding to the next migration.
+        // Try to complete rebalance tasks for the buckets in the response. A task completes only
+        // after the planned leader acknowledges its final state, so migrations run one by one.
         for (NotifyLeaderAndIsrResultForBucket notifyLeaderAndIsrResultForBucket :
                 notifyLeaderAndIsrResultForBuckets) {
-            tryToCompleteRebalanceTask(notifyLeaderAndIsrResultForBucket, serverId);
+            tryToCompleteRebalanceTask(
+                    notifyLeaderAndIsrResultForBucket,
+                    serverId,
+                    notifyLeaderAndIsrResponseReceivedEvent.getRequestBucketEpoch(
+                            notifyLeaderAndIsrResultForBucket.getTableBucket()));
         }
     }
 
@@ -1699,6 +1702,7 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) {
                     Collections.singleton(tableBucket),
                     OnlineBucket,
                     new ReassignmentLeaderElection(newReplicas));
+            awaitPlannedLeaderAck(tableBucket);
         } else {
             try {
                 LOG.info(
@@ -1716,7 +1720,8 @@ public void tryToExecuteRebalanceTask(RebalancePlanForBucket planForBucket) {
     /** try to finish rebalance tasks after receive notify leader and isr response. */
     private void tryToCompleteRebalanceTask(
             NotifyLeaderAndIsrResultForBucket notifyLeaderAndIsrResultForBucket,
-            int responseServerId) {
+            int responseServerId,
+            int requestBucketEpoch) {
         TableBucket tableBucket = notifyLeaderAndIsrResultForBucket.getTableBucket();
         RebalancePlanForBucket planForBucket =
                 rebalanceManager.getRebalancePlanForBucket(tableBucket);
@@ -1725,9 +1730,16 @@ private void tryToCompleteRebalanceTask(
                     ReplicaReassignment.build(
                             planForBucket.getOriginReplicas(), planForBucket.getNewReplicas());
             try {
-                if (planForBucket.isLeaderChanged() && !reassignment.isBeingReassigned()) {
-                    if (!isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader(
-                            notifyLeaderAndIsrResultForBucket, responseServerId, planForBucket)) {
+                Integer awaitedBucketEpoch =
+                        rebalanceManager.getAwaitedLeaderAckBucketEpoch(tableBucket);
+                if (awaitedBucketEpoch != null) {
+                    // A response to an earlier request (older bucket epoch) must not complete
+                    // the task.
+                    if (requestBucketEpoch < awaitedBucketEpoch
+                            || !isSuccessfulRebalanceResponseFromNewLeader(
+                                    notifyLeaderAndIsrResultForBucket,
+                                    responseServerId,
+                                    planForBucket)) {
                         return;
                     }
                     LeaderAndIsr leaderAndIsr = zooKeeperClient.getLeaderAndIsr(tableBucket).get();
@@ -1755,7 +1767,9 @@ private void tryToCompleteRebalanceTaskOnLeaderAndIsrChange(TableBucket tableBuc
             ReplicaReassignment reassignment =
                     ReplicaReassignment.build(
                             planForBucket.getOriginReplicas(), planForBucket.getNewReplicas());
-            if (planForBucket.isLeaderChanged() && !reassignment.isBeingReassigned()) {
+            // The task only waits for the planned leader to acknowledge its final state, running
+            // the reassignment again would repeat its last phase.
+            if (rebalanceManager.getAwaitedLeaderAckBucketEpoch(tableBucket) != null) {
                 return;
             }
             try {
@@ -1781,7 +1795,7 @@ private void tryToCompleteReassignmentTask(
     }
 
     @VisibleForTesting
-    static boolean isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader(
+    static boolean isSuccessfulRebalanceResponseFromNewLeader(
             NotifyLeaderAndIsrResultForBucket notifyLeaderAndIsrResultForBucket,
             int responseServerId,
             RebalancePlanForBucket planForBucket) {
@@ -1833,7 +1847,8 @@ static boolean isSuccessfulLeaderOnlyRebalanceResponseFromNewLeader(
      *   <li>B7. Update ZK with RS=TRS, AR=[], RR=[].
      *   <li>B8. After electing leader, the replicas and isr information changes. So resend the
      *       update metadata request to every tabletServer.
-     *   <li>B8. Mark the ongoing rebalance task to finish.
+     *   <li>B9. Wait for the new leader to acknowledge the final state, then mark the ongoing
+     *       rebalance task to finish.
      * </ul>
      *
      * <p>In general, there are two goals we want to aim for:
@@ -1918,9 +1933,38 @@ private void onBucketReassignment(
                     null,
                     null,
                     Collections.singleton(tableBucket));
-            // B8. Mark the ongoing rebalance task to finish.
-            rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.COMPLETED);
+            // B9. Finish the task only after the planned leader acknowledges this final state.
+            awaitPlannedLeaderAck(tableBucket);
+        }
+    }
+
+    private void awaitPlannedLeaderAck(TableBucket tableBucket) {
+        RebalancePlanForBucket planForBucket =
+                rebalanceManager.getRebalancePlanForBucket(tableBucket);
+        Optional<LeaderAndIsr> leaderAndIsr = coordinatorContext.getBucketLeaderAndIsr(tableBucket);
+        if (planForBucket == null) {
+            return;
+        }
+        // Without a leader the task can neither be acknowledged nor safely re-run, so fail it.
+        if (!leaderAndIsr.isPresent()) {
+            LOG.warn(
+                    "No leader for tableBucket {} after applying its rebalance task.", tableBucket);
+            rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED);
+            return;
+        }
+        // The election picks the first live in-sync target replica, which is not the planned
+        // leader when that one is offline. Its ACK can then never arrive, so fail the task now
+        // instead of blocking the next migrations until the task timeout.
+        if (leaderAndIsr.get().leader() != planForBucket.getNewLeader()) {
+            LOG.warn(
+                    "Planned leader {} for tableBucket {} was not elected, current leader is {}.",
+                    planForBucket.getNewLeader(),
+                    tableBucket,
+                    leaderAndIsr.get().leader());
+            rebalanceManager.finishRebalanceTask(tableBucket, RebalanceStatus.FAILED);
+            return;
         }
+        rebalanceManager.awaitLeaderAck(tableBucket, leaderAndIsr.get().bucketEpoch());
     }
 
     private boolean isReassignmentComplete(
@@ -1935,7 +1979,9 @@ private void maybeReassignedBucketLeaderIfRequired(
             TableBucket tableBucket, List<Integer> targetReplicas) throws Exception {
         LeaderAndIsr leaderAndIsr = coordinatorContext.getBucketLeaderAndIsr(tableBucket).get();
         int currentLeader = leaderAndIsr.leader();
-        if (!targetReplicas.contains(currentLeader)) {
+        // The first target replica is the planned leader (see GoalOptimizerUtils#getDiff), so
+        // re-elect also when the current leader stays in the target replicas but is not first.
+        if (currentLeader != targetReplicas.get(0)) {
             LOG.info(
                     "Leader {} for tableBucket {} being reassigned. Re-electing leader to {}",
                     currentLeader,

fluss-server/src/main/java/org/apache/fluss/server/coordinator/CoordinatorRequestBatch.java

@@ -460,8 +460,12 @@ private void sendNotifyLeaderAndIsrRequest(int coordinatorEpoch) {
             // Track exactly which buckets THIS request marked as pending leader activation. Only
             // those entries (where leader == serverId) need to be cleared if the request fails
             Set<TableBucket> addedToPendingLeaderActivation = new HashSet<>();
+            // The response carries no epoch, so remember what this request sent to tell a stale
+            // response apart from the one for the latest leader state.
+            Map<TableBucket, Integer> requestBucketEpochs = new HashMap<>();
             for (Map.Entry<TableBucket, PbNotifyLeaderAndIsrReqForBucket> entry :
                     notifyRequestEntry.getValue().entrySet()) {
+                requestBucketEpochs.put(entry.getKey(), entry.getValue().getBucketEpoch());
                 int leader = entry.getValue().getLeader();
                 if (leader == serverId) {
                     coordinatorContext.addPendingLeaderActivation(entry.getKey());
@@ -504,7 +508,9 @@ private void sendNotifyLeaderAndIsrRequest(int coordinatorEpoch) {
                         // put the response receive event into the event manager
                         eventManager.put(
                                 new NotifyLeaderAndIsrResponseReceivedEvent(
-                                        getNotifyLeaderAndIsrResponseData(response), serverId));
+                                        getNotifyLeaderAndIsrResponseData(response),
+                                        serverId,
+                                        requestBucketEpochs));
                     });
         }
         notifyLeaderAndIsrRequestMap.clear();

fluss-server/src/main/java/org/apache/fluss/server/coordinator/event/NotifyLeaderAndIsrResponseReceivedEvent.java

@@ -17,10 +17,12 @@
 
 package org.apache.fluss.server.coordinator.event;
 
+import org.apache.fluss.metadata.TableBucket;
 import org.apache.fluss.rpc.messages.NotifyLeaderAndIsrRequest;
 import org.apache.fluss.server.entity.NotifyLeaderAndIsrResultForBucket;
 
 import java.util.List;
+import java.util.Map;
 
 /** An event for receive the response of {@link NotifyLeaderAndIsrRequest} from tablet server. */
 public class NotifyLeaderAndIsrResponseReceivedEvent implements CoordinatorEvent {
@@ -30,11 +32,16 @@ public class NotifyLeaderAndIsrResponseReceivedEvent implements CoordinatorEvent
     // the server id that return the response
     private final int responseServerId;
 
+    // the bucket epoch sent in the request for each bucket, the response itself carries no epoch
+    private final Map<TableBucket, Integer> requestBucketEpochs;
+
     public NotifyLeaderAndIsrResponseReceivedEvent(
             List<NotifyLeaderAndIsrResultForBucket> notifyLeaderAndIsrResultForBuckets,
-            int responseServerId) {
+            int responseServerId,
+            Map<TableBucket, Integer> requestBucketEpochs) {
         this.notifyLeaderAndIsrResultForBuckets = notifyLeaderAndIsrResultForBuckets;
         this.responseServerId = responseServerId;
+        this.requestBucketEpochs = requestBucketEpochs;
     }
 
     public int getResponseServerId() {
@@ -44,4 +51,9 @@ public int getResponseServerId() {
     public List<NotifyLeaderAndIsrResultForBucket> getNotifyLeaderAndIsrResultForBuckets() {
         return notifyLeaderAndIsrResultForBuckets;
     }
+
+    /** Returns the bucket epoch the request sent for the bucket, or -1 if unknown. */
+    public int getRequestBucketEpoch(TableBucket tableBucket) {
+        return requestBucketEpochs.getOrDefault(tableBucket, -1);
+    }
 }

fluss-server/src/main/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManager.java

@@ -102,6 +102,13 @@ public class RebalanceManager {
     private final Map<TableBucket, RebalanceResultForBucket> finishedRebalanceTasks =
             new ConcurrentHashMap<>();
 
+    /**
+     * A mapping from table bucket to the bucket epoch its planned leader must acknowledge before
+     * the in-flight task completes.
+     */
+    private final Map<TableBucket, Integer> awaitedLeaderAckBucketEpochs =
+            new ConcurrentHashMap<>();
+
     private final GoalOptimizer goalOptimizer;
     private volatile long registerTime;
     private volatile @Nullable RebalanceStatus rebalanceStatus;
@@ -199,6 +206,7 @@ public void registerRebalance(
         inProgressRebalanceTasks.clear();
         inProgressRebalanceTasksQueue.clear();
         finishedRebalanceTasks.clear();
+        awaitedLeaderAckBucketEpochs.clear();
         // Clear gate (bucket) first, then data (startMs).
         inflightTaskBucket = null;
         inflightTaskStartMs = -1;
@@ -236,6 +244,7 @@ public void finishRebalanceTask(TableBucket tableBucket, RebalanceStatus statusF
         if (inProgressRebalanceTasksQueue.contains(tableBucket)) {
             inProgressRebalanceTasksQueue.remove(tableBucket);
             RebalanceResultForBucket resultForBucket = inProgressRebalanceTasks.remove(tableBucket);
+            awaitedLeaderAckBucketEpochs.remove(tableBucket);
             checkNotNull(resultForBucket, "RebalanceResultForBucket is null.");
             finishedRebalanceTasks.put(
                     tableBucket,
@@ -323,6 +332,7 @@ public void cancelRebalance(@Nullable String rebalanceId) {
         rebalanceStatus = CANCELED;
         inProgressRebalanceTasksQueue.clear();
         inProgressRebalanceTasks.clear();
+        awaitedLeaderAckBucketEpochs.clear();
         // Clear gate (bucket) first, then data (startMs).
         inflightTaskBucket = null;
         inflightTaskStartMs = -1;
@@ -375,6 +385,27 @@ public RebalanceTask generateRebalanceTask(List<Goal> goalsByPriority) {
         return null;
     }
 
+    /**
+     * Marks that the coordinator applied the final state of the in-flight task for the bucket. The
+     * task completes once the planned leader acknowledges a request with at least the given bucket
+     * epoch.
+     */
+    public void awaitLeaderAck(TableBucket tableBucket, int bucketEpoch) {
+        checkNotClosed();
+        if (inProgressRebalanceTasks.containsKey(tableBucket)) {
+            awaitedLeaderAckBucketEpochs.put(tableBucket, bucketEpoch);
+        }
+    }
+
+    /**
+     * Returns the bucket epoch the planned leader must acknowledge, or null if the task for the
+     * bucket is not waiting for the leader acknowledgement.
+     */
+    public @Nullable Integer getAwaitedLeaderAckBucketEpoch(TableBucket tableBucket) {
+        checkNotClosed();
+        return awaitedLeaderAckBucketEpochs.get(tableBucket);
+    }
+
     private void processNewRebalanceTask() {
         TableBucket tableBucket = inProgressRebalanceTasksQueue.peek();
         if (tableBucket != null && inProgressRebalanceTasks.containsKey(tableBucket)) {
@@ -410,6 +441,7 @@ private void completeRebalance() {
         rebalanceStatus = COMPLETED;
         inProgressRebalanceTasks.clear();
         inProgressRebalanceTasksQueue.clear();
+        awaitedLeaderAckBucketEpochs.clear();
 
         // Here, it will not clear finishedRebalanceTasks, because it will be used by
         // listRebalanceProgress. It will be cleared when next register.