[server] Support generate and execute reblance plan

Author: swuferhong

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

@@ -22,9 +22,10 @@
 import org.apache.fluss.client.metadata.KvSnapshots;
 import org.apache.fluss.client.metadata.LakeSnapshot;
 import org.apache.fluss.cluster.ServerNode;
-import org.apache.fluss.cluster.maintencance.GoalType;
-import org.apache.fluss.cluster.maintencance.RebalancePlanForBucket;
-import org.apache.fluss.cluster.maintencance.ServerTag;
+import org.apache.fluss.cluster.rebalance.GoalType;
+import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket;
+import org.apache.fluss.cluster.rebalance.RebalanceResultForBucket;
+import org.apache.fluss.cluster.rebalance.ServerTag;
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.exception.AuthorizationException;
 import org.apache.fluss.exception.DatabaseAlreadyExistException;
@@ -62,7 +63,6 @@
 import org.apache.fluss.metadata.TablePath;
 import org.apache.fluss.security.acl.AclBinding;
 import org.apache.fluss.security.acl.AclBindingFilter;
-import org.apache.fluss.shaded.netty4.io.netty.util.concurrent.CompleteFuture;
 
 import java.util.Collection;
 import java.util.List;
@@ -523,7 +523,7 @@ ListOffsetsResult listOffsets(
      *     it.
      * @return the generated rebalance plan for all the tableBuckets which need to do rebalance.
      */
-    CompleteFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
+    CompletableFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
             List<GoalType> priorityGoals, boolean dryRun);
 
     /**
@@ -537,7 +537,7 @@ CompleteFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
      *
      * @return the rebalance process for all the tableBuckets doing rebalance.
      */
-    CompleteFuture<Map<TableBucket, RebalanceResultForBucket>> listRebalanceProcess();
+    CompletableFuture<Map<TableBucket, RebalanceResultForBucket>> listRebalanceProcess();
 
     /**
      * Cannel the rebalance task.

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

@@ -24,9 +24,10 @@
 import org.apache.fluss.client.utils.ClientRpcMessageUtils;
 import org.apache.fluss.cluster.Cluster;
 import org.apache.fluss.cluster.ServerNode;
-import org.apache.fluss.cluster.maintencance.GoalType;
-import org.apache.fluss.cluster.maintencance.RebalancePlanForBucket;
-import org.apache.fluss.cluster.maintencance.ServerTag;
+import org.apache.fluss.cluster.rebalance.GoalType;
+import org.apache.fluss.cluster.rebalance.RebalancePlanForBucket;
+import org.apache.fluss.cluster.rebalance.RebalanceResultForBucket;
+import org.apache.fluss.cluster.rebalance.ServerTag;
 import org.apache.fluss.exception.LeaderNotAvailableException;
 import org.apache.fluss.metadata.DatabaseDescriptor;
 import org.apache.fluss.metadata.DatabaseInfo;
@@ -45,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;
@@ -69,13 +71,13 @@
 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;
 import org.apache.fluss.rpc.protocol.ApiError;
 import org.apache.fluss.security.acl.AclBinding;
 import org.apache.fluss.security.acl.AclBindingFilter;
-import org.apache.fluss.shaded.netty4.io.netty.util.concurrent.CompleteFuture;
 import org.apache.fluss.utils.MapUtils;
 
 import javax.annotation.Nullable;
@@ -486,19 +488,22 @@ public CompletableFuture<Void> removeServerTag(
     }
 
     @Override
-    public CompleteFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
+    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
-    public CompleteFuture<Map<TableBucket, RebalanceResultForBucket>> listRebalanceProcess() {
+    public CompletableFuture<Map<TableBucket, RebalanceResultForBucket>> listRebalanceProcess() {
         throw new UnsupportedOperationException("Support soon");
     }
 
     @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/admin/RebalanceResultForBucket.java

@@ -18,7 +18,7 @@
 package org.apache.fluss.client.admin;
 
 import org.apache.fluss.annotation.PublicEvolving;
-import org.apache.fluss.cluster.maintencance.RebalanceStatusForBucket;
+import org.apache.fluss.cluster.rebalance.RebalanceStatusForBucket;
 
 import java.util.List;
 

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-client/src/test/java/org/apache/fluss/client/admin/FlussAdminITCase.java

@@ -24,7 +24,7 @@
 import org.apache.fluss.client.table.Table;
 import org.apache.fluss.client.table.writer.UpsertWriter;
 import org.apache.fluss.cluster.ServerNode;
-import org.apache.fluss.cluster.maintencance.ServerTag;
+import org.apache.fluss.cluster.rebalance.ServerTag;
 import org.apache.fluss.config.AutoPartitionTimeUnit;
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.Configuration;

…fluss/cluster/maintencance/GoalType.java …he/fluss/cluster/rebalance/GoalType.javafluss-common/src/main/java/org/apache/fluss/cluster/maintencance/GoalType.java renamed to fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java fluss-common/src/main/java/org/apache/fluss/cluster/maintencance/GoalType.java renamed to fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/GoalType.java

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-package org.apache.fluss.cluster.maintencance;
+package org.apache.fluss.cluster.rebalance;
 
 import org.apache.fluss.annotation.PublicEvolving;
 
@@ -28,8 +28,20 @@
  */
 @PublicEvolving
 public enum GoalType {
+    /**
+     * Goal to generate replica movement tasks to ensure that the number of replicas on each
+     * tabletServer is near balanced.
+     */
     REPLICA_DISTRIBUTION_GOAL(0),
-    PREFERRED_LEADER_GOAL(1);
+
+    /**
+     * Goal to generate leadership movement and leader replica movement tasks to ensure that the
+     * number of leader replicas on each tabletServer is near balanced.
+     */
+    LEADER_REPLICA_DISTRIBUTION_GOAL(1),
+
+    /** Goal to move the leaders to the first replica of each tableBuckets. */
+    PREFERRED_LEADER_GOAL(2);
 
     public final int value;
 
@@ -40,6 +52,8 @@ public enum GoalType {
     public static GoalType valueOf(int value) {
         if (value == REPLICA_DISTRIBUTION_GOAL.value) {
             return REPLICA_DISTRIBUTION_GOAL;
+        } else if (value == LEADER_REPLICA_DISTRIBUTION_GOAL.value) {
+            return LEADER_REPLICA_DISTRIBUTION_GOAL;
         } else if (value == PREFERRED_LEADER_GOAL.value) {
             return PREFERRED_LEADER_GOAL;
         } else {

…maintencance/RebalancePlanForBucket.java …er/rebalance/RebalancePlanForBucket.javafluss-common/src/main/java/org/apache/fluss/cluster/maintencance/RebalancePlanForBucket.java renamed to fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalancePlanForBucket.java fluss-common/src/main/java/org/apache/fluss/cluster/maintencance/RebalancePlanForBucket.java renamed to fluss-common/src/main/java/org/apache/fluss/cluster/rebalance/RebalancePlanForBucket.java

@@ -15,41 +15,46 @@
  * limitations under the License.
  */
 
-package org.apache.fluss.cluster.maintencance;
+package org.apache.fluss.cluster.rebalance;
 
 import org.apache.fluss.annotation.PublicEvolving;
+import org.apache.fluss.metadata.TableBucket;
 
 import java.util.List;
 import java.util.Objects;
 
 /**
- * a Generated rebalance plan for a tableBucket.
+ * a generated rebalance plan for a tableBucket.
  *
  * @since 0.8
  */
 @PublicEvolving
 public class RebalancePlanForBucket {
-    private final int bucketId;
+    private final TableBucket tableBucket;
     private final int originalLeader;
     private final int newLeader;
     private final List<Integer> originReplicas;
     private final List<Integer> newReplicas;
 
     public RebalancePlanForBucket(
-            int bucketId,
+            TableBucket tableBucket,
             int originalLeader,
             int newLeader,
             List<Integer> originReplicas,
             List<Integer> newReplicas) {
-        this.bucketId = bucketId;
+        this.tableBucket = tableBucket;
         this.originalLeader = originalLeader;
         this.newLeader = newLeader;
         this.originReplicas = originReplicas;
         this.newReplicas = newReplicas;
     }
 
+    public TableBucket getTableBucket() {
+        return tableBucket;
+    }
+
     public int getBucketId() {
-        return bucketId;
+        return tableBucket.getBucket();
     }
 
     public Integer getOriginalLeader() {
@@ -68,11 +73,15 @@ public List<Integer> getNewReplicas() {
         return newReplicas;
     }
 
+    public boolean isLeaderAction() {
+        return originalLeader != newLeader;
+    }
+
     @Override
     public String toString() {
         return "RebalancePlanForBucket{"
-                + "bucketId="
-                + bucketId
+                + "tableBucket="
+                + tableBucket
                 + ", originalLeader="
                 + originalLeader
                 + ", newLeader="
@@ -93,7 +102,7 @@ public boolean equals(Object o) {
             return false;
         }
         RebalancePlanForBucket that = (RebalancePlanForBucket) o;
-        return bucketId == that.bucketId
+        return Objects.equals(tableBucket, that.tableBucket)
                 && originalLeader == that.originalLeader
                 && newLeader == that.newLeader
                 && Objects.equals(originReplicas, that.originReplicas)
@@ -102,6 +111,6 @@ public boolean equals(Object o) {
 
     @Override
     public int hashCode() {
-        return Objects.hash(bucketId, originalLeader, newLeader, originReplicas, newReplicas);
+        return Objects.hash(tableBucket, originalLeader, newLeader, originReplicas, newReplicas);
     }
 }