[server] Fix thread-safety bug in RebalanceManager ZooKeeper recovery path

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

@@ -86,6 +86,7 @@
 import org.apache.fluss.server.coordinator.event.NotifyLeaderAndIsrResponseReceivedEvent;
 import org.apache.fluss.server.coordinator.event.RebalanceEvent;
 import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent;
+import org.apache.fluss.server.coordinator.event.RecoverRebalanceEvent;
 import org.apache.fluss.server.coordinator.event.RemoveServerTagEvent;
 import org.apache.fluss.server.coordinator.event.ResumeDropEvent;
 import org.apache.fluss.server.coordinator.event.RetryOfflineLeaderEvent;
@@ -743,6 +744,13 @@ public void process(CoordinatorEvent event) {
             RebalanceEvent rebalanceEvent = (RebalanceEvent) event;
             completeFromCallable(
                     rebalanceEvent.getRespCallback(), () -> processRebalance(rebalanceEvent));
+        } else if (event instanceof RecoverRebalanceEvent) {
+            RecoverRebalanceEvent recoverRebalanceEvent = (RecoverRebalanceEvent) event;
+            RebalanceTask rebalanceTask = recoverRebalanceEvent.getRebalanceTask();
+            rebalanceManager.registerRebalance(
+                    rebalanceTask.getRebalanceId(),
+                    rebalanceTask.getExecutePlan(),
+                    rebalanceTask.getRebalanceStatus());
         } else if (event instanceof CancelRebalanceEvent) {
             CancelRebalanceEvent cancelRebalanceEvent = (CancelRebalanceEvent) event;
             completeFromCallable(

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

@@ -0,0 +1,39 @@
+/*
+ * 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.zk.data.RebalanceTask;
+
+/** Fired during startup to recover a rebalance task that was persisted in ZooKeeper. */
+public class RecoverRebalanceEvent implements CoordinatorEvent {
+
+    private final RebalanceTask rebalanceTask;
+
+    public RecoverRebalanceEvent(RebalanceTask rebalanceTask) {
+        this.rebalanceTask = rebalanceTask;
+    }
+
+    public RebalanceTask getRebalanceTask() {
+        return rebalanceTask;
+    }
+
+    @Override
+    public String toString() {
+        return "RecoverRebalanceEvent{rebalanceTask=" + rebalanceTask + "}";
+    }
+}

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

@@ -29,6 +29,7 @@
 import org.apache.fluss.server.coordinator.CoordinatorEventProcessor;
 import org.apache.fluss.server.coordinator.event.EventManager;
 import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent;
+import org.apache.fluss.server.coordinator.event.RecoverRebalanceEvent;
 import org.apache.fluss.server.coordinator.rebalance.goal.Goal;
 import org.apache.fluss.server.coordinator.rebalance.goal.GoalOptimizer;
 import org.apache.fluss.server.coordinator.rebalance.model.ClusterModel;
@@ -178,11 +179,8 @@ private void initialize() {
         try {
             zkClient.getRebalanceTask()
                     .ifPresent(
-                            rebalancePlan ->
-                                    registerRebalance(
-                                            rebalancePlan.getRebalanceId(),
-                                            rebalancePlan.getExecutePlan(),
-                                            rebalancePlan.getRebalanceStatus()));
+                            rebalanceTask ->
+                                    eventManager.put(new RecoverRebalanceEvent(rebalanceTask)));
         } catch (Exception e) {
             LOG.error(
                     "Failed to get rebalance plan from zookeeper, it will be treated as no"

fluss-server/src/test/java/org/apache/fluss/server/coordinator/rebalance/RebalanceManagerTest.java

@@ -34,6 +34,7 @@
 import org.apache.fluss.server.coordinator.event.CoordinatorEvent;
 import org.apache.fluss.server.coordinator.event.EventManager;
 import org.apache.fluss.server.coordinator.event.RebalanceTaskTimeoutEvent;
+import org.apache.fluss.server.coordinator.event.RecoverRebalanceEvent;
 import org.apache.fluss.server.coordinator.lease.KvSnapshotLeaseManager;
 import org.apache.fluss.server.coordinator.remote.RemoteDirDynamicLoader;
 import org.apache.fluss.server.metadata.CoordinatorMetadataCache;
@@ -177,6 +178,47 @@ void testRebalanceWithoutTask() throws Exception {
                 .hasValue(new RebalanceTask(rebalanceId, COMPLETED, new HashMap<>()));
     }
 
+    @Test
+    void testStartupQueuesRecoverRebalanceEvent() throws Exception {
+        ManualClock clock = new ManualClock(0L);
+        RecordingEventManager eventManager = new RecordingEventManager();
+        NoOpScheduledExecutor executor = new NoOpScheduledExecutor();
+        CoordinatorEventProcessor eventProcessor =
+                buildCoordinatorEventProcessor(new Configuration());
+
+        Map<TableBucket, RebalancePlanForBucket> plan = createRebalancePlan(2);
+        RebalanceTask rebalanceTask = new RebalanceTask("recover-test", NOT_STARTED, plan);
+        zookeeperClient.registerRebalanceTask(rebalanceTask);
+
+        RebalanceManager manager =
+                new RebalanceManager(
+                        eventProcessor, zookeeperClient, eventManager, clock, executor);
+        // If startup() finds a pending rebalance task in ZooKeeper, it should enqueue a
+        // RecoverRebalanceEvent to be processed by the coordinator event thread, instead of
+        // calling registerRebalance() directly on the startup thread.
+        manager.startup();
+
+        assertThat(eventManager.events).hasSize(1);
+        assertThat(eventManager.events.get(0)).isInstanceOf(RecoverRebalanceEvent.class);
+
+        RecoverRebalanceEvent recoverEvent = (RecoverRebalanceEvent) eventManager.events.get(0);
+        assertThat(recoverEvent.getRebalanceTask()).isEqualTo(rebalanceTask);
+
+        manager.close();
+    }
+
+    private Map<TableBucket, RebalancePlanForBucket> createRebalancePlan(int taskCount) {
+        Map<TableBucket, RebalancePlanForBucket> plan = new HashMap<>();
+        for (int i = 0; i < taskCount; i++) {
+            TableBucket tb = new TableBucket(1L, i);
+            plan.put(
+                    tb,
+                    new RebalancePlanForBucket(
+                            tb, 0, 0, Arrays.asList(0, 1, 2), Arrays.asList(0, 1, 2)));
+        }
+        return plan;
+    }
+
     @Test
     void testTimeoutEnqueuesEvent() throws Exception {
         ManualClock clock = new ManualClock(0L);