[server] Make timed-out rebalance tasks recoverable

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/CoordinatorRequestBatch.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/event/FinalizeRebalanceEvent.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/NotifyLeaderAndIsrRequestContext.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/NotifyLeaderAndIsrResponseReceivedEvent.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/RebalanceTaskTimeoutEvent.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/ReconcileRebalanceTaskEvent.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/RecoverRebalanceEvent.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) {