[server] Introduce new rebalance API

Author: swuferhong

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

@@ -22,7 +22,11 @@
 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.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 +36,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;
@@ -53,9 +62,11 @@
 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;
+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.
+     */
+    CompleteFuture<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.
+     */
+    CompleteFuture<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,9 @@
 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.exception.LeaderNotAvailableException;
 import org.apache.fluss.metadata.DatabaseDescriptor;
 import org.apache.fluss.metadata.DatabaseInfo;
@@ -70,6 +73,7 @@
 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;
@@ -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 CompleteFuture<Map<TableBucket, RebalancePlanForBucket>> rebalance(
+            List<GoalType> priorityGoals, boolean dryRun) {
+        throw new UnsupportedOperationException("Support soon");
+    }
+
+    @Override
+    public CompleteFuture<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-client/src/main/java/org/apache/fluss/client/admin/RebalanceResultForBucket.java

@@ -0,0 +1,69 @@
+/*
+ * 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.client.admin;
+
+import org.apache.fluss.annotation.PublicEvolving;
+import org.apache.fluss.cluster.maintencance.RebalanceStatusForBucket;
+
+import java.util.List;
+
+/**
+ * Result of rebalance process for a tabletBucket.
+ *
+ * @since 0.8
+ */
+@PublicEvolving
+public class RebalanceResultForBucket {
+
+    private final List<Integer> originReplicas;
+    private final List<Integer> newReplicas;
+    private final RebalanceStatusForBucket rebalanceStatusForBucket;
+
+    public RebalanceResultForBucket(
+            List<Integer> originReplicas,
+            List<Integer> newReplicas,
+            RebalanceStatusForBucket rebalanceStatusForBucket) {
+        this.originReplicas = originReplicas;
+        this.newReplicas = newReplicas;
+        this.rebalanceStatusForBucket = rebalanceStatusForBucket;
+    }
+
+    public List<Integer> replicas() {
+        return originReplicas;
+    }
+
+    public List<Integer> newReplicas() {
+        return newReplicas;
+    }
+
+    public RebalanceStatusForBucket rebalanceStatus() {
+        return rebalanceStatusForBucket;
+    }
+
+    @Override
+    public String toString() {
+        return "BucketReassignment{"
+                + "replicas="
+                + originReplicas
+                + ", newReplicas="
+                + newReplicas
+                + ", rebalanceStatus="
+                + rebalanceStatusForBucket
+                + '}';
+    }
+}

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

@@ -0,0 +1,51 @@
+/*
+ * 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.maintencance;
+
+import org.apache.fluss.annotation.PublicEvolving;
+
+import java.util.Arrays;
+
+/**
+ * The type of goal to optimize.
+ *
+ * @since 0.8
+ */
+@PublicEvolving
+public enum GoalType {
+    REPLICA_DISTRIBUTION_GOAL(0),
+    PREFERRED_LEADER_GOAL(1);
+
+    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 == 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())));
+        }
+    }
+}