[server] Introduce new rebalance API

Author: swuferhong

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

@@ -22,7 +22,12 @@
 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.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;
 import org.apache.fluss.exception.DatabaseNotEmptyException;
 import org.apache.fluss.exception.DatabaseNotExistException;
@@ -32,10 +37,15 @@
 import org.apache.fluss.exception.InvalidTableException;
 import org.apache.fluss.exception.KvSnapshotNotExistException;
 import org.apache.fluss.exception.LakeTableSnapshotNotExistException;
+import org.apache.fluss.exception.NoRebalanceInProgressException;
 import org.apache.fluss.exception.NonPrimaryKeyTableException;
 import org.apache.fluss.exception.PartitionAlreadyExistsException;
 import org.apache.fluss.exception.PartitionNotExistException;
+import org.apache.fluss.exception.RebalanceFailureException;
 import org.apache.fluss.exception.SchemaNotExistException;
+import org.apache.fluss.exception.ServerNotExistException;
+import org.apache.fluss.exception.ServerTagAlreadyExistException;
+import org.apache.fluss.exception.ServerTagNotExistException;
 import org.apache.fluss.exception.TableAlreadyExistException;
 import org.apache.fluss.exception.TableNotExistException;
 import org.apache.fluss.exception.TableNotPartitionedException;
@@ -56,6 +66,7 @@
 
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 
 /**
@@ -452,4 +463,90 @@ ListOffsetsResult listOffsets(
      * @return A CompletableFuture indicating completion of the operation.
      */
     DropAclsResult dropAcls(Collection<AclBindingFilter> filters);
+
+    /**
+     * Add server tag to the specified tabletServers, one tabletServer can only have one serverTag.
+     *
+     * <p>If one tabletServer failed adding tag, none of the tags will take effect.
+     *
+     * <ul>
+     *   <li>{@link AuthorizationException} If the authenticated user doesn't have reset config
+     *       access to the cluster.
+     *   <li>{@link ServerNotExistException} If the tabletServer in {@code tabletServers} does not
+     *       exist.
+     *   <li>{@link ServerTagAlreadyExistException} If the server tag already exists when {@code
+     *       overWriteIfExists} is false.
+     * </ul>
+     *
+     * @param tabletServers the tabletServers we want to add server tags.
+     * @param serverTag the server tag to be added.
+     */
+    CompletableFuture<Void> addServerTag(List<Integer> tabletServers, ServerTag serverTag);
+
+    /**
+     * Remove server tag from the specified tabletServers.
+     *
+     * <p>If one tabletServer failed removing tag, none of the tags will be removed.
+     *
+     * <ul>
+     *   <li>{@link AuthorizationException} If the authenticated user doesn't have reset config
+     *       access to the cluster.
+     *   <li>{@link ServerNotExistException} If the tabletServer in {@code tabletServers} does not
+     *       exist.
+     *   <li>{@link ServerTagNotExistException} If the server tag does not exist when {@code
+     *       overWriteIfExists} is false.
+     * </ul>
+     *
+     * @param tabletServers the tabletServers we want to remove server tags.
+     */
+    CompletableFuture<Void> removeServerTag(List<Integer> tabletServers, ServerTag serverTag);
+
+    /**
+     * Based on the provided {@code priorityGoals}, Fluss performs load balancing on the cluster's
+     * bucket load.
+     *
+     * <p>More details, Fluss collects the cluster's load information and optimizes to perform load
+     * 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.
+     *
+     * <ul>
+     *   <li>{@link AuthorizationException} If the authenticated user doesn't have reset config
+     *       access to the cluster.
+     *   <li>{@link RebalanceFailureException} If the rebalance failed. Such as there is an ongoing
+     *       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.
+     */
+    CompletableFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
+            List<GoalType> priorityGoals, boolean dryRun);
+
+    /**
+     * List the rebalance process.
+     *
+     * <ul>
+     *   <li>{@link AuthorizationException} If the authenticated user doesn't have reset config
+     *       access to the cluster.
+     *   <li>{@link NoRebalanceInProgressException} If there are no rebalance tasks in progress.
+     * </ul>
+     *
+     * @return the rebalance process for all the tableBuckets doing rebalance.
+     */
+    CompletableFuture<Map<TableBucket, RebalanceResultForBucket>> listRebalanceProcess();
+
+    /**
+     * Cannel the rebalance task.
+     *
+     * <ul>
+     *   <li>{@link AuthorizationException} If the authenticated user doesn't have reset config
+     *       access to the cluster.
+     *   <li>{@link NoRebalanceInProgressException} If there are no rebalance tasks in progress.
+     * </ul>
+     */
+    CompletableFuture<Void> cancelRebalance();
 }

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

@@ -24,6 +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.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;
@@ -464,6 +468,33 @@ public DropAclsResult dropAcls(Collection<AclBindingFilter> filters) {
         return result;
     }
 
+    @Override
+    public CompletableFuture<Void> addServerTag(List<Integer> tabletServers, ServerTag serverTag) {
+        throw new UnsupportedOperationException("Support soon");
+    }
+
+    @Override
+    public CompletableFuture<Void> removeServerTag(
+            List<Integer> tabletServers, ServerTag serverTag) {
+        throw new UnsupportedOperationException("Support soon");
+    }
+
+    @Override
+    public CompletableFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
+            List<GoalType> priorityGoals, boolean dryRun) {
+        throw new UnsupportedOperationException("Support soon");
+    }
+
+    @Override
+    public CompletableFuture<Map<TableBucket, RebalanceResultForBucket>> listRebalanceProcess() {
+        throw new UnsupportedOperationException("Support soon");
+    }
+
+    @Override
+    public CompletableFuture<Void> cancelRebalance() {
+        throw new UnsupportedOperationException("Support soon");
+    }
+
     @Override
     public void close() {
         // nothing to do yet

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

@@ -0,0 +1,65 @@
+/*
+ * 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.cluster.rebalance;
+
+import org.apache.fluss.annotation.PublicEvolving;
+
+import java.util.Arrays;
+
+/**
+ * The type of goal to optimize.
+ *
+ * @since 0.8
+ */
+@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),
+
+    /**
+     * 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;
+
+    GoalType(int value) {
+        this.value = value;
+    }
+
+    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 {
+            throw new IllegalArgumentException(
+                    String.format(
+                            "Value %s must be one of %s", value, Arrays.asList(GoalType.values())));
+        }
+    }
+}

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

@@ -0,0 +1,112 @@
+/*
+ * 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.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.
+ *
+ * @since 0.8
+ */
+@PublicEvolving
+public class RebalancePlanForBucket {
+    private final TableBucket tableBucket;
+    private final int originalLeader;
+    private final int newLeader;
+    private final List<Integer> originReplicas;
+    private final List<Integer> newReplicas;
+
+    public RebalancePlanForBucket(
+            TableBucket tableBucket,
+            int originalLeader,
+            int newLeader,
+            List<Integer> originReplicas,
+            List<Integer> newReplicas) {
+        this.tableBucket = tableBucket;
+        this.originalLeader = originalLeader;
+        this.newLeader = newLeader;
+        this.originReplicas = originReplicas;
+        this.newReplicas = newReplicas;
+    }
+
+    public TableBucket getTableBucket() {
+        return tableBucket;
+    }
+
+    public int getBucketId() {
+        return tableBucket.getBucket();
+    }
+
+    public Integer getOriginalLeader() {
+        return originalLeader;
+    }
+
+    public Integer getNewLeader() {
+        return newLeader;
+    }
+
+    public List<Integer> getOriginReplicas() {
+        return originReplicas;
+    }
+
+    public List<Integer> getNewReplicas() {
+        return newReplicas;
+    }
+
+    @Override
+    public String toString() {
+        return "RebalancePlanForBucket{"
+                + "tableBucket="
+                + tableBucket
+                + ", originalLeader="
+                + originalLeader
+                + ", newLeader="
+                + newLeader
+                + ", originReplicas="
+                + originReplicas
+                + ", newReplicas="
+                + newReplicas
+                + '}';
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        RebalancePlanForBucket that = (RebalancePlanForBucket) o;
+        return Objects.equals(tableBucket, that.tableBucket)
+                && originalLeader == that.originalLeader
+                && newLeader == that.newLeader
+                && Objects.equals(originReplicas, that.originReplicas)
+                && Objects.equals(newReplicas, that.newReplicas);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(tableBucket, originalLeader, newLeader, originReplicas, newReplicas);
+    }
+}