[server] Support generate and execute reblance plan (#1452)

Author: swuferhong

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

@@ -23,6 +23,7 @@
 import org.apache.fluss.client.metadata.LakeSnapshot;
 import org.apache.fluss.cluster.ServerNode;
 import org.apache.fluss.cluster.rebalance.GoalType;
+import org.apache.fluss.cluster.rebalance.RebalanceProgress;
 import org.apache.fluss.cluster.rebalance.ServerTag;
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.cluster.AlterConfig;
@@ -66,8 +67,11 @@
 import org.apache.fluss.security.acl.AclBinding;
 import org.apache.fluss.security.acl.AclBindingFilter;
 
+import javax.annotation.Nullable;
+
 import java.util.Collection;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -551,43 +555,53 @@ ListOffsetsResult listOffsets(
      * balancing according to the user-defined {@code priorityGoals}.
      *
      * <p>Currently, Fluss only supports one active rebalance task in the cluster. If an uncompleted
-     * rebalance task exists, an {@link RebalanceFailureException} will be thrown.
+     * rebalance task exists, Fluss will return the uncompleted rebalance task's progress.
+     *
+     * <p>If you want to cancel the rebalance task, you can use {@link #cancelRebalance(String)}
      *
      * <ul>
      *   <li>{@link AuthorizationException} If the authenticated user doesn't have cluster
      *       permissions.
-     *   <li>{@link RebalanceFailureException} If the rebalance failed. Such as there is an ongoing
-     *       execution.
+     *   <li>{@link RebalanceFailureException} If the rebalance failed. Such as there is an
+     *       inProgress execution.
      * </ul>
      *
      * @param priorityGoals the goals to be optimized.
-     * @param dryRun Calculate and return the rebalance optimization proposal, but do not execute
-     *     it.
-     * @return the generated rebalance plan for all the tableBuckets which need to do rebalance.
+     * @return the rebalance id. If there is no rebalance task in progress, it will trigger a new
+     *     rebalance task and return the rebalance id.
      */
-    CompletableFuture<RebalancePlan> rebalance(List<GoalType> priorityGoals, boolean dryRun);
+    CompletableFuture<String> rebalance(List<GoalType> priorityGoals);
 
     /**
      * List the rebalance progress.
      *
      * <ul>
      *   <li>{@link AuthorizationException} If the authenticated user doesn't have cluster
      *       permissions.
-     *   <li>{@link NoRebalanceInProgressException} If there are no rebalance tasks in progress.
+     *   <li>{@link NoRebalanceInProgressException} If there are no rebalance tasks in progress for
+     *       the input rebalanceId.
      * </ul>
      *
+     * @param rebalanceId the rebalance id to list progress, if it is null means list the in
+     *     progress rebalance task's.
      * @return the rebalance process.
      */
-    CompletableFuture<RebalanceProgress> listRebalanceProgress();
+    CompletableFuture<Optional<RebalanceProgress>> listRebalanceProgress(
+            @Nullable String rebalanceId);
 
     /**
      * Cannel the rebalance task.
      *
      * <ul>
      *   <li>{@link AuthorizationException} If the authenticated user doesn't have cluster
      *       permissions.
-     *   <li>{@link NoRebalanceInProgressException} If there are no rebalance tasks in progress.
+     *   <li>{@link NoRebalanceInProgressException} If there are no rebalance tasks in progress or
+     *       the rebalance id is not exists.
      * </ul>
+     *
+     * @param rebalanceId the rebalance id to cancel, if it is null means cancel the exists
+     *     rebalance task. If rebalanceId is not exists in server, {@link
+     *     NoRebalanceInProgressException} will be thrown.
      */
-    CompletableFuture<Void> cancelRebalance();
+    CompletableFuture<Void> cancelRebalance(@Nullable String rebalanceId);
 }

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

@@ -25,6 +25,7 @@
 import org.apache.fluss.cluster.Cluster;
 import org.apache.fluss.cluster.ServerNode;
 import org.apache.fluss.cluster.rebalance.GoalType;
+import org.apache.fluss.cluster.rebalance.RebalanceProgress;
 import org.apache.fluss.cluster.rebalance.ServerTag;
 import org.apache.fluss.config.cluster.AlterConfig;
 import org.apache.fluss.config.cluster.ConfigEntry;
@@ -49,6 +50,7 @@
 import org.apache.fluss.rpc.messages.AddServerTagRequest;
 import org.apache.fluss.rpc.messages.AlterClusterConfigsRequest;
 import org.apache.fluss.rpc.messages.AlterTableRequest;
+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,12 +71,15 @@
 import org.apache.fluss.rpc.messages.ListDatabasesResponse;
 import org.apache.fluss.rpc.messages.ListOffsetsRequest;
 import org.apache.fluss.rpc.messages.ListPartitionInfosRequest;
+import org.apache.fluss.rpc.messages.ListRebalanceProgressRequest;
 import org.apache.fluss.rpc.messages.ListTablesRequest;
 import org.apache.fluss.rpc.messages.ListTablesResponse;
 import org.apache.fluss.rpc.messages.PbAlterConfig;
 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.RebalanceResponse;
 import org.apache.fluss.rpc.messages.RemoveServerTagRequest;
 import org.apache.fluss.rpc.messages.TableExistsRequest;
 import org.apache.fluss.rpc.messages.TableExistsResponse;
@@ -91,6 +96,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 
 import static org.apache.fluss.client.utils.ClientRpcMessageUtils.makeAlterTableRequest;
@@ -551,19 +557,34 @@ public CompletableFuture<Void> removeServerTag(
     }
 
     @Override
-    public CompletableFuture<RebalancePlan> rebalance(
-            List<GoalType> priorityGoals, boolean dryRun) {
-        throw new UnsupportedOperationException("Support soon");
+    public CompletableFuture<String> rebalance(List<GoalType> priorityGoals) {
+        RebalanceRequest request = new RebalanceRequest();
+        priorityGoals.forEach(goal -> request.addGoal(goal.value));
+        return gateway.rebalance(request).thenApply(RebalanceResponse::getRebalanceId);
     }
 
     @Override
-    public CompletableFuture<RebalanceProgress> listRebalanceProgress() {
-        throw new UnsupportedOperationException("Support soon");
+    public CompletableFuture<Optional<RebalanceProgress>> listRebalanceProgress(
+            @Nullable String rebalanceId) {
+        ListRebalanceProgressRequest request = new ListRebalanceProgressRequest();
+
+        if (rebalanceId != null) {
+            request.setRebalanceId(rebalanceId);
+        }
+
+        return gateway.listRebalanceProgress(request)
+                .thenApply(ClientRpcMessageUtils::toRebalanceProgress);
     }
 
     @Override
-    public CompletableFuture<Void> cancelRebalance() {
-        throw new UnsupportedOperationException("Support soon");
+    public CompletableFuture<Void> cancelRebalance(@Nullable String rebalanceId) {
+        CancelRebalanceRequest request = new CancelRebalanceRequest();
+
+        if (rebalanceId != null) {
+            request.setRebalanceId(rebalanceId);
+        }
+
+        return gateway.cancelRebalance(request).thenApply(r -> null);
     }
 
     @Override

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

@@ -25,6 +25,10 @@
 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.cluster.rebalance.RebalanceProgress;
+import org.apache.fluss.cluster.rebalance.RebalanceResultForBucket;
+import org.apache.fluss.cluster.rebalance.RebalanceStatus;
 import org.apache.fluss.config.cluster.AlterConfigOpType;
 import org.apache.fluss.config.cluster.ColumnPositionType;
 import org.apache.fluss.config.cluster.ConfigEntry;
@@ -46,6 +50,7 @@
 import org.apache.fluss.rpc.messages.GetLatestLakeSnapshotResponse;
 import org.apache.fluss.rpc.messages.ListOffsetsRequest;
 import org.apache.fluss.rpc.messages.ListPartitionInfosResponse;
+import org.apache.fluss.rpc.messages.ListRebalanceProgressResponse;
 import org.apache.fluss.rpc.messages.LookupRequest;
 import org.apache.fluss.rpc.messages.MetadataRequest;
 import org.apache.fluss.rpc.messages.PbAddColumn;
@@ -61,6 +66,9 @@
 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.PbRebalanceProgressForBucket;
+import org.apache.fluss.rpc.messages.PbRebalanceProgressForTable;
 import org.apache.fluss.rpc.messages.PbRemotePathAndLocalFile;
 import org.apache.fluss.rpc.messages.PbRenameColumn;
 import org.apache.fluss.rpc.messages.PrefixLookupRequest;
@@ -77,10 +85,13 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import static org.apache.fluss.cluster.rebalance.RebalanceStatus.FINAL_STATUSES;
 import static org.apache.fluss.rpc.util.CommonRpcMessageUtils.toResolvedPartitionSpec;
+import static org.apache.fluss.utils.Preconditions.checkArgument;
 import static org.apache.fluss.utils.Preconditions.checkState;
 
 /**
@@ -370,6 +381,65 @@ public static AlterTableRequest makeAlterTableRequest(
         return request;
     }
 
+    public static Optional<RebalanceProgress> toRebalanceProgress(
+            ListRebalanceProgressResponse response) {
+        if (!response.hasRebalanceId()) {
+            return Optional.empty();
+        }
+
+        checkArgument(response.hasRebalanceStatus(), "Rebalance status is not set");
+        RebalanceStatus totalRebalanceStatus = RebalanceStatus.of(response.getRebalanceStatus());
+        int totalTask = 0;
+        int finishedTask = 0;
+        Map<TableBucket, RebalanceResultForBucket> rebalanceProgress = new HashMap<>();
+        for (PbRebalanceProgressForTable pbTable : response.getTableProgressesList()) {
+            long tableId = pbTable.getTableId();
+            for (PbRebalanceProgressForBucket pbBucket : pbTable.getBucketsProgressesList()) {
+                RebalanceStatus bucketStatus = RebalanceStatus.of(pbBucket.getRebalanceStatus());
+                RebalancePlanForBucket planForBucket =
+                        toRebalancePlanForBucket(tableId, pbBucket.getRebalancePlan());
+                rebalanceProgress.put(
+                        planForBucket.getTableBucket(),
+                        new RebalanceResultForBucket(planForBucket, bucketStatus));
+                if (FINAL_STATUSES.contains(bucketStatus)) {
+                    finishedTask++;
+                }
+                totalTask++;
+            }
+        }
+
+        // For these rebalance task without only bucket level rebalance tasks, we return -1 as
+        // progress.
+        double progress = -1d;
+        if (totalTask != 0) {
+            progress = (double) finishedTask / totalTask;
+        }
+
+        return Optional.of(
+                new RebalanceProgress(
+                        response.getRebalanceId(),
+                        totalRebalanceStatus,
+                        progress,
+                        rebalanceProgress));
+    }
+
+    private static RebalancePlanForBucket toRebalancePlanForBucket(
+            long tableId, PbRebalancePlanForBucket rebalancePlan) {
+        TableBucket tableBucket =
+                new TableBucket(
+                        tableId,
+                        rebalancePlan.hasPartitionId() ? rebalancePlan.getPartitionId() : null,
+                        rebalancePlan.getBucketId());
+        return new RebalancePlanForBucket(
+                tableBucket,
+                rebalancePlan.getOriginalLeader(),
+                rebalancePlan.getNewLeader(),
+                Arrays.stream(rebalancePlan.getOriginalReplicas())
+                        .boxed()
+                        .collect(Collectors.toList()),
+                Arrays.stream(rebalancePlan.getNewReplicas()).boxed().collect(Collectors.toList()));
+    }
+
     public static List<PartitionInfo> toPartitionInfos(ListPartitionInfosResponse response) {
         return response.getPartitionsInfosList().stream()
                 .map(