[server] Support generate and execute reblance plan

Author: swuferhong

fluss-client/src/main/java/org/apache/fluss/client/admin/FlussAdmin.java

@@ -46,6 +46,7 @@
 import org.apache.fluss.rpc.gateway.AdminReadOnlyGateway;
 import org.apache.fluss.rpc.gateway.TabletServerGateway;
 import org.apache.fluss.rpc.messages.AddServerTagRequest;
+import org.apache.fluss.rpc.messages.CancelRebalanceRequest;
 import org.apache.fluss.rpc.messages.CreateAclsRequest;
 import org.apache.fluss.rpc.messages.CreateDatabaseRequest;
 import org.apache.fluss.rpc.messages.CreateTableRequest;
@@ -70,6 +71,7 @@
 import org.apache.fluss.rpc.messages.PbListOffsetsRespForBucket;
 import org.apache.fluss.rpc.messages.PbPartitionSpec;
 import org.apache.fluss.rpc.messages.PbTablePath;
+import org.apache.fluss.rpc.messages.RebalanceRequest;
 import org.apache.fluss.rpc.messages.RemoveServerTagRequest;
 import org.apache.fluss.rpc.messages.TableExistsRequest;
 import org.apache.fluss.rpc.messages.TableExistsResponse;
@@ -488,7 +490,9 @@ public CompletableFuture<Void> removeServerTag(
     @Override
     public CompletableFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
             List<GoalType> priorityGoals, boolean dryRun) {
-        throw new UnsupportedOperationException("Support soon");
+        RebalanceRequest request = new RebalanceRequest().setDryRun(dryRun);
+        priorityGoals.forEach(goal -> request.addGoal(goal.value));
+        return gateway.rebalance(request).thenApply(ClientRpcMessageUtils::toRebalancePlan);
     }
 
     @Override
@@ -498,7 +502,8 @@ public CompletableFuture<Map<TableBucket, RebalanceResultForBucket>> listRebalan
 
     @Override
     public CompletableFuture<Void> cancelRebalance() {
-        throw new UnsupportedOperationException("Support soon");
+        CancelRebalanceRequest request = new CancelRebalanceRequest();
+        return gateway.cancelRebalance(request).thenApply(r -> null);
     }
 
     @Override

fluss-client/src/main/java/org/apache/fluss/client/utils/ClientRpcMessageUtils.java

@@ -25,6 +25,7 @@
 import org.apache.fluss.client.metadata.LakeSnapshot;
 import org.apache.fluss.client.write.KvWriteBatch;
 import org.apache.fluss.client.write.ReadyWriteBatch;
+import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket;
 import org.apache.fluss.fs.FsPath;
 import org.apache.fluss.fs.FsPathAndFileName;
 import org.apache.fluss.fs.token.ObtainedSecurityToken;
@@ -51,10 +52,14 @@
 import org.apache.fluss.rpc.messages.PbPrefixLookupReqForBucket;
 import org.apache.fluss.rpc.messages.PbProduceLogReqForBucket;
 import org.apache.fluss.rpc.messages.PbPutKvReqForBucket;
+import org.apache.fluss.rpc.messages.PbRebalancePlanForBucket;
+import org.apache.fluss.rpc.messages.PbRebalancePlanForPartition;
+import org.apache.fluss.rpc.messages.PbRebalancePlanForTable;
 import org.apache.fluss.rpc.messages.PbRemotePathAndLocalFile;
 import org.apache.fluss.rpc.messages.PrefixLookupRequest;
 import org.apache.fluss.rpc.messages.ProduceLogRequest;
 import org.apache.fluss.rpc.messages.PutKvRequest;
+import org.apache.fluss.rpc.messages.RebalanceResponse;
 
 import javax.annotation.Nullable;
 
@@ -323,6 +328,48 @@ public static DropPartitionRequest makeDropPartitionRequest(
         return dropPartitionRequest;
     }
 
+    public static Map<TableBucket, RebalancePlanForBucket> toRebalancePlan(
+            RebalanceResponse response) {
+        Map<TableBucket, RebalancePlanForBucket> rebalancePlan = new HashMap<>();
+        for (PbRebalancePlanForTable pbTable : response.getPlanForTablesList()) {
+            long tableId = pbTable.getTableId();
+            if (pbTable.getPartitionsPlansCount() == 0) {
+                // none-partition table.
+                for (PbRebalancePlanForBucket pbBucket : pbTable.getBucketsPlansList()) {
+                    int bucketId = pbBucket.getBucketId();
+                    rebalancePlan.put(
+                            new TableBucket(tableId, null, bucketId),
+                            toRebalancePlanForBucket(tableId, null, bucketId, pbBucket));
+                }
+            } else {
+                // partition table.
+                for (PbRebalancePlanForPartition pbPartition : pbTable.getPartitionsPlansList()) {
+                    long partitionId = pbPartition.getPartitionId();
+                    for (PbRebalancePlanForBucket pbBucket : pbPartition.getBucketsPlansList()) {
+                        int bucketId = pbBucket.getBucketId();
+                        rebalancePlan.put(
+                                new TableBucket(tableId, partitionId, bucketId),
+                                toRebalancePlanForBucket(tableId, partitionId, bucketId, pbBucket));
+                    }
+                }
+            }
+        }
+        return rebalancePlan;
+    }
+
+    private static RebalancePlanForBucket toRebalancePlanForBucket(
+            long tableId,
+            @Nullable Long partitionId,
+            int bucketId,
+            PbRebalancePlanForBucket pbBucket) {
+        return new RebalancePlanForBucket(
+                new TableBucket(tableId, partitionId, bucketId),
+                pbBucket.getOriginalLeader(),
+                pbBucket.getNewLeader(),
+                Arrays.stream(pbBucket.getOriginalReplicas()).boxed().collect(Collectors.toList()),
+                Arrays.stream(pbBucket.getNewReplicas()).boxed().collect(Collectors.toList()));
+    }
+
     public static List<PartitionInfo> toPartitionInfos(ListPartitionInfosResponse response) {
         return response.getPartitionsInfosList().stream()
                 .map(

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

@@ -73,6 +73,10 @@ public List<Integer> getNewReplicas() {
         return newReplicas;
     }
 
+    public boolean isLeaderAction() {
+        return originalLeader != newLeader;
+    }
+
     @Override
     public String toString() {
         return "RebalancePlanForBucket{"

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

@@ -165,7 +165,7 @@ public Map<Long, TablePath> allTables() {
         return tablePathById;
     }
 
-    public Set<TableBucket> allBuckets() {
+    public Set<TableBucket> getAllBuckets() {
         Set<TableBucket> allBuckets = new HashSet<>();
         for (Map.Entry<Long, Map<Integer, List<Integer>>> tableAssign :
                 tableAssignments.entrySet()) {