[server] Make timed-out rebalance tasks recoverable

Keep timed-out bucket migrations under bounded reconciliation while admitting
new work. Persist cancellation intent and final status, recover intermediate
assignments, and require a current leader acknowledgement before completion.
Fence stale responses and avoid repeated elections for applied leader moves.

Expose recovery limits and replan failed migrations using the configured
replica count. Document the state transitions and recovery settings, share
test setup, and retain targeted failure and restart coverage. Exempt only
ConfigOptions from the Checkstyle file length limit.

Skip startup recovery on ZooKeeper read failure to allow new rebalance
requests, and leave a TODO for recovery retries.

Validation: 89 targeted and integration tests passed, including five rebalance
integration cases. After the startup recovery adjustment, reran all 35
RebalanceManagerTest cases successfully. Checkstyle and Spotless checks passed.

Co-Authored-By: Codex <noreply@openai.com>
AI-Model: gpt-6
AI-Contributed/Feature: 7/7
AI-Contributed/UT: 0/0

Author: swuferhong

fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalanceStatus.java

@@ -38,7 +38,7 @@ public enum RebalanceStatus {
     TIMEOUT(5);
 
     public static final Set<RebalanceStatus> FINAL_STATUSES =
-            new HashSet<>(Arrays.asList(COMPLETED, CANCELED, FAILED, TIMEOUT));
+            new HashSet<>(Arrays.asList(COMPLETED, CANCELED, FAILED));
 
     private final int code;
 

fluss-common/src/main/java/org/apache/fluss/config/ConfigOptions.java

@@ -241,6 +241,39 @@ public class ConfigOptions {
                                             + "TableLifecycleThrottler scans in-flight drops for "
                                             + "timeouts.");
 
+    public static final ConfigOption<Duration> COORDINATOR_REBALANCE_TARGET_UNAVAILABLE_TIMEOUT =
+            key("coordinator.rebalance.target-unavailable-timeout")
+                    .durationType()
+                    .defaultValue(Duration.ofMinutes(30))
+                    .withDescription(
+                            "The time after which a timed-out bucket migration is marked FAILED "
+                                    + "if its target tablet servers remain unavailable and no "
+                                    + "bucket-state progress is observed. Observed progress or all "
+                                    + "target servers becoming live resets this timer. Must be at "
+                                    + "least 1 ms and fit in a signed 64-bit millisecond value.");
+
+    public static final ConfigOption<Duration> COORDINATOR_REBALANCE_NO_PROGRESS_TIMEOUT =
+            key("coordinator.rebalance.no-progress-timeout")
+                    .durationType()
+                    .defaultValue(Duration.ofHours(24))
+                    .withDescription(
+                            "The time after which a timed-out bucket migration is marked FAILED "
+                                    + "if no bucket-state progress is observed, even when its target "
+                                    + "tablet servers are live. Observed progress resets this timer; "
+                                    + "this does not limit the total migration duration. Must be at "
+                                    + "least 1 ms and fit in a signed 64-bit millisecond value.");
+
+    public static final ConfigOption<Integer> COORDINATOR_REBALANCE_MAX_TRACKED_TIMED_OUT_TASKS =
+            key("coordinator.rebalance.max-tracked-timed-out-tasks")
+                    .intType()
+                    .defaultValue(8)
+                    .withDescription(
+                            "The maximum number of non-final timed-out bucket migrations tracked "
+                                    + "before the coordinator stops admitting new bucket migrations. "
+                                    + "Admission resumes when a tracked migration reaches a final "
+                                    + "status. This limits concurrent migrations and reconciliation "
+                                    + "work on the coordinator and ZooKeeper. Must be at least 1.");
+
     public static final ConfigOption<Duration> COORDINATOR_OFFLINE_LEADER_RETRY_DELAY =
             key("coordinator.offline-leader.retry-delay")
                     .durationType()

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

@@ -38,7 +38,9 @@
 import org.apache.fluss.server.coordinator.event.AccessContextEvent;
 import org.apache.fluss.server.coordinator.event.DeleteReplicaResponseReceivedEvent;
 import org.apache.fluss.server.coordinator.event.EventManager;
+import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrRequestContext;
 import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent;
+import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey;
 import org.apache.fluss.server.entity.DeleteReplicaResultForBucket;
 import org.apache.fluss.server.entity.NotifyLeaderAndIsrData;
 import org.apache.fluss.server.metadata.BucketMetadata;
@@ -60,6 +62,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static org.apache.fluss.server.metadata.PartitionMetadata.DELETED_PARTITION_ID;
@@ -114,6 +117,8 @@ public class CoordinatorRequestBatch {
     private final CoordinatorChannelManager coordinatorChannelManager;
     private final EventManager eventManager;
     private final CoordinatorContext coordinatorContext;
+    private Function<TableBucket, RebalanceExecutionKey> rebalanceExecutionKeyResolver =
+            ignored -> null;
 
     public CoordinatorRequestBatch(
             CoordinatorChannelManager coordinatorChannelManager,
@@ -124,6 +129,11 @@ public CoordinatorRequestBatch(
         this.coordinatorContext = coordinatorContext;
     }
 
+    void setRebalanceExecutionKeyResolver(
+            Function<TableBucket, RebalanceExecutionKey> rebalanceExecutionKeyResolver) {
+        this.rebalanceExecutionKeyResolver = rebalanceExecutionKeyResolver;
+    }
+
     public void newBatch() {
         if (!notifyLeaderAndIsrRequestMap.isEmpty()) {
             throw new IllegalStateException(
@@ -456,6 +466,19 @@ private void sendNotifyLeaderAndIsrRequest(int coordinatorEpoch) {
             NotifyLeaderAndIsrRequest notifyLeaderAndIsrRequest =
                     makeNotifyLeaderAndIsrRequest(
                             coordinatorEpoch, notifyRequestEntry.getValue().values());
+            Map<TableBucket, NotifyLeaderAndIsrRequestContext> requestContexts = new HashMap<>();
+            for (Map.Entry<TableBucket, PbNotifyLeaderAndIsrReqForBucket> entry :
+                    notifyRequestEntry.getValue().entrySet()) {
+                PbNotifyLeaderAndIsrReqForBucket request = entry.getValue();
+                requestContexts.put(
+                        entry.getKey(),
+                        new NotifyLeaderAndIsrRequestContext(
+                                coordinatorEpoch,
+                                request.getLeader(),
+                                request.getLeaderEpoch(),
+                                request.getBucketEpoch(),
+                                rebalanceExecutionKeyResolver.apply(entry.getKey())));
+            }
 
             // 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
@@ -504,7 +527,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,
+                                        requestContexts));
                     });
         }
         notifyLeaderAndIsrRequestMap.clear();

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

@@ -0,0 +1,36 @@
+/*
+ * 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.fluss.server.coordinator.event;
+
+/** An event that retries persisting the final state of a rebalance. */
+public final class FinalizeRebalanceEvent implements CoordinatorEvent {
+    private final String rebalanceId;
+
+    public FinalizeRebalanceEvent(String rebalanceId) {
+        this.rebalanceId = rebalanceId;
+    }
+
+    public String getRebalanceId() {
+        return rebalanceId;
+    }
+
+    @Override
+    public String toString() {
+        return "FinalizeRebalanceEvent{rebalanceId='" + rebalanceId + "'}";
+    }
+}

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

@@ -0,0 +1,86 @@
+/*
+ * 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.fluss.server.coordinator.event;
+
+import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey;
+
+import javax.annotation.Nullable;
+
+/** The coordinator and bucket epochs of a sent NotifyLeaderAndIsr request. */
+public final class NotifyLeaderAndIsrRequestContext {
+
+    private final int coordinatorEpoch;
+    private final int leader;
+    private final int leaderEpoch;
+    private final int bucketEpoch;
+    private final @Nullable RebalanceExecutionKey rebalanceExecutionKey;
+
+    public NotifyLeaderAndIsrRequestContext(
+            int coordinatorEpoch, int leader, int leaderEpoch, int bucketEpoch) {
+        this(coordinatorEpoch, leader, leaderEpoch, bucketEpoch, null);
+    }
+
+    public NotifyLeaderAndIsrRequestContext(
+            int coordinatorEpoch,
+            int leader,
+            int leaderEpoch,
+            int bucketEpoch,
+            @Nullable RebalanceExecutionKey rebalanceExecutionKey) {
+        this.coordinatorEpoch = coordinatorEpoch;
+        this.leader = leader;
+        this.leaderEpoch = leaderEpoch;
+        this.bucketEpoch = bucketEpoch;
+        this.rebalanceExecutionKey = rebalanceExecutionKey;
+    }
+
+    public int getCoordinatorEpoch() {
+        return coordinatorEpoch;
+    }
+
+    public int getLeader() {
+        return leader;
+    }
+
+    public int getLeaderEpoch() {
+        return leaderEpoch;
+    }
+
+    public int getBucketEpoch() {
+        return bucketEpoch;
+    }
+
+    public @Nullable RebalanceExecutionKey getRebalanceExecutionKey() {
+        return rebalanceExecutionKey;
+    }
+
+    @Override
+    public String toString() {
+        return "NotifyLeaderAndIsrRequestContext{"
+                + "coordinatorEpoch="
+                + coordinatorEpoch
+                + ", leader="
+                + leader
+                + ", leaderEpoch="
+                + leaderEpoch
+                + ", bucketEpoch="
+                + bucketEpoch
+                + ", rebalanceExecutionKey="
+                + rebalanceExecutionKey
+                + '}';
+    }
+}

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

@@ -17,10 +17,14 @@
 
 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.Collections;
+import java.util.HashMap;
 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 +34,21 @@ public class NotifyLeaderAndIsrResponseReceivedEvent implements CoordinatorEvent
     // the server id that return the response
     private final int responseServerId;
 
+    private final Map<TableBucket, NotifyLeaderAndIsrRequestContext> requestContexts;
+
     public NotifyLeaderAndIsrResponseReceivedEvent(
             List<NotifyLeaderAndIsrResultForBucket> notifyLeaderAndIsrResultForBuckets,
             int responseServerId) {
+        this(notifyLeaderAndIsrResultForBuckets, responseServerId, Collections.emptyMap());
+    }
+
+    public NotifyLeaderAndIsrResponseReceivedEvent(
+            List<NotifyLeaderAndIsrResultForBucket> notifyLeaderAndIsrResultForBuckets,
+            int responseServerId,
+            Map<TableBucket, NotifyLeaderAndIsrRequestContext> requestContexts) {
         this.notifyLeaderAndIsrResultForBuckets = notifyLeaderAndIsrResultForBuckets;
         this.responseServerId = responseServerId;
+        this.requestContexts = Collections.unmodifiableMap(new HashMap<>(requestContexts));
     }
 
     public int getResponseServerId() {
@@ -44,4 +58,8 @@ public int getResponseServerId() {
     public List<NotifyLeaderAndIsrResultForBucket> getNotifyLeaderAndIsrResultForBuckets() {
         return notifyLeaderAndIsrResultForBuckets;
     }
+
+    public NotifyLeaderAndIsrRequestContext getRequestContext(TableBucket tableBucket) {
+        return requestContexts.get(tableBucket);
+    }
 }

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

@@ -17,23 +17,23 @@
 
 package org.apache.fluss.server.coordinator.event;
 
-import org.apache.fluss.metadata.TableBucket;
+import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey;
 
 /** An event fired when a rebalance task exceeds the timeout without completing. */
 public class RebalanceTaskTimeoutEvent implements CoordinatorEvent {
 
-    private final TableBucket tableBucket;
+    private final RebalanceExecutionKey executionKey;
 
-    public RebalanceTaskTimeoutEvent(TableBucket tableBucket) {
-        this.tableBucket = tableBucket;
+    public RebalanceTaskTimeoutEvent(RebalanceExecutionKey executionKey) {
+        this.executionKey = executionKey;
     }
 
-    public TableBucket getTableBucket() {
-        return tableBucket;
+    public RebalanceExecutionKey getExecutionKey() {
+        return executionKey;
     }
 
     @Override
     public String toString() {
-        return "RebalanceTaskTimeoutEvent{tableBucket=" + tableBucket + "}";
+        return "RebalanceTaskTimeoutEvent{executionKey=" + executionKey + "}";
     }
 }

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

@@ -0,0 +1,38 @@
+/*
+ * 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.fluss.server.coordinator.event;
+
+import org.apache.fluss.server.coordinator.rebalance.RebalanceExecutionKey;
+
+/** An event that reconciles a timed-out rebalance task against coordinator state. */
+public final class ReconcileRebalanceTaskEvent implements CoordinatorEvent {
+    private final RebalanceExecutionKey executionKey;
+
+    public ReconcileRebalanceTaskEvent(RebalanceExecutionKey executionKey) {
+        this.executionKey = executionKey;
+    }
+
+    public RebalanceExecutionKey getExecutionKey() {
+        return executionKey;
+    }
+
+    @Override
+    public String toString() {
+        return "ReconcileRebalanceTaskEvent{executionKey=" + executionKey + "}";
+    }
+}

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

@@ -19,9 +19,8 @@
 
 import org.apache.fluss.server.zk.data.RebalanceTask;
 
-/** Fired during startup to recover a rebalance task that was persisted in ZooKeeper. */
-public class RecoverRebalanceEvent implements CoordinatorEvent {
-
+/** An event that recovers and reconciles a persisted rebalance task. */
+public final class RecoverRebalanceEvent implements CoordinatorEvent {
     private final RebalanceTask rebalanceTask;
 
     public RecoverRebalanceEvent(RebalanceTask rebalanceTask) {